Flying Spaghetti Mass Optimisation¶

This example uses the beam dynamic adjoint to drive an optimisation. The design variable is the per-element cross-section mass per unit length, m_bar. The cost is a scalar measure of how straight the spaghetti is at the final timestep, taken as the inner product of the strains, which we aim to minimise. Stiffness is reduced by a factor of three relative to the reference case so that the beam undergoes large deformations and the optimisation problem is non-trivial.

The outer loop uses SLSQP, and is provided with both the primal value and the gradient.

This script can take a few minutes to run, depending on the number of iterations and the number of timesteps.

Imports¶

In [ ]:

from pathlib import Path

import jax
import jax.numpy as jnp
import numpy as np
from jax import Array
from matplotlib import pyplot as plt
from scipy.optimize import minimize

from flapjax.models.flying_spaghetti.flying_spaghetti import generate_flying_spaghetti
from flapjax.structure import StructureDesignVariables, StructureFullStates
from flapjax.structure.gradients.data_structures import StructureGradsToCompute
from flapjax.utils.print_utils import set_verbosity

set_verbosity("warning")

from pathlib import Path import jax import jax.numpy as jnp import numpy as np from jax import Array from matplotlib import pyplot as plt from scipy.optimize import minimize from flapjax.models.flying_spaghetti.flying_spaghetti import generate_flying_spaghetti from flapjax.structure import StructureDesignVariables, StructureFullStates from flapjax.structure.gradients.data_structures import StructureGradsToCompute from flapjax.utils.print_utils import set_verbosity set_verbosity("warning")

Case parameters¶

We integrate for 10 s with 0.01 s steps, giving 1001 timesteps. The first sample is deliberately placed one step behind zero so that the forcing history is well-defined at the initial condition.

For the optimisation, we apply bounds of [0.3, 3.0] kg/m on the mass per unit length, constraining the optimiser to physically plausible values.

In [ ]:

# case parameters
n_nodes = 21
n_elem = n_nodes - 1
dt = 0.01
t_end = 10.0
n_tstep = int(jnp.ceil(t_end / dt)) + 1

# optimisation parameters
n_iter = 10  # number of optimisation iterations
m_min = 0.3  # mass per unit length lower bound, kg/m
m_max = 3.0  # mass per unit length upper bound, kg/m

# start one timestep behind to prevent issues with initial condition
t = jnp.arange(n_tstep) * dt - dt

# plot to save to paraview if desired
out_dir = Path("./flying_spaghetti_optimise/")

# case parameters n_nodes = 21 n_elem = n_nodes - 1 dt = 0.01 t_end = 10.0 n_tstep = int(jnp.ceil(t_end / dt)) + 1 # optimisation parameters n_iter = 10 # number of optimisation iterations m_min = 0.3 # mass per unit length lower bound, kg/m m_max = 3.0 # mass per unit length upper bound, kg/m # start one timestep behind to prevent issues with initial condition t = jnp.arange(n_tstep) * dt - dt # plot to save to paraview if desired out_dir = Path("./flying_spaghetti_optimise/")

Cost function¶

The dynamic adjoint takes a callable of signature (states, dv, i_ts) -> Array, as in the static case. For dynamic cases, the cost is the sum of this function across timesteps all.

Here we only want the final-time straightness, so jax.lax.select returns the strain inner product when i_ts == n_tstep - 1 and zero otherwise. Using lax.select is a neccesity for JIT compilation, as a python if statement would break the JIT trace.

In [ ]:

def cost(
        states: StructureFullStates,
        _: StructureDesignVariables,
        i_ts: int,
) -> Array:
    # inner product of strains as a measure for how "straight" the spaghetti is
    return jax.lax.select(i_ts == n_tstep - 1, (states.eps ** 2).sum(), 0.0)

def cost( states: StructureFullStates, _: StructureDesignVariables, i_ts: int, ) -> Array: # inner product of strains as a measure for how "straight" the spaghetti is return jax.lax.select(i_ts == n_tstep - 1, (states.eps ** 2).sum(), 0.0)

Inner loop: primal primal and adjoint solves¶

make_sol is called once per SLSQP iteration with the current mass distribution. It creates a beam with a given mass distribution, solves the dynamic problem, and then runs the adjoint to get the gradient of the cost with respect to m_bar. It also saves VTK files and PNG snapshots of the beam shape, mass distribution, and gradient for each iteration.

A few points worth calling out:

• We use a module-level counter dict iter_counter instead of nonlocal so make_sol can be reused when the notebook is re-run.
• Each iteration writes VTK files and PNG snapshots to flying_spaghetti_optimise/iteration_N/ so the optimisation trajectory can be inspected after the fact.

In [ ]:

iter_counter = {"i": 0}
beam_snapshots: dict[int, np.ndarray] = {}


