Skip to content

Equilibrium¤

The heart of JAX FDM: compute states of static equilibrium with the force density method, differentiably.

Form-finding¤

The two entry points. fdm computes an equilibrium state for fixed parameters; constrained_fdm searches for the parameters whose equilibrium state minimizes a loss function, subject to constraints.

fdm ¤

fdm(datastructure: AnyFDDatastructure, sparse: bool = True, is_load_local: bool = False, tmax: int = 1, eta: float = 1e-06, itersolve_fn: Callable | None = None, iterload_fn: Callable | None = None, implicit_diff: bool = True, verbose: bool = False) -> AnyFDDatastructure

Compute a datastructure in static equilibrium with the force density method.

Parameters:

  • datastructure (AnyFDDatastructure) –

    The network or mesh to form-find.

  • sparse (bool, default: True ) –

    If True, assemble and solve the equilibrium system with a sparse solver.

  • is_load_local (bool, default: False ) –

    If True, apply edge and face loads in their local coordinate systems (follower loads).

  • tmax (int, default: 1 ) –

    The maximum number of iterations. With tmax=1 a single linear FDM step is taken and edge and face loads are ignored.

  • eta (float, default: 1e-06 ) –

    The convergence tolerance for the iterative solve.

  • itersolve_fn (Callable | None, default: None ) –

    The iterative equilibrium solver. If None, forward fixed-point iteration is used.

  • iterload_fn (Callable | None, default: None ) –

    A load callback invoked once before iterative equilibrium starts.

  • implicit_diff (bool, default: True ) –

    If True, apply implicit differentiation for the backward pass.

  • verbose (bool, default: False ) –

    Whether to print calculation info to the terminal.

Returns:

  • datastructure ( AnyFDDatastructure ) –

    A copy of the input datastructure in static equilibrium.

Notes

In the returned copy each node load is the model-aggregated tributary load (node, edge, and face) at equilibrium, while the original edge and face loads are retained. Form-finding the output again therefore redistributes those tributary loads a second time; always solve from the original datastructure.

constrained_fdm ¤

constrained_fdm(datastructure: AnyFDDatastructure, optimizer: Optimizer, loss: Loss, parameters: Sequence[Parameter] | None = None, constraints: Sequence[Constraint] | None = None, maxiter: int = 100, tol: float = 1e-06, tmax: int = 1, eta: float = 1e-06, callback: Callable | None = None, sparse: bool = True, is_load_local: bool = False, itersolve_fn: Callable | None = None, iterload_fn: Callable | None = None, implicit_diff: bool = True, nd: bool = False, verbose: bool = False, jit: bool = True) -> AnyFDDatastructure

Form-find a datastructure in constrained static equilibrium via optimization.

Parameters:

  • datastructure (AnyFDDatastructure) –

    The network or mesh to form-find.

  • optimizer (Optimizer) –

    The optimizer that minimizes the loss over the design parameters.

  • loss (Loss) –

    The loss function assembled from goals to minimize.

  • parameters (Sequence[Parameter] | None, default: None ) –

    The optimization parameters. If None, the optimizer uses its defaults.

  • constraints (Sequence[Constraint] | None, default: None ) –

    The constraints to enforce during optimization. If None, the problem is unconstrained.

  • maxiter (int, default: 100 ) –

    The maximum number of optimization iterations.

  • tol (float, default: 1e-06 ) –

    The convergence tolerance of the optimizer.

  • tmax (int, default: 1 ) –

    The maximum number of equilibrium iterations per optimization step. With tmax=1 a single linear FDM step is taken.

  • eta (float, default: 1e-06 ) –

    The convergence tolerance for the iterative equilibrium solve.

  • callback (Callable | None, default: None ) –

    A function invoked once per optimization iteration.

  • sparse (bool, default: True ) –

    If True, assemble and solve the equilibrium system with a sparse solver. Forced to False when constraints are given, which sparse does not support.

  • is_load_local (bool, default: False ) –

    If True, apply edge and face loads in their local coordinate systems (follower loads).

  • itersolve_fn (Callable | None, default: None ) –

    The iterative equilibrium solver. If None, forward fixed-point iteration is used.

  • iterload_fn (Callable | None, default: None ) –

    A load callback invoked once before iterative equilibrium starts.

  • implicit_diff (bool, default: True ) –

    If True, apply implicit differentiation for the backward pass.

  • nd (bool, default: False ) –

    Unused; kept for backward compatibility.

  • verbose (bool, default: False ) –

    Whether to print calculation info to the terminal.

  • jit (bool, default: True ) –

    If True, just-in-time compile the optimization problem.

Returns:

  • datastructure ( AnyFDDatastructure ) –

    A copy of the input datastructure in constrained static equilibrium.

Notes

Constraints are not yet supported for sparse inputs; passing constraints with sparse=True switches the solve to dense and prints a warning.

In the returned copy each node load is the model-aggregated tributary load (node, edge, and face) at equilibrium, while the original edge and face loads are retained. Form-finding the output again therefore redistributes those tributary loads a second time; always solve from the original datastructure.

Models¤

EquilibriumModel ¤

EquilibriumModel(tmax: int = 100, eta: float = 1e-06, is_load_local: bool = False, itersolve_fn: Callable | None = None, iterload_fn: Callable | None = None, implicit_diff: bool = True, verbose: bool = False)

A FDM model to calculate equilibrium states with shape-dependent loads.

