JAX FDM: A differentiable solver for inverse form-finding

Rafael Pastrana 1 2 Deniz Oktay 3 Ryan P. Adams 3 Sigrid Adriaenssens 2

Abstract
We introduce JAX FDM, a differentiable solver
to design mechanically efficient shapes for 3D
structures conditioned on target architectural, fabrication and structural properties. Examples ofsuch structures are domes, cable nets and towers. JAX FDM solves these inverse form-finding
problems by combining the force density method,
differentiable sparsity and gradient-based optimization. Our solver can be paired with otherlibraries in the JAX ecosystem to facilitate theintegration of form-finding simulations with neural networks. We showcase the features of JAX
FDM with two design examples. JAX FDM
is available as an open-source library at this
URL: https://github.com/arpastrana/jax fdm

Figure 1: Plan view of cable net with equalized edge lengths.

variety of shapes in static equilibrium (Fig. 3). However,in engineering practice, it is necessary to generate not any
mechanically efficient 3D shape, but a feasible one thatsatisfies constraints arising from architectural, fabrication,or other structural requirements.
Consider the case shown in Fig. 4 where it is of interest
to find a shape in equilibrium that is as close as possible
to a target surface Û . This surface may express architectural intent for a new roof or can represent the geometryof a historical masonry vault that needs to be analyzed forrestoration purposes (Panozzo et al., 2013; Marmo et al.,2019). Practical structural design, therefore, requires thesolution of an inverse form-finding problem, a mappingfrom Û → θ, where the goal is to estimate adequate values
for the parameters θ⋆ that are conducive to an equilibrium
state with prescribed characteristics Û .

1. Introduction
The force density method (FDM) (Schek, 1974) is aform-finding method that generates shapes in static equilibrium for meter-scale 3D structures, such as masonryvaults (Panozzo et al., 2013), cable nets (Veenendaal et al.,2017) and tensegrity systems (Zhang & Ohsaki, 2006). A
structure in static equilibrium carries loads like its selfweight or wind pressure only through internal tension andcompression forces (Adriaenssens et al., 2014). This axialdominant mechanical behavior enables a form-found structure to span long distances with low material usage compared to a structure that is not form-found (Schlaich, 2018;
Rippmann et al., 2016).

The design space of all possible shapes in static equilibrium parametrized by θ is vast, particularly as the dimensionality of these parameters grows proportionally to thehundreds or thousands of cables, bricks and blocks that compose a real-world structure. Numerical approaches basedon geometric heuristics (Lee et al., 2016) or genetic algorithms (Koohestani, 2012) offer limited support to navigate this high-dimensional design space towards feasible
designs. The current surge of differentiable physics solvers
and physics-informed neural networks in structural engineering (Cuvilliers, 2020; Chang & Cheng, 2020; Xue et al.,2023; Wu, 2023; Pastrana et al., 2023) provide insights todevelop new approaches to tackle inverse form-finding.

The FDM is a forward physics solver expressed as a functionf (θ, G) = U . Given a structure modeled as a sparse graph
G and a set of continuous design parameters θ, the FDM
computes a state of static equilibrium U for G (Fig. 2).
By inputting different values of θ, the FDM generates a1{School of Architecture, 2 Department of Civil and Environmental Engineering, 3 Department of Computer Science}, Princeton University, USA. Correspondence to: Rafael Pastrana <arpastrana@princeton.edu>.

In this paper, we present JAX FDM, a differentiable solver
to perform inverse form-finding on 3D structures modeled as
pin-jointed bar networks. JAX FDM implements the FDM in
JAX (Bradbury et al., 2018) and solves inverse form-finding
problems by estimating adequate inputs to the FDM via

Published at the Differentiable Almost Everything Workshop of the
40 th International Conference on Machine Learning, Honolulu,
Hawaii, USA. July 2023. Copyright 2023 by the author(s).

1JAX FDM: A differentiable solver for inverse form-finding

Figure 2: The FDM is a forward form-finding method that computes a state of static equilibrium U on an input graph G given
input parameters θ = (q, P, Xs ).