def make_sol(m_bar: np.ndarray) -> tuple[float, Array]:
    i_iter = iter_counter["i"]
    iter_counter["i"] += 1

    m_bar_j = jnp.asarray(m_bar)

    struct_, f_dead_2d_, _ = generate_flying_spaghetti(
        n_nodes,
        t,
        spectral_radius=0.7,
        m_bar=m_bar_j,
    )

    # reduce stiffness to make the problem more interesting
    struct_.k_cs /= 3.0

    primal_sol = struct_.dynamic_solve(
        init_state=None,
        n_tstep=n_tstep,
        dt=dt,
        f_ext_follower=None,
        f_ext_dead=f_dead_2d_,
        f_ext_aero=None,
        prescribed_dofs=(),
    )

    grads, _ = struct_.dynamic_adjoint(
        structure=primal_sol,
        objective=cost,
        p_q0_p_x=None,
        # approx_grads reduces some of the cost of the gradients by removing the sensitivities of the mass and gyroscopic terms to the beam deformation
        approx_grads=True,
        # grads_to_compute allows for disabling sensitivities we don't need, which saves some time in the adjoint sweep
        grads_to_compute=StructureGradsToCompute(
            x0=False,
            k_cs=False,
            m_cs=True,
            m_lumped=False,
            f_ext_follower=False,
            f_ext_dead=False,
        ),
    )

    assert grads.m_cs is not None
    # grads.m_cs has shape (1, n_elem, 6, 6), where the leading axis would increase in size for a vector cost. As the cross-sectional mass matrix contains the translational mass multiplied by the identity in the upper-left 3x3 block, we sum these entries. We here assume that perturbing the mass does not affect the local rotational inertia.
    m_bar_grad = (
            grads.m_cs[0, :, 0, 0] + grads.m_cs[0, :, 1, 1] + grads.m_cs[0, :, 2, 2]
    )

    # compute primal cost
    cost_val = (primal_sol.eps[-1, ...] ** 2).sum()

    # snapshot the final-timestep beam shape for later comparison of the first and last iterations
    beam_snapshots[i_iter] = np.asarray(primal_sol.x[-1, :, (0, 2)])

    iter_path = out_dir.joinpath(f"iteration_{i_iter}/")
    iter_path.mkdir(parents=True, exist_ok=True)

    # uncomment for plotting to vtk
    # primal_sol.plot(iter_path.joinpath("vtk"))

    # plot the beam shape at the final timestep
    _, ax = plt.subplots()
    ax.set_title(f"Beam at {i_iter}, cost = {float(cost_val)}")
    ax.set_xlabel("x")
    ax.set_ylabel("z")
    ax.axis("equal")
    ax.plot(primal_sol.hg[-1, :, 0, 3], primal_sol.hg[-1, :, 2, 3])
    plt.savefig(iter_path.joinpath("end_coords.png"))

    # plot the mass distribution
    _, ax = plt.subplots()
    ax.plot(jnp.arange(n_elem), m_bar_j)
    ax.set_title(f"Element m_bar at {i_iter}")
    ax.set_xlabel("Element number")
    ax.set_ylabel("m_bar, kg/m")
    plt.savefig(iter_path.joinpath("m_bar.png"))

    plt.close("all")

    print(f"iteration {i_iter}: cost = {float(cost_val):.6e}")

    return float(cost_val), m_bar_grad

iter_counter = {"i": 0} beam_snapshots: dict[int, np.ndarray] = {} def make_sol(m_bar: np.ndarray) -> tuple[float, Array]: i_iter = iter_counter["i"] iter_counter["i"] += 1 m_bar_j = jnp.asarray(m_bar) struct_, f_dead_2d_, _ = generate_flying_spaghetti( n_nodes, t, spectral_radius=0.7, m_bar=m_bar_j, ) # reduce stiffness to make the problem more interesting struct_.k_cs /= 3.0 primal_sol = struct_.dynamic_solve( init_state=None, n_tstep=n_tstep, dt=dt, f_ext_follower=None, f_ext_dead=f_dead_2d_, f_ext_aero=None, prescribed_dofs=(), ) grads, _ = struct_.dynamic_adjoint( structure=primal_sol, objective=cost, p_q0_p_x=None, # approx_grads reduces some of the cost of the gradients by removing the sensitivities of the mass and gyroscopic terms to the beam deformation approx_grads=True, # grads_to_compute allows for disabling sensitivities we don't need, which saves some time in the adjoint sweep grads_to_compute=StructureGradsToCompute( x0=False, k_cs=False, m_cs=True, m_lumped=False, f_ext_follower=False, f_ext_dead=False, ), ) assert grads.m_cs is not None # grads.m_cs has shape (1, n_elem, 6, 6), where the leading axis would increase in size for a vector cost. As the cross-sectional mass matrix contains the translational mass multiplied by the identity in the upper-left 3x3 block, we sum these entries. We here assume that perturbing the mass does not affect the local rotational inertia. m_bar_grad = ( grads.m_cs[0, :, 0, 0] + grads.m_cs[0, :, 1, 1] + grads.m_cs[0, :, 2, 2] ) # compute primal cost cost_val = (primal_sol.eps[-1, ...] ** 2).sum() # snapshot the final-timestep beam shape for later comparison of the first and last iterations beam_snapshots[i_iter] = np.asarray(primal_sol.x[-1, :, (0, 2)]) iter_path = out_dir.joinpath(f"iteration_{i_iter}/") iter_path.mkdir(parents=True, exist_ok=True) # uncomment for plotting to vtk # primal_sol.plot(iter_path.joinpath("vtk")) # plot the beam shape at the final timestep _, ax = plt.subplots() ax.set_title(f"Beam at {i_iter}, cost = {float(cost_val)}") ax.set_xlabel("x") ax.set_ylabel("z") ax.axis("equal") ax.plot(primal_sol.hg[-1, :, 0, 3], primal_sol.hg[-1, :, 2, 3]) plt.savefig(iter_path.joinpath("end_coords.png")) # plot the mass distribution _, ax = plt.subplots() ax.plot(jnp.arange(n_elem), m_bar_j) ax.set_title(f"Element m_bar at {i_iter}") ax.set_xlabel("Element number") ax.set_ylabel("m_bar, kg/m") plt.savefig(iter_path.joinpath("m_bar.png")) plt.close("all") print(f"iteration {i_iter}: cost = {float(cost_val):.6e}") return float(cost_val), m_bar_grad

