jax-fdm 0.4.4

Auto-differentiable and hardware-accelerated force density method

JAX FDM

A differentiable, hardware-accelerated framework for constrained form-finding in structural design.

Crafted with care in the Form-Finding Lab at Princeton University ❤️🇺🇸

JAX FDM enables the solution of inverse form-finding problems for discrete force networks using the force density method (FDM) and gradient-based optimization. It streamlines the integration of form-finding simulations into deep learning models for machine learning research.

Key features

• Derivatives, JIT compilation and parallelization. JAX FDM is written in JAX, a library for high-performance numerical computing and machine learning research, and it thus inherits many of JAX's perks: calculate derivatives, parallelize, and just-in-time (JIT) compile entire form-finding simulations written in Python code, and run them on a CPU, a GPU, or a TPU 🤯.

• Legendary form-finding solver. JAX FDM computes static equilibrium states for discrete force networks with the force density method (FDM), the time-tested form-finding solver backed up by over 50 years of peer-reviewed research 📚.

• Autotune those force densities. A form-found structure should fulfill additional design requirements to become a feasible structure in the real world. Formulate an inverse form-finding scenario like this as an optimization problem with JAX FDM. Then, let one of its gradient-based optimizers solve this readme by automatically tweaking the network's force densities 🕺🏻. (Coming soon: let the optimizer tune the support positions and the applied loads too!).

• A rich bank of goals, constraints and loss functions. No two structures are alike. JAX FDM allows you to model a custom inverse form-finding problem with its (growing!) collection of goals, constraints, and loss functions via a simple, object-oriented API. The available goals and constraints in the framework are granular and applicable to an entire network; to a subset of its nodes, edges, and combinations thereof 💡.

• Form-finding simulations as another layer in a neural network. Form-finding solver as another layer in a neural network. As an automatically differentiable library, JAX FDM can be seamlessly added as a module of a differentiable function approximator (like a neural network) that can be then trained end-to-end. Let the neural network learn the underlying physics of static equilibrium directly rom a form-finding solver instead of resorting to laborious techniques like data augmentation 🤖.

JAX FDM is a research project under development. Expect sharp edges and possibly some API breaking changes as we continue to support a broader set of features.

Installation

First, create a new Anaconda environment and then activate it:

conda create -n jaxenv
conda activate jaxenv

Next, install COMPAS and COMPAS VIEW2 via conda:

conda install -c conda-forge compas compas_view2==0.7.0 

Finally, install JAX FDM with a one-liner via pip:

pip install jax-fdm

JAX FDM requires Python 3.7+, JAX 0.3.17+, Numpy 1.23.3+, Scipy 1.9.1+, and COMPAS 1.16.0+. For visualization, use COMPAS_VIEW2 0.7.0.

Disclaimer

JAX is officially supported on Linux (Ubuntu 16.04 or later) and macOS (10.12 or later) platforms. This is the case with JAX FDM too. Consequently, installing JAX FDM on Windows may require a different approach from that given by the instructitons listed above. One alternative for Windows users is to use JAX and JAX FDM on CPU and GPU via the Windows Subsystem for Linux. Please refer to JAX's installation instructions for other alternatives.

Documentation

Work in progress! Expect a release soon.

Quick example

Suppose you are interested in generating a form in static equilibrium for a 10-meter span arch subjected to vertical point loads of 0.3 kN. The arch has to be a compression-only structure. You model the arch as a jax_fdm network, apply a force density of -1 to all of its edges, and compute the required shape with the force density method.

from jax_fdm.datastructures import FDNetwork
from jax_fdm.equilibrium import fdm


network = FDNetwork.from_json("data/json/arch.json")
network.edges_forcedensities(q=-1.0)
network.nodes_supports(keys=[node for node in network.nodes() if network.is_leaf(node)])
network.nodes_loads([0.0, 0.0, -0.3])

f_network = fdm(network)

You now wish to find a new form for this arch that minimizes the total Maxwell´s load path, while keeping the length of the arch segments between 0.75 and 1 meters. You solve this constrained form-finding problem with the SLSQP gradient-based optimizer.

from jax_fdm.equilibrium import constrained_fdm
from jax_fdm.optimization import SLSQP
from jax_fdm.constraints import EdgeLengthConstraint
from jax_fdm.goals import NetworkLoadPathGoal
from jax_fdm.losses import PredictionError
from jax_fdm.losses import Loss


loss = Loss(PredictionError(goals=[NetworkLoadPathGoal]))
constraints = [EdgeLengthConstraint(edge, 0.75, 1.0) for edge in network.edges()]
optimizer = SLSQP()

c_network = constrained_fdm(network, optimizer, loss, constraints=constraints)

You finally visualize the unconstrained form-found arch f_network (gray) and the constrained one, c_network (in teal) with the Viewer.

from jax_fdm.visualization import Viewer


viewer = Viewer(width=1600, height=900)
viewer.add(c_network)
viewer.add(f_network, as_wireframe=True)
viewer.show()

The constrained form is shallower than the unconstrained one as a result of the optimization process. The length of the arch segments also varies within the prescribe bounds to minimize the load path: segments are the longest where the arch's internal forces are lower (1.0 meter, at the appex); and conversely, the segments are shorter where the arch's internal forces are higher (0.75 m, at the base).

More examples

• Creased masonry vault: Best-fit a target surface (Panozzo, et al. 2013).
• Pointy dome: Control the tilt and the coarse width of a brick dome.
• Triple-branching saddle: Design the distribution of thrusts at the supports of a monkey saddle network while constraining the edge lengths.
• Saddle bridge: Create a crease in the middle of the bridge while constraining to transversal edges of the network to a target plane.

Citation

If you found this library to be useful in academic or industry work, please consider 1) starring the project on Github, and 2) citing it:

@software{pastrana_jaxfdm,
          title={{JAX~FDM}: Auto-differentiable and hardware-accelerated force density method},
          author={Rafael Pastrana and Sigrid Adriaenssens},
          year={2022},
          doi={10.5281/zenodo.7258292},
          url={https://github.com/arpastrana/jax\_fdm}}

Acknowledgements

This work has been supported by the U.S. National Science Foundation under grant OAC-2118201 and the Institute for Data Driven Dynamical Design.

See also

COMPAS CEM: Inverse design of 3D trusses with the extended Combinatorial Equilibrium Modeling framework.

JAX: Composable transformations of Python+NumPy programs.

License

MIT