Parameters:

  • tmax (int, default: 100 ) –

    The maximum number of iterations to calculate an equilibrium state. If tmax=1, the model does one linear FDM step and the remaining parameters are ignored; the edge and face loads are discarded too.

  • eta (float, default: 1e-06 ) –

    The convergence tolerance for calculating an equilibrium state.

  • is_load_local (bool, default: False ) –

    If True, the face and edge loads are applied in their local coordinate system at every iteration (follower loads).

  • itersolve_fn (Callable | None, default: None ) –

    The function that calculates an equilibrium state iteratively. If None, defaults to forward fixed-point iteration. The solver must be consistent with the choice of residual function.

  • iterload_fn (Callable | None, default: None ) –

    A load callback invoked once before iterative equilibrium starts. If None, the load state is left untouched.

  • implicit_diff (bool, default: True ) –

    If True, apply implicit differentiation to speed up backpropagation.

  • verbose (bool, default: False ) –

    Whether to print calculation info to the terminal.

edges_vectors staticmethod ¤

edges_vectors(xyz: Float[Array, 'nodes 3'], connectivity: Float[Array, 'edges nodes'] | Float[BCOO, 'edges nodes']) -> Float[Array, 'edges 3']

Calculate the unnormalized edge vectors as nodal coordinate differences.

Parameters:

  • xyz (Float[Array, 'nodes 3']) –

    The coordinates of the nodes.

  • connectivity (Float[Array, 'edges nodes'] | Float[BCOO, 'edges nodes']) –

    The signed edge-node incidence matrix of the structure.

Returns:

  • vectors ( Float[Array, 'edges 3'] ) –

    The edge vectors pointing from tail to head node.

edges_lengths staticmethod ¤

edges_lengths(vectors: Float[Array, 'edges 3']) -> Float[Array, 'edges 1']

Compute the length of the edges.

Parameters:

  • vectors (Float[Array, 'edges 3']) –

    The edge vectors.

Returns:

  • lengths ( Float[Array, 'edges 1'] ) –

    The Euclidean length of each edge.

edges_forces staticmethod ¤

edges_forces(q: Float[Array, edges], lengths: Float[Array, 'edges 1']) -> Float[Array, 'edges 1']

Calculate the axial force in the edges.

Parameters:

  • q (Float[Array, edges]) –

    The force densities of the edges.

  • lengths (Float[Array, 'edges 1']) –

    The lengths of the edges.

Returns:

  • forces ( Float[Array, 'edges 1'] ) –

    The axial force in each edge, the product of force density and length.

nodes_residuals staticmethod ¤

nodes_residuals(q: Float[Array, edges], loads: Float[Array, 'nodes 3'], vectors: Float[Array, 'edges 3'], connectivity: Float[Array, 'edges nodes'] | Float[BCOO, 'edges nodes']) -> Float[Array, 'nodes 3']

Compute the residual forces on the nodes of the structure.

Parameters:

  • q (Float[Array, edges]) –

    The force densities of the edges.

  • loads (Float[Array, 'nodes 3']) –

    The loads applied to the nodes.

  • vectors (Float[Array, 'edges 3']) –

    The edge vectors.

  • connectivity (Float[Array, 'edges nodes'] | Float[BCOO, 'edges nodes']) –

    The signed edge-node incidence matrix of the structure.

Returns:

  • residuals ( Float[Array, 'nodes 3'] ) –

    The residual force at each node, zero at nodes in equilibrium.

nodes_positions staticmethod ¤

nodes_positions(xyz_free: Float[Array, 'nodes_free 3'], xyz_fixed: Float[Array, 'nodes_fixed 3'], structure: EquilibriumStructure) -> Float[Array, 'nodes 3']

Concatenate the free and fixed node positions back into node order.

Parameters:

  • xyz_free (Float[Array, 'nodes_free 3']) –

    The coordinates of the free nodes.

  • xyz_fixed (Float[Array, 'nodes_fixed 3']) –

    The coordinates of the fixed (supported) nodes.

  • structure (EquilibriumStructure) –

    The structure whose free/fixed index map reorders the nodes.

Returns:

  • xyz ( Float[Array, 'nodes 3'] ) –

    The coordinates of all nodes, restored to the structure's node order.

nodes_free_positions ¤

nodes_free_positions(q: Float[Array, edges], xyz_fixed: Float[Array, 'nodes_fixed 3'], loads: Float[Array, 'nodes 3'], structure: EquilibriumStructure) -> Float[Array, 'nodes_free 3']

Calculate the coordinates of the free nodes by solving the FDM system.

Parameters:

  • q (Float[Array, edges]) –

    The force densities of the edges.

  • xyz_fixed (Float[Array, 'nodes_fixed 3']) –

    The coordinates of the fixed (supported) nodes.

  • loads (Float[Array, 'nodes 3']) –

    The loads applied to the nodes.

  • structure (EquilibriumStructure) –

    The structure that provides the connectivity matrices.

Returns:

  • xyz_free ( Float[Array, 'nodes_free 3'] ) –

    The coordinates of the free nodes at equilibrium.

Notes

Solves K @ xyz_free = P, the linear FDM equilibrium system, with the stiffness matrix K and the load matrix P.

nodes_equilibrium ¤

nodes_equilibrium(q: Float[Array, edges], xyz_fixed: Float[Array, 'nodes_fixed 3'], loads_nodes: Float[Array, 'nodes 3'], structure: EquilibriumStructure) -> Float[Array, 'nodes_free 3']

Calculate static equilibrium on the nodes of a structure.

