Truss analysis

In this notebook, we solve a truss analysis problem: computing forces and deformations in a 2D pin-jointed frame under load.

Features used:

• Var subclassing for node displacement variables

• @jaxls.Cost.factory for bar element strain energy

• Equality constraints for fixed supports

• Batched cost construction for all members

This is a classic introductory finite element analysis (FEA) problem using 1D bar elements.

Show code cell source

Hide code cell source

import sys
from loguru import logger

logger.remove()
logger.add(sys.stdout, format="<level>{level: <8}</level> | {message}");

import jax
import jax.numpy as jnp
import jaxls

Truss element theory

A truss is a structure of bar elements connected at pin joints (nodes). Each bar:

• Carries only axial force (tension or compression)

• Has stiffness \(k = \frac{EA}{L}\) where \(E\) is Young’s modulus, \(A\) is cross-sectional area, \(L\) is length

The strain energy in a bar element is

\[U = \frac{1}{2} k (\Delta L)^2 = \frac{1}{2} \frac{EA}{L} (L' - L)^2\]

where \(L'\) is the deformed length.

class NodeVar(jaxls.Var[jax.Array], default_factory=lambda: jnp.zeros(2)):
    """2D node displacement variable [dx, dy] in meters."""

Cost functions

1. Bar strain energy: Penalizes elongation/compression of each member

2. Support constraints: Fix displacements at support nodes

3. Load application: Prescribe displacement at load point (equilibrium is automatic)

@jaxls.Cost.factory
def bar_strain_energy(
    vals: jaxls.VarValues,
    node_i: NodeVar,
    node_j: NodeVar,
    pos_i: jax.Array,
    pos_j: jax.Array,
    EA: float,
) -> jax.Array:
    """Strain energy in a bar element: (1/2) * EA/L * (delta_L)^2.

    Args:
        node_i, node_j: Displacement variables at each end.
        pos_i, pos_j: Initial (undeformed) positions.
        EA: Axial stiffness (Young's modulus × area).
    """
    # Initial geometry.
    L0_vec = pos_j - pos_i
    L0 = jnp.sqrt(jnp.sum(L0_vec**2))

    # Deformed geometry.
    disp_i = vals[node_i]
    disp_j = vals[node_j]
    L_vec = L0_vec + (disp_j - disp_i)
    L = jnp.sqrt(jnp.sum(L_vec**2))

    # Return 2D residual instead of scalar: ||r||^2 = (EA/L0) * (L - L0)^2.
    # Using a 2D residual gives a rank-2 contribution to J^T J (the Gauss-Newton.
    # Hessian approximation), rather than rank-1 from a scalar residual.
    return jnp.sqrt(EA / L0) * (1 - L0 / L) * L_vec


@jaxls.Cost.factory(kind="constraint_eq_zero")
def pin_support(
    vals: jaxls.VarValues,
    node: NodeVar,
) -> jax.Array:
    """Pin support: both displacement components are zero."""
    return vals[node]


@jaxls.Cost.factory(kind="constraint_eq_zero")
def prescribed_displacement(
    vals: jaxls.VarValues,
    node: NodeVar,
    target_displacement: jax.Array,
) -> jax.Array:
    """Prescribe displacement at a node."""
    return vals[node] - target_displacement

Truss geometry

We model a Warren truss, a common bridge structure with diagonal members:

   5-----6-----7-----8
  /\    /\    /\    /\
 /  \  /  \  /  \  /  \
/    \/    \/    \/    \
0-----1-----2-----3-----4

• Nodes 0-4: Bottom chord

• Nodes 5-8: Top chord

• Nodes 0 and 4 are pinned (fixed in x and y)

• Load applied at center bottom node (node 2)

# Geometry: Warren truss bridge.
num_panels = 4  # Number of triangular panels
panel_width = 3.0  # [m] width of each panel
height = 2.0  # [m] truss height
span = num_panels * panel_width  # Total span

# Build node positions.
bottom_nodes = [[i * panel_width, 0.0] for i in range(num_panels + 1)]
top_nodes = [[(i + 0.5) * panel_width, height] for i in range(num_panels)]
node_positions = jnp.array(bottom_nodes + top_nodes)
num_nodes = len(node_positions)

# Node indices.
bottom_ids = list(range(num_panels + 1))  # 0, 1, 2, 3, 4
top_ids = list(range(num_panels + 1, num_nodes))  # 5, 6, 7, 8

# Build member connectivity.
member_list = []
# Bottom chord.
for i in range(num_panels):
    member_list.append([bottom_ids[i], bottom_ids[i + 1]])
# Top chord.
for i in range(num_panels - 1):
    member_list.append([top_ids[i], top_ids[i + 1]])
# Diagonals (left and right of each top node)
for i in range(num_panels):
    member_list.append([bottom_ids[i], top_ids[i]])  # Left diagonal
    member_list.append([top_ids[i], bottom_ids[i + 1]])  # Right diagonal

members = jnp.array(member_list)
num_members = len(members)

# Material properties.
EA = 50000.0  # [N] axial stiffness

# Load node.
load_node_id = 2  # Center bottom node

# Prescribed displacement (downward)
load_displacement = jnp.array([0.0, -0.02])  # 20 mm downward

print("Warren Truss Bridge:")
print(f"  Span: {span} m, Height: {height} m")
print(f"  Nodes: {num_nodes}, Members: {num_members}")
print(f"  Member stiffness EA = {EA:.0f} N")
print(
    f"  Prescribed displacement at node {load_node_id}: {float(load_displacement[1]) * 1000:.1f} mm (vertical)"
)

Warren Truss Bridge:
  Span: 12.0 m, Height: 2.0 m
  Nodes: 9, Members: 15
  Member stiffness EA = 50000 N
  Prescribed displacement at node 2: -20.0 mm (vertical)

Problem construction

# Create node displacement variables.
node_vars = NodeVar(id=jnp.arange(num_nodes))

# Support nodes (both pinned).
left_pin = 0
right_pin = num_panels

# Build costs.
costs: list[jaxls.Cost] = [
    # Strain energy in all members (batched).
    bar_strain_energy(
        NodeVar(id=members[:, 0]),
        NodeVar(id=members[:, 1]),
        node_positions[members[:, 0]],
        node_positions[members[:, 1]],
        EA,
    ),
    # Boundary conditions.
    pin_support(NodeVar(id=left_pin)),
    pin_support(NodeVar(id=right_pin)),
    # Applied load via prescribed displacement.
    prescribed_displacement(NodeVar(id=load_node_id), load_displacement),
]

print(f"Created {len(costs)} cost objects")
print(f"Load applied at node {load_node_id}")

Created 4 cost objects
Load applied at node 2

Solving

# Initial values: zero displacement.
initial_displacements = jnp.zeros((num_nodes, 2))
initial_vals = jaxls.VarValues.make([node_vars.with_value(initial_displacements)])

# Build the problem.
problem = jaxls.LeastSquaresProblem(costs, [node_vars])

# Visualize the problem structure structure.
problem.show()

# Analyze and solve.
problem = problem.analyze()
solution = problem.solve(initial_vals)

INFO     | Building optimization problem with 18 terms and 9 variables: 15 costs, 3 eq_zero, 0 leq_zero, 0 geq_zero
INFO     | Vectorizing constraint group with 2 constraints (constraint_eq_zero), 1 variables each: augmented_pin_support
INFO     | Vectorizing group with 15 costs, 2 variables each: bar_strain_energy
INFO     | Vectorizing constraint group with 1 constraints (constraint_eq_zero), 1 variables each: augmented_prescribed_displacement
INFO     | Augmented Lagrangian: initial snorm=2.0000e-02, csupn=2.0000e-02, max_rho=1.0000e+01, constraint_dim=6
INFO     |  step #0: cost=0.0000 lambd=0.0005 inexact_tol=1.0e-02
INFO     |      - augmented_pin_support(2): 0.00000 (avg 0.00000)
INFO     |      - bar_strain_energy(15): 0.00000 (avg 0.00000)
INFO     |      - augmented_prescribed_displacement(1): 0.00400 (avg 0.00200)
INFO     |      accepted=True ATb_norm=2.01e-01 cost_prev=0.0040 cost_new=0.0027
INFO     |  step #1: cost=0.0000 lambd=0.0003 inexact_tol=1.0e-02
INFO     |      - augmented_pin_support(2): 0.00088 (avg 0.00022)
INFO     |      - bar_strain_energy(15): 0.00001 (avg 0.00000)
INFO     |      - augmented_prescribed_displacement(1): 0.00177 (avg 0.00088)
INFO     |      accepted=True ATb_norm=1.28e-02 cost_prev=0.0027 cost_new=0.0027
INFO     |  AL update: snorm=1.3288e-02, csupn=1.3288e-02, max_rho=4.0000e+01
INFO     |  step #2: cost=0.0000 lambd=0.0001 inexact_tol=3.6e-03
INFO     |      - augmented_pin_support(2): 0.00552 (avg 0.00138)
INFO     |      - bar_strain_energy(15): 0.00001 (avg 0.00000)
INFO     |      - augmented_prescribed_displacement(1): 0.01104 (avg 0.00552)
INFO     |      accepted=True ATb_norm=6.59e-01 cost_prev=0.0166 cost_new=0.0164
INFO     |  step #3: cost=0.0002 lambd=0.0001 inexact_tol=3.6e-03
INFO     |      - augmented_pin_support(2): 0.00540 (avg 0.00135)
INFO     |      - bar_strain_energy(15): 0.00023 (avg 0.00001)
INFO     |      - augmented_prescribed_displacement(1): 0.01079 (avg 0.00540)
INFO     |      accepted=True ATb_norm=1.42e-02 cost_prev=0.0164 cost_new=0.0164
INFO     |  AL update: snorm=1.3103e-02, csupn=1.3103e-02, max_rho=1.6000e+02
INFO     |  step #4: cost=0.0002 lambd=0.0000 inexact_tol=4.2e-04
INFO     |      - augmented_pin_support(2): 0.02370 (avg 0.00592)
INFO     |      - bar_strain_energy(15): 0.00023 (avg 0.00001)
INFO     |      - augmented_prescribed_displacement(1): 0.04739 (avg 0.02369)
INFO     |      accepted=True ATb_norm=2.58e+00 cost_prev=0.0713 cost_new=0.0691
INFO     |  step #5: cost=0.0036 lambd=0.0000 inexact_tol=4.2e-04
INFO     |      - augmented_pin_support(2): 0.02184 (avg 0.00546)
INFO     |      - bar_strain_energy(15): 0.00365 (avg 0.00012)
INFO     |      - augmented_prescribed_displacement(1): 0.04364 (avg 0.02182)
INFO     |  step #6: cost=0.0036 lambd=0.0001 inexact_tol=4.2e-04
INFO     |      - augmented_pin_support(2): 0.02184 (avg 0.00546)
INFO     |      - bar_strain_energy(15): 0.00365 (avg 0.00012)
INFO     |      - augmented_prescribed_displacement(1): 0.04364 (avg 0.02182)
INFO     |  step #7: cost=0.0036 lambd=0.0005 inexact_tol=4.2e-04
INFO     |      - augmented_pin_support(2): 0.02184 (avg 0.00546)
INFO     |      - bar_strain_energy(15): 0.00365 (avg 0.00012)
INFO     |      - augmented_prescribed_displacement(1): 0.04364 (avg 0.02182)
INFO     |  step #8: cost=0.0036 lambd=0.0080 inexact_tol=4.2e-04
INFO     |      - augmented_pin_support(2): 0.02184 (avg 0.00546)
INFO     |      - bar_strain_energy(15): 0.00365 (avg 0.00012)
INFO     |      - augmented_prescribed_displacement(1): 0.04364 (avg 0.02182)
INFO     |  step #9: cost=0.0036 lambd=0.2560 inexact_tol=4.2e-04
INFO     |      - augmented_pin_support(2): 0.02184 (avg 0.00546)
INFO     |      - bar_strain_energy(15): 0.00365 (avg 0.00012)
INFO     |      - augmented_prescribed_displacement(1): 0.04364 (avg 0.02182)
INFO     |  step #10: cost=0.0036 lambd=16.3840 inexact_tol=4.2e-04
INFO     |      - augmented_pin_support(2): 0.02184 (avg 0.00546)
INFO     |      - bar_strain_energy(15): 0.00365 (avg 0.00012)
INFO     |      - augmented_prescribed_displacement(1): 0.04364 (avg 0.02182)
INFO     |  step #11: cost=0.0036 lambd=2097.1521 inexact_tol=4.2e-04
INFO     |      - augmented_pin_support(2): 0.02184 (avg 0.00546)
INFO     |      - bar_strain_energy(15): 0.00365 (avg 0.00012)
INFO     |      - augmented_prescribed_displacement(1): 0.04364 (avg 0.02182)
INFO     |      accepted=True ATb_norm=2.10e-02 cost_prev=0.0691 cost_new=0.0691
INFO     |  AL update: snorm=1.2409e-02, csupn=1.2409e-02, max_rho=6.4000e+02
INFO     |  step #12: cost=0.0036 lambd=1048.5760 inexact_tol=5.9e-05
INFO     |      - augmented_pin_support(2): 0.08763 (avg 0.02191)
INFO     |      - bar_strain_energy(15): 0.00365 (avg 0.00012)
INFO     |      - augmented_prescribed_displacement(1): 0.17504 (avg 0.08752)
INFO     |      accepted=True ATb_norm=9.79e+00 cost_prev=0.2663 cost_new=0.2436
INFO     |  step #13: cost=0.0240 lambd=524.2880 inexact_tol=5.9e-05
INFO     |      - augmented_pin_support(2): 0.06792 (avg 0.01698)
INFO     |      - bar_strain_energy(15): 0.02404 (avg 0.00080)
INFO     |      - augmented_prescribed_displacement(1): 0.15161 (avg 0.07581)
INFO     |      accepted=True ATb_norm=2.45e+00 cost_prev=0.2436 cost_new=0.2403
INFO     |  step #14: cost=0.0373 lambd=262.1440 inexact_tol=5.9e-05
INFO     |      - augmented_pin_support(2): 0.06203 (avg 0.01551)
INFO     |      - bar_strain_energy(15): 0.03730 (avg 0.00124)
INFO     |      - augmented_prescribed_displacement(1): 0.14100 (avg 0.07050)
INFO     |      accepted=True ATb_norm=5.84e-01 cost_prev=0.2403 cost_new=0.2400
INFO     |  step #15: cost=0.0414 lambd=131.0720 inexact_tol=5.9e-05
INFO     |      - augmented_pin_support(2): 0.06288 (avg 0.01572)
INFO     |      - bar_strain_energy(15): 0.04138 (avg 0.00138)
INFO     |      - augmented_prescribed_displacement(1): 0.13571 (avg 0.06786)
INFO     |      accepted=True ATb_norm=1.57e-01 cost_prev=0.2400 cost_new=0.2399
INFO     |  step #16: cost=0.0423 lambd=65.5360 inexact_tol=5.9e-05
INFO     |      - augmented_pin_support(2): 0.06489 (avg 0.01622)
INFO     |      - bar_strain_energy(15): 0.04227 (avg 0.00141)
INFO     |      - augmented_prescribed_displacement(1): 0.13273 (avg 0.06637)
INFO     |      accepted=True ATb_norm=5.83e-02 cost_prev=0.2399 cost_new=0.2399
INFO     |  step #17: cost=0.0425 lambd=32.7680 inexact_tol=5.9e-05
INFO     |      - augmented_pin_support(2): 0.06599 (avg 0.01650)
INFO     |      - bar_strain_energy(15): 0.04249 (avg 0.00142)
INFO     |      - augmented_prescribed_displacement(1): 0.13139 (avg 0.06570)
INFO     |      accepted=True ATb_norm=2.00e-02 cost_prev=0.2399 cost_new=0.2399
INFO     |  step #18: cost=0.0425 lambd=16.3840 inexact_tol=5.9e-05
INFO     |      - augmented_pin_support(2): 0.06630 (avg 0.01657)
INFO     |      - bar_strain_energy(15): 0.04254 (avg 0.00142)
INFO     |      - augmented_prescribed_displacement(1): 0.13103 (avg 0.06551)
INFO     |      accepted=True ATb_norm=2.04e-02 cost_prev=0.2399 cost_new=0.2399
INFO     |  AL update: snorm=1.0176e-02, csupn=1.0176e-02, max_rho=2.5600e+03
INFO     |  step #19: cost=0.0426 lambd=8.1920 inexact_tol=5.9e-05
INFO     |      - augmented_pin_support(2): 0.24684 (avg 0.06171)
INFO     |      - bar_strain_energy(15): 0.04255 (avg 0.00142)
INFO     |      - augmented_prescribed_displacement(1): 0.48419 (avg 0.24209)
INFO     |      accepted=True ATb_norm=3.23e+01 cost_prev=0.7736 cost_new=0.5982
INFO     |  step #20: cost=0.2573 lambd=4.0960 inexact_tol=5.9e-05
INFO     |      - augmented_pin_support(2): 0.12377 (avg 0.03094)
INFO     |      - bar_strain_energy(15): 0.25728 (avg 0.00858)
INFO     |      - augmented_prescribed_displacement(1): 0.21712 (avg 0.10856)
INFO     |      accepted=True ATb_norm=1.94e-01 cost_prev=0.5982 cost_new=0.5982
INFO     |  step #21: cost=0.2584 lambd=2.0480 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.12395 (avg 0.03099)
INFO     |      - bar_strain_energy(15): 0.25836 (avg 0.00861)
INFO     |      - augmented_prescribed_displacement(1): 0.21584 (avg 0.10792)
INFO     |  step #22: cost=0.2584 lambd=8.1920 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.12395 (avg 0.03099)
INFO     |      - bar_strain_energy(15): 0.25836 (avg 0.00861)
INFO     |      - augmented_prescribed_displacement(1): 0.21584 (avg 0.10792)
INFO     |  step #23: cost=0.2584 lambd=65.5360 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.12395 (avg 0.03099)
INFO     |      - bar_strain_energy(15): 0.25836 (avg 0.00861)
INFO     |      - augmented_prescribed_displacement(1): 0.21584 (avg 0.10792)
INFO     |  step #24: cost=0.2584 lambd=1048.5760 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.12395 (avg 0.03099)
INFO     |      - bar_strain_energy(15): 0.25836 (avg 0.00861)
INFO     |      - augmented_prescribed_displacement(1): 0.21584 (avg 0.10792)
INFO     |  step #25: cost=0.2584 lambd=33554.4336 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.12395 (avg 0.03099)
INFO     |      - bar_strain_energy(15): 0.25836 (avg 0.00861)
INFO     |      - augmented_prescribed_displacement(1): 0.21584 (avg 0.10792)
INFO     |      accepted=True ATb_norm=1.27e-02 cost_prev=0.5982 cost_new=0.5981
INFO     |  AL update: snorm=5.6059e-03, csupn=5.6059e-03, max_rho=1.0240e+04
INFO     |  step #26: cost=0.2583 lambd=16777.2168 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.40212 (avg 0.10053)
INFO     |      - bar_strain_energy(15): 0.25835 (avg 0.00861)
INFO     |      - augmented_prescribed_displacement(1): 0.63932 (avg 0.31966)
INFO     |      accepted=True ATb_norm=7.44e+01 cost_prev=1.2998 cost_new=1.0933
INFO     |  step #27: cost=0.3971 lambd=8388.6084 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.26187 (avg 0.06547)
INFO     |      - bar_strain_energy(15): 0.39713 (avg 0.01324)
INFO     |      - augmented_prescribed_displacement(1): 0.43430 (avg 0.21715)
INFO     |      accepted=True ATb_norm=3.44e+01 cost_prev=1.0933 cost_new=0.9951
INFO     |  step #28: cost=0.5128 lambd=4194.3042 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.17627 (avg 0.04407)
INFO     |      - bar_strain_energy(15): 0.51278 (avg 0.01709)
INFO     |      - augmented_prescribed_displacement(1): 0.30602 (avg 0.15301)
INFO     |      accepted=True ATb_norm=1.68e+01 cost_prev=0.9951 cost_new=0.9579
INFO     |  step #29: cost=0.5915 lambd=2097.1521 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.13929 (avg 0.03482)
INFO     |      - bar_strain_energy(15): 0.59153 (avg 0.01972)
INFO     |      - augmented_prescribed_displacement(1): 0.22710 (avg 0.11355)
INFO     |      accepted=True ATb_norm=6.83e+00 cost_prev=0.9579 cost_new=0.9481
INFO     |  step #30: cost=0.6313 lambd=1048.5760 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.13388 (avg 0.03347)
INFO     |      - bar_strain_energy(15): 0.63133 (avg 0.02104)
INFO     |      - augmented_prescribed_displacement(1): 0.18289 (avg 0.09144)
INFO     |      accepted=True ATb_norm=2.39e+00 cost_prev=0.9481 cost_new=0.9461
INFO     |  step #31: cost=0.6455 lambd=524.2880 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.13851 (avg 0.03463)
INFO     |      - bar_strain_energy(15): 0.64549 (avg 0.02152)
INFO     |      - augmented_prescribed_displacement(1): 0.16213 (avg 0.08106)
INFO     |      accepted=True ATb_norm=7.15e-01 cost_prev=0.9461 cost_new=0.9458
INFO     |  step #32: cost=0.6492 lambd=262.1440 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.14164 (avg 0.03541)
INFO     |      - bar_strain_energy(15): 0.64917 (avg 0.02164)
INFO     |      - augmented_prescribed_displacement(1): 0.15502 (avg 0.07751)
INFO     |      accepted=True ATb_norm=1.45e-01 cost_prev=0.9458 cost_new=0.9458
INFO     |  step #33: cost=0.6499 lambd=131.0720 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.14244 (avg 0.03561)
INFO     |      - bar_strain_energy(15): 0.64986 (avg 0.02166)
INFO     |      - augmented_prescribed_displacement(1): 0.15349 (avg 0.07675)
INFO     |  step #34: cost=0.6499 lambd=524.2880 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.14244 (avg 0.03561)
INFO     |      - bar_strain_energy(15): 0.64986 (avg 0.02166)
INFO     |      - augmented_prescribed_displacement(1): 0.15349 (avg 0.07675)
INFO     |  step #35: cost=0.6499 lambd=4194.3042 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.14244 (avg 0.03561)
INFO     |      - bar_strain_energy(15): 0.64986 (avg 0.02166)
INFO     |      - augmented_prescribed_displacement(1): 0.15349 (avg 0.07675)
INFO     |  step #36: cost=0.6499 lambd=67108.8672 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.14244 (avg 0.03561)
INFO     |      - bar_strain_energy(15): 0.64986 (avg 0.02166)
INFO     |      - augmented_prescribed_displacement(1): 0.15349 (avg 0.07675)
INFO     |      accepted=True ATb_norm=2.45e-02 cost_prev=0.9458 cost_new=0.9458
INFO     |  AL update: snorm=1.5760e-03, csupn=1.5760e-03, max_rho=1.0240e+04
INFO     |  step #37: cost=0.6499 lambd=33554.4336 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.35392 (avg 0.08848)
INFO     |      - bar_strain_energy(15): 0.64987 (avg 0.02166)
INFO     |      - augmented_prescribed_displacement(1): 0.30388 (avg 0.15194)
INFO     |      accepted=True ATb_norm=2.80e+01 cost_prev=1.3077 cost_new=1.2883
INFO     |  step #38: cost=0.6996 lambd=16777.2168 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.31348 (avg 0.07837)
INFO     |      - bar_strain_energy(15): 0.69964 (avg 0.02332)
INFO     |      - augmented_prescribed_displacement(1): 0.27515 (avg 0.13758)
INFO     |      accepted=True ATb_norm=1.55e+01 cost_prev=1.2883 cost_new=1.2768
INFO     |  step #39: cost=0.7474 lambd=8388.6084 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.27987 (avg 0.06997)
INFO     |      - bar_strain_energy(15): 0.74742 (avg 0.02491)
INFO     |      - augmented_prescribed_displacement(1): 0.24956 (avg 0.12478)
INFO     |      accepted=True ATb_norm=8.37e+00 cost_prev=1.2768 cost_new=1.2712
INFO     |  step #40: cost=0.7861 lambd=4194.3042 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.25598 (avg 0.06400)
INFO     |      - bar_strain_energy(15): 0.78606 (avg 0.02620)
INFO     |      - augmented_prescribed_displacement(1): 0.22919 (avg 0.11460)
INFO     |      accepted=True ATb_norm=3.92e+00 cost_prev=1.2712 cost_new=1.2695
INFO     |  step #41: cost=0.8105 lambd=2097.1521 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.24390 (avg 0.06097)
INFO     |      - bar_strain_energy(15): 0.81049 (avg 0.02702)
INFO     |      - augmented_prescribed_displacement(1): 0.21507 (avg 0.10753)
INFO     |      accepted=True ATb_norm=1.42e+00 cost_prev=1.2695 cost_new=1.2691
INFO     |  step #42: cost=0.8212 lambd=1048.5760 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.24112 (avg 0.06028)
INFO     |      - bar_strain_energy(15): 0.82121 (avg 0.02737)
INFO     |      - augmented_prescribed_displacement(1): 0.20679 (avg 0.10340)
INFO     |      accepted=True ATb_norm=4.32e-01 cost_prev=1.2691 cost_new=1.2691
INFO     |  step #43: cost=0.8245 lambd=524.2880 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.24179 (avg 0.06045)
INFO     |      - bar_strain_energy(15): 0.82450 (avg 0.02748)
INFO     |      - augmented_prescribed_displacement(1): 0.20283 (avg 0.10142)
INFO     |      accepted=True ATb_norm=1.20e-01 cost_prev=1.2691 cost_new=1.2690
INFO     |  step #44: cost=0.8251 lambd=262.1440 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.24241 (avg 0.06060)
INFO     |      - bar_strain_energy(15): 0.82513 (avg 0.02750)
INFO     |      - augmented_prescribed_displacement(1): 0.20147 (avg 0.10074)
INFO     |  step #45: cost=0.8251 lambd=1048.5760 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.24241 (avg 0.06060)
INFO     |      - bar_strain_energy(15): 0.82513 (avg 0.02750)
INFO     |      - augmented_prescribed_displacement(1): 0.20147 (avg 0.10074)
INFO     |      accepted=True ATb_norm=2.92e-02 cost_prev=1.2690 cost_new=1.2690
INFO     |  AL update: snorm=8.3996e-04, csupn=8.3996e-04, max_rho=4.0960e+04
INFO     |  step #46: cost=0.8252 lambd=524.2880 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.24110 (avg 0.06028)
INFO     |      - bar_strain_energy(15): 0.82523 (avg 0.02751)
INFO     |      - augmented_prescribed_displacement(1): 0.25546 (avg 0.12773)
INFO     |      accepted=True ATb_norm=5.20e+01 cost_prev=1.3218 cost_new=1.2709
INFO     |  step #47: cost=0.9460 lambd=262.1440 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.09125 (avg 0.02281)
INFO     |      - bar_strain_energy(15): 0.94603 (avg 0.03153)
INFO     |      - augmented_prescribed_displacement(1): 0.23362 (avg 0.11681)
INFO     |      accepted=True ATb_norm=6.46e-01 cost_prev=1.2709 cost_new=1.2709
INFO     |  step #48: cost=0.9478 lambd=131.0720 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.08923 (avg 0.02231)
INFO     |      - bar_strain_energy(15): 0.94778 (avg 0.03159)
INFO     |      - augmented_prescribed_displacement(1): 0.23390 (avg 0.11695)
INFO     |      accepted=True ATb_norm=2.02e-02 cost_prev=1.2709 cost_new=1.2709
INFO     |  AL update: snorm=3.4721e-04, csupn=3.4721e-04, max_rho=4.0960e+04
INFO     |  step #49: cost=0.9477 lambd=65.5360 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.12438 (avg 0.03109)
INFO     |      - bar_strain_energy(15): 0.94771 (avg 0.03159)
INFO     |      - augmented_prescribed_displacement(1): 0.09743 (avg 0.04872)
INFO     |      accepted=True ATb_norm=1.85e+01 cost_prev=1.1695 cost_new=1.1627
INFO     |  step #50: cost=1.0019 lambd=32.7680 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.09875 (avg 0.02469)
INFO     |      - bar_strain_energy(15): 1.00190 (avg 0.03340)
INFO     |      - augmented_prescribed_displacement(1): 0.06202 (avg 0.03101)
INFO     |  step #51: cost=1.0019 lambd=131.0720 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.09875 (avg 0.02469)
INFO     |      - bar_strain_energy(15): 1.00190 (avg 0.03340)
INFO     |      - augmented_prescribed_displacement(1): 0.06202 (avg 0.03101)
INFO     |  step #52: cost=1.0019 lambd=1048.5760 inexact_tol=3.3e-05
INFO     |      - augmented_pin_support(2): 0.09875 (avg 0.02469)
INFO     |      - bar_strain_energy(15): 1.00190 (avg 0.03340)
INFO     |      - augmented_prescribed_displacement(1): 0.06202 (avg 0.03101)
INFO     |      accepted=True ATb_norm=3.90e-02 cost_prev=1.1627 cost_new=1.1626
INFO     |  AL update: snorm=5.4603e-05, csupn=5.4603e-05, max_rho=4.0960e+04
INFO     |  step #53: cost=1.0020 lambd=524.2880 inexact_tol=4.0e-06
INFO     |      - augmented_pin_support(2): 0.10886 (avg 0.02722)
INFO     |      - bar_strain_energy(15): 1.00196 (avg 0.03340)
INFO     |      - augmented_prescribed_displacement(1): 0.06540 (avg 0.03270)
INFO     |      accepted=True ATb_norm=3.62e+00 cost_prev=1.1762 cost_new=1.1760
INFO     |  step #54: cost=1.0123 lambd=262.1440 inexact_tol=4.0e-06
INFO     |      - augmented_pin_support(2): 0.10119 (avg 0.02530)
INFO     |      - bar_strain_energy(15): 1.01229 (avg 0.03374)
INFO     |      - augmented_prescribed_displacement(1): 0.06257 (avg 0.03128)
INFO     |      accepted=True ATb_norm=4.24e-02 cost_prev=1.1760 cost_new=1.1760
INFO     |  AL update: snorm=1.3334e-05, csupn=1.3334e-05, max_rho=4.0960e+04
INFO     |  step #55: cost=1.0124 lambd=131.0720 inexact_tol=4.0e-06
INFO     |      - augmented_pin_support(2): 0.10344 (avg 0.02586)
INFO     |      - bar_strain_energy(15): 1.01240 (avg 0.03375)
INFO     |      - augmented_prescribed_displacement(1): 0.06315 (avg 0.03157)
INFO     |  step #56: cost=1.0124 lambd=524.2880 inexact_tol=4.0e-06
INFO     |      - augmented_pin_support(2): 0.10344 (avg 0.02586)
INFO     |      - bar_strain_energy(15): 1.01240 (avg 0.03375)
INFO     |      - augmented_prescribed_displacement(1): 0.06315 (avg 0.03157)
INFO     |  step #57: cost=1.0124 lambd=4194.3042 inexact_tol=4.0e-06
INFO     |      - augmented_pin_support(2): 0.10344 (avg 0.02586)
INFO     |      - bar_strain_energy(15): 1.01240 (avg 0.03375)
INFO     |      - augmented_prescribed_displacement(1): 0.06315 (avg 0.03157)
INFO     |  step #58: cost=1.0124 lambd=67108.8672 inexact_tol=4.0e-06
INFO     |      - augmented_pin_support(2): 0.10344 (avg 0.02586)
INFO     |      - bar_strain_energy(15): 1.01240 (avg 0.03375)
INFO     |      - augmented_prescribed_displacement(1): 0.06315 (avg 0.03157)
INFO     |  step #59: cost=1.0124 lambd=1000000.0000 inexact_tol=4.0e-06
INFO     |      - augmented_pin_support(2): 0.10344 (avg 0.02586)
INFO     |      - bar_strain_energy(15): 1.01240 (avg 0.03375)
INFO     |      - augmented_prescribed_displacement(1): 0.06315 (avg 0.03157)
INFO     |      accepted=True ATb_norm=8.40e-01 cost_prev=1.1790 cost_new=1.1790
INFO     |  AL update: snorm=1.2817e-05, csupn=1.2817e-05, max_rho=1.6384e+05
INFO     |  step #60: cost=1.0125 lambd=500000.0000 inexact_tol=4.0e-06
INFO     |      - augmented_pin_support(2): 0.02818 (avg 0.00704)
INFO     |      - bar_strain_energy(15): 1.01248 (avg 0.03375)
INFO     |      - augmented_prescribed_displacement(1): 0.01639 (avg 0.00819)
INFO     |      accepted=True ATb_norm=4.02e+00 cost_prev=1.0570 cost_new=1.0570
INFO     |  AL update: snorm=8.9874e-06, csupn=8.9874e-06, max_rho=6.5536e+05
INFO     | Terminated @ iteration #61: cost=1.0133 criteria=[0 0 1], term_deltas=8.3e-04,2.6e+00,4.1e-07 (solved in 0.2201 sec)

Results and visualization

Show code cell source

Hide code cell source

# Extract displacements.
displacements = solution[node_vars]
deformed_positions = node_positions + displacements


def compute_member_force(i: int, j: int, disp: jax.Array) -> jax.Array:
    """Compute axial force in member (positive = tension).

    Args:
        i: Start node index
        j: End node index
        disp: Node displacements array (num_nodes, 2)

    Returns:
        Axial force in the member (scalar)
    """
    L0_vec = node_positions[j] - node_positions[i]
    L0 = jnp.sqrt(jnp.sum(L0_vec**2))
    L_vec = L0_vec + (disp[j] - disp[i])
    L = jnp.sqrt(jnp.sum(L_vec**2))
    strain = (L - L0) / L0
    return EA * strain  # Force = EA * strain


member_forces = jax.vmap(lambda m: compute_member_force(m[0], m[1], displacements))(
    members
)

# Print results.
print("Node Displacements:")
print(f"{'Node':>4} {'dx [mm]':>10} {'dy [mm]':>10}")
print("-" * 26)
for i in range(num_nodes):
    dx, dy = displacements[i] * 1000  # Convert to mm
    print(f"{i:>4} {float(dx):>10.3f} {float(dy):>10.3f}")

print("\nMember Forces:")
print(f"{'Member':>6} {'Force [kN]':>12} {'Type':>10}")
print("-" * 30)
for idx, m in enumerate(members):
    f = member_forces[idx]
    f_kN = float(f) / 1000
    typ = "Tension" if f > 0 else "Compression"
    print(f"{int(m[0])}-{int(m[1]):>2} {f_kN:>12.2f} {typ:>10}")

Node Displacements:
Node    dx [mm]    dy [mm]
--------------------------
   0     -0.009     -0.002
   1     -1.156    -11.739
   2     -0.000    -19.996
   3      1.156    -11.739
   4      0.009     -0.002
   5      4.593     -5.457
   6      2.284    -16.314
   7     -2.284    -16.314
   8     -4.593     -5.457

Member Forces:
Member   Force [kN]       Type
------------------------------
0- 1        -0.02 Compression
1- 2         0.02    Tension
2- 3         0.02    Tension
3- 4        -0.02 Compression
5- 6        -0.04 Compression
6- 7        -0.08 Compression
7- 8        -0.04 Compression
0- 5        -0.03 Compression
5- 1         0.03    Tension
1- 6        -0.03 Compression
6- 2         0.03    Tension
2- 7         0.03    Tension
7- 3        -0.03 Compression
3- 8         0.03    Tension
8- 4        -0.03 Compression

Show code cell source

Hide code cell source

import plotly.graph_objects as go
from IPython.display import HTML

# Visualization.
scale = 20  # Displacement magnification for visibility
scaled_deformed = node_positions + scale * displacements

fig = go.Figure()

# Original structure (gray)
for m in members:
    i, j = int(m[0]), int(m[1])
    fig.add_trace(
        go.Scatter(
            x=[float(node_positions[i, 0]), float(node_positions[j, 0])],
            y=[float(node_positions[i, 1]), float(node_positions[j, 1])],
            mode="lines",
            line=dict(color="lightgray", width=6),
            showlegend=False,
        )
    )

# Deformed structure (colored by force)
max_force = float(jnp.max(jnp.abs(member_forces))) + 1e-6
for idx, m in enumerate(members):
    i, j = int(m[0]), int(m[1])
    force = float(member_forces[idx])
    # Color: blue for compression, red for tension.
    intensity = min(abs(force) / max_force, 1.0)
    if force > 0:
        color = f"rgba(220, {int(80 + 175 * (1 - intensity))}, {int(80 + 175 * (1 - intensity))}, 1)"
    else:
        color = f"rgba({int(80 + 175 * (1 - intensity))}, {int(80 + 175 * (1 - intensity))}, 220, 1)"

    fig.add_trace(
        go.Scatter(
            x=[float(scaled_deformed[i, 0]), float(scaled_deformed[j, 0])],
            y=[float(scaled_deformed[i, 1]), float(scaled_deformed[j, 1])],
            mode="lines",
            line=dict(color=color, width=5),
            showlegend=False,
        )
    )

# Nodes.
fig.add_trace(
    go.Scatter(
        x=[float(p) for p in scaled_deformed[:, 0]],
        y=[float(p) for p in scaled_deformed[:, 1]],
        mode="markers+text",
        marker=dict(size=10, color="steelblue"),
        text=[str(i) for i in range(num_nodes)],
        textposition="top center",
        textfont=dict(size=9),
        showlegend=False,
    )
)

# Load arrow at load node (pointing downward).
fig.add_annotation(
    x=float(scaled_deformed[load_node_id, 0]),
    y=float(scaled_deformed[load_node_id, 1]) - 0.8,
    ax=0,
    ay=-40,
    xref="x",
    yref="y",
    axref="pixel",
    ayref="pixel",
    showarrow=True,
    arrowhead=2,
    arrowsize=1.5,
    arrowwidth=3,
    arrowcolor="red",
)
fig.add_annotation(
    x=float(scaled_deformed[load_node_id, 0]),
    y=float(scaled_deformed[load_node_id, 1]) - 1.3,
    text="Load",
    showarrow=False,
    font=dict(size=12, color="red"),
)

fig.update_layout(
    title=f"Warren Truss Bridge Analysis (displacements x{scale})",
    xaxis=dict(title="x [m]", scaleanchor="y", scaleratio=1),
    yaxis=dict(title="y [m]"),
    height=400,
    showlegend=False,
    margin=dict(t=80, b=50, l=50, r=50),
)

# Add color legend.
fig.add_annotation(
    x=0.02,
    y=0.98,
    xref="paper",
    yref="paper",
    text="Red = Tension, Blue = Compression",
    showarrow=False,
    font=dict(size=10),
    align="left",
    bgcolor="white",
    bordercolor="gray",
    borderwidth=1,
)

HTML(fig.to_html(full_html=False, include_plotlyjs="cdn"))

The solver found the equilibrium configuration of the Warren truss bridge under a prescribed displacement:

• Deformed shape: Shown with exaggerated displacements for visibility

• Member colors: Red indicates tension, blue indicates compression

Key observations:

• Bottom chord is in tension (red) - it resists the spreading of the supports

• Top chord is in compression (blue) - it’s being squeezed as the bridge sags

• Diagonal members alternate between tension and compression

• Maximum deflection occurs at the center where the load is applied

This is the classic behavior of a simply-supported truss bridge under a center point load.

Varying displacements

We can animate the truss response to different prescribed displacements. As the displacement increases, the internal forces grow proportionally (since we’re in the linear elastic regime).

Using jax.vmap, we solve for all displacement magnitudes in parallel.

Show code cell source

Hide code cell source

# Solve for multiple displacement magnitudes using vmap.
displacement_magnitudes = jnp.linspace(0, 0.05, 21)  # 0 to 50 mm.


def solve_for_displacement(disp_y: jax.Array) -> jax.Array:
    """Solve truss for a given prescribed vertical displacement."""
    target_disp = jnp.array([0.0, -disp_y])
    costs_d: list[jaxls.Cost] = [
        bar_strain_energy(
            NodeVar(id=members[:, 0]),
            NodeVar(id=members[:, 1]),
            node_positions[members[:, 0]],
            node_positions[members[:, 1]],
            EA,
        ),
        pin_support(NodeVar(id=left_pin)),
        pin_support(NodeVar(id=right_pin)),
        prescribed_displacement(NodeVar(id=load_node_id), target_disp),
    ]
    problem_d = jaxls.LeastSquaresProblem(costs_d, [node_vars]).analyze()
    sol = problem_d.solve(verbose=False)
    return sol[node_vars]


# Use vmap to solve for all displacement values in parallel.
all_displacements = jax.vmap(solve_for_displacement)(displacement_magnitudes)
print(
    f"Solved for {len(displacement_magnitudes)} displacement values in parallel using vmap"
)

INFO     | Building optimization problem with 18 terms and 9 variables: 15 costs, 3 eq_zero, 0 leq_zero, 0 geq_zero
INFO     | Vectorizing group with 15 costs, 2 variables each: bar_strain_energy
INFO     | Vectorizing constraint group with 2 constraints (constraint_eq_zero), 1 variables each: augmented_pin_support
INFO     | Vectorizing constraint group with 1 constraints (constraint_eq_zero), 1 variables each: augmented_prescribed_displacement
Solved for 21 displacement values in parallel using vmap

Show code cell source

Hide code cell source

# Create animated visualization with force coloring.
scale_anim = 10  # Displacement magnification


def get_member_color(force: float, max_force: float) -> str:
    """Get color for a member based on its force (red=tension, blue=compression)."""
    intensity = min(abs(force) / (max_force + 1e-6), 1.0)
    if force > 0:  # Tension
        return f"rgba(220, {int(80 + 175 * (1 - intensity))}, {int(80 + 175 * (1 - intensity))}, 1)"
    else:  # Compression
        return f"rgba({int(80 + 175 * (1 - intensity))}, {int(80 + 175 * (1 - intensity))}, 220, 1)"


# Compute member forces for all displacement configurations.
def compute_forces_for_disp(disp: jax.Array) -> jax.Array:
    """Compute all member forces for a given displacement field."""
    return jax.vmap(lambda m: compute_member_force(m[0], m[1], disp))(members)


all_forces = jax.vmap(compute_forces_for_disp)(all_displacements)
global_max_force = float(jnp.max(jnp.abs(all_forces)))

# Build frames for animation.
frames = []
for i, (disp_mag, disp, forces) in enumerate(
    zip(displacement_magnitudes, all_displacements, all_forces)
):
    scaled_pos = node_positions + scale_anim * disp

    # Create individual traces for each member (for per-member coloring)
    member_traces = []
    for idx, m in enumerate(members):
        mi, mj = int(m[0]), int(m[1])
        color = get_member_color(float(forces[idx]), global_max_force)
        member_traces.append(
            go.Scatter(
                x=[float(scaled_pos[mi, 0]), float(scaled_pos[mj, 0])],
                y=[float(scaled_pos[mi, 1]), float(scaled_pos[mj, 1])],
                mode="lines",
                line=dict(color=color, width=4),
                showlegend=False,
            )
        )

    # Node markers.
    node_trace = go.Scatter(
        x=[float(p) for p in scaled_pos[:, 0]],
        y=[float(p) for p in scaled_pos[:, 1]],
        mode="markers",
        marker=dict(size=8, color="steelblue"),
        showlegend=False,
    )

    deflection = float(-disp[load_node_id, 1]) * 1000
    frames.append(
        go.Frame(
            data=member_traces + [node_trace],
            name=str(i),
            layout=go.Layout(
                title=f"Prescribed Displacement: {float(disp_mag) * 1000:.1f} mm, "
                f"Actual: {deflection:.1f} mm"
            ),
        )
    )

# Initial frame data (first displacement = 0, so all gray)
init_pos = node_positions + scale_anim * all_displacements[0]
init_forces = all_forces[0]

init_traces = []
for idx, m in enumerate(members):
    mi, mj = int(m[0]), int(m[1])
    color = get_member_color(float(init_forces[idx]), global_max_force)
    init_traces.append(
        go.Scatter(
            x=[float(init_pos[mi, 0]), float(init_pos[mj, 0])],
            y=[float(init_pos[mi, 1]), float(init_pos[mj, 1])],
            mode="lines",
            line=dict(color=color, width=4),
            showlegend=False,
        )
    )

init_traces.append(
    go.Scatter(
        x=[float(p) for p in init_pos[:, 0]],
        y=[float(p) for p in init_pos[:, 1]],
        mode="markers",
        marker=dict(size=8, color="steelblue"),
        showlegend=False,
    )
)

fig_anim = go.Figure(
    data=init_traces,
    frames=frames,
    layout=go.Layout(
        title="Prescribed Displacement: 0.0 mm, Actual: 0.0 mm",
        xaxis=dict(title="x [m]", range=[-1, span + 1], scaleanchor="y", scaleratio=1),
        yaxis=dict(title="y [m]", range=[-1, height + 2]),
        height=450,
        margin=dict(b=100),
        updatemenus=[
            dict(
                type="buttons",
                showactive=False,
                y=0,
                x=0.1,
                xanchor="right",
                buttons=[
                    dict(
                        label="Play",
                        method="animate",
                        args=[
                            None,
                            dict(
                                frame=dict(duration=100, redraw=True), fromcurrent=True
                            ),
                        ],
                    ),
                    dict(
                        label="Pause",
                        method="animate",
                        args=[
                            [None],
                            dict(
                                frame=dict(duration=0, redraw=False), mode="immediate"
                            ),
                        ],
                    ),
                ],
            )
        ],
        sliders=[
            dict(
                active=0,
                steps=[
                    dict(
                        args=[
                            [str(i)],
                            dict(frame=dict(duration=0, redraw=True), mode="immediate"),
                        ],
                        label=f"{float(d) * 1000:.0f}",
                        method="animate",
                    )
                    for i, d in enumerate(displacement_magnitudes)
                ],
                x=0.1,
                len=0.8,
                xanchor="left",
                y=-0.15,
                currentvalue=dict(
                    prefix="Displacement (mm): ", visible=True, xanchor="center"
                ),
                transition=dict(duration=0),
            )
        ],
    ),
)

# Add color legend.
fig_anim.add_annotation(
    x=0.02,
    y=0.98,
    xref="paper",
    yref="paper",
    text="Red = Tension, Blue = Compression",
    showarrow=False,
    font=dict(size=10),
    align="left",
    bgcolor="white",
    bordercolor="gray",
    borderwidth=1,
)

HTML(fig_anim.to_html(full_html=False, include_plotlyjs="cdn", auto_play=False))