gradient-based optimization. The required forward and backward calculations are executed efficiently by running a differentiable sparse solver on a CPU or a GPU. After presenting
the theory behind our work in Section 2, we use our solver
to address two inverse form-finding problems: the design ofa shell structure that matches an arbitrary target shape (Section 3.1) and the design of cable net with prescribed edge
lengths (Section 3.2). JAX FDM is open-source software accessible at this URL: https://github.com/arpastrana/jax fdm.

Figure 4: JAX FDM solves inverse form-finding problems estimating parameters θ⋆ that fit a prescribed equilibrium state Û via
gradient-based optimization. Here, we calculate the force densitiesq⋆ needed to solve a shape approximation problem on a creased
shell modeled as graphs G and GII .

The FDM is a function f that computes a state of static
equilibrium U on a fixed graph G given input parameters θ.
The input parameters θ = (q, P, Xs ) are features definedon the elements of G:

2. Auto-differentiable and sparsified FDM
2.1. The force density method (FDM)
The FDM models a structure as a pin-jointed, force network (Schek, 1974). Let G = (V, E) be a graph with nvertices V connected by m edges E encoding this network.
One portion of V of size ns is defined as the supported
vertices of the structure Vs , i.e., the locations in the structurethat are fixed and transfer reaction forces to its anchors. The
remaining nu unsupported vertices are denoted Vu .

• A diagonal matrix Q ∈ Rm×m with the force densities of the edges, q ∈ Rm×1 . The force density qi ofedge i is the ratio between the internal force ti and thelength li of the edge, qi = ti /li . A negative qi indicatescompression while a positive one indicates tension.
• A matrix P ∈ Rn×3 with the 3D vectors denoting theexternal loads applied to all the vertices of G. Submatrices Pu ∈ Rnu ×3 and Ps ∈ Rns ×3 correspond to
the rows of P with the loads applied to the vertices Vu
and Vs , respectively.
• A matrix Xs ∈ Rns ×3 containing the 3D coordinatesof the supported vertices, Vs .

A connectivity matrix C ∈ {−1, 0, 1}m×n encodes therelationship between the edges and the vertices of G. Entry
cij of C is equal to 1 if vertex j is the start node of edge i andequal to −1 if vertex j is the end node of edge i. Otherwise,
cij = 0. Submatrices Cu and Cs are formed by the columnsof C corresponding to the unsupported and the supported
vertices of G, respectively.

The state U = (Xu , Rs , t, l) characterizes the static equilibrium configuration of G:• A matrix Xu ∈ Rnu ×3 containing the 3D coordinatesin static equilibrium of the unsupported vertices, Vu .
• A matrix Rs ∈ Rns ×3 with the reaction forces incident
to the supported vertices, Vs .
• A vector t ∈ Rm×1 containing the tensile or compressive internal force of the edges.
• A vector l ∈ Rm×1 with the edge lengths.
The key step in the FDM is to find the 3D coordinates instatic equilibrium of the free vertices Xu with Eq. 1:

Figure 3: The FDM generates different static equilibrium configurations for variations of θ. From left to right: q ∈ {−0.1, −1},
q = 1, and q ∼ U(−0.1, −1). Colors indicate the internal axial
forces: blue denotes compression and red tension.

Xu = (CTu QCu )−1 (Pu − CTu QCs Xs )2

(1)

JAX FDM: A differentiable solver for inverse form-finding

The remaining components of U are computed as:
Rs = Ps − CTs QCX
T

t=q L

(2)
(3)

The matrix of 3D coordinates X results from concatenating
Xu and Xs . The diagonal matrix of edge lengths L can becalculated taking the row-wise L2 norm of the inner productof the connectivity matrix C and X, L = diag(∥C X∥2 ).
2.2. Solving inverse form-finding problems
The desideratum is to design structures in static equilibrium that attain additional architectural, fabrication, or othertarget properties.

Figure 5: Inverse form-found structures generated with JAX FDM.
Left: two compression-only structures, a monkey saddle (top) anda dome (bottom). Right: a tensegrity tower.

The FDM parametrizes form-finding in terms of θ, simplifying the computation of a state of static equilibrium U to
the solution of a linear system. However, the relationshipbetween θ and U is non-linear as linear perturbations in θ do
not correspond to linear changes in U . Moreover, the force
densities q are not interpretable quantities: they express aratio between the expected forces and lengths in the edges
of a structure, but neither of them concretely. Both issues
complicate tackling inverse form-finding problems without
an automated approach.