Parameters:

  • q (Float[Array, edges]) –

    The force densities of the edges.

  • xyz_fixed (Float[Array, 'nodes_fixed 3']) –

    The coordinates of the fixed (supported) nodes.

  • loads_nodes (Float[Array, 'nodes 3']) –

    The loads applied to the nodes.

  • structure (EquilibriumStructure) –

    The structure that provides the connectivity matrices.

Returns:

  • xyz_free ( Float[Array, 'nodes_free 3'] ) –

    The coordinates of the free nodes at equilibrium.

nodes_load ¤

nodes_load(xyz: Float[Array, 'nodes 3'], load_state: LoadState, structure: EquilibriumStructure) -> Float[Array, 'nodes 3']

Calculate the total load applied to the nodes of the structure.

Parameters:

  • xyz (Float[Array, 'nodes 3']) –

    The coordinates of all nodes.

  • load_state (LoadState) –

    The nodal, edge, and face loads to aggregate onto the nodes.

  • structure (EquilibriumStructure) –

    The structure that provides the connectivity used to distribute edge and face loads to the nodes.

Returns:

  • loads ( Float[Array, 'nodes 3'] ) –

    The load at each node, including tributary edge and face loads.

Notes

Edge and face loads are added only when present (non-scalar). A non-scalar face load occurs only for meshes, so structure is a mesh structure in that branch.

faces_load staticmethod ¤

faces_load(xyz: Float[Array, 'nodes 3'], faces_load: Float[Array, 'faces 3'], structure: EquilibriumMeshStructure, is_local: bool = False) -> Float[Array, 'nodes 3']

Distribute face area loads to the nodes of a structure.

Parameters:

  • xyz (Float[Array, 'nodes 3']) –

    The coordinates of all nodes.

  • faces_load (Float[Array, 'faces 3']) –

    The area load on each face.

  • structure (EquilibriumMeshStructure) –

    The mesh structure that maps faces to their nodes.

  • is_local (bool, default: False ) –

    If True, the face load is applied in the face's local coordinate system (a follower load that tracks the deformed geometry).

Returns:

  • loads ( Float[Array, 'nodes 3'] ) –

    The tributary face load carried by each node.

edges_load staticmethod ¤

edges_load(xyz: Float[Array, 'nodes 3'], edges_load: Float[Array, 'edges 3'], structure: EquilibriumStructure, is_local: bool = False) -> Float[Array, 'nodes 3']

Distribute edge line loads to the nodes of a structure.

Parameters:

  • xyz (Float[Array, 'nodes 3']) –

    The coordinates of all nodes.

  • edges_load (Float[Array, 'edges 3']) –

    The line load on each edge.

  • structure (EquilibriumStructure) –

    The structure that maps edges to their nodes.

  • is_local (bool, default: False ) –

    If True, the edge load is applied in the edge's local coordinate system (a follower load that tracks the deformed geometry).

Returns:

  • loads ( Float[Array, 'nodes 3'] ) –

    The tributary edge load carried by each node.

__call__ ¤

__call__(params: EquilibriumParametersState, structure: EquilibriumStructure) -> EquilibriumState

Compute an equilibrium state using the force density method (FDM).

Parameters:

Returns:

  • eq_state ( EquilibriumState ) –

    The equilibrium state with node coordinates, residuals, edge lengths, edge forces, node loads, and edge vectors.

Notes

A single linear FDM step is taken first. When tmax > 1 the solution is refined iteratively to account for shape-dependent loads, optionally with implicit differentiation for the backward pass.

equilibrium ¤

equilibrium(q: Float[Array, edges], xyz_fixed: Float[Array, 'nodes_fixed 3'], loads_nodes: Float[Array, 'nodes 3'], structure: EquilibriumStructure) -> Float[Array, 'nodes_free 3']

Calculate one linear force density step on a structure.

Parameters:

  • q (Float[Array, edges]) –

    The force densities of the edges.

  • xyz_fixed (Float[Array, 'nodes_fixed 3']) –

    The coordinates of the fixed (supported) nodes.

  • loads_nodes (Float[Array, 'nodes 3']) –

    The loads applied to the nodes.

  • structure (EquilibriumStructure) –

    The structure that provides the connectivity matrices.

Returns:

  • xyz_free ( Float[Array, 'nodes_free 3'] ) –

    The coordinates of the free nodes at equilibrium.

equilibrium_iterative_xyz ¤

equilibrium_iterative_xyz(q: Float[Array, edges], xyz_fixed: Float[Array, 'nodes_fixed 3'], load_state: LoadState, structure: EquilibriumStructure, xyz_free_init: Float[Array, 'nodes_free 3'] | None = None, tmax: int = 100, eta: float = 1e-06, solver: Callable | None = None, implicit_diff: bool = True, verbose: bool = False) -> Float[Array, 'nodes_free 3']

Calculate static equilibrium iteratively via fixed-point iteration on xyz.

Parameters:

  • q (Float[Array, edges]) –

    The force densities of the edges.

  • xyz_fixed (Float[Array, 'nodes_fixed 3']) –

    The coordinates of the fixed (supported) nodes.

  • load_state (LoadState) –

    The nodal, edge, and face loads.

  • structure (EquilibriumStructure) –

    The structure that provides the connectivity matrices.

  • xyz_free_init (Float[Array, 'nodes_free 3'] | None, default: None ) –

    The initial guess for the free node coordinates. If None, it is seeded with one linear FDM step.

  • tmax (int, default: 100 ) –

    The maximum number of fixed-point iterations.

  • eta (float, default: 1e-06 ) –

    The convergence tolerance on the iterates.

  • solver (Callable | None, default: None ) –

    The fixed-point solver. If None, the model's default solver is used.

  • implicit_diff (bool, default: True ) –

    If True, differentiate through the fixed point implicitly.

  • verbose (bool, default: False ) –

    Whether to print convergence info to the terminal.

