[Diagrams] Implemented FormDiagram.from_equilibrium_state

CHANGELOG.rst

@@ -11,6 +11,8 @@ Unreleased
 
 **Added**
 
+- Implemented `FormDiagram.from_equilibrium_state` constructor for compatibility with `jax_cem`.
+
 **Changed**
 
 **Fixed**

src/compas_cem/diagrams/form.py

@@ -1,3 +1,6 @@
+from compas_cem.elements import Node
+from compas_cem.elements import Edge
+
 from compas_cem.diagrams import Diagram
 
 
@@ -34,7 +37,7 @@ def __init__(self, *args, **kwargs):
     @classmethod
     def from_topology_diagram(cls, topology):
         """
-        Construct the base of a form diagram from a topology diagram.
+        Construct a form diagram from a topology diagram.
         """
         form = topology.copy(cls=cls)
 
@@ -45,6 +48,66 @@ def from_topology_diagram(cls, topology):
 
         return form
 
+    @classmethod
+    def from_equilibrium_state(cls, eq_state, structure):
+        """
+        Build a form diagram from an equilibrium state.
+        """
+        return form_from_eqstate(eq_state, structure, cls)
+
+
+# ==============================================================================
+# Helpers
+# ==============================================================================
+
+def form_from_eqstate(eqstate, structure, cls=None):
+    """
+    Generate a form diagram from an equilibrium state calculated with JAX CEM.
+    """
+    if cls is None:
+        cls = FormDiagram
+    form = cls()
+
+    # Add nodes
+    for node in structure.nodes:
+        form.add_node(Node(int(node)))
+
+    # Assign support nodes
+    for node in structure.support_nodes:
+        form.node_attribute(int(node), "type", "support")
+
+    # Add edges
+    for u, v in structure.edges:
+        form.add_edge(Edge(int(u), int(v), {}))
+
+    # Update form attributes
+    form_update(form, eqstate, structure)
+
+    return form
+
+
+def form_update(form, eqstate, structure):
+    """
+    Update in-place the attributes of a form diagram with an equilibrium state.
+    """
+    xyz = eqstate.xyz.tolist()
+    loads = eqstate.loads.tolist()
+    reactions = eqstate.reactions.tolist()
+    lengths = eqstate.lengths.tolist()
+    forces = eqstate.forces.tolist()
+
+    # Update q values and lengths on edges
+    for edge in structure.edges:
+        idx = structure.edge_index[tuple(edge)]
+        form.edge_attribute(edge, name="force", value=forces[idx].pop())
+        form.edge_attribute(edge, name="lengths", value=lengths[idx].pop())
+
+    # Update residuals on nodes
+    for node in structure.nodes:
+        form.node_attributes(node, "xyz", xyz[node])
+        form.node_attributes(node, ["rx", "ry", "rz"], reactions[node])
+        form.node_attributes(node, ["qx", "qy", "qz"], loads[node])
+
 # ==============================================================================
 # Main
 # ==============================================================================