map from force densities into the entries of the coefficientmatrix in compressed sparse-row (CSR) format. We thenuse the adjoint method to implement a custom gradientfor scipy.sparse.linalg.spsolve on CPU, andjax.experimental.sparse.linalg.spsolveon GPU. On CPU, we use a jax.pure callback toensure the sparse solve is compatible with JIT compilation.

To address these challenges, we solve an unconstrained
optimization problem w.r.t. parameters θ. Let g(f (θ, G))
be a non-linear goal function that computes a property ofinterest. One inverse form-finding problem may contain K
different goal functions that are individually scaled by aweight factor wk and aggregated in a loss function L(θ):L(θ) =K
X

wk gk (f (θ, G)))

3. Examples
JAX FDM features a rich bank of goal functions that simplify the modeling of inverse-form-finding problems onvarious structural systems (Fig. 5). Here, we present two
specific use cases with the current version of the library.

(4)

3.1. Shape approximation for shell structures

k=1

We want to calculate a network in static equilibrium
for a shell that approximates the geometry Û pictured
in Fig. 4 (Panozzo et al., 2013). This target shape is
supplied by the project architect as a COMPAS network (Van Mele & many others, 2017). JAX FDM offers
functions to convert such a network into a JAX-friendly
jfe.EquilibriumStructure. This structure encodes the graph representation G of the shell and its connectivity matrix C (Section 2.1).

We minimize Eq. 4 by estimating optimal parameters θ⋆
via gradient descent, iteratively updating θ in the negative
direction of the gradient ∇θ L. We conveniently estimate
the required value of ∇θ L with reverse-mode automaticdifferentiation (Bradbury et al., 2018).
2.3. Differentiable sparse solver
A bottleneck in our solver is the solution of the linear system
in Eq. 1. Although for small problems we can materialize
and invert the full dense coefficient matrix CTu QCu , forlarger problems we want to take advantage of the inherent
sparsity of Cu to get computational speedup, especially as
the linear solve is called many times during inverse design.

import jax_fdm.equilibrium as jfe
from jax_fdm.datastructures import FDNetwork# load a structure from a COMPAS network
net = FDNetwork.from_json("shell.json")eqs = jfe.EquilibriumStructure
structure = eqs.from_network(net)As sparse solvers have limited support on JAX at the timeof writing, we use implicit differentiation to derive a custom
differentiable sparse linear solver. Given that the sparsity
pattern only depends on C, which is fixed from the beginning for a given graph G, we implement a differentiable

To compute a state of static equilibrium eq state with
the FDM, we instantiate an jfe.EquilibriumModeland define θ as a tuple of design parameters, params.
3

JAX FDM: A differentiable solver for inverse form-finding
import jax.numpy as jnp# set the initial force densities
q = jnp.ones(structure.num_edges) * -1.0
# compute an equilibrium state
fdm = jfe.EquilibriumModel()params = (q, xyz_fixed, loads)
eq_state = fdm(params, structure)

The fdm model is a callable object that expresses f (θ, G)
and implements Eqs. 1-3. The initial vector of force densities is set to q = −1. The negative values denote compressive internal forces in the edges of G. The other arrays,
xyz fixed and loads, store the 3D coordinates of the
supports Xs , and the loads P applied to the vertices of G, respectively. Next, we set up an inverse form-finding problem
in terms of q with two functions:Figure 6: A cable net roof. The cables of this inverse form-found,tension-only structure have a target length of 0.15 m.

Post-optimization, the distance between the solution provided by q⋆ and the target shape decreases by four ordersof magnitude. The fit is comparable with an input graph
GII that has three times more edges and design parameters(Fig. 4). This example is available in a Colab notebook at
https://tinyurl.com/25czahvh.

def goal_fn(eq_state):
dist = (eq_state.xyz - xyz_target)**2
return jnp.sum(dist)
def loss_fn(q):params = (q, xyz_fixed, loads)
eq_state = fdm(params, structure)
return goal_fn(eq_state)