Returns:

  • xyz_free ( Float[Array, 'nodes_free 3'] ) –

    The coordinates of the free nodes at the converged fixed point.

Notes

This function only supports reverse mode auto-differentiation. To support forward-mode, we should define a custom jvp using implicit differentiation.

equilibrium_iterative_residual ¤

equilibrium_iterative_residual(q: Float[Array, edges], xyz_fixed: Float[Array, 'nodes_fixed 3'], load_state: LoadState, structure: EquilibriumStructure, xyz_free_init: Float[Array, 'nodes_free 3'] | None = None, tmax: int = 100, eta: float = 1e-06, solver: Callable | None = None, implicit_diff: bool = True, verbose: bool = False) -> Float[Array, 'nodes_free 3']

Calculate static equilibrium iteratively by driving nodal residuals to zero.

Parameters:

  • q (Float[Array, edges]) –

    The force densities of the edges.

  • xyz_fixed (Float[Array, 'nodes_fixed 3']) –

    The coordinates of the fixed (supported) nodes.

  • load_state (LoadState) –

    The nodal, edge, and face loads.

  • structure (EquilibriumStructure) –

    The structure that provides the connectivity matrices.

  • xyz_free_init (Float[Array, 'nodes_free 3'] | None, default: None ) –

    The initial guess for the free node coordinates. If None, it is seeded with one linear FDM step.

  • tmax (int, default: 100 ) –

    The maximum number of solver iterations.

  • eta (float, default: 1e-06 ) –

    The convergence tolerance on the residuals.

  • solver (Callable | None, default: None ) –

    The least-squares or root-finding solver. If None, the model's default solver is used.

  • implicit_diff (bool, default: True ) –

    If True, differentiate through the solution implicitly.

  • verbose (bool, default: False ) –

    Whether to print convergence info to the terminal.

Returns:

  • xyz_free ( Float[Array, 'nodes_free 3'] ) –

    The coordinates of the free nodes where the residuals vanish.

Notes

This function only supports reverse mode auto-differentiation. To support forward-mode, we should define a custom jvp using implicit differentiation.

select_equilibrium_iterative_fn ¤

select_equilibrium_iterative_fn(solver: Callable) -> Callable

Pick the iterative equilibrium function compatible with a solver.

Parameters:

  • solver (Callable) –

    The iterative solver function to match.

Returns:

  • equilibrium_fn ( Callable ) –

    The fixed-point (xyz) function for fixed-point solvers, or the residual function for least-squares and root-finding solvers.

Raises:

  • ValueError –

    If the solver does not belong to a supported family.

equilibrium_state ¤

equilibrium_state(q: Float[Array, edges], xyz: Float[Array, 'nodes 3'], loads_nodes: Float[Array, 'nodes 3'], structure: EquilibriumStructure) -> EquilibriumState

Assemble an equilibrium state from the equilibrated geometry.

Parameters:

  • q (Float[Array, edges]) –

    The force densities of the edges.

  • xyz (Float[Array, 'nodes 3']) –

    The coordinates of all nodes.

  • loads_nodes (Float[Array, 'nodes 3']) –

    The loads applied to the nodes.

  • structure (EquilibriumStructure) –

    The structure that provides the connectivity matrix.

Returns:

  • eq_state ( EquilibriumState ) –

    The equilibrium state bundling node coordinates, residuals, edge lengths, edge forces, node loads, and edge vectors.

stiffness_matrix staticmethod ¤

stiffness_matrix(q: Float[Array, edges], structure: EquilibriumStructure) -> Float[Array, 'nodes_free nodes_free']

Assemble the force density stiffness matrix of the free nodes.

Parameters:

  • q (Float[Array, edges]) –

    The force densities of the edges.

  • structure (EquilibriumStructure) –

    The structure that provides the free-node connectivity matrix.

Returns:

  • stiffness ( Float[Array, 'nodes_free nodes_free'] ) –

    The stiffness matrix Cf.T @ diag(q) @ Cf of the free nodes.

load_matrix ¤

load_matrix(q: Float[Array, edges], xyz_fixed: Float[Array, 'nodes_fixed 3'], load_nodes: Float[Array, 'nodes 3'], structure: EquilibriumStructure) -> Float[Array, 'nodes_free 3']

Assemble the right-hand side load matrix of the FDM system.

Parameters:

  • q (Float[Array, edges]) –

    The force densities of the edges.

  • xyz_fixed (Float[Array, 'nodes_fixed 3']) –

    The coordinates of the fixed (supported) nodes.

  • load_nodes (Float[Array, 'nodes 3']) –

    The loads applied to the nodes.

  • structure (EquilibriumStructure) –

    The structure that provides the connectivity matrices.

Returns:

  • load_matrix ( Float[Array, 'nodes_free 3'] ) –

    The free-node loads minus the fixed nodes' residual contribution.

load_xyz_matrix ¤

load_xyz_matrix(params: EquilibriumParametersState, xyz_free: Float[Array, 'nodes_free 3'], structure: EquilibriumStructure) -> Float[Array, 'nodes_free 3']

Assemble the free-node load matrix for shape-dependent loads.

