diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 2d82c6fb..e4da6335 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -11,6 +11,8 @@ Unreleased **Added** +- Implemented `FormDiagram.from_equilibrium_state` constructor for compatibility with `jax_cem`. + **Changed** **Fixed** diff --git a/src/compas_cem/diagrams/form.py b/src/compas_cem/diagrams/form.py index 3af61fbc..519915ea 100644 --- a/src/compas_cem/diagrams/form.py +++ b/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 # ==============================================================================