In [ ]:

m_bar0 = np.ones(n_elem)  # initial guess

total_mass_target = float(m_bar0.sum())  # total mass of beam

# mass constraint: total mass must remain constant, so the optimiser can only redistribute mass along the beam
constraints = (
    {
        "type": "eq",
        "fun": lambda x: x.sum() - total_mass_target,
        "jac": lambda x: np.ones_like(x),
    },
)

result = minimize(  # type: ignore
    fun=make_sol,
    x0=m_bar0,
    jac=True,
    bounds=[(m_min, m_max)] * n_elem,  # bounds on mass per unit length
    constraints=constraints,
    method="SLSQP",
    options={"maxiter": n_iter, "ftol": 1e-8},
)

print(result)

m_bar0 = np.ones(n_elem) # initial guess total_mass_target = float(m_bar0.sum()) # total mass of beam # mass constraint: total mass must remain constant, so the optimiser can only redistribute mass along the beam constraints = ( { "type": "eq", "fun": lambda x: x.sum() - total_mass_target, "jac": lambda x: np.ones_like(x), }, ) result = minimize( # type: ignore fun=make_sol, x0=m_bar0, jac=True, bounds=[(m_min, m_max)] * n_elem, # bounds on mass per unit length constraints=constraints, method="SLSQP", options={"maxiter": n_iter, "ftol": 1e-8}, ) print(result)

Inspecting the optimised design¶

Each iteration saved its own PNGs (and optionally VTK files) to flying_spaghetti_optimise/iteration_N/, the trajectory of the objective, mass distribution, and beam shape can be inspected. Below we plot the final mass distribution against the initial uniform guess to see how the optimiser has redistributed mass along the spaghetti.

In [ ]:

fig, ax_ = plt.subplots()
ax_.plot(np.arange(n_elem), m_bar0, label="Initial", linestyle="--")
ax_.plot(np.arange(n_elem), result.x, label="Optimised")
ax_.set_xlabel("Element number")
ax_.set_ylabel("Mass per unit length, kg/m")
ax_.set_title("Mass distribution")
ax_.legend()
plt.show()

fig, ax_ = plt.subplots() ax_.plot(np.arange(n_elem), m_bar0, label="Initial", linestyle="--") ax_.plot(np.arange(n_elem), result.x, label="Optimised") ax_.set_xlabel("Element number") ax_.set_ylabel("Mass per unit length, kg/m") ax_.set_title("Mass distribution") ax_.legend() plt.show()

Optimisation loop¶

The optimiser starts from a uniform mass distribution of 1.0 kg/m per element. As well as having bounds on the mass per unit length, we also constrain the total mass of the beam to remain constant.

The outer loop terminates after n_iter SLSQP iterations or when the cost converges to within ftol=1e-8.

Beam shape: first vs final iteration¶

To visualise what the optimiser has actually achieved, we compare the final-timestep beam shape from the very first optimisation call (uniform mass) against the shape from the last call (converged mass distribution).

In [ ]:

i_first = 0
i_last = max(beam_snapshots)
x_first, z_first = beam_snapshots[i_first]
x_last, z_last = beam_snapshots[i_last]

fig, ax_ = plt.subplots()
ax_.plot(x_first, z_first, label=f"iteration {i_first} (initial)", linestyle="--")
ax_.plot(x_last, z_last, label=f"iteration {i_last} (final)")
ax_.set_xlabel("x")
ax_.set_ylabel("z")
ax_.axis("equal")
ax_.set_title("Beam shape at final timestep: initial vs optimised mass distribution")
ax_.legend()
plt.show()

i_first = 0 i_last = max(beam_snapshots) x_first, z_first = beam_snapshots[i_first] x_last, z_last = beam_snapshots[i_last] fig, ax_ = plt.subplots() ax_.plot(x_first, z_first, label=f"iteration {i_first} (initial)", linestyle="--") ax_.plot(x_last, z_last, label=f"iteration {i_last} (final)") ax_.set_xlabel("x") ax_.set_ylabel("z") ax_.axis("equal") ax_.set_title("Beam shape at final timestep: initial vs optimised mass distribution") ax_.legend() plt.show()