Parameters:

  • params (EquilibriumParametersState) –

    The force densities, fixed node coordinates, and load state.

  • xyz_free (Float[Array, 'nodes_free 3']) –

    The current coordinates of the free nodes.

  • structure (EquilibriumStructure) –

    The structure that provides the connectivity matrices.

Returns:

  • load_matrix ( Float[Array, 'nodes_free 3'] ) –

    The right-hand side load matrix evaluated at the current geometry.

Notes

The full node coordinates are reassembled first so that edge and face loads can be recomputed against the current shape.

load_xyz_matrix_from_r_fixed ¤

load_xyz_matrix_from_r_fixed(params: tuple[StiffnessMatrix, Float[Array, 'nodes_free 3'], Float[Array, 'nodes_fixed 3'], LoadState], xyz_free: Float[Array, 'nodes_free 3'], structure: EquilibriumStructure) -> Float[Array, 'nodes_free 3']

Assemble the free-node load matrix reusing a precomputed fixed residual.

Parameters:

  • params (tuple[StiffnessMatrix, Float[Array, 'nodes_free 3'], Float[Array, 'nodes_fixed 3'], LoadState]) –

    A tuple of the stiffness matrix, the precomputed fixed-node residual matrix, the fixed node coordinates, and the load state.

  • xyz_free (Float[Array, 'nodes_free 3']) –

    The current coordinates of the free nodes.

  • structure (EquilibriumStructure) –

    The structure that provides the connectivity matrices.

Returns:

  • load_matrix ( Float[Array, 'nodes_free 3'] ) –

    The right-hand side load matrix evaluated at the current geometry.

Notes

This is the iteration hot path: the fixed-node residual is passed in rather than recomputed each step, unlike load_xyz_matrix.

residual_fixed_matrix staticmethod ¤

residual_fixed_matrix(q: Float[Array, edges], xyz_fixed: Float[Array, 'nodes_fixed 3'], structure: EquilibriumStructure) -> Float[Array, 'nodes_free 3']

Compute the fixed nodes' residual contribution to the free-node loads.

Parameters:

  • q (Float[Array, edges]) –

    The force densities of the edges.

  • xyz_fixed (Float[Array, 'nodes_fixed 3']) –

    The coordinates of the fixed (supported) nodes.

  • structure (EquilibriumStructure) –

    The structure that provides the free and fixed connectivity matrices.

Returns:

  • residual ( Float[Array, 'nodes_free 3'] ) –

    The term Cf.T @ diag(q) @ Cb @ xyz_fixed subtracted from the loads to form the right-hand side of the FDM system.

residual_free_matrix ¤

residual_free_matrix(params: EquilibriumParametersState, xyz_free: Float[Array, 'nodes_free 3'], structure: EquilibriumStructure) -> Float[Array, 'nodes_free 3']

Compute the residual forces at the free nodes for a given geometry.

Parameters:

  • params (EquilibriumParametersState) –

    The force densities, fixed node coordinates, and load state.

  • xyz_free (Float[Array, 'nodes_free 3']) –

    The current coordinates of the free nodes.

  • structure (EquilibriumStructure) –

    The structure that provides the connectivity matrices.

Returns:

  • residuals ( Float[Array, 'nodes_free 3'] ) –

    The out-of-balance force K @ xyz_free - P at each free node, zero at equilibrium.

EquilibriumModelSparse ¤

EquilibriumModelSparse(*args, **kwargs)

An FDM model that solves the equilibrium system with a sparse linear solver.

Notes

Identical to EquilibriumModel except the stiffness matrix is assembled in sparse format and the linear system is solved with a sparse solver, which scales to larger structures.

stiffness_matrix staticmethod ¤

stiffness_matrix(q: Float[Array, edges], structure: EquilibriumStructureSparse) -> Float[CSC, 'nodes_free nodes_free']

Assemble the force density stiffness matrix in sparse format.

Parameters:

  • q (Float[Array, edges]) –

    The force densities of the edges.

  • structure (EquilibriumStructureSparse) –

    The sparse structure that provides the sparse index array, diagonal indices, and the node-edge incidence used for the diagonal.

Returns:

  • stiffness ( Float[CSC, 'nodes_free nodes_free'] ) –

    The stiffness matrix of the free nodes, stored in sparse format.

Notes

The off-diagonal entries are the negated force densities gathered through the precomputed index array; the diagonal holds the per-node sum of incident force densities.

States¤

EquilibriumState ¤

The static equilibrium of a structure under load.

Attributes:

  • xyz (Float[Array, 'nodes 3']) –

    The coordinates of the nodes.

  • residuals (Float[Array, 'nodes 3']) –

    The residual force at each node, zero at nodes in equilibrium.

  • lengths (Float[Array, 'edges 1']) –

    The length of each edge.

  • forces (Float[Array, 'edges 1']) –

    The axial force in each edge.

  • loads (Float[Array, 'nodes 3']) –

    The load applied to each node.

  • vectors (Float[Array, 'edges 3']) –

    The edge vectors pointing from tail to head node.

from_datastructure classmethod ¤

from_datastructure(datastructure: FDNetwork | FDMesh, structure: EquilibriumStructure, dtype: DTypeLike | None = None) -> EquilibriumState

Read an equilibrium state off a datastructure's stored attributes.

Parameters:

  • datastructure (FDNetwork | FDMesh) –

    The network or mesh to read coordinates, residuals, lengths, forces, and loads from.

  • structure (EquilibriumStructure) –

    The structure whose connectivity derives the edge vectors, the one state field the datastructure does not store.

  • dtype (DTypeLike | None, default: None ) –

    The floating-point dtype for the state arrays. If None, uses DTYPE_JAX.