3.2. Equalizing edge lengths in a cable net
We design a self-stressed cable net inspired by the roofof the Rhön Klinikum (Oval, 2019). Building cable nets
from standardized components is important for fabrication
efficiency. Therefore, we calculate a tension-only shape forthe net that has a target edge length of 0.15 m. Fig. 6 displays
the solution to the inverse form-finding problem. The goal
function g(f (θ, G)) in this problem is:

The first one is a goal function g(f (θ, G)), which quantifies
the fitness of the shape approximation by measuring the
cumulative distance between the xyz coordinates in static
equilibrium of the vertices V produced by fdm, and thexyz target coordinates on the objective surface. The
second function represents Eq. 4, which we minimize with
an optax optimizer (Babuschkin et al., 2020):

def goal_fn(eq_state):diff = (eq_state.lengths - 0.15)**2
std = jnp.std(eq_state.lengths)
return jnp.sum(diff) + 0.01 * std

import optax
from jax import jit
from jax import value_and_grad@jit
def opt_step(q, o_state):loss, grad = value_and_grad(loss_fn)(q)upd, o_state = opt.update(grad, o_state)q = optax.apply_updates(q, upd)
return q, o_state, loss

After modeling the connectivity of the cable net as astructure, we can reuse the code blocks presented in
Section 3.1 to obtain q⋆ . The only requirement is to swap
goal fn in the body of the loss function loss fn. The
composition and interchangeability of such atomic goal
functions simplify the formulation of custom inverse formfinding problems with JAX FDM. We applied a similar
approach to generate the planar cable net shown in Fig. 1.

# optimization loop
opt = optax.adam(learning_rate=0.01)
o_state = opt.init(q)for i in range(5000):q, o_state, loss = opt_step(q, o_state)

4. Conclusion
We presented JAX FDM, an open-source solver that streamlines the design of 3D structures in static equilibrium conditioned on target properties. Future work will include oursolver as a differentiable layer in neural networks to build
accurate surrogate models that further accelerate inverse
form-finding.

The object abstractions and equilibrium calculations in JAX
FDM are compatible with JAX transformations, such as
jit and value and grad. This compatibility allows us
to write the optimization step for q with the same code
blocks conventionally used to train neural networks.
4

JAX FDM: A differentiable solver for inverse form-finding

Acknowledgements

ISSN 0266-3511. URL http://doi.org/10.1177/
0266351116660798.

This work has been supported by the U.S. National Science
Foundation under grant OAC-2118201, and by the Institutefor Data-Driven Dynamical Design (ID4).

Marmo, F., Masi, D., Mase, D., and Rosati, L. Thrust network analysis of masonry vaults. International Journal of
Masonry Research and Innovation, 4(1/2):64, 2019. ISSN
2056-9459, 2056-9467. URL http://www.doi.org/
10.1504/IJMRI.2019.096828.

References
Adriaenssens, S., Block, P., Veenendaal, D., and Williams,
C. (eds.). Shell structures for architecture: form finding
and optimization. Routledge/ Taylor & Francis Group,
London ; New York, 2014. ISBN 978-0-415-84059-0 9780-415-84060-6. URL https://doi.org/10.4324/
9781315849270.

Oval, R. Topology Finding of Patterns for Structural Design.
PhD Thesis, École des Ponts ParisTech, Paris, France,
2019. URL https://theses.hal.science/tel02917467.
Panozzo, D., Block, P., and Sorkine-Hornung, O. Designing Unreinforced Masonry Models. ACM Transactions on Graphics - SIGGRAPH 2013, 32(4):91:1–
91:12, 2013. URL http://doi.org/10.1145/
2461912.2461958.

Babuschkin, I., Baumli, K., Bell, A., Bhupatiraju, S., Bruce,J., Buchlovsky, P., Budden, D., Cai, T., Clark, A., Danihelka, I., Dedieu, A., Fantacci, C., Godwin, J., Jones, C.,
Hemsley, R., Hennigan, T., Hessel, M., Hou, S., Kapturowski, S., Keck, T., Kemaev, I., King, M., Kunesch,M., Martens, L., Merzic, H., Mikulik, V., Norman, T.,
Papamakarios, G., Quan, J., Ring, R., Ruiz, F., Sanchez,A., Schneider, R., Sezener, E., Spencer, S., Srinivasan,
S., Stokowiec, W., Wang, L., Zhou, G., and Viola, F.
The DeepMind JAX Ecosystem, 2020. URL http://github.com/deepmind.

