Static Beam Adjoint Gradients¶

This example computes the gradient of a scalar objective (in this case, chosen as the tip vertical displacement of the Geradin cantilever) with respect to the tip load magnitude and the beam bending stiffness. Gradients are obtained analytically via the adjoint of the nonlinear static solve and verified against central-free finite differences.

Imports¶

In [ ]:

from jax import Array
from jax import numpy as jnp

from flapjax.models.geradin_beam.geradin_beam import generate_geradin_beam
from flapjax.structure import BeamStructure, StructureCase, StructureFullStates
from flapjax.utils.data_structures import ConvergenceSettings
from flapjax.utils.print_utils import set_verbosity

set_verbosity("silent")  # silence the solver output for this example

from jax import Array from jax import numpy as jnp from flapjax.models.geradin_beam.geradin_beam import generate_geradin_beam from flapjax.structure import BeamStructure, StructureCase, StructureFullStates from flapjax.utils.data_structures import ConvergenceSettings from flapjax.utils.print_utils import set_verbosity set_verbosity("silent") # silence the solver output for this example

Case parameters¶

We use the same 20-node Geradin cantilever as the primal example, oriented along the x-axis, with a downward dead load applied at the tip node.

In [ ]:

n_nodes = 20
struct = generate_geradin_beam(n_nodes, "x")

load = 600000.0
f_ext = jnp.zeros((n_nodes, 6)).at[-1, 2].set(-load)

n_nodes = 20 struct = generate_geradin_beam(n_nodes, "x") load = 600000.0 f_ext = jnp.zeros((n_nodes, 6)).at[-1, 2].set(-load)

Tightening convergence for adjoint accuracy¶

The adjoint solution linearises about the converged primal residual. We here tighten the tolerance to improve the accuracy of the adjoint gradient, but also to improve the finite difference comparison.

In [ ]:

struct.struct_convergence_settings = ConvergenceSettings(
    max_n_iter=50,
    rel_force_tol=0.0,
    rel_disp_tol=0.0,
    abs_force_tol=0.0,
    abs_disp_tol=0.0,
)

struct.struct_convergence_settings = ConvergenceSettings( max_n_iter=50, rel_force_tol=0.0, rel_disp_tol=0.0, abs_force_tol=0.0, abs_disp_tol=0.0, )

Primal solve and objective¶

We wrap static_solve in a function so that the same call signature is used for the reference solve and the two perturbed solves.

The objective returns the vertical position of the tip node, read as the last entry of the x array in the StructureFullStates container.

Note that the objective function must match the type signature of StructureObjectiveFunction, being a function that takes arguments of:

• StructureFullStates: Data structure which holds the solution of the static problem, such as strains and displacements.
• StructureDesignVariables: Data structure which holds the design variables of the structure, such as the cross-sectional properties.
• int, Array or None: The current time step index, which is unused in static problems.

and this function should return an Array, which is the value of the objective function for which we want to compute the gradient. Whilst this objective is generally chosen to be a scalar, it can also be a vector-valued function.

In [ ]:

def solve(struct_: BeamStructure, f_ext_: Array) -> StructureCase:
    return struct_.static_solve(
        f_ext_follower=None,
        f_ext_dead=f_ext_,
        f_ext_aero=None,
        prescribed_dofs=jnp.arange(6),
        load_steps=3,
    )


def objective(states: StructureFullStates, *_) -> Array:
    # note that we here omit the unused arguments as the objective is only a function of the states
    # however, the states themselves are a function of the design variables, which means the gradient exists
    return states.x[-1, 2]  # tip vertical displacement


base_result = solve(struct, f_ext)  # solve for the deflected state

obj_base = objective(base_result.get_full_states())  # compute the scalar objective

def solve(struct_: BeamStructure, f_ext_: Array) -> StructureCase: return struct_.static_solve( f_ext_follower=None, f_ext_dead=f_ext_, f_ext_aero=None, prescribed_dofs=jnp.arange(6), load_steps=3, ) def objective(states: StructureFullStates, *_) -> Array: # note that we here omit the unused arguments as the objective is only a function of the states # however, the states themselves are a function of the design variables, which means the gradient exists return states.x[-1, 2] # tip vertical displacement base_result = solve(struct, f_ext) # solve for the deflected state obj_base = objective(base_result.get_full_states()) # compute the scalar objective

Solving the adjoint system¶

BeamStructure.static_adjoint linearises the primal residual and objective about the converged state and returns:

1. A StructureDesignVariables container grads_adj populated with the gradient of the objective with respect to every recognised design variable (nodal reference coordinates x0, orientation Euler angles, cross-sectional stiffness k_cs, cross-sectional inertia m_cs, lumped mass m_lumped, follower forces f_ext_follower, dead forces f_ext_dead, and any thrust histories). Fields that are not exposed on the case are left as None.
2. The adjoint state vector, which we do not use here but is returned for downstream chained-adjoint use-cases.

One call therefore gives us the full gradient with respect to all design variables simultaneously — the cost is independent of the number of design variables and dominated by a single linear solve at the converged state.

In [ ]:

grads_adj, adj = struct.static_adjoint(structure=base_result, objective=objective)

grads_adj, adj = struct.static_adjoint(structure=base_result, objective=objective)

Verification against finite differences: tip load¶

We compare the adjoint gradient against a forward finite-difference approximation obtained by perturbing the tip z-force and re-solving.

In [ ]:

f_eps = 1.0
f_pert = f_ext.at[-1, 2].add(f_eps)
f_pert_obj = objective(solve(struct, f_pert).get_full_states())

assert grads_adj.f_ext_dead is not None
f_fd_grad = (f_pert_obj - obj_base) / f_eps
print("Tip forcing gradient")
print(f"Adjoint: {float(grads_adj.f_ext_dead[-1, 2]):.04e}")
print(f"Finite difference: {float(f_fd_grad):.04e}")

f_eps = 1.0 f_pert = f_ext.at[-1, 2].add(f_eps) f_pert_obj = objective(solve(struct, f_pert).get_full_states()) assert grads_adj.f_ext_dead is not None f_fd_grad = (f_pert_obj - obj_base) / f_eps print("Tip forcing gradient") print(f"Adjoint: {float(grads_adj.f_ext_dead[-1, 2]):.04e}") print(f"Finite difference: {float(f_fd_grad):.04e}")

Verification against finite differences: bending stiffness¶

The cross-sectional stiffness tensor struct.k_cs has shape (n_elem_types, 6, 6); the (4, 4) entry corresponds to the bending stiffness about the local y-axis. We perturb it by a small increment and re-solve.

In [ ]:

k_eps = 10.0
struct.k_cs = struct.k_cs.at[0, 4, 4].add(k_eps)
k_pert_obj = objective(solve(struct, f_ext).get_full_states())
k_fd_grad = (k_pert_obj - obj_base) / k_eps

if grads_adj.k_cs is None:
    raise ValueError("k_cs is None")
print("Beam bending stiffness gradient")
print(f"Adjoint: {float(grads_adj.k_cs[0, 4, 4]):.04e}")
print(f"Finite difference: {float(k_fd_grad):.04e}")

k_eps = 10.0 struct.k_cs = struct.k_cs.at[0, 4, 4].add(k_eps) k_pert_obj = objective(solve(struct, f_ext).get_full_states()) k_fd_grad = (k_pert_obj - obj_base) / k_eps if grads_adj.k_cs is None: raise ValueError("k_cs is None") print("Beam bending stiffness gradient") print(f"Adjoint: {float(grads_adj.k_cs[0, 4, 4]):.04e}") print(f"Finite difference: {float(k_fd_grad):.04e}")