Returns:

  • eq_state ( EquilibriumState ) –

    The equilibrium state mirroring the datastructure's attributes.

Raises:

  • ValueError –

    If the datastructure is neither a network nor a mesh.

Notes

The state is read as-is, without form-finding: every field but the edge vectors comes straight from a stored attribute, so it is only a genuine equilibrium once the datastructure has been solved. Run fdm (or constrained_fdm) on the datastructure first; on a fresh, never-solved input the lengths, forces, and residuals default to zero. The edge vectors are derived as connectivity @ xyz since the datastructure never persists them.

LoadState ¤

The loads applied to a structure, grouped by the element they act on.

Attributes:

  • nodes (Float[Array, 'nodes 3']) –

    The load applied directly to each node.

  • edges (Float[Array, 'edges 3'] | float) –

    The line load on each edge, or 0.0 when there is no edge load.

  • faces (Float[Array, 'faces 3'] | float) –

    The area load on each face, or 0.0 when there is no face load.

Notes

Edge and face loads collapse to the scalar 0.0 when all their entries vanish, which lets the model skip distributing them to the nodes. Networks always carry faces=0.0.

from_datastructure classmethod ¤

from_datastructure(datastructure: FDNetwork | FDMesh, dtype: DTypeLike | None = None) -> LoadState

Create a load state from a network or mesh.

Parameters:

  • datastructure (FDNetwork | FDMesh) –

    The network or mesh to read the node, edge, and face loads from.

  • dtype (DTypeLike | None, default: None ) –

    The floating-point dtype for the load arrays. If None, uses the array library default.

Returns:

  • load_state ( LoadState ) –

    The load state, with edge and face loads collapsed to 0.0 when they are all zero.

EquilibriumParametersState ¤

The independent parameters that define an equilibrium problem.

Attributes:

  • q (Float[Array, edges]) –

    The force density of each edge.

  • xyz_fixed (Float[Array, 'nodes_fixed 3']) –

    The coordinates of the fixed (supported) nodes.

  • loads (LoadState) –

    The load state applied to the structure.

from_datastructure classmethod ¤

from_datastructure(datastructure: FDNetwork | FDMesh, dtype: DTypeLike | None = None) -> EquilibriumParametersState

Create a parameter state from a network or mesh.

Parameters:

  • datastructure (FDNetwork | FDMesh) –

    The network or mesh to read force densities, fixed coordinates, and loads from.

  • dtype (DTypeLike | None, default: None ) –

    The floating-point dtype for the parameter arrays. If None, uses DTYPE_JAX.

Returns:

Raises:

  • ValueError –

    If the datastructure is neither a network nor a mesh.

Structures¤

The static topology and indexing of a structure, precomputed once so that the equilibrium model can run shape updates as pure array operations.

EquilibriumStructure ¤

EquilibriumStructure(nodes: Int[ndarray, nodes], edges: Int[ndarray, 'edges 2'], supports: Int[ndarray, nodes], **kwargs)

A graph with supports and the connectivity matrices the FDM solve needs.

Notes

Extends Graph with a support mask and the free/fixed node partition, precomputing the free and fixed connectivity submatrices and the index maps that reorder nodes between the free-fixed and native orderings.

num_supports property ¤

num_supports: Int[Array, '']

The number of supports.

num_free property ¤

num_free: Int[Array, '']

The number of free (unsupported) nodes.

support_index property ¤

support_index: dict[int, int]

A mapping from support keys to indices.

nodes_free property ¤

nodes_free: Int[ndarray, nodes_free]

The free nodes.

nodes_fixed property ¤

nodes_fixed: Int[ndarray, nodes_fixed]

The fixed nodes.

from_network classmethod ¤

from_network(network: FDNetwork) -> EquilibriumStructure

Build an equilibrium structure from a force density network.

Parameters:

  • network (FDNetwork) –

    The network to read nodes, edges, and supports from.

Returns:

  • structure ( EquilibriumStructure ) –

    The structure with the network's connectivity and support mask.

EquilibriumStructureSparse ¤

EquilibriumStructureSparse(nodes: Int[ndarray, nodes], edges: Int[ndarray, 'edges 2'], supports: Int[ndarray, nodes], **kwargs)

An equilibrium structure that precomputes sparse stiffness-assembly helpers.

Notes

Precomputes an index array, the diagonal data positions, and a diagonal selector matrix so the sparse model can assemble the stiffness matrix by indexing rather than rebuilding its structure at every solve.

EquilibriumMeshStructure ¤

EquilibriumMeshStructure(vertices: Int[ndarray, vertices], faces: Int[ndarray, 'faces vertices'], edges: Int[ndarray, 'edges 2'], supports: Int[ndarray, vertices], **kwargs)

An equilibrium structure that also carries mesh face topology.

Notes

Adds face connectivity to EquilibriumStructure, so mesh vertices play the role of nodes while faces enable tributary face-load distribution.

num_free property ¤

num_free: Int[Array, '']

The number of free (unsupported) vertices.

support_index property ¤

support_index: dict[int, int]

A mapping from support vertices keys to indices.

vertices_free property ¤

vertices_free: Int[ndarray, vertices_free]

The free (unsupported) vertices.

vertices_fixed property ¤

vertices_fixed: Int[ndarray, vertices_fixed]

The fixed (supported) vertices.

from_mesh classmethod ¤