Pastrana, R., Ohlbrock, P. O., Oberbichler, T., D’Acunto,
P., and Parascho, S. Constrained Form-Finding of Tension–Compression Structures using Automatic Differentiation. Computer-Aided Design, 155:103435, 2023.
ISSN 0010-4485. URL http://doi.org/10.1016/
j.cad.2022.103435.
Rippmann, M., Popescu, M., Augustynowicz, E.,
Méndez Echenagucia, T., Calvo Barentin, C., Frick,U., Van Mele, T., and Block, P. The Armadillo Vault:
Computational Design and Digital Fabrication of a
Freeform Stone Shell. In Advances in Architectural
Geometry 2016, pp. 344–363, Zurich, Switzerland, 2016.
vdf Hochschulverlag AG an der ETH Zürich. URL
http://doi.org/10.3218/3778-4 23.

Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary,C., Maclaurin, D., and Wanderman-Milne, S. JAX: composable transformations of Python+NumPy programs,
2018. URL http://github.com/google/jax.
Chang, K.-H. and Cheng, C.-Y. Learning to simulate and design for structural engineering. In III, H. D. and Singh, A.
(eds.), Proceedings of the 37th international conferenceon machine learning, volume 119 of Proceedings of machine learning research, pp. 1426–1436. PMLR, July
2020. URL http://proceedings.mlr.press/
v119/chang20a.html.

Schek, H.-J. The force density method for form finding
and computation of general networks. Computer Methods in Applied Mechanics and Engineering, 3(1):115–
134, January 1974. ISSN 0045-7825. URL http://doi.org/10.1016/0045-7825(74)90045-0.

Cuvilliers, P. The constrained geometry of structures:
Optimization methods for inverse form-finding design.
PhD Thesis, Massachusetts Institute of Technology.,
2020. URL http://dspace.mit.edu/handle/
1721.1/7582.

Schlaich, M. Shell Bridges - and a New Specimen Madeof Stainless Steel. Journal of the International Association for Shell and Spatial Structures, 59(3):215–
224, September 2018. ISSN 1028-365X. URL http://doi.org/10.20898/j.iass.2018.197.027.

Koohestani, K. Form-finding of tensegrity structures
via genetic algorithm.
International Journal of
Solids and Structures, 49(5):739–747, March 2012.
ISSN 0020-7683. URL http://doi.org/10.1016/
j.ijsolstr.2011.11.015.

Van Mele, T. and many others. COMPAS: A framework
for computational research in architecture and structures., 2017. URL https://doi.org/10.5281/
zenodo.2594510.

Lee, J., Mueller, C., and Fivet, C. Automatic generation of diverse equilibrium structures through shapegrammars and graphic statics. International Journal of Space Structures, 31(2-4):147–164, June 2016.

Veenendaal, D., Bakker, J., and Block, P. Structural Design of the Flexibly Formed, Mesh-Reinforced Concrete Sandwich Shell Roof of NEST HiLo. Journal of the International Association for Shell and5

JAX FDM: A differentiable solver for inverse form-finding

Spatial Structures, 58(1):23–38, March 2017. ISSN
1028365X, 19969015.
URL http://doi.org/
10.20898/j.iass.2017.191.847.
Wu, G. A framework for structural shape optimization basedon automatic differentiation, the adjoint method and accelerated linear algebra. Structural and Multidisciplinary
Optimization, 66(7):151, June 2023. ISSN 1615-1488.
URL https://doi.org/10.1007/s00158-02303601-0.
Xue, T., Liao, S., Gan, Z., Park, C., Xie, X., Liu,
W. K., and Cao, J. JAX-FEM: A differentiable GPUaccelerated 3D finite element solver for automatic inverse design and mechanistic data science. Computer
Physics Communications, 291:108802, October 2023.
ISSN 0010-4655. URL http://doi.org/10.1016/
j.cpc.2023.108802.
Zhang, J. Y. and Ohsaki, M. Adaptive force density method
for form-finding problem of tensegrity structures. International Journal of Solids and Structures, 43(18):5658–
5673, September 2006. ISSN 0020-7683. URL http://doi.org/10.1016/j.ijsolstr.2005.10.011.

6