from_mesh(mesh: FDMesh) -> EquilibriumMeshStructure

Build an equilibrium mesh structure from a force density mesh.

Parameters:

  • mesh (FDMesh) –

    The mesh to read vertices, faces, edges, and supports from.

Returns:

Raises:

Notes

Faces are padded with -1 to a common length so they fit a rectangular index array; the padding is masked out downstream.

EquilibriumMeshStructureSparse ¤

EquilibriumMeshStructureSparse(vertices: Int[ndarray, vertices], faces: Int[ndarray, 'faces vertices'], edges: Int[ndarray, 'edges 2'], supports: Int[ndarray, vertices], **kwargs)

A mesh equilibrium structure with sparse stiffness-assembly helpers.

Notes

Combines the face topology of EquilibriumMeshStructure with the sparse assembly precomputation of EquilibriumStructureSparse.

Graph ¤

Graph(nodes: Int[ndarray, nodes], edges: Int[ndarray, 'edges 2'])

An immutable graph holding node and edge connectivity as dense matrices.

Notes

An equinox Module, so instances are registered pytrees. Nodes and edges are kept as static NumPy index arrays, while the connectivity and adjacency matrices are JAX arrays used in the differentiable equilibrium computation.

num_nodes property ¤

num_nodes: int

The number of nodes.

num_edges property ¤

num_edges: int

The number of edges.

node_index property ¤

node_index: dict[int, int]

A dictionary between node keys and their enumeration indices.

edge_index property ¤

edge_index: dict[tuple[int, int], int]

A dictionary between edge keys and their enumeration indices.

GraphSparse ¤

GraphSparse(nodes: Int[ndarray, nodes], edges: Int[ndarray, 'edges 2'])

A graph that keeps its connectivity and adjacency matrices in sparse format.

connectivity_scipy property ¤

connectivity_scipy: csc_matrix

The signed edge-node incidence matrix as a scipy sparse matrix.

This is the assembly substrate the sparse structures build from: the free and fixed column submatrices and the stiffness precomputation slice and factor it before converting the results to JAX sparse format.

Notes

Kept in scipy format on purpose. Slicing columns out of a JAX sparse matrix materializes a stored entry for every edge-column pair rather than only the nonzeros, and the stiffness precomputation relies on structural operations (transpose products, diagonal rewrites, index pointer reads) that JAX sparse arrays do not provide. Rebuilt on every access, which costs fractions of a millisecond.

Iterative solvers¤

Solver functions for equilibrium states with shape-dependent loads, passed to fdm and constrained_fdm via itersolve_fn.

solver_forward ¤

solver_forward(f: Callable, a: Any, x_init: Float[Array, ...], solver_config: dict[str, Any]) -> Float[Array, ...]

Find a fixed point of f(a, x) by forward iteration until convergence.

Parameters:

  • f (Callable) –

    The function to iterate upon.

  • a (Any) –

    The function parameters.

  • x_init (Float[Array, ...]) –

    The initial guess for the solution vector.

  • solver_config (dict[str, Any]) –

    The configuration options of the solver, read for tmax, eta, and verbose.

Returns:

  • x_star ( Float[Array, ...] ) –

    The solution vector at the fixed point.

Notes

Iteration stops when the mean nodal move between successive iterates drops below eta or after tmax steps, whichever comes first.

solver_fixedpoint ¤

solver_fixedpoint(f: Callable, a: SolverIterParams, x_init: Float[Array, 'nodes_free 3'], solver_config: dict[str, Any]) -> Float[Array, 'nodes_free 3']

Find a fixed point of f(a, x) using plain fixed-point iteration.

Parameters:

  • f (Callable) –

    The function to iterate upon.

  • a (SolverIterParams) –

    The function parameters.

  • x_init (Float[Array, 'nodes_free 3']) –

    The initial guess for the solution vector.

  • solver_config (dict[str, Any]) –

    The configuration options of the solver.

Returns:

  • x_star ( Float[Array, 'nodes_free 3'] ) –

    The solution vector at the fixed point.

solver_anderson ¤

solver_anderson(f: Callable, a: SolverIterParams, x_init: Float[Array, 'nodes_free 3'], solver_config: dict[str, Any]) -> Float[Array, 'nodes_free 3']

Find a fixed point of f(a, x) using Anderson acceleration.

Parameters:

  • f (Callable) –

    The function to iterate upon.

  • a (SolverIterParams) –

    The function parameters.

  • x_init (Float[Array, 'nodes_free 3']) –

    The initial guess for the solution vector.

  • solver_config (dict[str, Any]) –

    The configuration options of the solver.

Returns:

  • x_star ( Float[Array, 'nodes_free 3'] ) –

    The solution vector at the fixed point.

solver_newton ¤

solver_newton(fn: Callable, theta: SolverIterParams, x_init: Float[Array, nodes_free_flat], solver_config: dict[str, Any]) -> Float[Array, nodes_free_flat]

Find a root of fn(theta, x) = 0 using Newton's method.

Parameters:

  • fn (Callable) –

    The residual function whose root is sought.

  • theta (SolverIterParams) –

    The function parameters.

  • x_init (Float[Array, nodes_free_flat]) –

    The initial guess for the flattened solution vector.

  • solver_config (dict[str, Any]) –

    The configuration options of the solver.

Returns:

  • x_star ( Float[Array, nodes_free_flat] ) –

    The flattened solution vector at the root.

Notes

Uses an SVD linear solver inside Newton's method to stay robust when the Jacobian is singular or ill-conditioned.

solver_gauss_newton ¤

solver_gauss_newton(fn: Callable, theta: SolverIterParams, x_init: Float[Array, nodes_free_flat], solver_config: dict[str, Any]) -> Float[Array, nodes_free_flat]

Drive the residual fn(theta, x) to zero with the Gauss-Newton algorithm.

Parameters:

  • fn (Callable) –

    The residual function to drive to zero.

  • theta (SolverIterParams) –

    The function parameters.

  • x_init (Float[Array, nodes_free_flat]) –

    The initial guess for the flattened solution vector.

  • solver_config (dict[str, Any]) –

    The configuration options of the solver.

Returns:

  • x_star ( Float[Array, nodes_free_flat] ) –

    The flattened solution vector at the residual minimum.

solver_levenberg_marquardt ¤

solver_levenberg_marquardt(fn: Callable, theta: SolverIterParams, x_init: Float[Array, nodes_free_flat], solver_config: dict[str, Any]) -> Float[Array, nodes_free_flat]

Drive the residual fn(theta, x) to zero with the Levenberg-Marquardt method.

Parameters:

  • fn (Callable) –

    The residual function to drive to zero.

  • theta (SolverIterParams) –

    The function parameters.

  • x_init (Float[Array, nodes_free_flat]) –

    The initial guess for the flattened solution vector.

  • solver_config (dict[str, Any]) –

    The configuration options of the solver.

Returns:

  • x_star ( Float[Array, nodes_free_flat] ) –

    The flattened solution vector at the residual minimum.

Notes

Incompatible with the sparse equilibrium model because jax.experimental.sparse.csr_matmat does not implement a batching rule yet.

solver_dogleg ¤

solver_dogleg(fn: Callable, theta: SolverIterParams, x_init: Float[Array, nodes_free_flat], solver_config: dict[str, Any]) -> Float[Array, nodes_free_flat]

Drive the residual fn(theta, x) to zero with the Dogleg trust-region method.

Parameters:

  • fn (Callable) –

    The residual function to drive to zero.

  • theta (SolverIterParams) –

    The function parameters.

  • x_init (Float[Array, nodes_free_flat]) –

    The initial guess for the flattened solution vector.

  • solver_config (dict[str, Any]) –

    The configuration options of the solver.

Returns:

  • x_star ( Float[Array, nodes_free_flat] ) –

    The flattened solution vector at the residual minimum.

Implicit differentiation¤

The custom_vjp-wrapped solver functions that backpropagate through an equilibrium state via the implicit function theorem instead of unrolling solver iterations.

solver_fixedpoint_implicit ¤

solver_fixedpoint_implicit(solver: Callable, solver_config: dict[str, Any], f: Callable, a: SolverIterParams, x_init: Float[Array, 'nodes_free 3']) -> Float[Array, 'nodes_free 3']

Solve for a fixed point of f(a, x) with implicit differentiation.

Parameters:

  • solver (Callable) –

    The function that runs the fixed-point solve.

  • solver_config (dict[str, Any]) –

    The configuration options of the solver.

  • f (Callable) –

    The function to iterate upon.

  • a (SolverIterParams) –

    The function parameters, differentiated through implicitly.

  • x_init (Float[Array, 'nodes_free 3']) –

    The initial guess for the solution vector.

Returns:

  • x_star ( Float[Array, 'nodes_free 3'] ) –

    The solution vector at the fixed point.

Notes

Wrapped in a custom VJP so the backward pass differentiates through the fixed point implicitly rather than unrolling the solver iterations.

solver_nonlinear_implicit ¤

solver_nonlinear_implicit(solver: Callable, solver_config: dict[str, Any], fn: Callable, theta: SolverIterParams, x_init: Float[Array, nodes_free_flat]) -> Float[Array, nodes_free_flat]

Solve a nonlinear system for x with implicit differentiation.

Parameters:

  • solver (Callable) –

    The function that runs the nonlinear solve.

  • solver_config (dict[str, Any]) –

    The configuration options of the solver.

  • fn (Callable) –

    The residual function to drive to zero.

  • theta (SolverIterParams) –

    The function parameters, differentiated through implicitly.

  • x_init (Float[Array, nodes_free_flat]) –

    The initial guess for the flattened solution vector.

Returns:

  • x_star ( Float[Array, nodes_free_flat] ) –

    The flattened solution vector.

Notes

Wrapped in a custom VJP so the backward pass differentiates through the solution implicitly rather than unrolling the solver iterations.

Sparse solver¤

sparse_solve ¤

sparse_solve(A: SystemMatrixLHS, b: SystemMatrixRHS) -> SystemSolution

Solve the sparse linear system with a custom-differentiable solver.

Parameters:

  • A (SystemMatrixLHS) –

    The left-hand side matrix, assumed symmetric.

  • b (SystemMatrixRHS) –

    The right-hand side, one column per coordinate.

Returns:

  • x ( SystemSolution ) –

    The solution, one column per right-hand side.

Notes

Dispatches to the backend solver registered by :func:register_sparse_solver. A custom VJP supplies the backward pass, so only reverse-mode differentiation is supported.

register_sparse_solver ¤

register_sparse_solver(solvers: dict[str, Callable]) -> Callable

Pick the sparse solver matching the active JAX backend.

Parameters:

  • solvers (dict[str, Callable]) –

    A mapping from backend name ("cpu" or "gpu") to solver function.

Returns:

  • solver ( Callable ) –

    The solver registered for the current default backend.

Raises:

  • ValueError –

    If no solver is registered for the active backend.