Structure

flapjax.structure

StructureCase

StructureCase(
    hg: Array,
    conn: tuple[tuple[int, int], ...],
    o0: Array,
    d: Array,
    eps: Array,
    varphi: Array,
    f_ext_follower: Array | None,
    f_ext_dead: Array | None,
    f_ext_aero: Array | None,
    f_grav: Array | None,
    f_int: Array,
    f_elem: Array,
    f_res: Array,
    t: Array,
    thrust: dict[str, Array],
    thrust_nodes: tuple[tuple[str, int], ...],
    thrust_direction: tuple[
        tuple[str, tuple[float, float, float]], ...
    ],
    prescribed_dofs: tuple[int, ...] | Array,
    v: Array | None = None,
    v_dot: Array | None = None,
    a: Array | None = None,
    f_iner_gyr: Array | None = None,
    i_ts: int | None = None,
    local: bool = True,
    constraint_data: dict[str, dict[str, Array]]
    | None = None,
)

Object to hold the full state and forces of a structure analysis.

A single instance may represent any of three flavours:

• Static: dynamic-only fields (v, v_dot, a, f_iner_gyr) are not set — their public properties return an all-zeros array of appropriate shape. Array shapes are (n_nodes, ...).
• Dynamic snapshot: all dynamic fields populated for a single timestep. Array shapes are (n_nodes, ...).
• Dynamic trajectory: all dynamic fields populated for multiple timesteps. t is a (n_tstep,) array and i_ts is None. Array shapes are (n_tstep, n_nodes, ...).

Use :attr:is_dynamic and :attr:is_batched to distinguish at runtime.

Source code in src/flapjax/structure/data_structures.py

def __init__(
    self,
    hg: Array,
    conn: tuple[tuple[int, int], ...],
    o0: Array,
    d: Array,
    eps: Array,
    varphi: Array,
    f_ext_follower: Array | None,
    f_ext_dead: Array | None,
    f_ext_aero: Array | None,
    f_grav: Array | None,
    f_int: Array,
    f_elem: Array,
    f_res: Array,
    t: Array,
    thrust: dict[str, Array],
    thrust_nodes: tuple[tuple[str, int], ...],
    thrust_direction: tuple[tuple[str, tuple[float, float, float]], ...],
    prescribed_dofs: tuple[int, ...] | Array,
    v: Array | None = None,
    v_dot: Array | None = None,
    a: Array | None = None,
    f_iner_gyr: Array | None = None,
    i_ts: int | None = None,
    local: bool = True,
    constraint_data: dict[str, dict[str, Array]] | None = None,
):
    self.hg: Array = hg
    self.conn: tuple[tuple[int, int], ...] = conn
    self.o0: Array = o0
    self.d: Array = d
    self.eps: Array = eps
    self.varphi: Array = varphi
    self._v: Array | None = v
    self._v_dot: Array | None = v_dot
    self._a: Array | None = a
    self.f_ext_follower: Array | None = f_ext_follower
    self.f_ext_dead: Array | None = f_ext_dead
    self.f_ext_aero: Array | None = f_ext_aero
    self.f_grav: Array | None = f_grav
    self.f_int: Array = f_int
    self.f_elem: Array = f_elem
    self._f_iner_gyr: Array | None = f_iner_gyr
    self.f_res: Array = f_res
    self.thrust: dict[str, Array] = thrust
    self.thrust_nodes: tuple[tuple[str, int], ...] = thrust_nodes
    self.thrust_direction: tuple[tuple[str, tuple[float, float, float]], ...] = (
        thrust_direction
    )
    self.t: Array = t
    self.i_ts: int | None = i_ts
    self.prescribed_dofs: tuple[int, ...] = input_dof_index_to_tuple(
        prescribed_dofs
    )
    self.free_dofs: tuple[int, ...] = get_solve_dofs(
        n_dof=varphi.shape[-2] * 6, prescribed_dofs=self.prescribed_dofs
    )
    self.local: bool = local
    self.constraint_data: dict[str, dict[str, Array]] = constraint_data if constraint_data is not None else {}

to_dynamic

to_dynamic() -> StructureCase

to_dynamic(t: None) -> StructureCase

to_dynamic(t: Array) -> StructureCase

to_dynamic(t: Array | None = None) -> StructureCase

Convert static structure results to a dynamic snapshot (t=None) or a batched trajectory (t provided), zeroing velocity/acceleration fields. Calling on a Structure that is already dynamic returns self.

Source code in src/flapjax/structure/data_structures.py

def to_dynamic(self, t: Array | None = None) -> StructureCase:
    """Convert static structure results to a dynamic snapshot (``t=None``) or
    a batched trajectory (``t`` provided), zeroing velocity/acceleration
    fields. Calling on a Structure that is already dynamic returns ``self``.
    """
    if self.is_dynamic:
        return self

    dyn_snapshot = StructureCase(
        hg=self.hg,
        conn=self.conn,
        o0=self.o0,
        d=self.d,
        eps=self.eps,
        varphi=self.varphi,
        v=self.v,
        v_dot=self.v_dot,
        a=self.a,
        f_ext_follower=self.f_ext_follower,
        f_ext_dead=self.f_ext_dead,
        f_ext_aero=self.f_ext_aero,
        f_grav=self.f_grav,
        f_int=self.f_int,
        f_elem=self.f_elem,
        f_iner_gyr=self.f_iner_gyr,
        f_res=self.f_res,
        thrust=self.thrust,
        thrust_nodes=self.thrust_nodes,
        thrust_direction=self.thrust_direction,
        t=jnp.array(0.0),
        i_ts=-1,
        prescribed_dofs=self.prescribed_dofs,
        constraint_data=self.constraint_data,
    )

    if t is None:
        return dyn_snapshot
    return StructureCase.initialise(
        initial_snapshot=dyn_snapshot,
        t=t,
        use_f_ext_aero=self.f_ext_aero is not None,
        use_f_ext_follower=self.f_ext_follower is not None,
        use_f_ext_dead=self.f_ext_dead is not None,
    )

to_static

to_static() -> StructureCase

Return a static Structure, dropping velocity/acceleration fields. If already static, returns self. For a batched trajectory, raises: use self[i_ts].to_static() to extract a single time step first.

Source code in src/flapjax/structure/data_structures.py

def to_static(self) -> StructureCase:
    """Return a static Structure, dropping velocity/acceleration
    fields. If already static, returns ``self``. For a batched trajectory,
    raises: use ``self[i_ts].to_static()`` to extract a single time step first.
    """
    if not self.is_dynamic:
        return self
    if self.is_batched:
        raise ValueError(
            "to_static() on a batched Structure is ambiguous; index a "
            "single time step first (e.g. `structure[i_ts].to_static()`)."
        )
    return StructureCase(
        hg=self.hg,
        conn=self.conn,
        o0=self.o0,
        d=self.d,
        eps=self.eps,
        varphi=self.varphi,
        f_ext_follower=self.f_ext_follower,
        f_ext_dead=self.f_ext_dead,
        f_ext_aero=self.f_ext_aero,
        f_grav=self.f_grav,
        f_int=self.f_int,
        f_elem=self.f_elem,
        f_res=self.f_res,
        t=self.t,
        thrust=self.thrust,
        thrust_nodes=self.thrust_nodes,
        thrust_direction=self.thrust_direction,
        prescribed_dofs=self.prescribed_dofs,
        constraint_data=self.constraint_data,
    )

initialise classmethod

initialise(
    initial_snapshot: StructureCase,
    t: Array,
    use_f_ext_follower: bool,
    use_f_ext_dead: bool,
    use_f_ext_aero: bool,
) -> StructureCase

Initialise a batched dynamic Structure from a single dynamic snapshot.

Parameters:

Name	Type	Description	Default
initial_snapshot	StructureCase	Snapshot at initial time step (must be dynamic, i.e. is_dynamic == True and not batched).	required
t	Array	Time step array, (n_tstep, )	required
use_f_ext_follower	bool	Whether to include follower force array	required
use_f_ext_dead	bool	Whether to include dead force array	required
use_f_ext_aero	bool	Whether to include aero force array	required

Returns:

Type	Description
StructureCase	Batched Structure with arrays initialised to zero except for the first time step.

Source code in src/flapjax/structure/data_structures.py

@classmethod
def initialise(
    cls,
    initial_snapshot: StructureCase,
    t: Array,
    use_f_ext_follower: bool,
    use_f_ext_dead: bool,
    use_f_ext_aero: bool,
) -> StructureCase:
    r"""
    Initialise a batched dynamic Structure from a single dynamic snapshot.
    :param initial_snapshot: Snapshot at initial time step (must be dynamic,
    i.e. ``is_dynamic == True`` and not batched).
    :param t: Time step array, ``(n_tstep, )``
    :param use_f_ext_follower: Whether to include follower force array
    :param use_f_ext_dead: Whether to include dead force array
    :param use_f_ext_aero: Whether to include aero force array
    :return: Batched Structure with arrays initialised to zero except for
    the first time step.
    """
    if not initial_snapshot.is_dynamic:
        raise ValueError(
            "initial_snapshot must be dynamic; call to_dynamic() first"
        )
    if initial_snapshot.is_batched:
        raise ValueError("initial_snapshot must be a single snapshot, not batched")

    n_node = initial_snapshot.hg.shape[0]
    n_elem = initial_snapshot.d.shape[0]
    n_tstep = t.shape[0]

    hg = jnp.zeros((n_tstep, n_node, 4, 4)).at[0, ...].set(initial_snapshot.hg)
    d = jnp.zeros((n_tstep, n_elem, 6)).at[0, ...].set(initial_snapshot.d)
    eps = jnp.zeros((n_tstep, n_elem, 6)).at[0, ...].set(initial_snapshot.eps)
    varphi = jnp.zeros((n_tstep, n_node, 6)).at[0, ...].set(initial_snapshot.varphi)
    v = jnp.zeros((n_tstep, n_node, 6)).at[0, ...].set(initial_snapshot.v)
    v_dot = jnp.zeros((n_tstep, n_node, 6)).at[0, ...].set(initial_snapshot.v_dot)
    a = jnp.zeros((n_tstep, n_node, 6)).at[0, ...].set(initial_snapshot.a)

    if use_f_ext_follower:
        f_ext_follower = jnp.zeros((n_tstep, n_node, 6))
        if initial_snapshot.f_ext_follower is not None:
            f_ext_follower = f_ext_follower.at[0, ...].set(
                initial_snapshot.f_ext_follower
            )
    else:
        f_ext_follower = None

    if use_f_ext_dead:
        f_ext_dead = jnp.zeros((n_tstep, n_node, 6))
        if initial_snapshot.f_ext_dead is not None:
            f_ext_dead = f_ext_dead.at[0, ...].set(initial_snapshot.f_ext_dead)
    else:
        f_ext_dead = None

    if use_f_ext_aero:
        f_ext_aero = jnp.zeros((n_tstep, n_node, 6))
        if initial_snapshot.f_ext_aero is not None:
            f_ext_aero = f_ext_aero.at[0, ...].set(initial_snapshot.f_ext_aero)
    else:
        f_ext_aero = None

    f_grav = (
        jnp.zeros((n_tstep, n_node, 6)).at[0, ...].set(initial_snapshot.f_grav)
        if initial_snapshot.f_grav is not None
        else None
    )
    f_int = jnp.zeros((n_tstep, n_node, 6)).at[0, ...].set(initial_snapshot.f_int)
    f_elem = jnp.zeros((n_tstep, n_elem, 6)).at[0, ...].set(initial_snapshot.f_elem)
    f_iner_gyr = (
        jnp.zeros((n_tstep, n_node, 6)).at[0, ...].set(initial_snapshot.f_iner_gyr)
    )
    f_res = jnp.zeros((n_tstep, n_node, 6)).at[0, ...].set(initial_snapshot.f_res)

    thrust = {k: jnp.full(n_tstep, v) for k, v in initial_snapshot.thrust.items()}
    return cls(
        hg=hg,
        conn=initial_snapshot.conn,
        o0=initial_snapshot.o0,
        d=d,
        eps=eps,
        varphi=varphi,
        v=v,
        v_dot=v_dot,
        a=a,
        f_ext_follower=f_ext_follower,
        f_ext_dead=f_ext_dead,
        f_ext_aero=f_ext_aero,
        f_grav=f_grav,
        f_int=f_int,
        f_elem=f_elem,
        f_iner_gyr=f_iner_gyr,
        f_res=f_res,
        thrust=thrust,
        thrust_nodes=initial_snapshot.thrust_nodes,
        thrust_direction=initial_snapshot.thrust_direction,
        t=t,
        prescribed_dofs=initial_snapshot.prescribed_dofs,
    )

to_global

to_global() -> None

Convert local structure results to global frame.

Source code in src/flapjax/structure/data_structures.py

def to_global(self) -> None:
    """Convert local structure results to global frame."""
    if not self.local:
        warn("Results already in global frame, skipping conversion.")
        return
    self.local = False
    self._transform(rmat=self.rmat)

to_local

to_local() -> None

Convert global structure results to local frame.

Source code in src/flapjax/structure/data_structures.py

def to_local(self) -> None:
    """Convert global structure results to local frame."""
    if self.local:
        warn("Results already in local frame, skipping conversion.")
        return
    self.local = True
    if self.is_batched:
        rmat_t = jnp.transpose(self.rmat, (0, 1, 3, 2))
    else:
        rmat_t = jnp.transpose(self.rmat, (0, 2, 1))
    self._transform(rmat=rmat_t)

plot

plot(directory: PathLike | str, n_interp: int = 0) -> Path

plot(
    directory: PathLike | str,
    n_interp: int = 0,
    *,
    index: slice
    | Sequence[int]
    | int
    | Array
    | None = None,
) -> Path

plot(
    directory: PathLike | str,
    n_interp: int = 0,
    *,
    index: slice
    | Sequence[int]
    | int
    | Array
    | None = None,
) -> Path

Plot beam results to VTK/VTU files in the specified directory. For a batched Structure, a PVD is written alongside per-timestep VTUs.

Parameters:

Name	Type	Description	Default
directory	PathLike | str	Path to write files to.	required
n_interp	int	Number of interpolation points to add between each element for smoother visualisation.	0
index	slice | Sequence[int] | int | Array | None	For batched Structures only, time step indices to plot.	None

Source code in src/flapjax/structure/data_structures.py

def plot(
    self,
    directory: os.PathLike | str,
    n_interp: int = 0,
    *,
    index: slice | Sequence[int] | int | Array | None = None,
) -> Path:
    r"""
    Plot beam results to VTK/VTU files in the specified directory. For a
    batched Structure, a PVD is written alongside per-timestep VTUs.
    :param directory: Path to write files to.
    :param n_interp: Number of interpolation points to add between each element for smoother visualisation.
    :param index: For batched Structures only, time step indices to plot.
    """
    if self.is_batched:
        index_ = index_to_arr(index=index, n_entries=self.n_tstep)
        directory_path = Path(directory).resolve()
        directory_path.mkdir(parents=True, exist_ok=True)

        paths = [self[i_ts]._plot_single(directory, n_interp) for i_ts in index_]

        assert self.t is not None
        return write_pvd(directory, "beam_dynamic_ts", paths, list(self.t[index_]))

    if index is not None:
        raise ValueError("`index` is only used for batched Structure")

    if not self.is_dynamic:
        return self.to_dynamic()._plot_single(directory, n_interp)
    return self._plot_single(directory, n_interp)

BeamStructure

BeamStructure(
    num_nodes: int,
    connectivity: Array,
    y_vector: Array,
    k_cs_index: Array | None = None,
    m_cs_index: Array | None = None,
    m_lumped_index: Array | None = None,
    gravity: Array | Sequence[float] | None = None,
    thrust_nodes: dict[str, int] | None = None,
    thrust_direction: dict[str, Array] | None = None,
    optional_jacobians: OptionalJacobians | None = None,
    relaxation_factor: float = 1.0,
    spectral_radius: float = 0.9,
    alpha_m: float = 0.0,
    beta_k: float = 0.0,
    struct_convergence_settings: ConvergenceSettings = DEFAULT_STRUCT_CONVERGENCE_SETTINGS,
    constraints: dict[str, SoftConstraint | HardConstraint]
    | None = None,
)

Bases: BaseBeamStructure

Initialise BaseBeamStructure class with all non-design parameters.

Parameters:

Name	Type	Description	Default
num_nodes	int	Number of nodes in the structure.	required
connectivity	Array	Connectivity array, `(n_elem, 2)``.	required
y_vector	Array	Vector defining the y direction for each element, (n_elem, 3).	required
k_cs_index	Array | None	Array defining the index from the library of k_cs to use for each element, (n_elem, ). If None, all elements will use the first entry in the k_cs library.	None
m_cs_index	Array | None	Array defining the index from the library of m_cs to use for each element, (n_elem, ). If None, all elements will use the first entry in the m_cs library.	None
m_lumped_index	Array | None	Node index for nodes which are to have a lumped mass attached. The order is the same as that for the lumped mass data (n_lumped_mass, ).	None
gravity	Array | Sequence[float] | None	Gravity vector in global reference frame, or None for no gravity_vec, (3, ).	None
thrust_nodes	dict[str, int] | None	Dictionary of thrust node names and their corresponding node indices, {keys, int}.	None
thrust_direction	dict[str, Array] | None	Dictionary of thrust node names and their corresponding thrust direction vectors, {keys, (3, )}.	None
optional_jacobians	OptionalJacobians | None	Define which Jacobians contributions are to be used for solution.	None
relaxation_factor	float	Relaxation factor which reduces the displacement update at each iteration. A value of 1 is no relaxation, and a value of 0 is no update.	1.0
spectral_radius	float	Spectral radius for structural time integrator, where a value of 0 is highly damped and a value of 1 is undamped.	0.9
alpha_m	float	Mass-proportional Rayleigh damping coefficient.	0.0
beta_k	float	Stiffness-proportional Rayleigh damping coefficient.	0.0
struct_convergence_settings	ConvergenceSettings	Structure convergence settings.	DEFAULT_STRUCT_CONVERGENCE_SETTINGS
constraints	dict[str, SoftConstraint | HardConstraint] | None	Named dict {name: constraint}, or None, which add	None

Source code in src/flapjax/structure/beam.py

def __init__(
    self,
    num_nodes: int,
    connectivity: Array,
    y_vector: Array,
    k_cs_index: Array | None = None,
    m_cs_index: Array | None = None,
    m_lumped_index: Array | None = None,
    gravity: Array | Sequence[float] | None = None,
    thrust_nodes: dict[str, int] | None = None,
    thrust_direction: dict[str, Array] | None = None,
    optional_jacobians: OptionalJacobians | None = None,
    relaxation_factor: float = 1.0,
    spectral_radius: float = 0.9,
    alpha_m: float = 0.0,
    beta_k: float = 0.0,
    struct_convergence_settings: ConvergenceSettings = DEFAULT_STRUCT_CONVERGENCE_SETTINGS,
    constraints: (dict[str, SoftConstraint | HardConstraint] | None) = None,
) -> None:
    r"""
    Initialise BaseBeamStructure class with all non-design parameters.
    :param num_nodes: Number of nodes in the structure.
    :param connectivity: Connectivity array, `(n_elem, 2)``.
    :param y_vector: Vector defining the y direction for each element, ``(n_elem, 3)``.
    :param k_cs_index: Array defining the index from the library of k_cs to use for each element, ``(n_elem, )``.
    If ``None``, all elements will use the first entry in the k_cs library.
    :param m_cs_index: Array defining the index from the library of m_cs to use for each element, ``(n_elem, )``.
    If ``None``, all elements will use the first entry in the m_cs library.
    :param m_lumped_index: Node index for nodes which are to have a lumped mass attached. The order is the same as
    that for the lumped mass data ``(n_lumped_mass, )``.
    :param gravity: Gravity vector in global reference frame, or None for no gravity_vec, ``(3, )``.
    :param thrust_nodes: Dictionary of thrust node names and their corresponding node indices, {keys, int}.
    :param thrust_direction: Dictionary of thrust node names and their corresponding thrust direction vectors,
    ``{keys, (3, )}``.
    :param optional_jacobians: Define which Jacobians contributions are to be used for solution.
    :param relaxation_factor: Relaxation factor which reduces the displacement update at each iteration. A value of
    1 is no relaxation, and a value of 0 is no update.
    :param spectral_radius: Spectral radius for structural time integrator, where a value of 0 is highly damped and
    a value of 1 is undamped.
    :param alpha_m: Mass-proportional Rayleigh damping coefficient.
    :param beta_k: Stiffness-proportional Rayleigh damping coefficient.
    :param struct_convergence_settings: Structure convergence settings.
    :param constraints: Named dict ``{name: constraint}``, or None, which add
    """

    check_type(num_nodes, int)

    if constraints is None:
        named_constraints: dict[str, SoftConstraint | HardConstraint] = {}
    elif isinstance(constraints, dict):
        named_constraints = constraints
    else:
        raise ValueError("Invalid constraint input")

    hard_constraints = tuple(
        c for c in named_constraints.values() if isinstance(c, HardConstraint)
    )

    check_arr_shape(connectivity, (None, 2), "connectivity")
    check_arr_dtype(connectivity, int, "connectivity")
    _check_connectivity(connectivity, num_nodes)

    auto_node_sources: list[int] = []
    conn_list: list[list[int]] = connectivity.tolist()

    # add extra nodes and alter connectivity to account for constraints with the auto-generate node behaviour
    for con in hard_constraints:
        if con.node_j is None and not con.is_grounded:
            node_j_new = num_nodes + len(auto_node_sources)
            con.node_j = node_j_new
            auto_node_sources.append(con.node_i)
            conn_list = _split_connectivity(conn_list, con.node_i, node_j_new)
    self._auto_node_sources: tuple[int, ...] = tuple(auto_node_sources)

    num_nodes += len(auto_node_sources)
    connectivity = jnp.array(conn_list, dtype=int)
    self.n_nodes: int = num_nodes
    self.n_dof: int = num_nodes * 6

    self.connectivity: tuple[tuple[int, int], ...] = nested_list_to_tuple(
        connectivity.tolist()
    )  # (n_elem, 2)
    self.n_elem_per_node: tuple[int] = tuple(
        _n_elem_per_node(connectivity=connectivity, n_nodes=num_nodes).tolist()
    )  # (n_nodes, )
    self.n_elem: int = connectivity.shape[0]

    self.dof_per_elem: tuple[tuple[float]] = nested_list_to_tuple(
        jnp.zeros((self.n_elem, 12), dtype=int)
        .at[:, :6]
        .set(6 * self.connectivity_arr[:, [0]] + jnp.arange(6)[None, :])
        .at[:, 6:]
        .set(6 * self.connectivity_arr[:, [1]] + jnp.arange(6)[None, :])
        .tolist()
    )

    # allow for a single y_vector to be broadcast to all elements
    if y_vector.shape == (3,):
        y_vector = y_vector[None, :]
    if y_vector.shape == (1, 3):
        y_vector = jnp.broadcast_to(y_vector, (self.n_elem, 3))

    # y vectors in reference unoriented configuration, and placeholder for oriented equivalent.
    check_arr_shape(y_vector, (self.n_elem, 3), "y_vector")
    self.y_vector_reference: tuple[tuple[tuple[float]]] = nested_list_to_tuple(
        y_vector.tolist()
    )
    self.y_vector: Array = jnp.zeros_like(jnp.array(y_vector))

    # initialise design variables with default values
    self.x0_reference: Array = jnp.zeros((num_nodes, 3))  # unoriented
    self.x0: Array = jnp.zeros((num_nodes, 3))  # oriented

    self.m_cs = None
    self.k_cs = None
    self.m_lumped = None
    self.use_lumped_mass: bool = m_lumped_index is not None

    # initialise auxiliary arrays
    self.o0: Array = jnp.zeros((self.n_elem, 3, 3))
    self.l0: Array = jnp.zeros(self.n_elem)
    self.d0: Array = jnp.zeros((self.n_elem, 6))

    # initialise undeformed algebra and group
    self.hg0_reference: Array = jnp.zeros((self.n_nodes, 4, 4))  # unoriented
    self.hg0: Array = jnp.zeros((self.n_nodes, 4, 4))  # oriented

    # grads inverse action for the reference rotations
    self.ad_inv_o0: Array = jnp.zeros((self.n_elem, 6, 6))

    # gravity settings
    if not isinstance(gravity, jnp.ndarray) and gravity is not None:
        gravity = jnp.array(gravity)
    self.use_gravity: bool = gravity is not None and bool(jnp.any(gravity))
    if self.use_gravity:
        assert gravity is not None
        check_arr_shape(gravity, (3,), "gravity")
        self.gravity_vec: tuple[float, float, float] = tuple(gravity.tolist())
    else:
        self.gravity_vec = (0.0, 0.0, 0.0)

    # indexing
    if k_cs_index is None:
        k_cs_index_ = jnp.zeros(self.n_elem, dtype=int)
    else:
        check_arr_shape(k_cs_index, (self.n_elem,), "k_cs_index")
        check_arr_dtype(k_cs_index, int, "k_cs_index")
        k_cs_index_ = k_cs_index
    self.k_cs_index: tuple[int] = tuple(k_cs_index_.tolist())

    if m_cs_index is None:
        m_cs_index_ = jnp.zeros(self.n_elem, dtype=int)
    else:
        check_arr_shape(m_cs_index, (self.n_elem,), "m_cs_index")
        check_arr_dtype(m_cs_index, int, "m_cs_index")
        m_cs_index_ = m_cs_index
    self.m_cs_index: tuple[int] = tuple(m_cs_index_.tolist())

    self.m_lumped_index: tuple[int] | None = None
    if m_lumped_index is not None:
        check_arr_dtype(m_lumped_index, int, "m_lumped_index")
        if m_lumped_index.ndim not in (0, 1):
            raise ValueError("m_lumped_index.ndim must be 0 or 1.")
        self.m_lumped_index = tuple(jnp.atleast_1d(m_lumped_index).tolist())

    # add thrust
    self.thrust_nodes: tuple[tuple[str, int], ...] = ()
    self.thrust_direction: tuple[tuple[str, tuple[float, float, float]], ...] = ()
    if thrust_nodes is not None and thrust_direction is not None:
        if thrust_nodes.keys() != thrust_direction.keys():
            raise ValueError(
                f"Mismatch in keys of thrust_nodes ({thrust_nodes.keys()}) and thrust_direction ({thrust_direction.keys()}))."
            )

        for k, v in thrust_direction.items():
            check_arr_shape(v, (3,), f"thrust_direction[{k}]")

        self.thrust_nodes = tuple([(k, v) for k, v in thrust_nodes.items()])
        self.thrust_direction = tuple(
            [
                (k, nested_list_to_tuple((v / jnp.linalg.norm(v)).tolist()))
                for k, v in thrust_direction.items()
            ]
        )  # make unit vectors
    elif thrust_nodes is not None or thrust_direction is not None:
        warn(
            "One of thrust_nodes or thrust_direction has not been passed. Running with no thrust nodes."
        )

    # set the reference thrust to be zero, which can be overwritten later
    self.thrust_reference: dict[str, Array] = {
        k: jnp.atleast_1d(1) for k in [k_ for k_, v in self.thrust_nodes]
    }

    # set the reference orientation, which can be overwritten later.
    self.orientation_euler: Array = jnp.zeros(3)
    self.orientation: Array = jnp.eye(3)

    self.optional_jacobians: OptionalJacobians = (
        optional_jacobians
        if optional_jacobians is not None
        else OptionalJacobians()
    )
    self.struct_convergence_settings: ConvergenceSettings = (
        struct_convergence_settings
    )
    self.relaxation_factor: float = relaxation_factor
    self.spectral_radius: float = spectral_radius
    self.alpha_m: float = float(alpha_m)
    self.beta_k: float = float(beta_k)

    self.time_integrator = None

    self.constraints: dict[str, SoftConstraint | HardConstraint] = named_constraints

n_multibody_constraints property

n_multibody_constraints: int

Total number of scalar Lagrange-multiplier constraints.

n_holonomic_constraints property

n_holonomic_constraints: int

Number of scalar holonomic (position-level) Lagrange-multiplier constraints.

n_nonholonomic_constraints property

n_nonholonomic_constraints: int

Number of scalar non-holonomic (velocity-level) Lagrange-multiplier constraints.

set_design_variables

set_design_variables(
    coords: Array,
    k_cs: Array,
    m_cs: Array | None,
    m_lumped: Array | None = None,
    orientation_euler: Array | None = None,
    thrust_reference: dict[str, Array | float]
    | None = None,
    *,
    remove_checks: bool = False,
) -> None

Set design variables and compute initial configuration dependent quantities.

Parameters:

Name	Type	Description	Default
coords	Array	Node coordinates in the reference configuration, (n_nodes, 3).	required
k_cs	Array	Cross-section stiffness matrices, (n_entry, 6, 6) or (6, 6).	required
m_cs	Array | None	Cross-section mass matrices, (n_entry, 6, 6) or (6, 6).	required
m_lumped	Array | None	Lumped mass matrices at nodes, (n_entry, 6, 6).	None
orientation_euler	Array | None	Euler angles in radians which to rotate the reference configuration by, (3, ). This rotation is performed about the origin, and will default to the identity is no Array is passed. These are rotated in z-y-x order.	None
thrust_reference	dict[str, Array | float] | None	Reference thrust magnitude, {keys, (1, )}.	None
remove_checks	bool	Flag to ignore input checks, used when function is JIT compiled.	False

Source code in src/flapjax/structure/beam.py

def set_design_variables(
    self,
    coords: Array,
    k_cs: Array,
    m_cs: Array | None,
    m_lumped: Array | None = None,
    orientation_euler: Array | None = None,
    thrust_reference: dict[str, Array | float] | None = None,
    *,
    remove_checks: bool = False,
) -> None:
    r"""
    Set design variables and compute initial configuration dependent quantities.
    :param coords: Node coordinates in the reference configuration, ``(n_nodes, 3)``.
    :param k_cs: Cross-section stiffness matrices, ``(n_entry, 6, 6)`` or ``(6, 6)``.
    :param m_cs: Cross-section mass matrices, ``(n_entry, 6, 6)`` or ``(6, 6)``.
    :param m_lumped: Lumped mass matrices at nodes, ``(n_entry, 6, 6)``.
    :param orientation_euler: Euler angles in radians which to rotate the reference configuration by, ``(3, )``. This rotation
    is performed about the origin, and will default to the identity is no Array is passed. These are rotated in
    z-y-x order.
    :param thrust_reference: Reference thrust magnitude, ``{keys, (1, )}``.
    :param remove_checks: Flag to ignore input checks, used when function is JIT compiled.
    """

    # orientation
    if orientation_euler is not None:
        check_arr_shape(orientation_euler, (3,), "orientation_euler")
        self.orientation_euler = orientation_euler
        self.orientation = Rotation.from_euler(
            seq="zyx", angles=orientation_euler
        ).as_matrix()

    # rotate the y vectors
    self.y_vector = jnp.einsum(
        "jk,ik->ij",
        self.orientation,
        self.y_vector_reference_arr,
    )

    # coordinates — auto-extend for nodes created by multibody constraints
    if self._auto_node_sources and coords.shape[0] == self.n_nodes - len(
        self._auto_node_sources
    ):
        coords = jnp.concatenate(
            [coords, coords[jnp.array(self._auto_node_sources)]],
            axis=0,
        )
    check_arr_shape(coords, (self.n_nodes, 3), "coords")
    self.x0_reference = coords
    self.x0 = jnp.einsum("jk,ik->ij", self.orientation, coords)

    # populate arrays
    if k_cs.ndim == 2:
        k_cs = k_cs[None, ...]
    check_arr_shape(k_cs, (None, 6, 6), "k_cs")

    if (
        not remove_checks
        and k_cs.shape[0] != jnp.unique_values(jnp.array(self.k_cs_index)).size
    ):
        warn(
            "Redundant values in k_cs which are not used for solution due to no corresponding entry in k_cs_index."
        )

    self.k_cs = k_cs
    if m_cs is None:
        if not remove_checks and self.use_gravity and m_lumped is None:
            warn(
                "No mass matrices provided, but gravity is enabled. Assuming zero mass.",
            )
        m_cs_ = jnp.zeros((6, 6))
    else:
        m_cs_ = m_cs

    if m_cs_.ndim == 2:
        m_cs_ = m_cs_[None, ...]

    check_arr_shape(m_cs_, (None, 6, 6), "m_cs")

    if (
        not remove_checks
        and m_cs_.shape[0] != jnp.unique_values(jnp.array(self.m_cs_index)).size
        and m_cs is not None
    ):
        warn(
            "Redundant values in m_cs which are not used for solution due to no corresponding entry in "
            "m_cs_index."
        )

    self.m_cs = m_cs_

    # thrust
    if thrust_reference is not None:
        self.thrust_reference = {
            k: jnp.atleast_1d(v) for k, v in thrust_reference.items()
        }

        for k, v in self.thrust_reference.items():
            check_arr_shape(v, (1,), f"thrust_reference[{k}]")

    if m_lumped is not None:
        if not remove_checks:
            check_arr_shape(m_lumped, (None, 6, 6), "m_lumped")

            if self.m_lumped_index is None:
                raise ValueError("m_lumped_index has not been set")

            if m_lumped.shape[0] != len(self.m_lumped_index):
                raise ValueError(
                    "Number of entries in m_lumped does not match number of indices in m_lumped_index."
                )

        self.m_lumped = m_lumped

    # obtain initial orientation and length
    x_elem = jnp.take(
        self.x0_reference, self.connectivity_arr, axis=0
    )  # (n_elem, 2, 3)
    dx = x_elem[:, 1, :] - x_elem[:, 0, :]  # (n_elem, 3)

    # ensure out-of-plane vector and beam vector are not collinear
    if not remove_checks and jnp.any(
        jnp.linalg.norm(jnp.cross(dx, self.y_vector_reference_arr, 1, 1), axis=-1)
        < 1e-6
    ):
        raise ValueError(
            "y_vector is collinear with beam element direction for at least one element. "
            "Please provide a different y_vector."
        )

    self.l0 = jnp.linalg.norm(dx, axis=-1)  # (n_elem,)
    self.d0 = self.d0.at[:, 0].set(self.l0)

    dx_unit = dx / self.l0[:, None]  # unit vector in beam direction, (n_elem, 3)
    dz = jnp.cross(
        dx_unit, self.y_vector_reference_arr, axis=-1
    )  # vector in plane(n_elem, 3)
    dz_unit = dz / jnp.linalg.norm(dz, axis=-1)[:, None]  # (n_elem, 3)

    dy_unit = jnp.cross(dz_unit, dx_unit)

    self.o0 = self.o0.at[..., 0].set(dx_unit)
    self.o0 = self.o0.at[..., 1].set(dy_unit)
    self.o0 = self.o0.at[..., 2].set(dz_unit)

    self.ad_inv_o0 = vmap(rmat_to_ha_hat)(jnp.transpose(self.o0, (0, 2, 1)))

    # set unoriented initial coordinates
    self.hg0_reference = jnp.broadcast_to(
        jnp.eye(4)[None, ...], (self.n_nodes, 4, 4)
    )  # (n_nodes, 4, 4)
    self.hg0_reference = self.hg0_reference.at[:, :3, 3].set(self.x0_reference)

    # set oriented initial coordinates
    self.hg0 = self.hg0.at[:, :3, :3].set(
        jnp.broadcast_to(self.orientation[None, ...], (self.n_nodes, 3, 3))
    )  # (n_nodes, 4, 4)
    self.hg0 = self.hg0.at[:, :3, 3].set(self.x0)
    self.hg0 = self.hg0.at[:, 3, 3].set(1.0)

    # add reference frames to the nodal constraints
    for con in self.nodal_constraints:
        con.resolve_hg_ref(self.hg0)
    for con in self.multibody_constraints:
        con.resolve_hg_ref(self.hg0)

get_design_variables

get_design_variables(
    struct_case: StructureCase,
    thrust_t: dict[str, Array],
    grads_to_compute: StructureGradsToCompute | None,
) -> StructureDesignVariables

Obtain the design variables for the structural problem. As the external forcing is defined for each solve, the chosen forcing is required as input.

Parameters:

Name	Type	Description	Default
struct_case	StructureCase	Structural case	required
thrust_t	dict[str, Array]	Thrust time history, {keys, (n_tstep,)}.	required
grads_to_compute	StructureGradsToCompute | None	Data structure which describes which design variables should be obtained. If none, all variables are obtained.	required

Returns:

Type	Description
StructureDesignVariables	StructureDesignVariables dataclass containing design variables

Source code in src/flapjax/structure/beam.py

def get_design_variables(
    self,
    struct_case: StructureCase,
    thrust_t: dict[str, Array],
    grads_to_compute: StructureGradsToCompute | None,
) -> StructureDesignVariables:
    r"""
    Obtain the design variables for the structural problem. As the external forcing is defined for each solve, the
    chosen forcing is required as input.
    :param struct_case: Structural case
    :param thrust_t: Thrust time history, {keys, ``(n_tstep,)``}.
    :param grads_to_compute: Data structure which describes which design variables should be obtained. If none, all
    variables are obtained.
    :return: StructureDesignVariables dataclass containing design variables
    """

    # struct_case.f_ext_dead is stored in local frame: f_local = R^T @ f_global,
    # so recover f_global = R @ f_local
    hg = struct_case.hg
    if hg.ndim == 4:  # batched case: (n_tstep, n_nodes, 4, 4)
        rmat = hg[:, :, :3, :3]
    else:  # snapshot: (n_nodes, 4, 4)
        rmat = hg[:, :3, :3]
    f_ext_dead_global = (
        transform_nodal_vect(struct_case.f_ext_dead, rmat)
        if struct_case.f_ext_dead is not None
        else None
    )
    if isinstance(grads_to_compute, StructureGradsToCompute):
        return StructureDesignVariables(
            x0=self.x0 if grads_to_compute.x0 else None,
            orientation_euler=self.orientation_euler
            if grads_to_compute.orientation_euler
            else None,
            m_cs=self.m_cs if grads_to_compute.m_cs else None,
            k_cs=self.k_cs if grads_to_compute.k_cs else None,
            m_lumped=self._m_lumped if grads_to_compute.m_lumped else None,
            f_ext_dead=f_ext_dead_global if grads_to_compute.f_ext_dead else None,
            f_ext_follower=struct_case.f_ext_follower
            if grads_to_compute.f_ext_follower
            else None,
            thrust_t=thrust_t if grads_to_compute.thrust_t else None,
            f_shape=(),
        )
    else:
        return StructureDesignVariables(
            x0=self.x0,
            orientation_euler=self.orientation_euler,
            m_cs=self.m_cs,
            k_cs=self.k_cs,
            m_lumped=self._m_lumped,
            f_ext_dead=f_ext_dead_global,
            f_ext_follower=struct_case.f_ext_follower,
            thrust_t=thrust_t,
            f_shape=(),
        )

reference_configuration

reference_configuration(
    prescribed_dofs: Sequence[int]
    | Array
    | slice
    | int = (),
    use_f_ext_follower: bool = True,
    use_f_ext_dead: bool = True,
    use_f_aero: bool = True,
    use_f_grav: bool = True,
) -> StructureCase

Get the reference configuration of the structure.

Parameters:

Name	Type	Description	Default
prescribed_dofs	Sequence[int] | Array | slice | int	Prescribed degrees of freedom, which are not solved for. Defaults to no prescribed DoFs.	()
use_f_ext_follower	bool	Whether to include follower forces in the reference configuration.	True
use_f_ext_dead	bool	Whether to include dead forces in the reference configuration.	True
use_f_aero	bool	Whether to include aerodynamic forces in the reference configuration.	True
use_f_grav	bool	Whether to include gravitational forces in the reference configuration.	True

Returns:

Type	Description
StructureCase	Structure dataclass containing reference configuration.

Source code in src/flapjax/structure/beam.py

def reference_configuration(
    self,
    prescribed_dofs: Sequence[int] | Array | slice | int = (),
    use_f_ext_follower: bool = True,
    use_f_ext_dead: bool = True,
    use_f_aero: bool = True,
    use_f_grav: bool = True,
) -> StructureCase:
    r"""
    Get the reference configuration of the structure.
    :param prescribed_dofs: Prescribed degrees of freedom, which are not solved for. Defaults to no prescribed DoFs.
    :param use_f_ext_follower: Whether to include follower forces in the reference configuration.
    :param use_f_ext_dead: Whether to include dead forces in the reference configuration.
    :param use_f_aero: Whether to include aerodynamic forces in the reference configuration.
    :param use_f_grav: Whether to include gravitational forces in the reference configuration.
    :return: Structure dataclass containing reference configuration.
    """
    prescribed_dofs = self.make_prescribed_dofs_tuple(prescribed_dofs)
    return StructureCase(
        hg=self.hg0,
        conn=self.connectivity,
        o0=self.o0,
        d=self.d0,
        eps=jnp.zeros((self.n_elem, 6)),
        varphi=jnp.zeros((self.n_nodes, 6)),
        f_ext_follower=jnp.zeros((self.n_nodes, 6)) if use_f_ext_follower else None,
        f_ext_dead=jnp.zeros((self.n_nodes, 6)) if use_f_ext_dead else None,
        f_ext_aero=jnp.zeros((self.n_nodes, 6)) if use_f_aero else None,
        f_grav=jnp.zeros((self.n_nodes, 6)) if use_f_grav else None,
        f_int=jnp.zeros((self.n_nodes, 6)),
        f_elem=jnp.zeros((self.n_elem, 6)),
        f_res=jnp.zeros((self.n_nodes, 6)),
        thrust=self.thrust_reference,
        thrust_direction=self.thrust_direction,
        thrust_nodes=self.thrust_nodes,
        local=True,
        prescribed_dofs=prescribed_dofs,
        t=jnp.zeros(1),
    )

compute_varphi_from_hg

compute_varphi_from_hg(hg: Array) -> Array

Calculate the twist vector from the reference configuration to hg

Parameters:

Name	Type	Description	Default
hg	Array	Deformed coordinates, (n_nodes, 4, 4)	required

Returns:

Type	Description
Array	Vector of twists, (n_nodes, 6)

Source code in src/flapjax/structure/beam.py

def compute_varphi_from_hg(self, hg: Array) -> Array:
    r"""
    Calculate the twist vector from the reference configuration to hg
    :param hg: Deformed coordinates, ``(n_nodes, 4, 4)``
    :return: Vector of twists, ``(n_nodes, 6)``
    """
    return vmap(hg_to_d, (0, 0), 0)(self.hg0, hg)

assemble_matrix_from_entries

assemble_matrix_from_entries(entries: Array) -> Array

Assemble global matrix from element entries

Parameters:

Name	Type	Description	Default
entries	Array	Array of element matrix entries, (n_elem, 12, 12)	required

Returns:

Type	Description
Array	System global matrix, (n_dof, n_dof)

Source code in src/flapjax/structure/beam.py

def assemble_matrix_from_entries(self, entries: Array) -> Array:
    r"""
    Assemble global matrix from element entries
    :param entries: Array of element matrix entries, ``(n_elem, 12, 12)``
    :return: System global matrix, ``(n_dof, n_dof)``
    """

    row_idx = jnp.broadcast_to(
        self.dof_per_elem_arr[:, :, None], (self.n_elem, 12, 12)
    )
    col_idx = jnp.broadcast_to(
        self.dof_per_elem_arr[:, None, :], (self.n_elem, 12, 12)
    )
    return (
        jnp.zeros((self.n_dof, self.n_dof))
        .at[row_idx.ravel(), col_idx.ravel()]
        .add(entries.ravel())
    )

assemble_vector_from_entries

assemble_vector_from_entries(entries: Array) -> Array

Assemble global vector from element entries

Parameters:

Name	Type	Description	Default
entries	Array	Array of element vector entries, (n_elem, 12)	required

Returns:

Type	Description
Array	System global vector, (n_dof, )

Source code in src/flapjax/structure/beam.py

def assemble_vector_from_entries(self, entries: Array) -> Array:
    r"""
    Assemble global vector from element entries
    :param entries: Array of element vector entries, ``(n_elem, 12)``
    :return: System global vector, ``(n_dof, )``
    """

    vect = jnp.zeros(self.n_dof)
    vect = vect.at[self.dof_per_elem_arr[:, :6]].add(entries[:, :6])
    return vect.at[self.dof_per_elem_arr[:, 6:]].add(entries[:, 6:])

add_lumped_contributions_to_arr

add_lumped_contributions_to_arr(
    arr: Array, lumped_arr: Array
) -> Array

Add lumped contributions to an array

Parameters:

Name	Type	Description	Default
arr	Array	Full array, (6*n_node, 6*n_node)	required
lumped_arr	Array	Lumped contributions, (n_lump, 6, 6)	required

Returns:

Type	Description
Array	In-place updated array, (6*n_node, 6*n_node)

Source code in src/flapjax/structure/beam.py

def add_lumped_contributions_to_arr(self, arr: Array, lumped_arr: Array) -> Array:
    r"""
    Add lumped contributions to an array
    :param arr: Full array, ``(6*n_node, 6*n_node)``
    :param lumped_arr: Lumped contributions, ``(n_lump, 6, 6)``
    :return: In-place updated array, ``(6*n_node, 6*n_node)``
    """

    assert self.m_lumped_index is not None

    def add_block(carry, x):
        node_idx, block = x
        dofs = node_idx * 6 + jnp.arange(6)
        return carry.at[jnp.ix_(dofs, dofs)].add(block), None

    arr, _ = jax.lax.scan(
        add_block, arr, (jnp.array(self.m_lumped_index), lumped_arr)
    )
    return arr

add_lumped_contributions_to_vec

add_lumped_contributions_to_vec(
    vec: Array, lumped_vec: Array
) -> Array

Add lumped contributions to an array

Parameters:

Name	Type	Description	Default
vec	Array	Full vector, (6*n_node, )	required
lumped_vec	Array	Lumped contributions, (n_lump, 6)	required

Returns:

Type	Description
Array	In-place updated vector, (6*n_node, ).

Source code in src/flapjax/structure/beam.py

def add_lumped_contributions_to_vec(self, vec: Array, lumped_vec: Array) -> Array:
    r"""
    Add lumped contributions to an array
    :param vec: Full vector, ``(6*n_node, )``
    :param lumped_vec: Lumped contributions, ``(n_lump, 6)``
    :return: In-place updated vector, ``(6*n_node, )``.
    """

    assert self.m_lumped_index is not None

    idx = (
        jnp.array(self.m_lumped_index)[:, None] * 6 + jnp.arange(6)[None, :]
    ).ravel()  # (n_lump * 6,)

    return vec.at[idx].add(lumped_vec)

make_k_t

make_k_t(d: Array, p_d: Array, eps: Array) -> Array

Assemble tangent stiffness matrix as a function of the element relative configuration vectors

Parameters:

Name	Type	Description	Default
d	Array	Element relative configuration, (n_elem, 6).	required
p_d	Array	P(d) operator, (n_elem, 6, 12).	required
eps	Array	Element strains, (n_elem, 6).	required

Returns:

Type	Description
Array	Elementwise stiffness matrix entries, (n_elem, 12, 12).

Source code in src/flapjax/structure/beam.py

def make_k_t(
    self,
    d: Array,
    p_d: Array,
    eps: Array,
) -> Array:
    r"""
    Assemble tangent stiffness matrix as a function of the element relative configuration vectors
    :param d: Element relative configuration, ``(n_elem, 6)``.
    :param p_d: P(d) operator, ``(n_elem, 6, 12)``.
    :param eps: Element strains, ``(n_elem, 6)``.
    :return: Elementwise stiffness matrix entries, ``(n_elem, 12, 12)``.
    """
    # compute stiffness matrix entries
    return vmap(
        partial(
            _k_t_entry,
            include_geometric=self.optional_jacobians.d_f_int_d_p_d,
        ),
        (0, 0, 0, 0, 0, 0),
        0,
    )(
        d,
        p_d,
        self.l0,
        eps,
        self.k_cs[self.k_cs_index, ...],
        self.ad_inv_o0,
    )  # (n_elem, 12, 12)

make_k_t_full

make_k_t_full(
    d: Array,
    p_d: Array,
    eps: Array,
    f_ext_dead: Array | None,
    rmat: Array,
    m_t: Array | None,
) -> Array

Compute the full tangent stiffness matrix, with contributions from stiffness, dead forces and gravity.

Parameters:

Name	Type	Description	Default
d	Array	Element relative configuration, (n_elem, 6).	required
p_d	Array	P(d) operator, (n_elem, 6, 12).	required
eps	Array	Strain vectors, (n_elem, 6).	required
f_ext_dead	Array | None	External dead forces in global reference, (n_node, 6).	required
rmat	Array	Nodal rotation matrices, (n_node, 3, 3).	required
m_t	Array | None	Disassembled system mass matrix, (n_elem, 12, 12).	required

Returns:

Type	Description
Array	Tangent stiffness matrix with all contributions, (n_dof, n_dof).

Source code in src/flapjax/structure/beam.py

def make_k_t_full(
    self,
    d: Array,
    p_d: Array,
    eps: Array,
    f_ext_dead: Array | None,
    rmat: Array,
    m_t: Array | None,
) -> Array:
    r"""
    Compute the full tangent stiffness matrix, with contributions from stiffness, dead forces and gravity.
    :param d: Element relative configuration, ``(n_elem, 6)``.
    :param p_d: P(d) operator, ``(n_elem, 6, 12)``.
    :param eps: Strain vectors, ``(n_elem, 6)``.
    :param f_ext_dead: External dead forces in global reference, ``(n_node, 6)``.
    :param rmat: Nodal rotation matrices, ``(n_node, 3, 3)``.
    :param m_t: Disassembled system mass matrix, ``(n_elem, 12, 12)``.
    :return: Tangent stiffness matrix with all contributions, ``(n_dof, n_dof)``.
    """

    k_t = self.assemble_matrix_from_entries(self.make_k_t(d, p_d, eps))
    if f_ext_dead is not None and self.optional_jacobians.d_f_ext_dead_d_n:
        k_t += block_diag(*self._make_k_t_dead(rmat, f_ext_dead))

    if self.use_gravity and self.optional_jacobians.d_f_grav_d_n:
        if m_t is None:
            raise ValueError("m_t needs to be provided")
        k_t += self.assemble_matrix_from_entries(
            self._make_k_t_grav(d, p_d, rmat, m_t)
        )
        if self.use_lumped_mass:
            k_t_lumped = self._make_k_t_grav_lumped(rmat)
            k_t = self.add_lumped_contributions_to_arr(
                arr=k_t, lumped_arr=k_t_lumped
            )
    return k_t

make_m_t

make_m_t(
    d: Array,
    int_order: Literal[3, 4, 5] = BASE_LOBATTO_ORDER,
) -> Array

Assemble tangent mass matrix as a function of the element relative configuration vectors. This does not include the lumped mass contribution.

Parameters:

Name	Type	Description	Default
d	Array	Element relative configuration, (n_elem, 6)	required
int_order	Literal[3, 4, 5]	Integration order for mass matrix computation	BASE_LOBATTO_ORDER

Returns:

Type	Description
Array	Elementwise mass matrix, (n_elem, 12, 12)

Source code in src/flapjax/structure/beam.py

def make_m_t(
    self, d: Array, int_order: Literal[3, 4, 5] = BASE_LOBATTO_ORDER
) -> Array:
    r"""
    Assemble tangent mass matrix as a function of the element relative configuration vectors. This does not include
    the lumped mass contribution.
    :param d: Element relative configuration, ``(n_elem, 6)``
    :param int_order: Integration order for mass matrix computation
    :return: Elementwise mass matrix, ``(n_elem, 12, 12)``
    """
    return vmap(partial(_integrate_m_l, int_order=int_order), (0, 0, 0, 0), 0)(
        self.m_cs[self.m_cs_index, ...], d, self.ad_inv_o0, self.l0
    )

make_nodal_m_k

make_nodal_m_k(
    case: StructureCase,
    int_order: Literal[3, 4, 5] = BASE_LOBATTO_ORDER,
) -> tuple[Array, Array]

Create the global mass and stiffness matrices for a given static structure case. These can be used for modal analysis or other purposes. These matrices are the Jacobians of the local forcing residual with respect to global perturbations in acceleration and displacement, respectively.

Parameters:

Name	Type	Description	Default
case	StructureCase	Static structure case for which to compute the global mass and stiffness matrices.	required
int_order	Literal[3, 4, 5]	Integration order for mass matrix computation.	BASE_LOBATTO_ORDER

Returns:

Type	Description
tuple[Array, Array]	Global mass and stiffness matrices, (n_free_dof, n_free_dof).

Source code in src/flapjax/structure/beam.py

def make_nodal_m_k(
    self,
    case: StructureCase,
    int_order: Literal[3, 4, 5] = BASE_LOBATTO_ORDER,
) -> tuple[Array, Array]:
    r"""
    Create the global mass and stiffness matrices for a given static structure case. These can be used for modal
    analysis or other purposes. These matrices are the Jacobians of the local forcing residual with respect to
    global perturbations in acceleration and displacement, respectively.
    :param case: Static structure case for which to compute the global mass and stiffness matrices.
    :param int_order: Integration order for mass matrix computation.
    :return: Global mass and stiffness matrices, ``(n_free_dof, n_free_dof)``.
    """
    # extract variables from case
    d = case.d
    eps = self.make_eps(d=d)
    p_d = self.make_p_d(d=d)
    t_varphi = vmap(t_se3)(case.varphi)  # (n_node, 6, 6)
    rmat = case.hg[:, :3, :3]  # (n_node, 3, 3)

    # get dead external forcing as this has a stiffness contribution
    f_ext_dead_local: Array | None
    if case.f_ext_dead is not None and case.f_ext_aero is not None:
        f_ext_dead_local = case.f_ext_dead + case.f_ext_aero
    else:
        f_ext_dead_local = (
            case.f_ext_dead if case.f_ext_dead is not None else case.f_ext_aero
        )

    # convert to global frame, as it required for creating the stiffness matrix
    f_ext_dead: Array | None = (
        transform_nodal_vect(f_ext_dead_local, rmat)
        if f_ext_dead_local is not None
        else None
    )
    free_dofs = jnp.array(
        get_solve_dofs(n_dof=self.n_dof, prescribed_dofs=case.prescribed_dofs)
    )

    def transform_mat_to_global(mat: Array) -> Array:
        # function to rotate a forcing Jacobian matrix from the local frame to the global frame.
        mat_reshaped = mat.reshape(self.n_nodes, 6, self.n_dof)
        m_lin = jnp.einsum("nij,njk->nik", rmat, mat_reshaped[:, :3, :])
        m_rot = jnp.einsum("nij,njk->nik", rmat, mat_reshaped[:, 3:, :])
        return jnp.concatenate((m_lin, m_rot), axis=1).reshape(
            self.n_dof, self.n_dof
        )

    # mass
    m_t = self.assemble_matrix_from_entries(
        self.make_m_t(d=d, int_order=int_order)
    )  # (n_dof, n_dof)
    if self.use_lumped_mass:
        m_t = self.add_lumped_contributions_to_arr(
            arr=m_t, lumped_arr=self.m_lumped
        )

    m_modal_full = transform_mat_to_global(
        mat=jnp.einsum("ijk,jkl->ijl", m_t.reshape(self.n_dof, -1, 6), t_varphi)
    )

    m_modal = m_modal_full.reshape(self.n_dof, self.n_dof)[
        jnp.ix_(free_dofs, free_dofs)
    ]

    # stiffness
    k_t = self.make_k_t_full(
        d=case.d, p_d=p_d, eps=eps, f_ext_dead=f_ext_dead, rmat=rmat, m_t=m_t
    )

    k_modal_full = transform_mat_to_global(
        mat=jnp.einsum("ijk,jkl->ijl", k_t.reshape(self.n_dof, -1, 6), t_varphi)
    )

    k_modal = k_modal_full.reshape(self.n_dof, self.n_dof)[
        jnp.ix_(free_dofs, free_dofs)
    ]

    return m_modal, k_modal

modal

modal(
    case: StructureCase,
    remove_complex_conjugate: bool = True,
    int_order: Literal[3, 4, 5] = BASE_LOBATTO_ORDER,
    n_modes: int = 20,
    freq_range: tuple[float | Array, float | Array] = (
        0.0,
        jnp.inf,
    ),
    damp_range: tuple[float | Array, float | Array] = (
        -jnp.inf,
        jnp.inf,
    ),
    vtu_directory: str | PathLike = "./modal",
    n_plot_vtu: int | None = None,
    aero: UVLM | None = None,
    n_phase: int = 8,
    n_interp: int = 0,
    max_disp: float = 0.2,
    max_ang: float = 0.2,
) -> tuple[Array, Array, Array]

Perform modal analysis on the structure.

Parameters:

Name	Type	Description	Default
case	StructureCase	The static structure case for which to perform modal analysis.	required
remove_complex_conjugate	bool	If true, keep only one mode from each complex conjugate pair.	True
int_order	Literal[3, 4, 5]	Integration order for mass matrix computation.	BASE_LOBATTO_ORDER
n_modes	int	Number of modes to preserve.	20
freq_range	tuple[float | Array, float | Array]	Frequency range for filtering out modes.	(0.0, inf)
damp_range	tuple[float | Array, float | Array]	Damping range for filtering out modes.	(-inf, inf)
vtu_directory	str | PathLike	Directory to for saving the mode shapes to vtu files.	'./modal'
n_plot_vtu	int | None	Number of modes to plot to vtu files. Will default to "./modal".	None
aero	UVLM | None	UVLM aerodynamic model. If passed, the vtu files will include the aerodynamic grid. If not, they will just be the beam structure.	None
n_phase	int	Number of phases to use when plotting the modes to vtu files.	8
n_interp	int	Number of times to interpolate between beam nodes for vtu plotting.	0
max_disp	float	Maximum displacement of structure for plotted modes, used for scaling.	0.2
max_ang	float	Maximum angle of structure for plotted modes in radians, used for scaling.	0.2

Returns:

Type	Description
tuple[Array, Array, Array]	Tuple of natural frequencies (n_free_dof), damping ratios (n_free_dof), and mode shapes with no normalisation (n_modes, n_free_dof).

Source code in src/flapjax/structure/beam.py

def modal(
    self,
    case: StructureCase,
    remove_complex_conjugate: bool = True,
    int_order: Literal[3, 4, 5] = BASE_LOBATTO_ORDER,
    n_modes: int = 20,
    freq_range: tuple[float | Array, float | Array] = (0.0, jnp.inf),
    damp_range: tuple[float | Array, float | Array] = (-jnp.inf, jnp.inf),
    vtu_directory: str | os.PathLike = "./modal",
    n_plot_vtu: int | None = None,
    aero: UVLM | None = None,
    n_phase: int = 8,
    n_interp: int = 0,
    max_disp: float = 0.2,
    max_ang: float = 0.2,
) -> tuple[Array, Array, Array]:
    r"""
    Perform modal analysis on the structure.
    :param case: The static structure case for which to perform modal analysis.
    :param remove_complex_conjugate: If true, keep only one mode from each complex conjugate pair.
    :param int_order: Integration order for mass matrix computation.
    :param n_modes: Number of modes to preserve.
    :param freq_range: Frequency range for filtering out modes.
    :param damp_range: Damping range for filtering out modes.
    :param vtu_directory: Directory to for saving the mode shapes to vtu files.
    :param n_plot_vtu: Number of modes to plot to vtu files. Will default to "./modal".
    :param aero: UVLM aerodynamic model. If passed, the vtu files will include the aerodynamic grid. If not, they
    will just be the beam structure.
    :param n_phase: Number of phases to use when plotting the modes to vtu files.
    :param n_interp: Number of times to interpolate between beam nodes for vtu plotting.
    :param max_disp: Maximum displacement of structure for plotted modes, used for scaling.
    :param max_ang: Maximum angle of structure for plotted modes in radians, used for scaling.
    :return: Tuple of natural frequencies (n_free_dof), damping ratios (n_free_dof), and mode shapes with no
     normalisation ``(n_modes, n_free_dof)``.
    """
    freqs, damping, modes, *_ = self.base_modal(
        case=case,
        freq_range=freq_range,
        damp_range=damp_range,
        int_order=int_order,
        n_modes=n_modes,
        remove_complex_conjugate=remove_complex_conjugate,
    )

    if n_plot_vtu is not None:
        q_full = (
            jnp.zeros((n_plot_vtu, self.n_nodes * 6))
            .at[:, case.free_dofs]
            .set(modes[:n_plot_vtu, :])
        )

        for _i_mode in range(n_plot_vtu):
            plot_modes_vtu(
                reference=case,
                directory=vtu_directory,
                q_full=q_full.reshape(n_plot_vtu, self.n_nodes, 6),
                freqs=freqs,
                dampings=damping,
                gamma_b_full=None,
                gamma_w_full=None,
                zeta_w_full=None,
                uvlm=aero,
                n_interp=n_interp,
                n_phase=n_phase,
                max_disp=max_disp,
                max_ang=max_ang,
                max_gamma=1e6,
            )

    return freqs, damping, modes

linearise

linearise(
    reference: StructureCase,
    dt: float,
    n_modes: int | None = None,
    modal_inputs: bool = False,
    modal_outputs: bool = False,
    prescribed_dofs: Sequence[int]
    | Array
    | slice
    | int
    | None = None,
) -> LinearBeam

Linearise the beam about a given static structure case. This creates a LinearBeam object which can be used for linear dynamic analysis.

Parameters:

Name	Type	Description	Default
reference	StructureCase	Static structure case about which to linearise the beam.	required
dt	float	Time step size, used for conversions between continuous and discrete time.	required
n_modes	int | None	If not None, the linearised system uses modal state coordinates truncated to this many modes.	None
modal_inputs	bool	If True, external forcing inputs are provided as modal forces (requires n_modes).	False
modal_outputs	bool	If True, outputs are exposed as modal coordinates (requires n_modes).	False
prescribed_dofs	Sequence[int] | Array | slice | int | None	If provided, overrides the prescribed DOFs from the reference case.	None

Returns:

Type	Description
LinearBeam	Continuous-time linearised beam object.

Source code in src/flapjax/structure/beam.py

def linearise(
    self,
    reference: StructureCase,
    dt: float,
    n_modes: int | None = None,
    modal_inputs: bool = False,
    modal_outputs: bool = False,
    prescribed_dofs: Sequence[int] | Array | slice | int | None = None,
) -> LinearBeam:
    r"""
    Linearise the beam about a given static structure case. This creates a LinearBeam object which can be used for
    linear dynamic analysis.
    :param reference: Static structure case about which to linearise the beam.
    :param dt: Time step size, used for conversions between continuous and discrete time.
    :param n_modes: If not None, the linearised system uses modal state coordinates truncated to this many modes.
    :param modal_inputs: If True, external forcing inputs are provided as modal forces (requires n_modes).
    :param modal_outputs: If True, outputs are exposed as modal coordinates (requires n_modes).
    :param prescribed_dofs: If provided, overrides the prescribed DOFs from the reference case.
    :return: Continuous-time linearised beam object.
    """
    return LinearBeam(
        beam=self,
        reference=reference,
        dt=dt,
        n_modes=n_modes,
        modal_inputs=modal_inputs,
        modal_outputs=modal_outputs,
        prescribed_dofs=prescribed_dofs,
    )

apply_nodal_constraint_tangent

apply_nodal_constraint_tangent(
    mat: Array,
    hg: Array,
    i_ts: int,
    gamma_prime: float | Array | None,
) -> Array

Add nodal constraint contributions to a system matrix.

Parameters:

Name	Type	Description	Default
mat	Array	System matrix to update, (n_dof, n_dof).	required
hg	Array	SE(3) coordiantes, (n_nodes, 4, 4).	required
i_ts	int	Time-step index (0 for static solves).	required
gamma_prime	float | Array | None	Time-integrator gamma_prime for damping scaling, or None to skip damping.	required

Returns:

Type	Description
Array	Updated system matrix.

Source code in src/flapjax/structure/beam.py

def apply_nodal_constraint_tangent(
    self,
    mat: Array,
    hg: Array,
    i_ts: int,
    gamma_prime: float | Array | None,
) -> Array:
    r"""
    Add nodal constraint contributions to a system matrix.
    :param mat: System matrix to update, ``(n_dof, n_dof)``.
    :param hg: SE(3) coordiantes, ``(n_nodes, 4, 4)``.
    :param i_ts: Time-step index (0 for static solves).
    :param gamma_prime: Time-integrator gamma_prime for damping scaling, or ``None`` to skip damping.
    :return: Updated system matrix.
    """
    for con in self.nodal_constraints:
        node = con.node_index
        hg_i = hg[node]
        dofs = node * 6 + jnp.arange(6)
        mat = mat.at[jnp.ix_(dofs, dofs)].add(con.k_tangent(hg_i, i_ts))
        if gamma_prime is not None:
            mat = mat.at[jnp.ix_(dofs, dofs)].add(
                gamma_prime * con.c_tangent(hg_i, i_ts)
            )

    # hard constraint tangent contributions (e.g. hinge spring-damper)
    for con in self.multibody_constraints:
        if con.has_f_res:
            dofs_i = con.node_i * 6 + jnp.arange(6)
            dofs_j = con.node_j * 6 + jnp.arange(6)
            dofs_ij = jnp.concatenate([dofs_i, dofs_j])
            k_12 = con.k_tangent(hg[con.node_i], hg[con.node_j])
            mat = mat.at[jnp.ix_(dofs_ij, dofs_ij)].add(k_12)
            if gamma_prime is not None:
                # add damping terms
                mat = mat.at[jnp.ix_(dofs_ij, dofs_ij)].add(
                    gamma_prime * con.c_tangent(hg[con.node_i], hg[con.node_j])
                )

    return mat

postprocess_constraints

postprocess_constraints(
    hg: Array,
) -> dict[str, dict[str, Array]]

Postprocess all constraints to extract derived quantities (e.g. hinge angles).

Parameters:

Name	Type	Description	Default
hg	Array	Nodal SE(3) frames, (n_nodes, 4, 4) or (n_tstep, n_nodes, 4, 4).	required

Returns:

Type	Description
dict[str, dict[str, Array]]	Nested dict {constraint_name: {quantity_name: Array}}.

Source code in src/flapjax/structure/beam.py

def postprocess_constraints(self, hg: Array) -> dict[str, dict[str, Array]]:
    r"""
    Postprocess all constraints to extract derived quantities (e.g. hinge angles).
    :param hg: Nodal SE(3) frames, ``(n_nodes, 4, 4)`` or ``(n_tstep, n_nodes, 4, 4)``.
    :return: Nested dict ``{constraint_name: {quantity_name: Array}}``.
    """
    data: dict[str, dict[str, Array]] = {}
    for name, con in self.constraints.items():
        pp = con.postprocess(hg)
        if pp:
            data[name] = pp
    return data

solve_constrained

solve_constrained(
    sys_mat_solve: Array,
    f_res_solve: Array,
    hg_eval: Array,
    solve_dofs: Array,
    hg_base: Array | None = None,
    phi: Array | None = None,
    v: Array | None = None,
    gamma_prime: float | Array | None = None,
) -> tuple[Array, Array, Array]

Solve the augmented system with Lagrange multipliers, supporting both holonomic and non-holonomic constraints.

Parameters:

Name	Type	Description	Default
sys_mat_solve	Array	System matrix at solve DOFs, (n_solve, n_solve).	required
f_res_solve	Array	Force residual at solve DOFs, (n_solve,).	required
hg_eval	Array	SE(3) frames for constraint evaluation, (n_nodes, 4, 4).	required
solve_dofs	Array	Free DOF indices, (n_solve,).	required
hg_base	Array | None	Base frames for Jacobian computation (defaults to hg_eval).	None
phi	Array | None	Accumulated configuration increment, (n_nodes, 6).	None
v	Array | None	Current nodal velocities for non-holonomic constraints, (n_nodes, 6).	None
gamma_prime	float | Array | None	Newmark parameter for non-holonomic constraints.	None

Returns:

Type	Description
tuple[Array, Array, Array]	(delta_phi, lagrange_multipliers, constraint_violation).

Source code in src/flapjax/structure/beam.py

def solve_constrained(
    self,
    sys_mat_solve: Array,
    f_res_solve: Array,
    hg_eval: Array,
    solve_dofs: Array,
    hg_base: Array | None = None,
    phi: Array | None = None,
    v: Array | None = None,
    gamma_prime: float | Array | None = None,
) -> tuple[Array, Array, Array]:
    r"""
    Solve the augmented system with Lagrange multipliers, supporting both holonomic and non-holonomic constraints.
    :param sys_mat_solve: System matrix at solve DOFs, ``(n_solve, n_solve)``.
    :param f_res_solve: Force residual at solve DOFs, ``(n_solve,)``.
    :param hg_eval: SE(3) frames for constraint evaluation, ``(n_nodes, 4, 4)``.
    :param solve_dofs: Free DOF indices, ``(n_solve,)``.
    :param hg_base: Base frames for Jacobian computation (defaults to ``hg_eval``).
    :param phi: Accumulated configuration increment, ``(n_nodes, 6)``.
    :param v: Current nodal velocities for non-holonomic constraints, ``(n_nodes, 6)``.
    :param gamma_prime: Newmark parameter for non-holonomic constraints.
    :return: ``(delta_phi, lagrange_multipliers, constraint_violation)``.
    """
    hg_jac = hg_base if hg_base is not None else hg_eval

    n_s = sys_mat_solve.shape[0]
    n_h = self.n_holonomic_constraints
    n_nh = self.n_nonholonomic_constraints
    n_c = n_h + n_nh

    aug = jnp.zeros((n_s + n_c, n_s + n_c))
    aug = aug.at[:n_s, :n_s].set(sys_mat_solve)

    rhs_parts: list[Array] = [f_res_solve]
    violation_parts: list[Array] = []

    if n_h > 0:
        viol_h = self._compute_holonomic_violation(hg_eval)
        jac_h = self._compute_holonomic_jacobian(hg_jac, solve_dofs, phi)
        aug = aug.at[:n_s, n_s : n_s + n_h].set(-jac_h.T)
        aug = aug.at[n_s : n_s + n_h, :n_s].set(jac_h)
        rhs_parts.append(-viol_h)
        violation_parts.append(viol_h)

    if n_nh > 0:
        assert v is not None
        vel_viol_nh = self._compute_nonholonomic_vel_violation(hg_eval, v)
        a_vel_solve = self._compute_nonholonomic_a_vel(hg_eval, v, solve_dofs)
        a_phi_solve = self._compute_nonholonomic_a_phi(hg_jac, v, solve_dofs, phi)
        aug = aug.at[:n_s, n_s + n_h : n_s + n_c].set(-a_vel_solve.T)
        aug = aug.at[n_s + n_h : n_s + n_c, :n_s].set(
            a_phi_solve + gamma_prime * a_vel_solve
        )
        rhs_parts.append(-vel_viol_nh)
        violation_parts.append(vel_viol_nh)

    rhs = jnp.concatenate(rhs_parts)
    sol = jnp.linalg.solve(aug, rhs)

    return sol[:n_s], sol[n_s:], jnp.concatenate(violation_parts)

compute_centre_of_mass

compute_centre_of_mass(hg: Array) -> Array

Compute the centre of mass for an arbitrary system.

Parameters:

Name	Type	Description	Default
hg	Array	Node SE(3) coordinates, (n_node, 4, 4) or (n_tstep, n_node, 4, 4).	required

Returns:

Type	Description
Array	Centre of mass, (3) or (n_tstep, 3).

Source code in src/flapjax/structure/beam.py

def compute_centre_of_mass(self, hg: Array) -> Array:
    r"""
    Compute the centre of mass for an arbitrary system.
    :param hg: Node SE(3) coordinates, ``(n_node, 4, 4)`` or ``(n_tstep, n_node, 4, 4)``.
    :return: Centre of mass, (3) or ``(n_tstep, 3)``.
    """

    def inner_func(hg_: Array) -> Array:
        d = self.make_d(hg=hg_)
        m = self.assemble_matrix_from_entries(self.make_m_t(d=d))  # (n_dof, n_dof)
        if self.use_lumped_mass:
            m = self.add_lumped_contributions_to_arr(
                arr=m, lumped_arr=self.m_lumped
            )
        m_lin = m[::6, ::6]
        return jnp.einsum("ij,jk->k", m_lin, hg_[:, :3, 3]) / m_lin.sum()  # (3, )

    if hg.ndim == 3:
        return inner_func(hg)  # single timestep, (3, ).
    elif hg.ndim == 4:
        return vmap(inner_func, 0, 0)(hg)  # multiple timesteps, (n_tstep, 3)
    else:
        raise ValueError("hg.ndim must be 3 or 4")

make_f_elem

make_f_elem(eps: Array) -> Array

Compute the forces within the elements as :math:\mathbf{f}_{elem} = \mathcal{K}_{cs} \epsilon.

Parameters:

Name	Type	Description	Default
eps	Array	Element strain vectors, (n_elem, 6).	required

Returns:

Type	Description
Array	Element forces, (n_elem, 6).

Source code in src/flapjax/structure/beam.py

def make_f_elem(self, eps: Array) -> Array:
    r"""
    Compute the forces within the elements as :math:`\mathbf{f}_{elem} = \mathcal{K}_{cs} \epsilon`.
    :param eps: Element strain vectors, ``(n_elem, 6)``.
    :return: Element forces, ``(n_elem, 6)``.
    """
    return jnp.einsum("ijk,ik->ij", self.k_cs[self.k_cs_index, ...], eps)

make_f_int

make_f_int(p_d: Array, eps: Array) -> Array

Assemble global internal force vector as a function of the element relative configuration vectors.

Parameters:

Name	Type	Description	Default
p_d	Array	P(d) operator, (n_elem, 6, 12).	required
eps	Array	Element strain vectors, (n_elem, 6).	required

Returns:

Type	Description
Array	Internal forces, (n_elem, 12).

Source code in src/flapjax/structure/beam.py

def make_f_int(self, p_d: Array, eps: Array) -> Array:
    r"""
    Assemble global internal force vector as a function of the element relative configuration vectors.
    :param p_d: P(d) operator, ``(n_elem, 6, 12)``.
    :param eps: Element strain vectors, ``(n_elem, 6)``.
    :return: Internal forces, ``(n_elem, 12)``.
    """

    return -jnp.einsum("ikj,ikl,il->ij", p_d, self.k_cs[self.k_cs_index, ...], eps)

make_f_dead_ext staticmethod

make_f_dead_ext(f_ext: Array, rmat: Array) -> Array

Compute the global external dead force vector.

Parameters:

Name	Type	Description	Default
f_ext	Array	External forces array of dead forces in global reference, (n_node, 6)	required
rmat	Array	Deformation rotation matrices, (n_node, 3, 3)	required

Returns:

Type	Description
Array	External forces, (n_node, 6)

Source code in src/flapjax/structure/beam.py

@staticmethod
def make_f_dead_ext(f_ext: Array, rmat: Array) -> Array:
    r"""
    Compute the global external dead force vector.
    :param f_ext: External forces array of dead forces in global reference, ``(n_node, 6)``
    :param rmat: Deformation rotation matrices, ``(n_node, 3, 3)``
    :return: External forces, ``(n_node, 6)``
    """

    return transform_nodal_vect(f_ext, jnp.swapaxes(rmat, -1, -2))

add_thrust_force

add_thrust_force(
    force: Array, thrust: dict[str, Array]
) -> Array

Add thrust acting at nodes onto full system forcing.

Parameters:

Name	Type	Description	Default
force	Array	Input forcing, (n_node, 6).	required
thrust	dict[str, Array]	Input thrust at the current step, {key: ()}.	required

Returns:

Type	Description
Array	Updated forcing, (n_node, 6).

Source code in src/flapjax/structure/beam.py

def add_thrust_force(self, force: Array, thrust: dict[str, Array]) -> Array:
    r"""
    Add thrust acting at nodes onto full system forcing.
    :param force: Input forcing, ``(n_node, 6)``.
    :param thrust: Input thrust at the current step, ``{key: ()}``.
    :return: Updated forcing, ``(n_node, 6)``.
    """

    for k, v in thrust.items():
        node = dict(self.thrust_nodes)[k]
        direction = jnp.array(dict(self.thrust_direction)[k])
        force = force.at[node, :3].add(v * direction)
    return force

make_eps

make_eps(d: Array) -> Array

Compute the element strain vectors as a function of the element relative configuration vectors. Formulation from Geometrically exact beam finite element formulated on the special Euclidean group SE(3), by Sonneville et al., 2013, Eq 64.

Parameters:

Name	Type	Description	Default
d	Array	Element relative configuration, (n_elem, 6)	required

Returns:

Type	Description
Array	Element strain vectors, (n_elem, 6)

Source code in src/flapjax/structure/beam.py

def make_eps(self, d: Array) -> Array:
    r"""
    Compute the element strain vectors as a function of the element relative configuration vectors. Formulation from
    Geometrically exact beam finite element formulated on the special Euclidean group SE(3), by Sonneville et al.,
    2013, Eq 64.
    :param d: Element relative configuration, ``(n_elem, 6)``
    :return: Element strain vectors, ``(n_elem, 6)``
    """

    return (d - self.d0) / self.l0[:, None]

make_p_d

make_p_d(d: Array) -> Array

Compute the P(d) operator as a function of the element relative configuration vectors.

Parameters:

Name	Type	Description	Default
d	Array	Relative configuration vectors, (n_elem, 6)	required

Returns:

Type	Description
Array	P(d) operator, (n_elem, 6, 12)

Source code in src/flapjax/structure/beam.py

def make_p_d(self, d: Array) -> Array:
    r"""
    Compute the P(d) operator as a function of the element relative configuration vectors.
    :param d: Relative configuration vectors, ``(n_elem, 6)``
    :return: P(d) operator, ``(n_elem, 6, 12)``
    """
    return vmap(p, (0, 0), 0)(d, self.ad_inv_o0)  # [n_elem, 6, 12]

make_d

make_d(hg: Array) -> Array

Compute the element relative configuration vectors from the nodal homogeneous transformation matrices

Parameters:

Name	Type	Description	Default
hg	Array	Nodal homogeneous transformation matrices, (n_nodes, 4, 4)	required

Returns:

Type	Description
Array	Element relative configuration vectors, (n_elem, 6)

Source code in src/flapjax/structure/beam.py

def make_d(self, hg: Array) -> Array:
    r"""
    Compute the element relative configuration vectors from the nodal homogeneous transformation matrices
    :param hg: Nodal homogeneous transformation matrices, ``(n_nodes, 4, 4)``
    :return: Element relative configuration vectors, ``(n_elem, 6)``
    """

    base_hg = jnp.zeros((self.n_elem, 4, 4))
    base_hg = base_hg.at[:, :3, :3].set(self.o0)
    base_hg = base_hg.at[:, 3, 3].set(1.0)

    haha0 = jnp.einsum(
        "ijk,ikl->ijl", hg[self.connectivity_arr[:, 0], :, :], base_hg
    )  # (n_elem, 4, 4)
    haha1 = jnp.einsum(
        "ijk,ikl->ijl", hg[self.connectivity_arr[:, 1], :, :], base_hg
    )  # (n_elem, 4, 4)

    return vmap(hg_to_d, (0, 0), 0)(haha0, haha1)  # (n_elem, 6)

make_hg_dot staticmethod

make_hg_dot(hg: Array, v: Array) -> Array

Obtain the time derivative of the nodal coordinates.

Parameters:

Name	Type	Description	Default
hg	Array	Node coordinates, (n_node, 4, 4).	required
v	Array	Node local velocities, (n_node, 6)	required

Returns:

Type	Description
Array	Coordinate time derivative, (n_node, 4, 4)

Source code in src/flapjax/structure/beam.py

@staticmethod
def make_hg_dot(hg: Array, v: Array) -> Array:
    r"""
    Obtain the time derivative of the nodal coordinates.
    :param hg: Node coordinates, ``(n_node, 4, 4)``.
    :param v: Node local velocities, ``(n_node, 6)``
    :return: Coordinate time derivative, ``(n_node, 4, 4)``
    """
    return jnp.einsum(
        "ijk,ikl->ijl", hg, vmap(ha_to_ha_tilde, 0, 0)(v)
    )  # (n_nodes, 4, 4)

resolve_forces

resolve_forces(
    hg: Array,
    dynamic: Literal[True],
    f_ext_follower: Array | None,
    f_ext_dead: Array | None,
    f_ext_aero: Array | None,
    thrust: dict[str, Array],
    v: Array,
    v_dot: Array,
    approx_gradients: bool = False,
) -> tuple[
    Array,
    Array,
    Array | None,
    Array | None,
    Array | None,
    Array,
    Array,
    Array,
    Array,
]

resolve_forces(
    hg: Array,
    dynamic: Literal[False],
    f_ext_follower: Array | None,
    f_ext_dead: Array | None,
    f_ext_aero: Array | None,
    thrust: dict[str, Array],
    v: None,
    v_dot: None,
    approx_gradients: bool = False,
) -> tuple[
    Array,
    Array,
    Array | None,
    Array | None,
    Array | None,
    Array,
    None,
    None,
    Array,
]

resolve_forces(
    hg: Array,
    dynamic: bool,
    f_ext_follower: Array | None,
    f_ext_dead: Array | None,
    f_ext_aero: Array | None,
    thrust: dict[str, Array],
    v: Array | None,
    v_dot: Array | None,
    approx_gradients: bool = False,
) -> tuple[
    Array,
    Array,
    Array | None,
    Array | None,
    Array | None,
    Array,
    Array | None,
    Array | None,
    Array,
]

Obtain all components of the force from a final solution.

Parameters:

Name	Type	Description	Default
hg	Array	Nodal homogeneous transformation matrices, (n_nodes, 4, 4).	required
dynamic	bool	Whether to compute dynamic forces.	required
f_ext_follower	Array | None	External follower forces in local reference, (n_node, 6).	required
f_ext_dead	Array | None	External dead forces in global reference, (n_node, 6).	required
f_ext_aero	Array | None	External aero forces in global reference, (n_node, 6).	required
thrust	dict[str, Array]	Thrust forces at current step, {keys, ()}.	required
v	Array | None	Nodal velocities in global frame, (n_node, 6).	required
v_dot	Array | None	Nodal accelerations in global frame, (n_node, 6).	required
approx_gradients	bool	Whether to stop computing gradients of the inertial and gyroscopic forces with respect to the node coordinates, as these are small but nonzero values in practice.	False

Returns:

Type	Description
tuple[Array, Array, Array | None, Array | None, Array | None, Array, Array | None, Array | None, Array]	Configuration vectors, strain vectors, Dead external forces, aero external forces, gravitational forces, internal forces, gyroscopic forces, inertial forces and residual forces.

Source code in src/flapjax/structure/beam.py

def resolve_forces(
    self,
    hg: Array,
    dynamic: bool,
    f_ext_follower: Array | None,
    f_ext_dead: Array | None,
    f_ext_aero: Array | None,
    thrust: dict[str, Array],
    v: Array | None,
    v_dot: Array | None,
    approx_gradients: bool = False,
) -> tuple[
    Array,
    Array,
    Array | None,
    Array | None,
    Array | None,
    Array,
    Array | None,
    Array | None,
    Array,
]:
    r"""
    Obtain all components of the force from a final solution.
    :param hg: Nodal homogeneous transformation matrices, ``(n_nodes, 4, 4)``.
    :param dynamic: Whether to compute dynamic forces.
    :param f_ext_follower: External follower forces in local reference, ``(n_node, 6)``.
    :param f_ext_dead: External dead forces in global reference, ``(n_node, 6)``.
    :param f_ext_aero: External aero forces in global reference, ``(n_node, 6)``.
    :param thrust: Thrust forces at current step, {keys, ``()``}.
    :param v: Nodal velocities in global frame, ``(n_node, 6)``.
    :param v_dot: Nodal accelerations in global frame, ``(n_node, 6)``.
    :param approx_gradients: Whether to stop computing gradients of the inertial and gyroscopic forces with respect to
    the node coordinates, as these are small but nonzero values in practice.
    :return: Configuration vectors, strain vectors, Dead external forces, aero external forces, gravitational forces, internal forces,
    gyroscopic forces, inertial forces and residual forces.
    """

    def prop_grad(x: Array) -> Array:
        return jax.lax.stop_gradient(x) if approx_gradients else x

    d = self.make_d(hg)
    eps = self.make_eps(d)
    p_d = self.make_p_d(d)

    if dynamic or self.use_gravity:
        m_t = self.make_m_t(prop_grad(d))
    else:
        m_t = None

    if dynamic:
        assert v is not None

        d_dot = self._make_d_dot(p_d, v)
        c_l = self._make_c_t(prop_grad(d), prop_grad(d_dot), v)[0]
        c_l_lumped = self._make_c_t_lumped(v)[0] if self.use_lumped_mass else None
    else:
        d_dot, c_l, c_l_lumped = None, None, None

    this_f_res = self.add_thrust_force(
        force=jnp.zeros((self.n_nodes, 6)), thrust=thrust
    )

    if f_ext_dead is not None:
        this_f_ext_dead = self.make_f_dead_ext(f_ext_dead, hg[:, :3, :3])
        this_f_res += this_f_ext_dead
    else:
        this_f_ext_dead = None

    if f_ext_aero is not None:
        this_f_ext_aero = self.make_f_dead_ext(f_ext_aero, hg[:, :3, :3])
        this_f_res += this_f_ext_aero
    else:
        this_f_ext_aero = None

    if self.use_gravity:
        assert m_t is not None
        this_f_grav = self.assemble_vector_from_entries(
            self._make_f_grav(m_t, hg[:, :3, :3])
        ).reshape(-1, 6)
        if self.use_lumped_mass:
            f_grav_lumped = self._make_f_grav_lumped(hg[:, :3, :3])
            this_f_grav = self.add_lumped_contributions_to_vec(
                vec=this_f_grav.ravel(), lumped_vec=f_grav_lumped.ravel()
            ).reshape(-1, 6)
        this_f_res += this_f_grav
    else:
        this_f_grav = None

    this_f_int = self.assemble_vector_from_entries(
        self.make_f_int(p_d, eps)
    ).reshape(-1, 6)
    this_f_res += this_f_int

    if dynamic:
        assert (
            m_t is not None
            and c_l is not None
            and v is not None
            and v_dot is not None
        )
        this_f_iner, this_f_gyr = self._make_f_iner_gyr(m_t, c_l, v, v_dot)
        this_f_iner = self.assemble_vector_from_entries(this_f_iner).reshape(-1, 6)
        this_f_gyr = self.assemble_vector_from_entries(this_f_gyr).reshape(-1, 6)

        if self.use_lumped_mass:
            assert c_l_lumped is not None
            f_iner_lumped, f_gyr_lumped = self._make_f_iner_gyr_lumped(
                c_l_lumped, v, v_dot
            )
            this_f_iner = self.add_lumped_contributions_to_vec(
                this_f_iner.ravel(), (f_iner_lumped + f_gyr_lumped).ravel()
            ).reshape(-1, 6)
        this_f_res += this_f_iner
    else:
        this_f_iner = None
        this_f_gyr = None

    if f_ext_follower is not None:
        this_f_res += f_ext_follower

    return (
        d,
        eps,
        this_f_ext_dead,
        this_f_ext_aero,
        this_f_grav,
        this_f_int,
        this_f_gyr,
        this_f_iner,
        this_f_res,
    )

make_f_res

make_f_res(
    solve_dofs: Array | None,
    p_d: Array,
    eps: Array,
    hg: Array,
    f_ext_follower_n: Array | None,
    f_ext_dead_n: Array | None,
    thrust_n: dict[str, Array],
    dynamic: Literal[True],
    m_t: Array,
    c_l: Array,
    c_l_lumped: Array | None,
    v: Array,
    v_dot: Array,
    i_ts: int = 0,
    k_t_assembled: Array | None = None,
) -> tuple[Array, Array]

make_f_res(
    solve_dofs: Array | None,
    p_d: Array,
    eps: Array,
    hg: Array,
    f_ext_follower_n: Array | None,
    f_ext_dead_n: Array | None,
    thrust_n: dict[str, Array],
    dynamic: Literal[False],
    m_t: Array | None,
    c_l: None,
    c_l_lumped: None,
    v: None,
    v_dot: None,
    i_ts: int = 0,
    k_t_assembled: Array | None = None,
) -> tuple[Array, Array]

make_f_res(
    solve_dofs: Array | None,
    p_d: Array,
    eps: Array,
    hg: Array,
    f_ext_follower_n: Array | None,
    f_ext_dead_n: Array | None,
    thrust_n: dict[str, Array],
    dynamic: bool,
    m_t,
    c_l,
    c_l_lumped,
    v,
    v_dot,
    i_ts: int = 0,
    k_t_assembled: Array | None = None,
) -> tuple[Array, Array]

Compute the residual force vector for a given configuration and external forces, used in the nonlinear solve. This is the force imbalance that the nonlinear solver will seek to drive to zero. Additionally, returns an "absolute sum" of all forces, used for relative convergence checks.

Parameters:

Name	Type	Description	Default
solve_dofs	Array | None	Optional array of degrees of freedom to solve for (n_solve_dofs, ).	required
p_d	Array	P(d) operator, (n_elem, 6, 12).	required
eps	Array	Element strain vectors, (n_elem, 6).	required
hg	Array	Nodal homogeneous transformation matrices, (n_nodes, 4, 4).	required
f_ext_follower_n	Array | None	Nodal follower forces, (n_nodes, 6).	required
f_ext_dead_n	Array | None	Nodal dead forces, (n_nodes, 6).	required
thrust_n	dict[str, Array]	Thrust magnitude, {key: ()}.	required
dynamic	bool	Flag for whether to compute dynamic entries.	required
m_t		Disassembled system mass matrix, (n_elem, 12, 12).	required
c_l		Dissembled system gyroscopic matrix, (n_elem, 12, 12).	required
c_l_lumped		Lumped gyroscopic matrix, (n_nodes, 6, 6).	required
v		Nodal velocities, (n_nodes, 6).	required
v_dot		Nodal accelerations, (n_node, 6).	required
i_ts	int	Time-step index (0 for static solves).	0
k_t_assembled	Array | None	Assembled global tangent stiffness matrix (n_dof, n_dof), required for Rayleigh damping.	None

Returns:

Type	Description
tuple[Array, Array]	Residual force vector, (n_dof, ), absolute sum of forces, (n_dof, ).

Source code in src/flapjax/structure/beam.py

def make_f_res(
    self,
    solve_dofs: Array | None,
    p_d: Array,
    eps: Array,
    hg: Array,
    f_ext_follower_n: Array | None,
    f_ext_dead_n: Array | None,
    thrust_n: dict[str, Array],
    dynamic: bool,
    m_t,
    c_l,
    c_l_lumped,
    v,
    v_dot,
    i_ts: int = 0,
    k_t_assembled: Array | None = None,
) -> tuple[Array, Array]:
    r"""
    Compute the residual force vector for a given configuration and external forces, used in the nonlinear solve.
    This is the force imbalance that the nonlinear solver will seek to drive to zero. Additionally, returns an
    "absolute sum" of all forces, used for relative convergence checks.
    :param solve_dofs: Optional array of degrees of freedom to solve for ``(n_solve_dofs, )``.
    :param p_d: P(d) operator, ``(n_elem, 6, 12)``.
    :param eps: Element strain vectors, ``(n_elem, 6)``.
    :param hg: Nodal homogeneous transformation matrices, ``(n_nodes, 4, 4)``.
    :param f_ext_follower_n: Nodal follower forces, ``(n_nodes, 6)``.
    :param f_ext_dead_n: Nodal dead forces, ``(n_nodes, 6)``.
    :param thrust_n: Thrust magnitude, ``{key: ()}``.
    :param dynamic: Flag for whether to compute dynamic entries.
    :param m_t: Disassembled system mass matrix, ``(n_elem, 12, 12)``.
    :param c_l: Dissembled system gyroscopic matrix, ``(n_elem, 12, 12)``.
    :param c_l_lumped: Lumped gyroscopic matrix, ``(n_nodes, 6, 6)``.
    :param v: Nodal velocities, ``(n_nodes, 6)``.
    :param v_dot: Nodal accelerations, ``(n_node, 6)``.
    :param i_ts: Time-step index (0 for static solves).
    :param k_t_assembled: Assembled global tangent stiffness matrix ``(n_dof, n_dof)``, required for Rayleigh
    damping.
    :return: Residual force vector, ``(n_dof, )``, absolute sum of forces, ``(n_dof, )``.
    """

    f_res = self.make_f_int(p_d, eps)  # (n_elem, 12)
    f_abs_sum = jnp.abs(f_res)

    if self.use_gravity:
        f_grav = self._make_f_grav(m_t, hg[:, :3, :3])
        f_res += f_grav
        f_abs_sum += jnp.abs(f_grav)

    if dynamic:
        f_iner, f_gyr = self._make_f_iner_gyr(m_t, c_l, v, v_dot)
        f_res += f_iner + f_gyr
        f_abs_sum += jnp.abs(f_iner + f_gyr)

    f_res_vect = self.assemble_vector_from_entries(f_res)
    f_abs_sum_vect = self.assemble_vector_from_entries(f_abs_sum)

    # add external forcing contributions
    if f_ext_follower_n is not None:
        f_res_vect += f_ext_follower_n.reshape(self.n_dof).ravel()
        f_abs_sum_vect += jnp.abs(f_ext_follower_n.reshape(self.n_dof).ravel())
    if f_ext_dead_n is not None:
        f_dead = self.make_f_dead_ext(f_ext_dead_n, hg[:, :3, :3]).ravel()
        f_res_vect += f_dead
        f_abs_sum_vect += jnp.abs(f_dead)

    f_thrust = self.add_thrust_force(
        force=jnp.zeros((self.n_nodes, 6)), thrust=thrust_n
    ).ravel()
    f_res_vect += f_thrust
    f_abs_sum_vect += jnp.abs(f_thrust)

    if self.use_lumped_mass:
        if dynamic:
            f_iner_lumped, f_gyr_lumped = self._make_f_iner_gyr_lumped(
                c_l_lumped, v, v_dot
            )
            f_iner_gyr_lumped = (f_iner_lumped + f_gyr_lumped).ravel()
            f_res_vect = self.add_lumped_contributions_to_vec(
                f_res_vect, f_iner_gyr_lumped
            )
            f_abs_sum_vect = self.add_lumped_contributions_to_vec(
                f_abs_sum_vect, jnp.abs(f_iner_gyr_lumped)
            )
        if self.use_gravity:
            f_grav_lumped = self._make_f_grav_lumped(hg[:, :3, :3]).ravel()
            f_res_vect = self.add_lumped_contributions_to_vec(
                vec=f_res_vect, lumped_vec=f_grav_lumped
            )
            f_abs_sum_vect = self.add_lumped_contributions_to_vec(
                vec=f_abs_sum_vect, lumped_vec=f_grav_lumped
            )

    # nodal constraint contributions
    for con in self.nodal_constraints:
        node = con.node_index
        v_node = v[node] if dynamic else jnp.zeros(6)
        f_constraint = con.f_res(hg[node], v_node, i_ts)
        dofs = node * 6 + jnp.arange(6)
        f_res_vect = f_res_vect.at[dofs].add(f_constraint)
        f_abs_sum_vect = f_abs_sum_vect.at[dofs].add(jnp.abs(f_constraint))

    # hard constraint force contributions (e.g. hinge spring-damper)
    for con in self.multibody_constraints:
        if con.has_f_res:
            v_i = v[con.node_i] if dynamic else jnp.zeros(6)
            v_j = v[con.node_j] if dynamic else jnp.zeros(6)
            f_i, f_j = con.f_res(hg[con.node_i], hg[con.node_j], v_i, v_j)
            dofs_i = con.node_i * 6 + jnp.arange(6)
            dofs_j = con.node_j * 6 + jnp.arange(6)
            f_res_vect = f_res_vect.at[dofs_i].add(f_i)
            f_res_vect = f_res_vect.at[dofs_j].add(f_j)
            f_abs_sum_vect = f_abs_sum_vect.at[dofs_i].add(jnp.abs(f_i))
            f_abs_sum_vect = f_abs_sum_vect.at[dofs_j].add(jnp.abs(f_j))

    # Rayleigh structural damping
    if dynamic and (self.alpha_m != 0.0 or self.beta_k != 0.0):
        if self.beta_k != 0.0 and k_t_assembled is None:
            raise ValueError(
                "k_t_assembled must be provided when beta_k != 0 for dynamic residual."
            )
        f_damp = self._make_f_rayleigh_damp(
            m_t=m_t,
            k_t_assembled=k_t_assembled
            if k_t_assembled is not None
            else jnp.zeros((self.n_dof, self.n_dof)),
            v=v,
        )
        f_res_vect += f_damp
        f_abs_sum_vect += jnp.abs(f_damp)

    if solve_dofs is not None:
        return f_res_vect[solve_dofs], f_abs_sum_vect[
            solve_dofs
        ]  # (n_solve_dof, ), (n_solve_dof, )
    else:
        return f_res_vect, f_abs_sum_vect  # (n_dof, ), (n_dof, )

update_hg staticmethod

update_hg(hg: Array, phi: Array) -> Array

Update the nodal homogeneous transformation matrices with the configuration increments.

Parameters:

Name	Type	Description	Default
hg	Array	Existing nodal homogeneous transformation matrices, (n_nodes, 4, 4)	required
phi	Array	Perturbation to the configuration vector, (n_nodes, 6)	required

Returns:

Type	Description
Array	Updated nodal homogeneous transformation matrices, (n_nodes, 4, 4)

Source code in src/flapjax/structure/beam.py

@staticmethod
def update_hg(hg: Array, phi: Array) -> Array:
    r"""
    Update the nodal homogeneous transformation matrices with the configuration increments.
    :param hg: Existing nodal homogeneous transformation matrices, ``(n_nodes, 4, 4)``
    :param phi: Perturbation to the configuration vector, ``(n_nodes, 6)``
    :return: Updated nodal homogeneous transformation matrices, ``(n_nodes, 4, 4)``
    """
    return jnp.einsum(
        "ijk,ikl->ijl",
        hg,
        vmap(exp_se3, 0, 0)(phi.reshape(-1, 6)),
    )

static_solve

static_solve(
    prescribed_dofs: Sequence[int] | Array | slice | int,
    f_ext_follower: Array | None = None,
    f_ext_dead: Array | None = None,
    f_ext_aero: Array | None = None,
    load_steps: int = 1,
    *,
    print_header: bool = True,
    postprocess_constraints: bool = True,
) -> StructureCase

Perform static solve of the structure under external loads.

Parameters:

Name	Type	Description	Default
f_ext_follower	Array | None	External forces array of follower forces (n_node, 6).	None
f_ext_dead	Array | None	External forces array of dead loads (n_node, 6).	None
f_ext_aero	Array | None	External forces array of aerodynamic loads (n_node, 6).	None
prescribed_dofs	Sequence[int] | Array | slice | int	Index of degrees of freedom which are prescribed (not solved for).	required
load_steps	int	Number of load steps to apply the external loads over.	1
print_header	bool	If False, suppress the "Static Solve" table header and trailing line.	True
postprocess_constraints	bool	If True, apply constraint postprocessing to the final solution.	True

Returns:

Type	Description
StructureCase	StructureCase object containing results of the static analysis.

Source code in src/flapjax/structure/beam.py

def static_solve(
    self,
    prescribed_dofs: Sequence[int] | Array | slice | int,
    f_ext_follower: Array | None = None,
    f_ext_dead: Array | None = None,
    f_ext_aero: Array | None = None,
    load_steps: int = 1,
    *,
    print_header: bool = True,
    postprocess_constraints: bool = True,
) -> StructureCase:
    r"""
    Perform static solve of the structure under external loads.
    :param f_ext_follower: External forces array of follower forces ``(n_node, 6)``.
    :param f_ext_dead: External forces array of dead loads ``(n_node, 6)``.
    :param f_ext_aero: External forces array of aerodynamic loads ``(n_node, 6)``.
    :param prescribed_dofs: Index of degrees of freedom which are prescribed (not solved for).
    :param load_steps: Number of load steps to apply the external loads over.
    :param print_header: If False, suppress the "Static Solve" table header and trailing line.
    :param postprocess_constraints: If True, apply constraint postprocessing to the final solution.
    :return: StructureCase object containing results of the static analysis.
    """

    if load_steps < 1:
        raise ValueError("load_steps must be at least 1")

    # check inputs
    if f_ext_follower is not None:
        check_arr_shape(f_ext_follower, (self.n_nodes, 6), "f_ext_follower")
    if f_ext_dead is not None:
        check_arr_shape(f_ext_dead, (self.n_nodes, 6), "f_ext_dead")

    if not (0.0 < self.relaxation_factor <= 1.0):
        raise ValueError("struct_relaxation_factor must be in the range (0, 1]")

    # degrees of freedom to solve for
    prescribed_dofs_: tuple[int, ...] = self.make_prescribed_dofs_tuple(
        prescribed_dofs
    )
    solve_dofs: Array = jnp.array(
        get_solve_dofs(n_dof=self.n_dof, prescribed_dofs=prescribed_dofs_)
    )

    # process external forces for load stepping
    load_step_weight: Array = jnp.linspace(0.0, 1.0, load_steps + 1)[
        1:
    ]  # (load_steps, )

    f_ext_follower_steps = self._make_load_steps_f(
        f_ext_follower, load_step_weight, apply_alpha_weighting=False
    )
    f_ext_dead_steps = self._make_load_steps_f(
        f_ext_dead, load_step_weight, apply_alpha_weighting=False
    )
    f_ext_aero_steps = self._make_load_steps_f(
        f_ext_aero, load_step_weight, apply_alpha_weighting=False
    )

    def _update(
        i_load_step: int,
        converge_status: ConvergenceStatus,
        hg_n: Array,
    ) -> tuple[int, ConvergenceStatus, Array]:
        # base parameters
        d_n = self.make_d(hg_n)  # (n_elem, 6)
        p_d_n = self.make_p_d(d_n)  # (n_elem, 6, 12)
        eps_n = self.make_eps(d_n)  # (n_elem, 6)
        m_t = self.make_m_t(d_n) if self.use_gravity else None  # (n_elem, 12, 12)

        # get total dead forces for this load step, (n_node, 6)
        total_f_ext_dead_step = self.make_f_ext_dead_tot(
            f_ext_dead_steps, f_ext_aero_steps, i_load_step
        )

        # assemble tangent stiffness matrix, (n_dof, n_dof)
        k_t_full_n = self.make_k_t_full(
            d=d_n,
            p_d=p_d_n,
            eps=eps_n,
            f_ext_dead=total_f_ext_dead_step,
            rmat=hg_n[:, :3, :3],
            m_t=m_t,
        )
        # apply nodal constraint contributions
        k_t_full_n = self.apply_nodal_constraint_tangent(
            mat=k_t_full_n, hg=hg_n, i_ts=0, gamma_prime=None
        )
        k_t_solve_n = k_t_full_n[jnp.ix_(solve_dofs, solve_dofs)]

        # compute residual forces, (n_solve_dofs, )
        f_res_solve_n, f_abs_sum_n = self.make_f_res(
            solve_dofs=solve_dofs,
            p_d=p_d_n,
            eps=eps_n,
            hg=hg_n,
            f_ext_follower_n=f_ext_follower_steps[i_load_step, ...]
            if f_ext_follower_steps is not None
            else None,
            f_ext_dead_n=total_f_ext_dead_step,
            thrust_n=self.thrust_reference,  # use reference thrust in static case
            dynamic=False,
            m_t=m_t,
            c_l=None,
            c_l_lumped=None,
            v=None,
            v_dot=None,
        )

        # solve for configuration increment, (n_solve_dofs, )
        if self.n_holonomic_constraints:
            d_varphi_np1, _, _ = self.solve_constrained(
                sys_mat_solve=k_t_solve_n,
                f_res_solve=f_res_solve_n,
                hg_eval=hg_n,
                solve_dofs=solve_dofs,
            )
            d_varphi_np1 *= self.relaxation_factor
        else:
            d_varphi_np1 = (
                jnp.linalg.solve(k_t_solve_n, f_res_solve_n)
                * self.relaxation_factor
            )

        # update configuration, (n_nodes, 4, 4)
        hg_np1_full = self.update_hg(
            hg_n, jnp.zeros(self.n_dof).at[solve_dofs].set(d_varphi_np1)
        )

        # algebra between undeformed and deformed shape, used to check relative convergence, (n_solve_dofs, )
        # this is relatively expensive to compute
        if self.struct_convergence_settings.rel_disp_tol is not None:
            h_full = vmap(hg_to_d, (0, 0), 0)(self.hg0, hg_np1_full).ravel()[
                solve_dofs
            ]
        else:
            h_full = None

        # update convergence status
        converge_status.update(
            delta_disp=d_varphi_np1,
            total_disp=h_full,
            delta_force=f_res_solve_n,
            total_force=f_abs_sum_n,
        )

        if map_verbosity_level(get_verbosity()) >= map_verbosity_level("verbose"):
            converge_status.print_struct_message(
                i_ts=None, t=None, i_load_step=i_load_step
            )

        return i_load_step, converge_status, hg_np1_full

    def convergence_loop(
        i_load_step: int,
        hg_init: Array,
    ) -> Array:
        r"""
        Convergence loop
        :param i_load_step: Index of load step.
        :param hg_init: Initial coordinates, ``(n_nodes, 4, 4)``.
        :return: Converged coordinates, ``(n_nodes, 4, 4)``.
        """
        _, convergence_status, hg_solve = eqxi.while_loop(
            lambda args_: ~args_[1].get_status(),
            lambda args_: _update(*args_),
            (
                i_load_step,
                ConvergenceStatus(
                    self.struct_convergence_settings,
                ),
                hg_init,
            ),
            max_steps=self.struct_convergence_settings.max_n_iter,
            kind="bounded",
        )

        if map_verbosity_level(get_verbosity()) >= map_verbosity_level("normal"):
            convergence_status.print_struct_message(
                i_ts=None, t=None, i_load_step=i_load_step
            )

        return hg_solve

    if print_header and map_verbosity_level(get_verbosity()) >= map_verbosity_level(
        "normal"
    ):
        ConvergenceStatus.print_header(dynamic=False)

    # solve for each load step
    hg = jax.lax.fori_loop(
        0,
        load_steps,
        lambda *args: convergence_loop(*args),
        self.hg0,
    )

    if print_header and map_verbosity_level(get_verbosity()) >= map_verbosity_level(
        "normal"
    ):
        ConvergenceStatus.print_line(dynamic=False)

    # postprocess final results
    d, eps, f_ext_dead_local, f_ext_aero_local, f_grav, f_int, _, _, f_res = (
        self.resolve_forces(
            hg=hg,
            dynamic=False,
            f_ext_dead=f_ext_dead,
            f_ext_follower=f_ext_follower,
            f_ext_aero=f_ext_aero,
            thrust=self.thrust_reference,
            v=None,
            v_dot=None,
        )
    )
    varphi = self.compute_varphi_from_hg(hg)
    f_elem = self.make_f_elem(eps=eps)  # compute loads in each element

    result = StructureCase(
        hg=hg,
        conn=self.connectivity,
        o0=self.o0,
        d=d,
        eps=eps,
        varphi=varphi,
        f_int=f_int,
        f_elem=f_elem,
        f_ext_follower=f_ext_follower,
        f_ext_dead=f_ext_dead_local,
        f_ext_aero=f_ext_aero_local,
        f_grav=f_grav,
        f_res=f_res,
        thrust=self.thrust_reference,
        thrust_nodes=self.thrust_nodes,
        thrust_direction=self.thrust_direction,
        prescribed_dofs=prescribed_dofs_,
        t=jnp.zeros(1),
    )
    if postprocess_constraints:
        result.constraint_data = self.postprocess_constraints(hg)
    return result

base_dynamic_solve

base_dynamic_solve(
    struct_case: StructureCase,
    struct_convergence_status: ConvergenceStatus,
    t: Array,
    solve_dofs: tuple[int, ...],
    load_steps: int,
    f_ext_dead: Array | None,
    f_ext_follower: Array | None,
    thrust_t: dict[str, Array],
    aero_obj: None,
    aero_case: None,
    fsi_convergence_status: None,
    cs_ang_t: None,
    cs_vel_t: None,
) -> StructureCase

base_dynamic_solve(
    struct_case: StructureCase,
    struct_convergence_status: ConvergenceStatus,
    t: Array,
    solve_dofs: tuple[int, ...],
    load_steps: int,
    f_ext_dead: Array | None,
    f_ext_follower: Array | None,
    thrust_t: dict[str, Array],
    aero_obj: DynamicAeroSolver,
    aero_case: AeroCase,
    fsi_convergence_status: ConvergenceStatus,
    cs_ang_t: dict[str, Array],
    cs_vel_t: dict[str, Array],
) -> AeroelasticCase

base_dynamic_solve(
    struct_case: StructureCase,
    struct_convergence_status: ConvergenceStatus,
    t: Array,
    solve_dofs: tuple[int, ...],
    load_steps: int,
    f_ext_dead: Array | None,
    f_ext_follower: Array | None,
    thrust_t: dict[str, Array],
    aero_obj: DynamicAeroSolver | None,
    aero_case: AeroCase | None,
    fsi_convergence_status: ConvergenceStatus | None,
    cs_ang_t: dict[str, Array] | None,
    cs_vel_t: dict[str, Array] | None,
) -> StructureCase | AeroelasticCase

Generic dynamic solver. Both the structural dynamic solve, and aeroelastic dynamic solve, are formed as wrappers of this

Source code in src/flapjax/structure/beam.py

def base_dynamic_solve(
    self,
    struct_case: StructureCase,
    struct_convergence_status: ConvergenceStatus,
    t: Array,
    solve_dofs: tuple[int, ...],
    load_steps: int,
    f_ext_dead: Array | None,
    f_ext_follower: Array | None,
    thrust_t: dict[str, Array],
    aero_obj: DynamicAeroSolver | None,
    aero_case: AeroCase | None,
    fsi_convergence_status: ConvergenceStatus | None,
    cs_ang_t: dict[str, Array] | None,
    cs_vel_t: dict[str, Array] | None,
) -> StructureCase | AeroelasticCase:
    r"""
    Generic dynamic solver. Both the structural dynamic solve, and aeroelastic dynamic solve, are formed as wrappers
    of this
    """

    if not (0.0 < self.relaxation_factor <= 1.0):
        raise ValueError("Relaxation factor must be in range (0, 1]")

    n_tstep = len(t)

    include_aero: bool = aero_obj is not None

    # process external forces for load stepping
    load_step_weight: Array = jnp.linspace(0.0, 1.0, load_steps + 1)[
        1:
    ]  # (load_steps, )
    f_ext_follower_alpha_steps = self._make_load_steps_f(
        f_ext_follower, load_step_weight, apply_alpha_weighting=True
    )
    f_ext_dead_alpha_steps = self._make_load_steps_f(
        f_ext_dead, load_step_weight, apply_alpha_weighting=True
    )

    solve_dofs_arr: Array = jnp.array(solve_dofs)
    prescribed_dofs_arr: Array = jnp.array(
        sorted(set(range(self.n_dof)) - set(solve_dofs)), dtype=int
    )

    def _update(
        i_load_step: int,
        i_ts: int,
        struct_convergence_status_: ConvergenceStatus,
        hg_n: Array,
        phi_alpha: Array,
        q_alpha: StructureMinimalStates,
        f_ext_aero_alpha_steps: Array | None,
        thrust_alpha: dict[str, Array],
    ) -> tuple[
        int,
        int,
        ConvergenceStatus,
        Array,
        Array,
        StructureMinimalStates,
        Array | None,
        dict[str, Array],
    ]:
        r"""
        Solution update for a single iteration of the nonlinear solver at a given time step and load step.
        :param i_load_step: Load step index.
        :param i_ts: Time step index.
        :param struct_convergence_status_: ConvergenceStatus object for the current iteration, used to track
        convergence and print messages.
        :param hg_n: Transformation matrices at iteration varphi, ``(n_nodes, 4, 4)``.
        :param phi_alpha: Timestep increment to the alpha step, ``(n_nodes, 6)``.
        :param f_ext_aero_alpha_steps: Load steps for the external aerodynamic forcing, ``(n_steps, n_nodes, 6)``.
        :param thrust_alpha: Thrust magnitude at the alpha step, ``{keys: ()}``.
        :return: Load and time step indices, updated ConvergenceStatus object, updated transformation matrices,
        configuration, velocities and accelerations for iteration n+1.
        """

        hg_update = self.update_hg(hg_n, phi_alpha)  # (n_node, 4, 4)

        # base parameters
        d_n = self.make_d(hg_update)  # (n_elem, 6)
        p_d_n = self.make_p_d(d_n)  # (n_elem, 6, 12)
        eps_n = self.make_eps(d_n)  # (n_elem, 6)
        d_dot_n = self._make_d_dot(p_d_n, q_alpha.v)  # (n_elem, 6)
        t_n = vmap(t_se3, 0, 0)(phi_alpha)  # (n_node, 6, 6)

        # tangent matrices
        m_t = self.make_m_t(d_n)  # (n_elem, 12, 12)
        c_l, c_t = self._make_c_t(
            d_n, d_dot_n, q_alpha.v
        )  # (n_elem, 12, 12), (n_elem, 12, 12)

        total_f_ext_dead = self.make_f_ext_dead_tot(
            f_ext_dead=f_ext_dead_alpha_steps[:, i_ts, :, :]
            if f_ext_dead_alpha_steps is not None
            else None,
            f_ext_aero=f_ext_aero_alpha_steps,
            i_load_step=i_load_step,
        )  # (n_node, 6)

        k_t = self.make_k_t_full(
            d_n,
            p_d_n,
            eps_n,
            total_f_ext_dead,
            hg_update[:, :3, :3],
            m_t,
        )  # (n_dof, n_dof)

        # add lumped mass contributions if applicable
        if self.use_lumped_mass:
            c_l_lumped, c_t_lumped = self._make_c_t_lumped(
                q_alpha.v
            )  # (n_node, 6, 6), (n_node, 6, 6)
        else:
            c_l_lumped, c_t_lumped = None, None

        # residual forces, (n_solve_dofs, )
        f_res_n_solve, f_abs_sum_n = self.make_f_res(
            solve_dofs=solve_dofs_arr,
            p_d=p_d_n,
            eps=eps_n,
            hg=hg_update,
            f_ext_follower_n=f_ext_follower_alpha_steps[i_load_step, i_ts, ...]
            if f_ext_follower_alpha_steps is not None
            else None,
            f_ext_dead_n=total_f_ext_dead,
            thrust_n=thrust_alpha,
            dynamic=True,
            m_t=m_t,
            c_l=c_l,
            c_l_lumped=c_l_lumped,
            v=q_alpha.v,
            v_dot=q_alpha.v_dot,
            i_ts=i_ts,
            k_t_assembled=k_t,
        )

        # system matrix, (n_dof, n_dof)
        sys_mat_full = self._make_sys_matrix(
            m_t=m_t,
            c_t=c_t,
            c_t_lumped=c_t_lumped,
            k_t=k_t,
            t_n=t_n,
            ti=self.time_integrator,
        )
        # add nodal constraint contributions
        sys_mat_full = self.apply_nodal_constraint_tangent(
            mat=sys_mat_full,
            hg=hg_update,
            i_ts=i_ts,
            gamma_prime=self.time_integrator.gamma_prime,
        )
        sys_mat = sys_mat_full[jnp.ix_(solve_dofs_arr, solve_dofs_arr)]

        # solve for configuration increment, (n_solve_dofs, )
        if self.multibody_constraints:
            d_n_np1, _, _ = self.solve_constrained(
                sys_mat_solve=sys_mat,
                f_res_solve=f_res_n_solve,
                hg_eval=hg_update,
                solve_dofs=solve_dofs_arr,
                hg_base=hg_n,
                phi=phi_alpha,
                v=q_alpha.v,
                gamma_prime=self.time_integrator.gamma_prime,
            )
            d_n_np1 *= self.relaxation_factor
        else:
            d_n_np1 = (
                jnp.linalg.solve(sys_mat, f_res_n_solve) * self.relaxation_factor
            )
        phi_np1 = phi_alpha.ravel().at[solve_dofs_arr].add(d_n_np1).reshape(-1, 6)

        # update configuration, velocities and accelerations
        v_np1 = (
            q_alpha.v.ravel()
            .at[solve_dofs_arr]
            .add(self.time_integrator.gamma_prime * d_n_np1)
            .reshape(-1, 6)
        )
        v_dot_np1 = (
            q_alpha.v_dot.ravel()
            .at[solve_dofs_arr]
            .add(self.time_integrator.beta_prime * d_n_np1)
            .reshape(-1, 6)
        )

        # update convergence status
        struct_convergence_status_.update(
            delta_disp=d_n_np1,
            total_disp=phi_np1,
            delta_force=f_res_n_solve,
            total_force=f_abs_sum_n,
        )

        if map_verbosity_level(get_verbosity()) >= map_verbosity_level("verbose"):
            struct_convergence_status_.print_struct_message(
                i_ts=i_ts, t=t[i_ts], i_load_step=i_load_step
            )

        q_alpha_update = StructureMinimalStates(
            varphi=None, v=v_np1, v_dot=v_dot_np1, a=q_alpha.a
        )

        return (
            i_load_step,
            i_ts,
            struct_convergence_status_,
            hg_n,
            phi_np1,
            q_alpha_update,
            f_ext_aero_alpha_steps,
            thrust_alpha,
        )

    @overload
    def time_step_loop(
        i_ts: int,
        struct_sol: StructureCase,
        struct_convergence_status_: ConvergenceStatus,
        aero_sol: None,
        fsi_convergence_status_: None,
        thrust_t_: dict[str, Array],
        cs_ang_t_: None,
        cs_vel_t_: None,
    ) -> tuple[
        StructureCase,
        ConvergenceStatus,
        None,
        None,
        dict[str, Array],
        None,
        None,
    ]: ...

    @overload
    def time_step_loop(
        i_ts: int,
        struct_sol: StructureCase,
        struct_convergence_status_: ConvergenceStatus,
        aero_sol: AeroCase,
        fsi_convergence_status_: ConvergenceStatus,
        thrust_t_: dict[str, Array],
        cs_ang_t_: dict[str, Array],
        cs_vel_t_: dict[str, Array],
    ) -> tuple[
        StructureCase,
        ConvergenceStatus,
        AeroCase,
        ConvergenceStatus,
        dict[str, Array],
        dict[str, Array],
        dict[str, Array],
    ]: ...

    def time_step_loop(
        i_ts: int,
        struct_sol: StructureCase,
        struct_convergence_status_: ConvergenceStatus,
        aero_sol: AeroCase | None,
        fsi_convergence_status_: ConvergenceStatus | None,
        thrust_t_: dict[str, Array],
        cs_ang_t_: dict[str, Array] | None,
        cs_vel_t_: dict[str, Array] | None,
    ) -> tuple[
        StructureCase,
        ConvergenceStatus,
        AeroCase | None,
        ConvergenceStatus | None,
        dict[str, Array],
        dict[str, Array] | None,
        dict[str, Array] | None,
    ]:
        r"""
        Performs analysis on a single time step, including load stepping
        :param i_ts: Index of time step to solve
        :param struct_sol: Solution object, with results up to time step i_ts-1.
        :param struct_convergence_status_: Convergence status object.
        :param aero_sol: Aero solution object, with results up to time step i_ts-1, if aero is included.
        :param fsi_convergence_status_: Convergence status object.
        :param thrust_t_: Thrust magnitude time history, ``{name: (n_tstep, )}``.
        :param cs_ang_t_: Control surface angle time history, ``{name: (n_tstep, )}``.
        :param cs_vel_t_: Control surface velocity time history, ``{name: (n_tstep, )}``.
        :return: Solution object with results up to time step i_ts.
        """

        # predictor step
        q_nm1 = struct_sol.get_minimal_states(i_ts - 1)
        phi_init, q_init = self.time_integrator.predict_q(q_nm1)
        phi_alpha_init, q_alpha_init = self.time_integrator.compute_q_alpha(
            q_nm1=q_nm1,
            q_n=q_init,
            phi_n=phi_init,
        )

        # prescribed DOFs should not be influenced by the time integration
        phi_alpha_init = (
            phi_alpha_init.ravel().at[prescribed_dofs_arr].set(0.0).reshape(-1, 6)
        )
        q_alpha_init.v = (
            q_alpha_init.v.ravel()
            .at[prescribed_dofs_arr]
            .set(q_nm1.v.ravel()[prescribed_dofs_arr])
            .reshape(-1, 6)
        )
        q_alpha_init.v_dot = (
            q_alpha_init.v_dot.ravel()
            .at[prescribed_dofs_arr]
            .set(q_nm1.v_dot.ravel()[prescribed_dofs_arr])
            .reshape(-1, 6)
        )
        q_alpha_init.a = (
            q_alpha_init.a.ravel()
            .at[prescribed_dofs_arr]
            .set(q_nm1.a.ravel()[prescribed_dofs_arr])
            .reshape(-1, 6)
        )

        q_alpha_init.varphi = None  # this value is not used during the loop

        # thrust force
        thrust_alpha: dict[str, Array] = {
            k: self.time_integrator.compute_f_alpha(f_nm1=v[i_ts - 1], f_n=v[i_ts])
            for k, v in thrust_t_.items()
        }
        thrust_n: dict[str, Array] = {k: v[i_ts] for k, v in thrust_t_.items()}

        if include_aero:
            assert (
                aero_sol is not None
                and fsi_convergence_status_ is not None
                and struct_sol.f_ext_aero is not None
                and cs_ang_t_ is not None
                and cs_vel_t_ is not None
            )

            fsi_convergence_status_.reset_status()

            # f_ext_aero is stored in local frame, so we convert back to global
            # so that both operands of the alpha blend are in the same (global) frame.
            f_aero_nm1 = jnp.concatenate(
                [
                    jnp.einsum(
                        "ijk,ik->ij",
                        struct_sol.hg[i_ts - 1, :, :3, :3],
                        struct_sol.f_ext_aero[i_ts - 1, :, :3],
                    ),
                    jnp.einsum(
                        "ijk,ik->ij",
                        struct_sol.hg[i_ts - 1, :, :3, :3],
                        struct_sol.f_ext_aero[i_ts - 1, :, 3:],
                    ),
                ],
                axis=-1,
            )

            # get control surface angles and velocities
            cs_ang_nm1 = {k: v[i_ts - 1] for k, v in cs_ang_t_.items()}
            cs_ang_n = {k: v[i_ts] for k, v in cs_ang_t_.items()}
            cs_vel_n = {k: v[i_ts] for k, v in cs_vel_t_.items()}
            assert fsi_convergence_status is not None
            (
                _,
                struct_sol,
                aero_sol,
                struct_convergence_status_,
                fsi_convergence_status_,
                phi_alpha,
                q_alpha,
                *_,
            ) = eqxi.while_loop(
                lambda args_: ~cast(ConvergenceStatus, args_[4]).get_status(),
                lambda args_: fsi_convergence_loop(*args_),
                (
                    i_ts,
                    struct_sol,
                    aero_sol,
                    struct_convergence_status_,
                    fsi_convergence_status_,
                    phi_alpha_init,
                    q_alpha_init,
                    f_aero_nm1,  # this value is for the previous timesteps force, and is propagated unaltered
                    f_aero_nm1,  # first guess for forcing at alpha is to use value from i_ts=n-1
                    thrust_alpha,
                    cs_ang_n,
                    cs_ang_nm1,
                    cs_vel_n,
                ),
                max_steps=fsi_convergence_status.convergence_settings.max_n_iter,
                kind="bounded",
            )

        else:
            # solve pure structural problem
            _, struct_convergence_status_, _, phi_alpha, q_alpha, *_ = (
                load_step_loop(
                    i_ts=i_ts,
                    struct_convergence_status_=struct_convergence_status_,
                    hg_alpha=struct_sol.hg[i_ts - 1, ...],
                    phi_alpha=phi_alpha_init,
                    q_alpha=q_alpha_init,
                    f_ext_aero_steps=None,
                    thrust_alpha=thrust_alpha,
                )
            )

        # print message where we only require one message per timestep
        if map_verbosity_level(get_verbosity()) == map_verbosity_level("normal"):
            struct_convergence_status_.print_struct_message(
                i_ts=i_ts, t=struct_sol.t[i_ts], i_load_step=load_steps - 1
            )
            if include_aero and fsi_convergence_status_ is not None:
                fsi_convergence_status_.print_fsi_message(
                    i_ts=i_ts, t=struct_sol.t[i_ts]
                )

        # postprocess results for time step and store in solution object
        q_n, phi_n = self.time_integrator.compute_q_n_from_q_alpha(
            q_alpha=q_alpha,
            q_nm1=struct_sol.get_minimal_states(i_ts - 1),
            phi_alpha=phi_alpha,
        )

        # update pseudo-acceleration
        q_n.a = self.time_integrator.compute_a_n(
            a_nm1=struct_sol.a[i_ts - 1, ...],
            v_dot_nm1=struct_sol.v_dot[i_ts - 1, ...],
            v_dot_n=q_n.v_dot,
        )

        # final node coordinates
        hg_n = self.update_hg(struct_sol.hg[i_ts - 1, ...], phi_n)

        if include_aero:
            if (
                aero_sol is None
                or aero_obj is None
                or fsi_convergence_status_ is None
            ):
                raise ValueError("Missing aero arguments")

            f_ext_aero = aero_sol.project_forcing_to_beam(
                i_ts=i_ts,
                rmat=hg_n[:, :3, :3],
                x0_aero=aero_obj.zeta_b0,
                include_unsteady=aero_obj.include_unsteady_force,
            )

        else:
            f_ext_aero = None

        (
            d,
            eps,
            f_ext_dead_local,
            f_ext_aero_local,
            f_grav,
            f_int,
            f_gyr,
            f_iner,
            f_res,
        ) = self.resolve_forces(
            hg=hg_n,
            dynamic=True,
            f_ext_dead=f_ext_dead[i_ts, ...] if f_ext_dead is not None else None,
            f_ext_follower=f_ext_follower[i_ts, ...]
            if f_ext_follower is not None
            else None,
            thrust=thrust_n,
            f_ext_aero=f_ext_aero,
            v=q_n.v,
            v_dot=q_n.v_dot,
        )
        struct_sol.d = struct_sol.d.at[i_ts, ...].set(d)
        struct_sol.eps = struct_sol.eps.at[i_ts, ...].set(eps)
        struct_sol.v = struct_sol.v.at[i_ts, ...].set(q_n.v)
        struct_sol.v_dot = struct_sol.v_dot.at[i_ts, ...].set(q_n.v_dot)
        struct_sol.a = struct_sol.a.at[i_ts, ...].set(q_n.a)
        struct_sol.hg = struct_sol.hg.at[i_ts, ...].set(hg_n)
        struct_sol.varphi = struct_sol.varphi.at[i_ts, ...].set(
            vmap(hg_to_d, (0, 0), 0)(self.hg0, hg_n)
        )

        if f_ext_follower is not None and struct_sol.f_ext_follower is not None:
            struct_sol.f_ext_follower = struct_sol.f_ext_follower.at[i_ts, ...].set(
                f_ext_follower[i_ts, ...]
            )
        if f_ext_dead is not None and struct_sol.f_ext_dead is not None:
            struct_sol.f_ext_dead = struct_sol.f_ext_dead.at[i_ts, ...].set(
                f_ext_dead_local
            )

        if f_ext_aero is not None and struct_sol.f_ext_aero is not None:
            struct_sol.f_ext_aero = struct_sol.f_ext_aero.at[i_ts, ...].set(
                f_ext_aero_local
            )

        if self.use_gravity:
            if struct_sol.f_grav is None:
                raise ValueError("struct_sol.f_grav is None")
            struct_sol.f_grav = struct_sol.f_grav.at[i_ts, ...].set(f_grav)
        struct_sol.f_int = struct_sol.f_int.at[i_ts, ...].set(f_int)
        struct_sol.f_elem = struct_sol.f_elem.at[i_ts, ...].set(
            self.make_f_elem(eps=eps)
        )
        struct_sol.f_iner_gyr = struct_sol.f_iner_gyr.at[i_ts, ...].set(
            f_iner + f_gyr
        )
        struct_sol.f_res = struct_sol.f_res.at[i_ts, ...].set(f_res)

        if include_aero and aero_sol is not None:
            assert cs_ang_t_ is not None and cs_vel_t_ is not None
            cs_ang_n = {k: v[i_ts] for k, v in cs_ang_t_.items()}
            cs_vel_n = {k: v[i_ts] for k, v in cs_vel_t_.items()}
            aero_sol.cs_ang = {
                k: v.at[i_ts].set(cs_ang_n[k]) for k, v in aero_sol.cs_ang.items()
            }
            aero_sol.cs_vel = {
                k: v.at[i_ts].set(cs_vel_n[k]) for k, v in aero_sol.cs_vel.items()
            }

        return (
            struct_sol,
            struct_convergence_status_,
            aero_sol,
            fsi_convergence_status_,
            thrust_t_,
            cs_ang_t_,
            cs_vel_t_,
        )

    def fsi_convergence_loop(
        i_ts: int,
        struct_sol: StructureCase,
        aero_sol: AeroCase,
        struct_convergence_status_: ConvergenceStatus,
        fsi_convergence_status_: ConvergenceStatus,
        phi_alpha_init: Array,
        q_alpha_init: StructureMinimalStates,
        f_aero_nm1: Array,
        f_aero_alpha_prev: Array,
        thrust_alpha: dict[str, Array],
        cs_ang_n: dict[str, Array],
        cs_ang_nm1: dict[str, Array],
        cs_vel_n: dict[str, Array],
    ) -> tuple[
        int,
        StructureCase,
        AeroCase,
        ConvergenceStatus,
        ConvergenceStatus,
        Array,
        StructureMinimalStates,
        Array,
        Array,
        dict[str, Array],
        dict[str, Array],
        dict[str, Array],
        dict[str, Array],
    ]:
        # obtain coordinates at timestep (not alpha)
        phi_n = self.time_integrator.compute_phi_from_phi_alpha(
            phi_alpha=phi_alpha_init
        )
        v_n = self.time_integrator.compute_v_from_v_alpha(
            v_alpha=q_alpha_init.v, v_nm1=struct_sol.v[i_ts - 1, ...]
        )

        hg_n = self.update_hg(hg=struct_sol.hg[i_ts - 1, ...], phi=phi_n)
        hg_dot = self.make_hg_dot(hg=hg_n, v=v_n)

        if aero_obj is None or struct_sol.f_ext_aero is None:
            raise ValueError("Missing aero parameters")

        # evaluate aerodynamic forcing on beam
        aero_sol = aero_obj.case_solve(
            case=aero_sol,
            i_ts=i_ts,
            hg_n=hg_n,
            hg_nm1=struct_sol.hg[i_ts - 1, ...],
            hg_dot_n=hg_dot,
            static=False,
            horseshoe=False,
            cs_ang_n=cs_ang_n,
            cs_ang_nm1=cs_ang_nm1,
            cs_vel_n=cs_vel_n,
        )

        f_aero_n = aero_sol.project_forcing_to_beam(
            i_ts=i_ts,
            rmat=hg_n[:, :3, :3],
            x0_aero=aero_obj.zeta_b0,
            include_unsteady=aero_obj.include_unsteady_force,
        )

        # aerodynamic force at alpha point, subsequently divided into load steps
        f_aero_alpha = self.time_integrator.compute_f_alpha(
            f_nm1=f_aero_nm1, f_n=f_aero_n
        )

        f_aero_alpha_steps = self._make_load_steps_f(
            f=f_aero_alpha, weighting=load_step_weight, apply_alpha_weighting=False
        )

        # reset convergence status
        struct_convergence_status_.reset_status()

        # solve structural problem for given aero load
        _, struct_convergence_status_, _, phi_alpha, q_alpha, *_ = load_step_loop(
            i_ts,
            struct_convergence_status_,
            struct_sol.hg[i_ts - 1, ...],
            phi_alpha_init,
            q_alpha_init,
            f_aero_alpha_steps,
            thrust_alpha,
        )

        # update the FSI convergence object
        # note that for convenience we use the alpha properties
        fsi_convergence_status_.update(
            delta_disp=(phi_alpha_init - phi_alpha).ravel()[solve_dofs_arr],
            total_disp=phi_alpha.ravel()[solve_dofs_arr],
            delta_force=(f_aero_alpha - f_aero_alpha_prev).ravel()[solve_dofs_arr],
            total_force=f_aero_alpha.ravel()[solve_dofs_arr],
        )

        if map_verbosity_level(get_verbosity()) >= map_verbosity_level("verbose"):
            fsi_convergence_status_.print_fsi_message(i_ts=i_ts, t=t[i_ts])

        return (
            i_ts,
            struct_sol,
            aero_sol,
            struct_convergence_status_,
            fsi_convergence_status_,
            phi_alpha,
            q_alpha,
            f_aero_nm1,
            f_aero_alpha,
            thrust_alpha,
            cs_ang_n,
            cs_ang_nm1,
            cs_vel_n,
        )

    def struct_convergence_loop(
        i_load_step: int,
        i_ts: int,
        struct_convergence_status_: ConvergenceStatus,
        hg_alpha: Array,
        phi_alpha: Array,
        q_alpha: StructureMinimalStates,
        f_ext_aero_steps: Array | None,
        thrust_alpha: dict[str, Array],
    ) -> tuple[
        int,
        ConvergenceStatus,
        Array,
        Array,
        StructureMinimalStates,
        Array | None,
        dict[str, Array],
    ]:
        r"""
        Convergence loop within each load step of a time step.
        :param i_load_step: Load step index.
        :param i_ts: Time step index.
        :param struct_convergence_status_: ConvergenceStatus object to update with convergence information during load
        stepping.
        :param hg_alpha: Node transformations at the beginning of the load step, ``(n_nodes, 4, 4)``.
        :param phi_alpha: Node configuration increments in algebra space, ``(n_nodes, 6)``.
        :param q_alpha: Minimal states at intermediate alpha step.
        :param f_ext_aero_steps: Optional aerodynamic forcing alpha load steps ``(n_steps, n_nodes, 6)``.
        :param thrust_alpha: Thrust at the alpha step, ``{key: ()}``.
        :return: Time step index, convergence status, and updated configuration, velocities, accelerations, and
        optional aerodynamic forcing.
        """

        struct_convergence_status_.reset_status()

        _, _, struct_convergence_status_, hg_solve, phi_alpha, q_alpha, _, _ = (
            eqxi.while_loop(
                lambda args_: ~args_[2].get_status(),
                lambda args_: _update(*args_),
                (
                    i_load_step,
                    i_ts,
                    struct_convergence_status_,
                    hg_alpha,
                    phi_alpha,
                    q_alpha,
                    f_ext_aero_steps,
                    thrust_alpha,
                ),
                max_steps=self.struct_convergence_settings.max_n_iter,
                kind="bounded",
            )
        )

        if map_verbosity_level(get_verbosity()) >= map_verbosity_level("verbose"):
            struct_convergence_status_.print_struct_message(
                i_ts=i_ts, t=t[i_ts], i_load_step=i_load_step
            )

        return (
            i_ts,
            struct_convergence_status_,
            hg_solve,
            phi_alpha,
            q_alpha,
            f_ext_aero_steps,
            thrust_alpha,
        )

    def load_step_loop(
        i_ts: int,
        struct_convergence_status_: ConvergenceStatus,
        hg_alpha: Array,
        phi_alpha: Array,
        q_alpha: StructureMinimalStates,
        f_ext_aero_steps: Array | None,
        thrust_alpha: dict[str, Array],
    ) -> tuple[
        int,
        ConvergenceStatus,
        Array,
        Array,
        StructureMinimalStates,
        Array | None,
    ]:
        r"""
        Performs load stepping iterations for a given time step. Load stepping is not performed for thrust.
        :param i_ts: Timestep index for which to perform load stepping.
        :param struct_convergence_status_: ConvergenceStatus object to update with load stepping convergence information.
        :param hg_alpha: SE(3) nodal transformation matrices at the beginning of the load step, ``(n_nodes, 4, 4)``.
        :param phi_alpha: Nodal updates to the configuration in the algebra space, ``(n_nodes, 6)``.
        :param q_alpha: Minimal states at intermediate alpha step.
        :param f_ext_aero_steps: Optional aerodynamic forcing alpha load steps ``(n_steps, n_nodes, 6)``.
        :param thrust_alpha: Thrust at the alpha step, ``{key: ()}``.
        :return: Time step index, updated ConvergenceStatus object, and updated configuration, velocities and accelerations after load stepping
        """
        return jax.lax.fori_loop(
            0,
            load_steps,
            lambda i_load_step, args: struct_convergence_loop(i_load_step, *args),
            (
                i_ts,
                struct_convergence_status_,
                hg_alpha,
                phi_alpha,
                q_alpha,
                f_ext_aero_steps,
                thrust_alpha,
            ),
        )

    def time_step_loop_checked(
        i_ts: int,
        struct_sol: StructureCase,
        struct_convergence_status_: ConvergenceStatus,
        aero_sol: AeroCase | None,
        fsi_convergence_status_: ConvergenceStatus | None,
        thrust_t_: dict[str, Array],
        cs_ang_t_: dict[str, Array] | None,
        cs_vel_t_: dict[str, Array] | None,
        diverged: Array,
    ) -> tuple[
        StructureCase,
        ConvergenceStatus,
        AeroCase | None,
        ConvergenceStatus | None,
        dict[str, Array],
        dict[str, Array] | None,
        dict[str, Array] | None,
        Array,
    ]:
        r"""
        Wraps ``time_step_loop`` with a check for solution divergence. Once a NaN is detected, this becomes a no-op
        for all remaining time steps. The corresponding time history entries are left at their initialised value
        (zero).
        """

        if include_aero:
            assert aero_sol is not None
            assert fsi_convergence_status_ is not None
            assert cs_ang_t_ is not None
            assert cs_vel_t_ is not None
            aero_sol_ok: AeroCase = aero_sol
            fsi_convergence_status_ok: ConvergenceStatus = fsi_convergence_status_
            cs_ang_t_ok: dict[str, Array] = cs_ang_t_
            cs_vel_t_ok: dict[str, Array] = cs_vel_t_
            false_branch = lambda: time_step_loop(
                i_ts,
                struct_sol,
                struct_convergence_status_,
                aero_sol_ok,
                fsi_convergence_status_ok,
                thrust_t_,
                cs_ang_t_ok,
                cs_vel_t_ok,
            )
        else:
            assert aero_sol is None
            assert fsi_convergence_status_ is None
            assert cs_ang_t_ is None
            assert cs_vel_t_ is None
            aero_sol_none: None = aero_sol
            fsi_convergence_status_none: None = fsi_convergence_status_
            cs_ang_t_none: None = cs_ang_t_
            cs_vel_t_none: None = cs_vel_t_
            false_branch = lambda: time_step_loop(
                i_ts,
                struct_sol,
                struct_convergence_status_,
                aero_sol_none,
                fsi_convergence_status_none,
                thrust_t_,
                cs_ang_t_none,
                cs_vel_t_none,
            )

        (
            struct_sol,
            struct_convergence_status_,
            aero_sol,
            fsi_convergence_status_,
            thrust_t_,
            cs_ang_t_,
            cs_vel_t_,
        ) = jax.lax.cond(
            diverged,
            lambda: (
                struct_sol,
                struct_convergence_status_,
                aero_sol,
                fsi_convergence_status_,
                thrust_t_,
                cs_ang_t_,
                cs_vel_t_,
            ),
            false_branch,
        )

        has_nan = struct_convergence_status_.has_nan
        if include_aero:
            assert fsi_convergence_status_ is not None
            has_nan = has_nan | fsi_convergence_status_.has_nan
        new_diverged = diverged | has_nan

        jax.lax.cond(
            new_diverged & ~diverged,
            lambda: warn(
                "NaN detected in dynamic solve at time step {i_ts} (t={t_val:.03e}) - skipping remaining time steps.",
                i_ts=i_ts,
                t_val=t[i_ts],
            ),
            lambda: None,
        )

        return (
            struct_sol,
            struct_convergence_status_,
            aero_sol,
            fsi_convergence_status_,
            thrust_t_,
            cs_ang_t_,
            cs_vel_t_,
            new_diverged,
        )

    struct_case, _, aero_case, *_ = jax.lax.fori_loop(
        1,
        n_tstep,
        lambda i_ts, args: time_step_loop_checked(i_ts, *args),
        (
            struct_case,
            struct_convergence_status,
            aero_case,
            fsi_convergence_status,
            thrust_t,
            cs_ang_t,
            cs_vel_t,
            jnp.zeros((), dtype=bool),
        ),
    )

    struct_case.constraint_data = self.postprocess_constraints(struct_case.hg)

    if include_aero:
        if aero_case is None:
            raise ValueError("aero_case cannot be None")

        from flapjax.coupled.data_structures import (
            AeroelasticCase,
        )  # import here to prevent circular references

        return AeroelasticCase(structure=struct_case, aero=aero_case)
    else:
        return struct_case

dynamic_solve

dynamic_solve(
    init_state: StructureCase | None,
    n_tstep: int,
    dt: Array | float,
    prescribed_dofs: Sequence[int]
    | Array
    | slice
    | int
    | None = None,
    f_ext_follower: Array | None = None,
    f_ext_dead: Array | None = None,
    f_ext_aero: Array | None = None,
    thrust_t: dict[str, Array] | None = None,
    load_steps: int = 1,
) -> StructureCase

Perform dynamic solve of the structure under external loads

Parameters:

Name	Type	Description	Default
init_state	StructureCase | None	Initial state of the structure, either static or a dynamic snapshot. If None, the reference configuration is used with zero velocities.	required
prescribed_dofs	Sequence[int] | Array | slice | int | None	Degrees of freedom which are prescribed (not solved for). If None, inherit from the initial state.	None
n_tstep	int	Number of time steps to simulate.	required
dt	Array | float	Time step length.	required
f_ext_follower	Array | None	Following external forces array, (n_tstep, n_node, 6), (n_node, 6) or None for zero external follower forces.	None
f_ext_dead	Array | None	Dead external forces array, (n_tstep, n_node, 6), (n_node, 6) or None for zero external dead forces.	None
f_ext_aero	Array | None	Aerodynamic external forces array, (n_tstep, n_node, 6), (n_node, 6) or None for zero external aerodynamic forces.	None
thrust_t	dict[str, Array] | None	Thrust time history, {key: (n_tstep, )}. If none, this will use the reference value.	None
load_steps	int	Number of load steps to apply the external loads over.	1

Returns:

Type	Description
StructureCase	Structure dataclass containing results of the dynamic analysis.

Source code in src/flapjax/structure/beam.py

def dynamic_solve(
    self,
    init_state: StructureCase | None,
    n_tstep: int,
    dt: Array | float,
    prescribed_dofs: Sequence[int] | Array | slice | int | None = None,
    f_ext_follower: Array | None = None,
    f_ext_dead: Array | None = None,
    f_ext_aero: Array | None = None,
    thrust_t: dict[str, Array] | None = None,
    load_steps: int = 1,
) -> StructureCase:
    r"""
    Perform dynamic solve of the structure under external loads
    :param init_state: Initial state of the structure, either static or a
    dynamic snapshot. If None, the reference configuration is used with zero
    velocities.
    :param prescribed_dofs: Degrees of freedom which are prescribed (not solved for). If None, inherit
    from the initial state.
    :param n_tstep: Number of time steps to simulate.
    :param dt: Time step length.
    :param f_ext_follower: Following external forces array, ``(n_tstep, n_node, 6)``, ``(n_node, 6)`` or None for zero external follower forces.
    :param f_ext_dead: Dead external forces array, ``(n_tstep, n_node, 6)``, ``(n_node, 6)`` or None for zero external dead forces.
    :param f_ext_aero: Aerodynamic external forces array, ``(n_tstep, n_node, 6)``, ``(n_node, 6)`` or None for zero external aerodynamic forces.
    :param thrust_t: Thrust time history, ``{key: (n_tstep, )}``. If none, this will use the reference value.
    :param load_steps: Number of load steps to apply the external loads over.
    :return: Structure dataclass containing results of the dynamic analysis.
    """

    if load_steps <= 0:
        raise ValueError("load_steps must be a positive integer")

    # set thrust if not provided
    thrust_t_: dict[str, Array] = (
        thrust_t
        if thrust_t is not None
        else {k: jnp.full(n_tstep, v) for k, v in self.thrust_reference.items()}
    )

    if prescribed_dofs is None:
        # inherit prescribed dofs from initial state
        if init_state is None:
            raise ValueError("prescribed_dofs cannot be None if init_state is None")
        prescribed_dofs = init_state.prescribed_dofs

    # degrees of freedom to solve for
    prescribed_dofs_arr = self.make_prescribed_dofs_tuple(prescribed_dofs)
    solve_dofs = get_solve_dofs(
        n_dof=self.n_dof, prescribed_dofs=prescribed_dofs_arr
    )

    # check and process external forces
    def check_force(arr: Array | None, name: str) -> Array | None:
        if arr is None:
            return None
        match arr.ndim:
            case 2:
                out_ = jnp.broadcast_to(arr[None, ...], (n_tstep, self.n_nodes, 6))
            case 3:
                out_ = arr
            case _:
                raise ValueError(
                    f"{name} must have shape [n_node, 6] or [n_tstep, n_node, 6]"
                )
        check_arr_shape(out_, (n_tstep, self.n_nodes, 6), name)
        return out_

    f_ext_dead = check_force(f_ext_dead, "f_ext_dead")  # (n_tstep, n_node, 6)
    f_ext_follower = check_force(
        f_ext_follower, "f_ext_follower"
    )  # (n_tstep, n_node, 6)
    f_ext_aero = check_force(f_ext_aero, "f_ext_aero")  # (n_tstep, n_node, 6)

    # time integration parameters
    self.time_integrator = TimeIntegrator(
        spectral_radius=self.spectral_radius, dt=jnp.array(dt)
    )

    def evaluate_initial_equilibrium(
        init_state__: StructureCase,
    ) -> StructureCase:
        r"""
        Evaluates the forces for a given initial state to check whether it is in equilibrium. If not, a warning is
        raised with the maximum residual force. This is important to ensure that the time integration starts from a
        consistent state.
        :param init_state__: Structure containing the initial state to evaluate.
        :return: Structure with the forces evaluated for the initial state.
        """
        d, eps, f_ext_dead_, f_ext_aero_, f_grav, f_int, f_gyr, f_iner, f_res = (
            self.resolve_forces(
                hg=init_state__.hg,
                dynamic=True,
                f_ext_dead=init_state__.f_ext_dead,
                f_ext_aero=init_state__.f_ext_aero,
                thrust=init_state__.thrust,
                f_ext_follower=init_state__.f_ext_follower,
                v=init_state__.v,
                v_dot=init_state__.v_dot,
            )
        )

        max_res = jnp.max(jnp.abs(f_res))
        jax_print(
            "Initial state maximum residual force: {max_res:.3e}",
            max_res=max_res,
            verbose_level="normal",
        )

        f_elem = self.make_f_elem(eps=eps)

        return StructureCase(
            hg=init_state__.hg,
            conn=self.connectivity,
            o0=self.o0,
            d=d,
            eps=eps,
            varphi=init_state__.varphi,
            v=init_state__.v,
            v_dot=init_state__.v_dot,
            a=init_state__.v_dot,  # initial pseudo-acceleration set equal to initial acceleration
            f_ext_follower=init_state__.f_ext_follower,
            f_ext_dead=f_ext_dead_,
            f_ext_aero=f_ext_aero_,
            f_grav=f_grav,
            f_int=f_int,
            f_elem=f_elem,
            f_iner_gyr=f_iner + f_gyr,  # type: ignore
            f_res=f_res,
            thrust=init_state__.thrust,
            thrust_nodes=self.thrust_nodes,
            thrust_direction=self.thrust_direction,
            t=init_state__.t,
            i_ts=init_state__.i_ts,
            prescribed_dofs=prescribed_dofs_arr,
        )

    # time steps
    t = jnp.arange(n_tstep) * dt
    if init_state is not None and init_state.is_dynamic:
        if init_state.is_batched:
            t += init_state.t[0]
        else:
            t += init_state.t

    # set up initial state
    if init_state is None:
        init_state_: StructureCase = self.reference_configuration(
            use_f_aero=f_ext_aero is not None,
            use_f_ext_dead=f_ext_dead is not None,
            use_f_ext_follower=f_ext_follower is not None,
            prescribed_dofs=tuple(prescribed_dofs_arr),
        ).to_dynamic(t=None)
    elif not init_state.is_dynamic:
        init_state_ = init_state.to_dynamic(t=None)
    elif not init_state.is_batched:
        init_state_ = init_state
    else:
        raise TypeError(
            "dynamic_solve init_state cannot be a batched Structure; pass a "
            "snapshot or static state"
        )

    # check if initial state satisfies equilibrium
    init_state_eval = evaluate_initial_equilibrium(init_state_)
    dynamic_struct = StructureCase.initialise(
        initial_snapshot=init_state_eval,
        t=t,
        use_f_ext_follower=f_ext_follower is not None,
        use_f_ext_dead=f_ext_dead is not None,
        use_f_ext_aero=False,
    )
    converge_status = ConvergenceStatus(
        convergence_settings=self.struct_convergence_settings
    )

    ConvergenceStatus.print_header(dynamic=True)

    out = self.base_dynamic_solve(
        struct_case=dynamic_struct,
        struct_convergence_status=converge_status,
        t=t,
        solve_dofs=solve_dofs,
        load_steps=load_steps,
        f_ext_dead=f_ext_dead,
        f_ext_follower=f_ext_follower,
        aero_obj=None,
        aero_case=None,
        fsi_convergence_status=None,
        thrust_t=thrust_t_,
        cs_ang_t=None,
        cs_vel_t=None,
    )

    ConvergenceStatus.print_line(dynamic=True)
    return out

case_from_dv

case_from_dv(dv: StructureDesignVariables) -> BeamStructure

Obtain a structural object as a function of design variables, allowing it to have defined gradients w.r.t. design variables.

Parameters:

Name	Type	Description	Default
dv	StructureDesignVariables	Design variables.	required

Returns:

Type	Description
BeamStructure	Beam structure object with the same functionality as self.

Source code in src/flapjax/structure/gradients/beam.py

def case_from_dv(self, dv: StructureDesignVariables) -> BeamStructure:
    r"""
    Obtain a structural object as a function of design variables, allowing it to have defined gradients w.r.t. design variables.
    :param dv: Design variables.
    :return: Beam structure object with the same functionality as self.
    """
    inner_case = pytree_clone(self)
    inner_case.set_design_variables(
        coords=dv_or(dv.x0, self.x0),
        k_cs=dv_or(dv.k_cs, self.k_cs),
        m_cs=dv_or(dv.m_cs, self.m_cs),
        m_lumped=dv_or(dv.m_lumped, self._m_lumped),
        remove_checks=True,
    )

    return inner_case

minimal_states_to_full_states

minimal_states_to_full_states(
    i_ts: int,
    q: StructureMinimalStates,
    dv: StructureDesignVariables,
    dv_full: StructureDesignVariables,
) -> StructureFullStates

Obtain the full set of states from the minimal states and the design variables.

Parameters:

Name	Type	Description	Default
i_ts	int	Index of the time step.	required
q	StructureMinimalStates	Minimal dynamic structure states.	required
dv	StructureDesignVariables	Design variables, where entries for gradients which aren't needed are set to None.	required
dv_full	StructureDesignVariables	Design variables, without omissions. These values are fallen back to when an entry in dv is none, to give an equivalent with zero gradient.	required

Returns:

Type	Description
StructureFullStates	Full set of structural states used inside objective function.

Source code in src/flapjax/structure/gradients/beam.py

def minimal_states_to_full_states(
    self,
    i_ts: int,
    q: StructureMinimalStates,
    dv: StructureDesignVariables,
    dv_full: StructureDesignVariables,
) -> StructureFullStates:
    r"""
    Obtain the full set of states from the minimal states and the design variables.
    :param i_ts: Index of the time step.
    :param q: Minimal dynamic structure states.
    :param dv: Design variables, where entries for gradients which aren't needed are set to None.
    :param dv_full: Design variables, without omissions. These values are fallen back to when an entry in ``dv`` is
    none, to give an equivalent with zero gradient.
    :return: Full set of structural states used inside objective function.
    """
    struct = self.case_from_dv(dv)
    hg = struct.compute_hg_from_varphi(q.varphi)
    d = struct.make_d(hg=hg)
    p_d = struct.make_p_d(d=d)
    eps = struct.make_eps(d=d)
    f_elem = struct.make_f_elem(eps=eps)

    f_ext_dead_i = (
        dv.f_ext_dead[i_ts, ...]
        if dv.f_ext_dead is not None
        else dv_full.f_ext_dead[i_ts, ...]
        if dv_full.f_ext_dead is not None
        else None
    )
    m_t = struct.make_m_t(d=d)

    # k_t_assembled is only required when Rayleigh damping is active. Skip the extra assembly when unused.
    k_t_assembled = (
        struct.make_k_t_full(
            d=d,
            p_d=p_d,
            eps=eps,
            f_ext_dead=f_ext_dead_i,
            rmat=hg[:, :3, :3],
            m_t=m_t,
        )
        if struct.beta_k != 0.0
        else None
    )

    assert dv_full.thrust_t is not None
    f_res = struct.make_f_res(
        solve_dofs=None,
        p_d=p_d,
        eps=eps,
        hg=hg,
        f_ext_follower_n=dv.f_ext_follower[i_ts, ...]
        if dv.f_ext_follower is not None
        else dv_full.f_ext_follower[i_ts, ...]
        if dv_full.f_ext_follower is not None
        else None,
        f_ext_dead_n=f_ext_dead_i,
        thrust_n={k: (v[i_ts] if v.ndim > 0 else v) for k, v in dv.thrust_t.items()}
        if dv.thrust_t is not None
        else {
            k: (v[i_ts] if v.ndim > 0 else v) for k, v in dv_full.thrust_t.items()
        },
        dynamic=True,
        m_t=m_t,
        c_l=self._make_c_t(d=d, d_dot=self._make_d_dot(p_d=p_d, v=q.v), v=q.v)[0],
        c_l_lumped=self._make_c_t_lumped(v=q.v)[0]
        if self.use_lumped_mass
        else None,
        v=q.v,
        v_dot=q.v_dot,
        i_ts=i_ts,
        k_t_assembled=k_t_assembled,
    )[0]
    return StructureFullStates(
        v=q.v,
        v_dot=q.v_dot,
        eps=eps,
        varphi=q.varphi,
        hg=hg,
        f_elem=f_elem,
        f_res=f_res,
    )

static_adjoint

static_adjoint(
    structure: StructureCase,
    objective: StructureObjectiveFunction,
    optional_jacobians: OptionalJacobians
    | None = OPTIONAL_JACOBIANS_DEFAULT,
    ad_mode: ADMode = "reverse",
) -> tuple[StructureDesignVariables, Array]

Computes the static grads of the structure, which is used to compute gradients of the loss with respect to the structure's parameters.

Parameters:

Name	Type	Description	Default
structure	StructureCase	StructureCase containing the current state of the structure.	required
objective	StructureObjectiveFunction	Objective function that takes the structure and design variables and returns an array	required
optional_jacobians	OptionalJacobians | None	OptionalJacobians object specifying which Jacobians to compute.	OPTIONAL_JACOBIANS_DEFAULT
ad_mode	ADMode	Flag on which to use of the forward or reverse adjoint.	'reverse'

Returns:

Type	Description
tuple[StructureDesignVariables, Array]	Gradient of objective function output with respect to design variables, and adjoint states.

Source code in src/flapjax/structure/gradients/beam.py

def static_adjoint(
    self,
    structure: StructureCase,
    objective: StructureObjectiveFunction,
    optional_jacobians: OptionalJacobians | None = OPTIONAL_JACOBIANS_DEFAULT,
    ad_mode: ADMode = "reverse",
) -> tuple[StructureDesignVariables, Array]:
    r"""
    Computes the static grads of the structure, which is used to compute gradients of the loss with respect to
    the structure's parameters.
    :param structure: StructureCase containing the current state of the structure.
    :param objective: Objective function that takes the structure and design variables and returns an array
    :param optional_jacobians: OptionalJacobians object specifying which Jacobians to compute.
    :param ad_mode: Flag on which to use of the forward or reverse adjoint.
    :return: Gradient of objective function output with respect to design variables, and adjoint states.
    """

    solve_dofs = jnp.array(
        get_solve_dofs(n_dof=self.n_dof, prescribed_dofs=structure.prescribed_dofs)
    )

    if optional_jacobians is not None:
        self.optional_jacobians = optional_jacobians

    # Recover original global dead force: structure.f_ext_dead is stored in local frame as
    # f_local = R^T @ f_global, so f_global = R @ f_local
    rmat = structure.hg[:, :3, :3]
    f_ext_dead_global = (
        transform_nodal_vect(structure.f_ext_dead, rmat)
        if structure.f_ext_dead is not None
        else None
    )

    # make design variables for current state of structure
    dv = StructureDesignVariables(
        x0=self.x0,
        orientation_euler=self.orientation_euler,
        k_cs=self.k_cs,
        m_cs=self._m_cs,
        m_lumped=self.m_lumped if self.use_lumped_mass else None,
        f_ext_follower=structure.f_ext_follower,
        f_ext_dead=f_ext_dead_global,
        thrust_t=structure.thrust,
        f_shape=(),
    )

    struct_states = structure.get_full_states()

    # find shape of objective function output without evaluating function
    f_properties = jax.eval_shape(lambda: objective(struct_states, dv, None))
    f_shape = f_properties.shape
    n_f = f_properties.size
    n_x = dv.n_x
    n_u = len(solve_dofs)
    n_u_full = self.n_dof

    # gradient of objective w.r.t. minimal states
    p_f_p_n, p_f_p_x = jax.jacrev(
        lambda varphi_, dv_: objective(
            self._structural_states_res_from_dv_varphi(
                dv=dv_, varphi=varphi_, thrust=structure.thrust
            ),
            dv_,
            None,
        ),
        argnums=(0, 1),
        allow_int=True,
    )(structure.varphi, dv)

    p_f_p_n = p_f_p_n.reshape(n_f, n_u_full)[:, solve_dofs]  # (n_f, n_u)
    p_f_p_x = p_f_p_x.ravel_jacobian(n_f, n_x)  # (n_f, n_x)

    # gradient of residual w.r.t. design variables and minimal states
    p_res_p_x, p_res_p_varphi = (jax.jacfwd if n_u > n_x else jax.jacrev)(
        lambda dv_, varphi_: (
            self._structural_states_res_from_dv_varphi(
                dv=dv_, varphi=varphi_, thrust=structure.thrust
            ).f_res
        ),
        argnums=(0, 1),
        allow_int=True,
    )(dv, structure.varphi)

    p_res_p_x = p_res_p_x.ravel_jacobian(n_u_full, n_x)[solve_dofs, :]  # (n_u, n_x)
    p_res_p_varphi = p_res_p_varphi.reshape(n_u_full, n_u_full)[
        jnp.ix_(solve_dofs, solve_dofs)
    ]  # (n_u, n_u)

    if ad_mode == "forward":
        # forward mode
        adj = jnp.linalg.solve(p_res_p_varphi, p_res_p_x)  # (n_u, n_x)
        rhs = p_f_p_n @ adj  # (n_f, n_x)
    elif ad_mode == "reverse":
        # reverse mode
        adj = jnp.linalg.solve(p_res_p_varphi.T, p_f_p_n.T).T  # (n_f, n_u)
        rhs = adj @ p_res_p_x  # (n_f, n_x)
    else:
        raise ValueError("AD mode must be either 'forward' or 'reverse'")

    return StructureDesignVariables(
        **dv.from_adjoint(f_shape, p_f_p_x - rhs), f_shape=f_shape
    ), adj

timestep_residual

timestep_residual(
    i_ts: int | Array,
    q_nm1: StructureMinimalStates,
    q_n: StructureMinimalStates,
    dv_: StructureDesignVariables,
    thrust_t: dict[str, Array],
    solve_dofs: tuple[int, ...],
    approx_grads: bool,
) -> Array

Routine to compute the full residual for the structural dynamic problem.

Parameters:

Name	Type	Description	Default
i_ts	int | Array	Time step index.	required
q_nm1	StructureMinimalStates	Previous minimal state.	required
q_n	StructureMinimalStates	Current minimal state.	required
dv_	StructureDesignVariables	Design variables.	required
thrust_t	dict[str, Array]	Thrust time history, {key, (n_tstep, )}.	required
solve_dofs	tuple[int, ...]	Solve degrees of freedom.	required
approx_grads	bool	If true, block gradients from some parts of the solution.	required

Returns:

Type	Description
Array	Residual vector, (4 * n_solve_dof,).

Source code in src/flapjax/structure/gradients/beam.py

def timestep_residual(
    self,
    i_ts: int | Array,
    q_nm1: StructureMinimalStates,
    q_n: StructureMinimalStates,
    dv_: StructureDesignVariables,
    thrust_t: dict[str, Array],
    solve_dofs: tuple[int, ...],
    approx_grads: bool,
) -> Array:
    r"""
    Routine to compute the full residual for the structural dynamic problem.
    :param i_ts: Time step index.
    :param q_nm1: Previous minimal state.
    :param q_n: Current minimal state.
    :param dv_: Design variables.
    :param thrust_t: Thrust time history, ``{key, (n_tstep, )}``.
    :param solve_dofs: Solve degrees of freedom.
    :param approx_grads: If true, block gradients from some parts of the solution.
    :return: Residual vector, ``(4 * n_solve_dof,)``.
    """
    return jnp.stack(
        (
            self.varphi_res_func(
                varphi_nm1=q_nm1.varphi,
                varphi_n=q_n.varphi,
                v_nm1=q_nm1.v,
                a_nm1=q_nm1.a,
                a_n=q_n.a,
                solve_dofs=solve_dofs,
            ),
            self.v_res_func(
                v_nm1=q_nm1.v,
                v_n=q_n.v,
                a_nm1=q_nm1.a,
                a_n=q_n.a,
                solve_dofs=solve_dofs,
            ),
            self.v_dot_res_func(
                i_ts=i_ts,
                varphi_nm1=q_nm1.varphi,
                varphi_n=q_n.varphi,
                v_nm1=q_nm1.v,
                v_n=q_n.v,
                v_dot_nm1=q_nm1.v_dot,
                v_dot_n=q_n.v_dot,
                approx_grads=approx_grads,
                f_aero_nm1=q_nm1.f_ext_aero,
                f_aero_n=q_n.f_ext_aero,
                thrust_t=thrust_t,
                dv=dv_,
                solve_dofs=solve_dofs,
            ),
            self.a_res_func(
                v_dot_nm1=q_nm1.v_dot,
                v_dot_n=q_n.v_dot,
                a_nm1=q_nm1.a,
                a_n=q_n.a,
                solve_dofs=solve_dofs,
            ),
        ),
        axis=0,
    ).ravel()  # [4*n_free_dof]

timestep_residual_jacobians

timestep_residual_jacobians(
    i_ts: int | Array,
    q_nm1: StructureMinimalStates,
    q_n: StructureMinimalStates,
    f_ext_aero_nm1: Array | None,
    f_ext_aero_n: Array | None,
    dv: StructureDesignVariables,
    thrust_t: dict[str, Array],
    solve_dofs: tuple[int, ...],
    approx_grads: bool,
    n_profile_loops: int | None,
    jac_options: dict[
        str, dict[str, Callable[..., Any] | None]
    ],
    mode: ADMode = "reverse",
) -> tuple[
    Array,
    Array,
    StructureDesignVariables,
    Array | None,
    Array | None,
    dict[str, dict[str, float]] | None,
    dict[str, dict[str, float]] | None,
]

Obtain the Jacobians of the structural residual with respect to the current states and previous states.

Parameters:

Name	Type	Description	Default
i_ts	int | Array	Time step index.	required
q_nm1	StructureMinimalStates	Previous minimal states.	required
q_n	StructureMinimalStates	Current minimal states.	required
f_ext_aero_nm1	Array | None	Optional aerodynamic forcing for previous time step, (n_nodes, 6).	required
f_ext_aero_n	Array | None	Optional aerodynamic forcing for current time step, (n_nodes, 6).	required
dv	StructureDesignVariables	Design variables.	required
thrust_t	dict[str, Array]	Thrust time history, {key, (n_tstep, )}.	required
solve_dofs	tuple[int, ...]	Index of degrees of freedom to solve for.	required
approx_grads	bool	If True, remove some gradient terms which are generally small.	required
n_profile_loops	int | None	Number of profile loops to run for timing function. If None, no profiling is done.	required
jac_options	dict[str, dict[str, Callable[..., Any] | None]]	Input which passes functions which can be used to approximate the Jacobians. If entries are None, AD is used.	required
mode	ADMode	AD mode used for Jacobian construction.	'reverse'

Returns:

Type	Description
tuple[Array, Array, StructureDesignVariables, Array | None, Array | None, dict[str, dict[str, float]] | None, dict[str, dict[str, float]] | None]	Jacobians with respect to previous state and current state, gradients with respect to design variables, previous, and current aerodynamic forces respectively, and profiling times for compilation and run time.

Source code in src/flapjax/structure/gradients/beam.py

def timestep_residual_jacobians(
    self,
    i_ts: int | Array,
    q_nm1: StructureMinimalStates,
    q_n: StructureMinimalStates,
    f_ext_aero_nm1: Array | None,
    f_ext_aero_n: Array | None,
    dv: StructureDesignVariables,
    thrust_t: dict[str, Array],
    solve_dofs: tuple[int, ...],
    approx_grads: bool,
    n_profile_loops: int | None,
    jac_options: dict[str, dict[str, Callable[..., Any] | None]],
    mode: ADMode = "reverse",
) -> tuple[
    Array,
    Array,
    StructureDesignVariables,
    Array | None,
    Array | None,
    dict[str, dict[str, float]] | None,
    dict[str, dict[str, float]] | None,
]:
    r"""
    Obtain the Jacobians of the structural residual with respect to the current states and previous states.
    :param i_ts: Time step index.
    :param q_nm1: Previous minimal states.
    :param q_n: Current minimal states.
    :param f_ext_aero_nm1: Optional aerodynamic forcing for previous time step, ``(n_nodes, 6)``.
    :param f_ext_aero_n: Optional aerodynamic forcing for current time step, ``(n_nodes, 6)``.
    :param dv: Design variables.
    :param thrust_t: Thrust time history, ``{key, (n_tstep, )}``.
    :param solve_dofs: Index of degrees of freedom to solve for.
    :param approx_grads: If True, remove some gradient terms which are generally small.
    :param n_profile_loops: Number of profile loops to run for timing function. If None, no profiling is done.
    :param jac_options: Input which passes functions which can be used to approximate the Jacobians. If entries are
    None, AD is used.
    :param mode: AD mode used for Jacobian construction.
    :return: Jacobians with respect to previous state and current state, gradients with respect to design variables,
    previous, and current aerodynamic forces respectively, and profiling times for compilation and run time.
    """

    compile_time: dict[str, dict[str, float]] = {}
    run_time: dict[str, dict[str, float]] = {}

    compute_f_aero_grads = f_ext_aero_nm1 is not None and f_ext_aero_n is not None
    if compute_f_aero_grads:
        jac_options["v_dot"].update({"f_aero_nm1": None, "f_aero_n": None})

    # varphi
    d_varphi, compile_time["varphi"], run_time["varphi"] = jacrev_custom(
        func=self.varphi_res_func,
        jac_options=jac_options["varphi"],
        n_profile_loops=n_profile_loops,
        func_name="varphi",
        static_argnames=("solve_dofs",),
        mode=mode,
    )(
        varphi_nm1=q_nm1.varphi.ravel(),
        varphi_n=q_n.varphi.ravel(),
        v_nm1=q_nm1.v.ravel(),
        a_nm1=q_nm1.a.ravel(),
        a_n=q_n.a.ravel(),
        solve_dofs=solve_dofs,
    )

    # velocity
    d_v, compile_time["v"], run_time["v"] = jacrev_custom(
        func=self.v_res_func,
        jac_options=jac_options["v"],
        n_profile_loops=n_profile_loops,
        func_name="v",
        static_argnames=("solve_dofs",),
        mode=mode,
    )(
        v_nm1=q_nm1.v.ravel(),
        v_n=q_n.v.ravel(),
        a_nm1=q_nm1.a.ravel(),
        a_n=q_n.a.ravel(),
        solve_dofs=solve_dofs,
    )

    # acceleration
    d_v_dot, compile_time["v_dot"], run_time["v_dot"] = jacrev_custom(
        func=self.v_dot_res_func,
        jac_options=jac_options["v_dot"],
        n_profile_loops=n_profile_loops,
        func_name="v_dot",
        static_argnames=("solve_dofs", "approx_grads"),
        mode=mode,
    )(
        i_ts=i_ts,
        varphi_nm1=q_nm1.varphi.ravel(),
        varphi_n=q_n.varphi.ravel(),
        v_nm1=q_nm1.v.ravel(),
        v_n=q_n.v.ravel(),
        v_dot_nm1=q_nm1.v_dot.ravel(),
        v_dot_n=q_n.v_dot.ravel(),
        dv=dv,
        f_aero_nm1=f_ext_aero_nm1.ravel() if compute_f_aero_grads else None,  # type: ignore
        f_aero_n=f_ext_aero_n.ravel() if compute_f_aero_grads else None,  # type: ignore
        thrust_t=thrust_t,
        solve_dofs=solve_dofs,
        approx_grads=approx_grads,
    )

    if not compute_f_aero_grads:
        # no Jacobians for aero case, but include a None to keep the keys consistent
        d_v_dot.update({"f_aero_nm1": None, "f_aero_n": None})

    # pseudo-acceleration
    d_a, compile_time["a"], run_time["a"] = jacrev_custom(
        func=self.a_res_func,
        jac_options=jac_options["a"],
        n_profile_loops=n_profile_loops,
        func_name="a",
        static_argnames=("solve_dofs",),
        mode=mode,
    )(
        v_dot_nm1=q_nm1.v_dot.ravel(),
        v_dot_n=q_n.v_dot.ravel(),
        a_nm1=q_nm1.a.ravel(),
        a_n=q_n.a.ravel(),
        solve_dofs=solve_dofs,
    )

    struct_sizes = (
        len(solve_dofs),
        len(solve_dofs),
        len(solve_dofs),
        len(solve_dofs),
    )

    nm1_keys = ("varphi_nm1", "v_nm1", "v_dot_nm1", "a_nm1")
    p_r_n_p_q_nm1 = construct_named_block_jacobian(
        entries=tuple(
            [
                {k: v[:, solve_dofs] for k, v in jacs.items() if k in nm1_keys}
                for jacs in (d_varphi, d_v, d_v_dot, d_a)
            ]
        ),
        keys=nm1_keys,
        widths=struct_sizes,
        heights=struct_sizes,
    )

    n_keys = ("varphi_n", "v_n", "v_dot_n", "a_n")
    p_r_n_p_q_n = construct_named_block_jacobian(
        entries=tuple(
            [
                {k: v[:, solve_dofs] for k, v in jacs.items() if k in n_keys}
                for jacs in (d_varphi, d_v, d_v_dot, d_a)
            ]
        ),
        keys=n_keys,
        widths=struct_sizes,
        heights=struct_sizes,
    )

    return (
        p_r_n_p_q_nm1,
        p_r_n_p_q_n,
        d_v_dot["dv"],
        d_v_dot["f_aero_nm1"],
        d_v_dot["f_aero_n"],
        compile_time if n_profile_loops is not None else None,
        run_time if n_profile_loops is not None else None,
    )

j_from_q_x

j_from_q_x(
    q_n_mat: Array,
    dv: StructureDesignVariables,
    dv_full: StructureDesignVariables,
    objective: StructureObjectiveFunction,
    i_ts: int,
) -> Array

Obtain the objective as a function of the minimal states and design variables.

Parameters:

Name	Type	Description	Default
q_n_mat	Array	Matrix representation of the minimal states.	required
dv	StructureDesignVariables	Design variables, with unwanted entries replaced with None.	required
dv_full	StructureDesignVariables	Design variables which are defined for all entries.	required
objective	StructureObjectiveFunction	Objective function.	required
i_ts	int	Time step index.	required

Returns:

Type	Description
Array	Objective value.

Source code in src/flapjax/structure/gradients/beam.py

def j_from_q_x(
    self,
    q_n_mat: Array,
    dv: StructureDesignVariables,
    dv_full: StructureDesignVariables,
    objective: StructureObjectiveFunction,
    i_ts: int,
) -> Array:
    r"""
    Obtain the objective as a function of the minimal states and design variables.
    :param q_n_mat: Matrix representation of the minimal states.
    :param dv: Design variables, with unwanted entries replaced with None.
    :param dv_full: Design variables which are defined for all entries.
    :param objective: Objective function.
    :param i_ts: Time step index.
    :return: Objective value.
    """
    full_states = self.minimal_states_to_full_states(
        i_ts=i_ts,
        q=StructureMinimalStates.from_mat(q_n_mat),
        dv=dv,
        dv_full=dv_full,
    )
    return jnp.atleast_1d(objective(full_states, dv, i_ts))

p_j

p_j(
    objective: StructureObjectiveFunction,
    i_ts: int,
    dv: StructureDesignVariables,
    dv_full: StructureDesignVariables,
    q_n: StructureMinimalStates,
) -> tuple[Array, StructureDesignVariables]

Obtains Jacobians of the objective function.

Parameters:

Name	Type	Description	Default
objective	StructureObjectiveFunction	Objective function.	required
i_ts	int	Time step index.	required
dv	StructureDesignVariables	Design variables.	required
dv_full	StructureDesignVariables	Design variables which are defined for all entries.	required
q_n	StructureMinimalStates	Current minimal states.	required

Returns:

Type	Description
tuple[Array, StructureDesignVariables]	Jacobian with respect to minimal states and design variables.

Source code in src/flapjax/structure/gradients/beam.py

@jax.jit(static_argnums=(0, 1, 3, 4))
def p_j(
    self,
    objective: StructureObjectiveFunction,
    i_ts: int,
    dv: StructureDesignVariables,
    dv_full: StructureDesignVariables,
    q_n: StructureMinimalStates,
) -> tuple[Array, StructureDesignVariables]:
    r"""
    Obtains Jacobians of the objective function.
    :param objective: Objective function.
    :param i_ts: Time step index.
    :param dv: Design variables.
    :param dv_full: Design variables which are defined for all entries.
    :param q_n: Current minimal states.
    :return: Jacobian with respect to minimal states and design variables.
    """

    def _j(q_n_mat: Array, dv_: StructureDesignVariables) -> Array:
        return self.j_from_q_x(
            q_n_mat=q_n_mat, dv=dv_, dv_full=dv_full, objective=objective, i_ts=i_ts
        )

    p_j_n_p_q_n, p_j_n_p_x = jax.jacrev(_j, argnums=(0, 1), allow_int=True)(
        q_n.to_mat(), dv
    )

    return cast(Array, p_j_n_p_q_n), cast(StructureDesignVariables, p_j_n_p_x)

adjoint_time_loop

adjoint_time_loop(
    rev_i_ts: int,
    d_j_d_x_: StructureDesignVariables,
    adj_: Array,
    p_r_np1_p_q_n: Array | None,
    adj_t_p_r_np1_p_q_n: Array | None,
    q_n: StructureMinimalStates,
    structure: StructureCase,
    objective: StructureObjectiveFunction,
    dv: StructureDesignVariables,
    dv_full: StructureDesignVariables,
    thrust_t: dict[str, Array],
    solve_dofs: tuple[int, ...],
    approx_grads: bool,
    save_adjoint: bool,
    matrix_free: bool,
    n_j: int,
    jac_options: dict[
        str, dict[str, Callable[..., Any] | None]
    ],
    i_ts_end: int | None = None,
) -> tuple[
    StructureDesignVariables,
    Array,
    Array,
    StructureMinimalStates,
]

Function to obtain the grads states at timestep varphi, which is dependent on the grads at timestep varphi+1.

Parameters:

Name	Type	Description	Default
rev_i_ts	int	Reversed timestep index. JAX loop does not allow for reverse indexing, and so this is. explicitly reversed within the function body to obtain i_ts.	required
d_j_d_x_	StructureDesignVariables	Design gradient to accumulate.	required
adj_	Array	Full grads matrix which is updated inplace, (n_tstep, *j_shape, 5*n_dof).	required
p_r_np1_p_q_n	Array | None	Gradient of future step with respect to current state, used when computing the full Jacobian (5*n_dof, 5*n_dof).	required
adj_t_p_r_np1_p_q_n	Array | None	VJP of the future adjoint step and the Jacobian of the future residual with respect to the current state, (n_adj_dof, ).	required
q_n	StructureMinimalStates	Current minimal states.	required
structure	StructureCase	Dynamic structure solution.	required
objective	StructureObjectiveFunction	Objective function.	required
dv	StructureDesignVariables	Structure design variables.	required
dv_full	StructureDesignVariables	Structure design variables which are defined for all entries.	required
thrust_t	dict[str, Array]	Thrust time history, {key, (n_tstep, )}.	required
solve_dofs	tuple[int, ...]	Tuple of dof index to solve.	required
approx_grads	bool	Whether to approximate the gradient or not.	required
save_adjoint	bool	Whether to save the full adjoint time history.	required
matrix_free	bool	If False, solve the system using the residual Jacobian-vector product using GMRES.	required
n_j	int	Number of objective function outputs.	required
jac_options	dict[str, dict[str, Callable[..., Any] | None]]	Input which passes functions which can be used to approximate the Jacobians. If entries are None, AD is used.	required
i_ts_end	int | None	Largest time step index for which the adjoint is computed. Defaults to structure.n_tstep - 1 when None	None

Returns:

Type	Description
tuple[StructureDesignVariables, Array, Array, StructureMinimalStates]	Updated grads matrix, gradient of current step with respect to previous state and current state.

Source code in src/flapjax/structure/gradients/beam.py

def adjoint_time_loop(
    self,
    rev_i_ts: int,
    d_j_d_x_: StructureDesignVariables,
    adj_: Array,
    p_r_np1_p_q_n: Array | None,
    adj_t_p_r_np1_p_q_n: Array | None,
    q_n: StructureMinimalStates,
    structure: StructureCase,
    objective: StructureObjectiveFunction,
    dv: StructureDesignVariables,
    dv_full: StructureDesignVariables,
    thrust_t: dict[str, Array],
    solve_dofs: tuple[int, ...],
    approx_grads: bool,
    save_adjoint: bool,
    matrix_free: bool,
    n_j: int,
    jac_options: dict[str, dict[str, Callable[..., Any] | None]],
    i_ts_end: int | None = None,
) -> tuple[StructureDesignVariables, Array, Array, StructureMinimalStates]:
    r"""
    Function to obtain the grads states at timestep varphi, which is dependent on the grads at timestep varphi+1.
    :param rev_i_ts: Reversed timestep index. JAX loop does not allow for reverse indexing, and so this is.
    explicitly reversed within the function body to obtain i_ts.
    :param d_j_d_x_: Design gradient to accumulate.
    :param adj_: Full grads matrix which is updated inplace, ``(n_tstep, *j_shape, 5*n_dof)``.
    :param p_r_np1_p_q_n: Gradient of future step with respect to current state, used when computing the full
    Jacobian ``(5*n_dof, 5*n_dof)``.
    :param adj_t_p_r_np1_p_q_n: VJP of the future adjoint step and the Jacobian of the future residual with respect
    to the current state, ``(n_adj_dof, )``.
    :param q_n: Current minimal states.
    :param structure: Dynamic structure solution.
    :param objective: Objective function.
    :param dv: Structure design variables.
    :param dv_full: Structure design variables which are defined for all entries.
    :param thrust_t: Thrust time history, ``{key, (n_tstep, )}``.
    :param solve_dofs: Tuple of dof index to solve.
    :param approx_grads: Whether to approximate the gradient or not.
    :param save_adjoint: Whether to save the full adjoint time history.
    :param matrix_free: If False, solve the system using the residual Jacobian-vector product using GMRES.
    :param n_j: Number of objective function outputs.
    :param jac_options: Input which passes functions which can be used to approximate the Jacobians. If entries are
    None, AD is used.
    :param i_ts_end: Largest time step index for which the adjoint is computed. Defaults to
    ``structure.n_tstep - 1`` when ``None``
    :return: Updated grads matrix, gradient of current step with respect to previous state and current state.
    """

    i_ts_end_ = structure.n_tstep - 1 if i_ts_end is None else i_ts_end
    i_ts = i_ts_end_ - rev_i_ts  # index for timestep n, which decrements

    i_ts_nm1 = jnp.maximum(i_ts - 1, 0)  # index for timestep varphi-1

    solve_idx = jnp.array(solve_dofs)

    # find minimal states for timestep varphi-1
    q_nm1 = structure.get_minimal_states(i_ts_nm1)

    # Objective sensitivities
    p_j_n_p_q_n, p_j_n_p_x = self.p_j(
        objective=objective, i_ts=i_ts, dv=dv, dv_full=dv_full, q_n=q_n
    )

    if matrix_free:

        def _residual_states(
            q_n_: StructureMinimalStates, q_nm1_: StructureMinimalStates
        ):
            return self.timestep_residual(
                i_ts=i_ts,
                q_nm1=q_nm1_,
                q_n=q_n_,
                dv_=dv,
                thrust_t=thrust_t,
                solve_dofs=solve_dofs,
                approx_grads=approx_grads,
            )

        # Linearise the timestep residual around (q_n, q_nm1). This single VJP returns:
        # p_r_n_dot_v(v)[0] = (p_r_n/p_q_n).T @ v, p_r_n_dot_v(v)[1] = (p_r_n/p_q_nm1).T @ v
        _, p_r_n_dot_v = jax.vjp(_residual_states, q_n, q_nm1)

        def _cot_to_solve_vec(cot: StructureMinimalStates) -> Array:
            # collapse cotangent to [n_adj_dof]
            mat = cot.to_mat()  # [4, n_nodes, 6]
            return mat.reshape(mat.shape[0], -1)[:, solve_idx].ravel()

        def matvec_qn_t(v: Array) -> Array:
            # function to compute (p_r_n/p_q_n).T @ v for some vector v
            return _cot_to_solve_vec(p_r_n_dot_v(v)[0])

        # sensitivity of objective to degrees of freedom, (n_j, n_adj_dof)
        p_j_solve = (
            p_j_n_p_q_n.reshape(n_j, 4, -1, 6)
            .reshape(n_j, 4, -1)[..., solve_idx]
            .reshape(n_j, -1)
        )
        assert adj_t_p_r_np1_p_q_n is not None, (
            "The adjoint-Jacobian product has not been passed"
        )
        b_rhs = -(p_j_solve + adj_t_p_r_np1_p_q_n)  # (n_j, n_adj_dof)

        # solve for the adjoint vector at timestep n, batched along the size of the objective.
        def _solve_row(b_row: Array) -> Array:
            # noinspection PyTypeChecker
            x, _ = jax.scipy.sparse.linalg.gmres(
                matvec_qn_t,
                b_row,
                tol=1e-10,
                atol=1e-10,
                maxiter=50,
                solve_method="batched",
            )
            return x

        adj_n = jax.vmap(_solve_row)(b_rhs)  # (n_j, n_adj_dof)

        # Design gradient accumulation via a separate VJP to obtain adj.T @ p_r_v_dot_n_p_dv.
        def _residual_dv(dv_: StructureDesignVariables) -> Array:
            return self.timestep_residual(
                i_ts=i_ts,
                q_nm1=q_nm1,
                q_n=q_n,
                dv_=dv_,
                thrust_t=thrust_t,
                solve_dofs=solve_dofs,
                approx_grads=approx_grads,
            )

        _, pull_dv = jax.vjp(_residual_dv, dv)
        dv_grads = jax.vmap(pull_dv)(adj_n)[0]

        # accumulate with seperate statements as there is no __add__ member
        d_j_d_x_ += dv_grads
        d_j_d_x_ += p_j_n_p_x

        # compute adj_n @ p_r_n/p_q_nm1 for next iteration
        def _coupling_row(a: Array) -> Array:
            _, cot_qnm1 = p_r_n_dot_v(a)
            return _cot_to_solve_vec(cot_qnm1)

        adj_t_p_r_n_p_q_nm1 = jax.vmap(_coupling_row)(adj_n)  # (n_j, n_adj_dof)

        p_r_n_p_q_nm1: Array | None = None  # unused
    else:
        # find gradients of residual function (state Jacobians only)
        p_r_n_p_q_nm1, p_r_n_p_q_n, p_r_v_dot_n_p_dv, *_ = (
            self.timestep_residual_jacobians(
                i_ts=i_ts,
                q_n=q_n,
                q_nm1=q_nm1,
                dv=dv,
                solve_dofs=solve_dofs,
                approx_grads=approx_grads,
                f_ext_aero_nm1=None,
                f_ext_aero_n=None,
                thrust_t=thrust_t,
                n_profile_loops=None,
                jac_options=jac_options,
            )
        )

        # solve for adjoint at current timestep
        prev_adjoint = adj_[i_ts + 1, ...] if save_adjoint else adj_
        b: Array = -(p_j_n_p_q_n.reshape(n_j, -1) + prev_adjoint @ p_r_np1_p_q_n).T
        adj_n = jnp.linalg.solve(p_r_n_p_q_n.T, b).T

        # accumulate design derivative
        d_j_d_x_ += p_r_v_dot_n_p_dv.premultiply_adj(
            adj_n[:, solve_idx + 2 * len(solve_dofs)]
        )

        # add on direct contribution from objective
        d_j_d_x_ += p_j_n_p_x

        adj_t_p_r_n_p_q_nm1 = None  # unused

    # update adjoint vector time history if requested
    if save_adjoint:
        adj_ = adj_.at[i_ts, ...].set(adj_n)

    # print to console
    jax_print(
        "Adjoint step: {i_ts}",
        i_ts=i_ts,
        verbose_level="normal",
    )

    if matrix_free:
        assert adj_t_p_r_n_p_q_nm1 is not None
        return d_j_d_x_, adj_ if save_adjoint else adj_n, adj_t_p_r_n_p_q_nm1, q_nm1
    else:
        assert p_r_n_p_q_nm1 is not None
        return d_j_d_x_, adj_ if save_adjoint else adj_n, p_r_n_p_q_nm1, q_nm1

construct_approximate_jacobians

construct_approximate_jacobians(
    sol: StructureCase,
    jacobian_approximations: StructureJacobianApproximations,
) -> dict[str, dict[str, Callable[..., Any] | None]]

Compute approximations for Jacobians which are specified in the jacobian_approximations data structure.

Parameters:

Name	Type	Description	Default
sol	StructureCase	Solution for which approximations will be created for the initial time step.	required
jacobian_approximations	StructureJacobianApproximations	Data structure which defines which approximations to create.	required

Returns:

Type	Description
dict[str, dict[str, Callable[..., Any] | None]]	Dictionary of approximations.

Source code in src/flapjax/structure/gradients/beam.py

def construct_approximate_jacobians(
    self,
    sol: StructureCase,
    jacobian_approximations: StructureJacobianApproximations,
) -> dict[str, dict[str, Callable[..., Any] | None]]:
    r"""
    Compute approximations for Jacobians which are specified in the jacobian_approximations data structure.
    :param sol: Solution for which approximations will be created for the initial time step.
    :param jacobian_approximations: Data structure which defines which approximations to create.
    :return: Dictionary of approximations.
    """
    q_nm1 = sol.get_minimal_states(0)
    q_n = sol.get_minimal_states(1)
    dv = self.get_design_variables(
        struct_case=sol, thrust_t=sol.thrust, grads_to_compute=None
    )
    solve_dofs = tuple(
        int(i)
        for i in get_solve_dofs(
            n_dof=self.n_dof, prescribed_dofs=sol.prescribed_dofs
        )
    )
    if sol.f_ext_aero is not None:
        f_aero_nm1 = sol.f_ext_aero[0, ...].ravel()
        f_aero_n = sol.f_ext_aero[1, ...].ravel()
    else:
        f_aero_nm1 = None
        f_aero_n = None

    res_args: dict[
        str, tuple[Callable[..., Array], dict[str, Any], Sequence[str]]
    ] = {
        "varphi": (
            self.varphi_res_func,
            {
                "varphi_nm1": q_nm1.varphi.ravel(),
                "varphi_n": q_n.varphi.ravel(),
                "v_nm1": q_nm1.v.ravel(),
                "a_nm1": q_nm1.a.ravel(),
                "a_n": q_n.a.ravel(),
                "solve_dofs": solve_dofs,
            },
            [f.name for f in fields(VarphiApprox)],
        ),
        "v": (
            self.v_res_func,
            {
                "v_nm1": q_nm1.v.ravel(),
                "v_n": q_n.v.ravel(),
                "a_nm1": q_nm1.a.ravel(),
                "a_n": q_n.a.ravel(),
                "solve_dofs": solve_dofs,
            },
            [f.name for f in fields(VApprox)],
        ),
        "v_dot": (
            self.v_dot_res_func,
            {
                "i_ts": 1,
                "varphi_nm1": q_nm1.varphi.ravel(),
                "varphi_n": q_n.varphi.ravel(),
                "v_nm1": q_nm1.v.ravel(),
                "v_n": q_n.v.ravel(),
                "v_dot_nm1": q_nm1.v_dot.ravel(),
                "v_dot_n": q_n.v_dot.ravel(),
                "dv": dv,
                "f_aero_nm1": f_aero_nm1,
                "f_aero_n": f_aero_n,
                "thrust_t": sol.thrust,
                "solve_dofs": solve_dofs,
                "approx_grads": True,
            },
            [f.name for f in fields(VDotApprox)],
        ),
        "a": (
            self.a_res_func,
            {
                "v_dot_nm1": q_nm1.v_dot.ravel(),
                "v_dot_n": q_n.v_dot.ravel(),
                "a_nm1": q_nm1.a.ravel(),
                "a_n": q_n.a.ravel(),
                "solve_dofs": solve_dofs,
            },
            [f.name for f in fields(AApprox)],
        ),
    }

    return construct_approximation(
        res_args=res_args, jacobian_approximations=jacobian_approximations
    )

dynamic_adjoint

dynamic_adjoint(
    structure: StructureCase,
    objective: StructureObjectiveFunction,
    matrix_free: bool = False,
    jacobian_approximations: StructureJacobianApproximations = JACOBIAN_APPROXIMATIONS_DEFAULT,
    p_q0_p_x: StructureDesignVariables | None = None,
    save_adjoint: bool = False,
    approx_grads: bool = True,
    grads_to_compute: StructureGradsToCompute = GRADS_TO_COMPUTE_DEFAULT,
    i_ts_adjoint_range: tuple[int | None, int | None] = (
        None,
        None,
    ),
) -> tuple[StructureDesignVariables, Array | None]

Dynamic structure grads problem. This computes the gradient of the objective of the dynamic response with respect to design variables. The objective has structure :math:J = \sum_{i=1}^N \left(j(\mathbf{x}, \mathbf{y}_i)\right) where :math:\mathbf{x} are the design variables and :math:\mathbf{y} are the structural states at each timestep, which depend on the design variables through the dynamic structure equations. The gradient is computed by first solving a backward pass to obtain the grads states, and then using these to compute the gradient w.r.t. design variables in a forward pass.

Parameters:

Name	Type	Description	Default
structure	StructureCase	Dynamic structure solution object.	required
objective	StructureObjectiveFunction	Objective function :math:j(\mathbf{x}, \mathbf{y}_i).	required
matrix_free	bool	Whether to use matrix-free methods for solving the linear systems. Default is False, as structural problems generally do not benefit from this solve.	False
jacobian_approximations	StructureJacobianApproximations	Data structure which specifies Jacobian approximations to use for each part of the problem. The value can either be None for no approximation, constant for the assumption that the Jacobian does not vary with any variables. Alternatively, it can be tuple pairs with first entry being dense_linear or lazy_linear, with the second entry being a sequence of argument names for which to obtain the Hessian.	JACOBIAN_APPROXIMATIONS_DEFAULT
p_q0_p_x	StructureDesignVariables | None	Optional Jacobian used to describe the sensitivities of the initial structural degrees of freedom to the design variables.	None
save_adjoint	bool	Whether to save the full adjoint vectors.	False
approx_grads	bool	If true, some gradient contributions which are assumed to be near-zero are removed to decrease computational cost.	True
grads_to_compute	StructureGradsToCompute	Design variables with which to compute design gradients for.	GRADS_TO_COMPUTE_DEFAULT
i_ts_adjoint_range	tuple[int | None, int | None]	Optional (start, end) window of time step indices over which the objective contributes to the gradient. Either entry may be None to leave that side untruncated. Defining a start time step that is nonzero will skip the initial state adjoint contribution.	(None, None)

Returns:

Type	Description
tuple[StructureDesignVariables, Array | None]	Objective gradient :math:\frac{dJ}{d\mathbf{x}} and adjoint states

Source code in src/flapjax/structure/gradients/beam.py

def dynamic_adjoint(
    self,
    structure: StructureCase,
    objective: StructureObjectiveFunction,
    matrix_free: bool = False,
    jacobian_approximations: StructureJacobianApproximations = JACOBIAN_APPROXIMATIONS_DEFAULT,
    p_q0_p_x: StructureDesignVariables | None = None,
    save_adjoint: bool = False,
    approx_grads: bool = True,
    grads_to_compute: StructureGradsToCompute = GRADS_TO_COMPUTE_DEFAULT,
    i_ts_adjoint_range: tuple[int | None, int | None] = (None, None),
) -> tuple[StructureDesignVariables, Array | None]:
    r"""
    Dynamic structure grads problem. This computes the gradient of the objective of the dynamic response with
    respect to design variables. The objective has structure
    :math:`J = \sum_{i=1}^N \left(j(\mathbf{x}, \mathbf{y}_i)\right)` where :math:`\mathbf{x}` are the design variables
    and :math:`\mathbf{y}` are the structural states at each timestep, which depend on the design variables through
    the dynamic structure equations. The gradient is computed by first solving a backward pass to obtain the grads
    states, and then using these to compute the gradient w.r.t. design variables in a forward pass.
    :param structure: Dynamic structure solution object.
    :param objective: Objective function :math:`j(\mathbf{x}, \mathbf{y}_i)`.
    :param matrix_free: Whether to use matrix-free methods for solving the linear systems. Default is False, as
    structural problems generally do not benefit from this solve.
    :param jacobian_approximations: Data structure which specifies Jacobian approximations to use for each part of
    the problem. The value can either be None for no approximation, `constant` for the assumption that the Jacobian
    does not vary with any variables. Alternatively, it can be tuple pairs with first entry being `dense_linear` or
    `lazy_linear`, with the second entry being a sequence of argument names for which to obtain the Hessian.
    :param p_q0_p_x: Optional Jacobian used to describe the sensitivities of the initial structural degrees of
    freedom to the design variables.
    :param save_adjoint: Whether to save the full adjoint vectors.
    :param approx_grads: If true, some gradient contributions which are assumed to be near-zero are removed to
    decrease computational cost.
    :param grads_to_compute: Design variables with which to compute design gradients for.
    :param i_ts_adjoint_range: Optional ``(start, end)`` window of time step indices over which the objective
    contributes to the gradient. Either entry may be ``None`` to leave that side untruncated. Defining a start
    time step that is nonzero will  skip the initial state adjoint contribution.
    :return: Objective gradient :math:`\frac{dJ}{d\mathbf{x}}` and adjoint states
    """

    dv = self.get_design_variables(
        struct_case=structure,
        thrust_t=structure.thrust,
        grads_to_compute=grads_to_compute,
    )

    dv_full = self.get_design_variables(
        struct_case=structure, thrust_t=structure.thrust, grads_to_compute=None
    )

    struct_states_init = structure.get_full_states(i_ts=0)

    j_properties = jax.eval_shape(
        lambda: jnp.atleast_1d(objective(struct_states_init, dv, None))
    )
    j_shape = j_properties.shape
    n_j = j_properties.size

    # assemble
    solve_dofs: tuple[int, ...] = tuple(
        int(i)
        for i in get_solve_dofs(
            n_dof=self.n_dof, prescribed_dofs=structure.prescribed_dofs
        )
    )

    dv_grad_init = StructureDesignVariables(
        x0=jnp.zeros((*j_shape, *self.x0.shape)) if dv.x0 is not None else None,
        orientation_euler=jnp.zeros((*j_shape, 3))
        if dv.orientation_euler is not None
        else None,
        k_cs=jnp.zeros((*j_shape, *self.k_cs.shape))
        if dv.k_cs is not None
        else None,
        m_cs=jnp.zeros((*j_shape, *self.m_cs.shape))
        if dv.m_cs is not None
        else None,
        m_lumped=jnp.zeros((*j_shape, *self.m_lumped.shape))
        if self.use_lumped_mass and dv.m_lumped is not None
        else None,
        f_ext_dead=jnp.zeros((*j_shape, *structure.f_ext_dead.shape))
        if structure.f_ext_dead is not None and dv.f_ext_dead is not None
        else None,
        f_ext_follower=jnp.zeros((*j_shape, *structure.f_ext_follower.shape))
        if structure.f_ext_follower is not None and dv.f_ext_follower is not None
        else None,
        thrust_t={
            k: jnp.zeros((*j_shape, *v.shape)) for k, v in structure.thrust.items()
        }
        if dv.thrust_t is not None
        else None,
        f_shape=(),
    )

    n_adj_dof = 4 * (
        self.n_dof - len(structure.prescribed_dofs)
    )  # number of grads degrees of freedom

    # compute Jacobian approximations, if requested
    jac_options = self.construct_approximate_jacobians(
        sol=structure, jacobian_approximations=jacobian_approximations
    )

    # check adjoint window
    i_ts_start_adj, i_ts_end_adj = i_ts_adjoint_range
    i_ts_start_adj_: int = 1 if i_ts_start_adj is None else i_ts_start_adj
    i_ts_end_adj_: int = (
        structure.n_tstep - 1 if i_ts_end_adj is None else i_ts_end_adj
    )
    if i_ts_start_adj_ < 1:
        raise ValueError(
            f"i_ts_adjoint_range start must be >= 1, got {i_ts_start_adj_}"
        )
    if i_ts_end_adj_ > structure.n_tstep - 1:
        raise ValueError(
            f"i_ts_adjoint_range end must be <= n_tstep - 1 = "
            f"{structure.n_tstep - 1}, got {i_ts_end_adj_}"
        )
    if i_ts_end_adj_ < i_ts_start_adj_:
        raise ValueError(
            f"i_ts_adjoint_range end ({i_ts_end_adj_}) must be >= start "
            f"({i_ts_start_adj_})"
        )
    n_adj_iters: int = i_ts_end_adj_ - i_ts_start_adj_ + 1

    # wrap in a local JIT so structure/aero_dv become closure constants
    @jax.jit
    def adjoint_step(
        rev_i_ts_: int,
        d_j_d_x_: StructureDesignVariables,
        adj_: Array,
        coupling_arr: Array,
        q_n: StructureMinimalStates,
    ) -> tuple[StructureDesignVariables, Array, Array, StructureMinimalStates]:
        return self.adjoint_time_loop(
            rev_i_ts=rev_i_ts_,
            d_j_d_x_=d_j_d_x_,
            adj_=adj_,
            p_r_np1_p_q_n=None if matrix_free else coupling_arr,
            adj_t_p_r_np1_p_q_n=coupling_arr if matrix_free else None,
            q_n=q_n,
            structure=structure,
            objective=objective,
            dv=dv,
            dv_full=dv_full,
            thrust_t=structure.thrust,
            solve_dofs=solve_dofs,
            approx_grads=approx_grads,
            save_adjoint=save_adjoint,
            matrix_free=matrix_free,
            n_j=n_j,
            jac_options=jac_options,
            i_ts_end=i_ts_end_adj_,
        )

    # coupling array is either a Jacobian or a VJP depending on if using matrix free or not
    coupling_init = (
        jnp.zeros((n_j, n_adj_dof))
        if matrix_free
        else jnp.zeros((n_adj_dof, n_adj_dof))
    )

    # pass through time steps backwards to obtain adjoints
    # coupling0 is p_r1_p_q0 when matrix_free is False, and adj_1 @ p_r1_p_q0 when matrix_free is True
    d_j_d_x, adj, coupling0, _ = jax.lax.fori_loop(
        lower=0,
        upper=n_adj_iters,
        body_fun=lambda i_ts_, args: adjoint_step(i_ts_, *args),
        init_val=(
            dv_grad_init,
            jnp.zeros((structure.n_tstep + 1, n_j, n_adj_dof))
            if save_adjoint
            else jnp.zeros((n_j, n_adj_dof)),
            coupling_init,
            structure.get_minimal_states(i_ts_end_adj_),
        ),
    )

    # solve initial timestep adjoint, as there is no r0. Skipped when the adjoint window truncates early time steps
    if i_ts_start_adj_ <= 1:
        p_j0_p_q0, p_j0_p_x = self.p_j(
            objective=objective,
            i_ts=0,
            dv=dv,
            dv_full=dv_full,
            q_n=structure.get_minimal_states(0),
        )

        if matrix_free:
            adj0 = -p_j0_p_q0.reshape(n_j, -1) - coupling0
        else:
            adj0 = (
                -p_j0_p_q0.reshape(n_j, -1)
                - (adj[1, ...] if save_adjoint else adj) @ coupling0
            )

        # add initial direct sensitivity
        d_j_d_x += p_j0_p_x

        # include initial state sensitivity
        if p_q0_p_x is not None:
            d_j_d_x += p_q0_p_x.premultiply_adj(-adj0)
    else:
        adj0 = jnp.zeros((n_j, n_adj_dof))

    # restore original shape of j, and cut off zeros for past-end timestep
    if save_adjoint:
        adj = adj.at[0, ...].set(adj0)
        return d_j_d_x, adj.reshape(adj.shape[0], *j_shape, *adj.shape[2:])[:-1]
    else:
        return d_j_d_x, None

dynamic_adjoint_jacobian_profile

dynamic_adjoint_jacobian_profile(
    sol: StructureCase,
    approx_grads: bool,
    jacobian_approximations: StructureJacobianApproximations = JACOBIAN_APPROXIMATIONS_DEFAULT,
    grads_to_compute: StructureGradsToCompute | None = None,
    f_aero_nm1_n: tuple[Array, Array] | None = None,
    i_ts: int = 1,
    n_loop: int = 10,
    *,
    print_header: bool = True,
) -> tuple[
    dict[str, dict[str, float]], dict[str, dict[str, float]]
]

Function to time evaluation of the Jacobians used for the adjoint solution.

Parameters:

Name	Type	Description	Default
sol	StructureCase	Dynamic structural solution to extract states from.	required
approx_grads	bool	If True, neglect small gradient terms.	required
jacobian_approximations	StructureJacobianApproximations	Data structure which specifies Jacobian approximations to use for each part of the problem.	JACOBIAN_APPROXIMATIONS_DEFAULT
grads_to_compute	StructureGradsToCompute | None	StructureGradsToCompute object which describes which design gradients to compute. If None, all gradients will be computed.	None
f_aero_nm1_n	tuple[Array, Array] | None	Tuple of [f_aero_nm1, f_aero_n] which are passed from the aero problem. If None, no aerodynamic force gradients will be computed.	None
i_ts	int	Time step index where to evaluate residual Jacobians.	1
n_loop	int	Number of times to loop the Jacobian evaluation time for averaging the runtime.	10
print_header	bool	Flag used to prevent heading printer when called by the coupled profiler.	True

Returns:

Type	Description
tuple[dict[str, dict[str, float]], dict[str, dict[str, float]]]	Dictionary of {residual_name: {gradient_argument: val}} for compile time and run time respectively.

Source code in src/flapjax/structure/gradients/beam.py

def dynamic_adjoint_jacobian_profile(
    self,
    sol: StructureCase,
    approx_grads: bool,
    jacobian_approximations: StructureJacobianApproximations = JACOBIAN_APPROXIMATIONS_DEFAULT,
    grads_to_compute: StructureGradsToCompute | None = None,
    f_aero_nm1_n: tuple[Array, Array] | None = None,
    i_ts: int = 1,
    n_loop: int = 10,
    *,
    print_header: bool = True,
) -> tuple[dict[str, dict[str, float]], dict[str, dict[str, float]]]:
    r"""
    Function to time evaluation of the Jacobians used for the adjoint solution.
    :param sol: Dynamic structural solution to extract states from.
    :param approx_grads: If True, neglect small gradient terms.
    :param jacobian_approximations: Data structure which specifies Jacobian approximations to use for each part of
    the problem.
    :param grads_to_compute: StructureGradsToCompute object which describes which design gradients to compute. If
    None, all gradients will be computed.
    :param f_aero_nm1_n: Tuple of [f_aero_nm1, f_aero_n] which are passed from the aero problem. If None, no
    aerodynamic force gradients will be computed.
    :param i_ts: Time step index where to evaluate residual Jacobians.
    :param n_loop: Number of times to loop the Jacobian evaluation time for averaging the runtime.
    :param print_header: Flag used to prevent heading printer when called by the coupled profiler.
    :return: Dictionary of {residual_name: {gradient_argument: val}} for compile time and run time respectively.
    """

    if print_header:
        print_table_title(inner_width=95, title="Structure Adjoint Profile")

    # compute Jacobian approximations, if requested
    jac_options = self.construct_approximate_jacobians(
        sol=sol, jacobian_approximations=jacobian_approximations
    )

    common_kwargs = {
        "i_ts": i_ts,
        "q_nm1": sol.get_minimal_states(i_ts - 1),
        "q_n": sol.get_minimal_states(i_ts),
        "dv": self.get_design_variables(
            struct_case=sol, thrust_t=sol.thrust, grads_to_compute=grads_to_compute
        ),
        "thrust_t": sol.thrust,
        "solve_dofs": tuple(
            get_solve_dofs(n_dof=self.n_dof, prescribed_dofs=sol.prescribed_dofs)
        ),
        "approx_grads": approx_grads,
        "n_profile_loops": n_loop,
        "jac_options": jac_options,
    }

    if f_aero_nm1_n is not None:
        *_, compile_time, run_time = self.timestep_residual_jacobians(
            f_ext_aero_nm1=f_aero_nm1_n[0],
            f_ext_aero_n=f_aero_nm1_n[1],
            **common_kwargs,
        )
    else:
        *_, compile_time, run_time = self.timestep_residual_jacobians(
            f_ext_aero_nm1=None,
            f_ext_aero_n=None,
            **common_kwargs,
        )

    if print_header:
        print_table_line(inner_width=95)

    assert compile_time is not None and run_time is not None, (
        "No output timings passed"
    )

    return compile_time, run_time

LinearBeam

LinearBeam(
    beam: BaseBeamStructure,
    reference: StructureCase,
    n_modes: int | None,
    dt: float | Array,
    modal_inputs: bool = False,
    modal_outputs: bool = False,
    int_order: Literal[3, 4, 5] = BASE_LOBATTO_ORDER,
    prescribed_dofs: Sequence[int]
    | Array
    | slice
    | int
    | None = None,
)

Bases: LinearModel[StructureCase, StructureInputUnflattened, StructureStateUnflattened, StructureOutputUnflattened, StructureLinearResult]

Class to represent a linearised beam system about a reference state.

Source code in src/flapjax/structure/linear/linear_beam.py

def __init__(
    self,
    beam: BaseBeamStructure,
    reference: StructureCase,
    n_modes: int | None,
    dt: float | Array,
    modal_inputs: bool = False,
    modal_outputs: bool = False,
    int_order: Literal[3, 4, 5] = BASE_LOBATTO_ORDER,
    prescribed_dofs: Sequence[int] | Array | slice | int | None = None,
):
    if prescribed_dofs is not None:
        prescribed_dofs = beam.make_prescribed_dofs_tuple(prescribed_dofs)
        free_dofs = get_solve_dofs(
            n_dof=beam.n_dof, prescribed_dofs=prescribed_dofs
        )
    else:
        # inherit from the reference by default
        free_dofs = reference.free_dofs
    self.free_dofs: tuple = free_dofs
    self.n_free_dof: int = len(free_dofs)
    self.n_nodes: int = beam.n_nodes
    self.modal_states: bool = n_modes is not None
    self.modal_inputs: bool = modal_inputs
    self.modal_outputs: bool = modal_outputs
    self._n_modes: int | None = n_modes

    if not self.modal_states and (self.modal_inputs or self.modal_outputs):
        raise ValueError(
            "Modal projection for inputs or outputs requires the system states to be modal."
        )

    # compute m, k, mode_shapes so that the superclass has access when initialised
    if self.modal_states:
        assert n_modes is not None
        self.m, self.k, self.mode_shapes = beam.make_modal_m_k(
            case=reference,
            int_order=int_order,
            n_modes=n_modes,
        )
    else:
        self.m, self.k = beam.make_nodal_m_k(case=reference, int_order=int_order)
        self.mode_shapes = None

    # Rayleigh damping matrix
    self.alpha_m: float = beam.alpha_m
    self.beta_k: float = beam.beta_k
    if beam.alpha_m != 0.0 or beam.beta_k != 0.0:
        self.c = beam.alpha_m * self.m + beam.beta_k * self.k
    else:
        self.c = None

    super().__init__(reference=reference, dt=dt)

    self.sys: LinearSystem = self.linearise()

unpack_state_vector

unpack_state_vector(x: Array) -> S

Unpack a state vector into its components.

Parameters:

Name	Type	Description	Default
x	Array	State vector, (n_states, )	required

Returns:

Type	Description
S	StateUnflattened object.

Source code in src/flapjax/utils/linear.py

def unpack_state_vector(self, x: Array) -> S:
    r"""
    Unpack a state vector into its components.
    :param x: State vector, ``(n_states, )``
    :return: StateUnflattened object.
    """
    return self.state_object(**self._unpack_vector(x=x, slices=self.state_slices))

unpack_output_vector

unpack_output_vector(y: Array) -> O

Unpack an output vector into its components.

Parameters:

Name	Type	Description	Default
y	Array	Output vector, (n_outputs, )	required

Returns:

Type	Description
O	OutputUnflattened object.

Source code in src/flapjax/utils/linear.py

def unpack_output_vector(self, y: Array) -> O:
    r"""
    Unpack an output vector into its components.
    :param y: Output vector, ``(n_outputs, )``
    :return: OutputUnflattened object.
    """
    return self.output_object(**self._unpack_vector(x=y, slices=self.output_slices))

pack_input_vector

pack_input_vector(u_input: InputUnflattened) -> Array

Pack an input unflattened object into a vector.

Parameters:

Name	Type	Description	Default
u_input	InputUnflattened	InputUnflattened object.	required

Returns:

Type	Description
Array	Input vector, (n_inputs, ).

Source code in src/flapjax/utils/linear.py

def pack_input_vector(self, u_input: InputUnflattened) -> Array:
    r"""
    Pack an input unflattened object into a vector.
    :param u_input: InputUnflattened object.
    :return: Input vector, ``(n_inputs, )``.
    """
    return self._pack_vector(
        slices=self.input_slices,
        vec_length=self.n_inputs,
        arrs=shallow_as_dict(u_input),
    )

pack_state_vector

pack_state_vector(x_state: StateUnflattened) -> Array

Pack a state unflattened object into a vector.

Parameters:

Name	Type	Description	Default
x_state	StateUnflattened	StateUnflattened object.	required

Returns:

Type	Description
Array	State vector, (n_states, ).

Source code in src/flapjax/utils/linear.py

def pack_state_vector(self, x_state: StateUnflattened) -> Array:
    r"""
    Pack a state unflattened object into a vector.
    :param x_state: StateUnflattened object.
    :return: State vector, ``(n_states, )``.
    """
    return self._pack_vector(
        slices=self.state_slices,
        vec_length=self.n_states,
        arrs=shallow_as_dict(x_state),
    )

pack_output_vector

pack_output_vector(y_output: OutputUnflattened) -> Array

Pack an output unflattened object into a vector.

Parameters:

Name	Type	Description	Default
y_output	OutputUnflattened	OutputUnflattened object.	required

Returns:

Type	Description
Array	Output vector, (n_outputs, ).

Source code in src/flapjax/utils/linear.py

def pack_output_vector(self, y_output: OutputUnflattened) -> Array:
    r"""
    Pack an output unflattened object into a vector.
    :param y_output: OutputUnflattened object.
    :return: Output vector, ``(n_outputs, )``.
    """
    return self._pack_vector(
        slices=self.output_slices,
        vec_length=self.n_outputs,
        arrs=shallow_as_dict(y_output),
    )

get_total_input

get_total_input(u: InputUnflattened) -> InputUnflattened

Get the total input by adding the reference to the input perturbation.

Parameters:

Name	Type	Description	Default
u	InputUnflattened	InputUnflattened perturbation object.	required

Returns:

Type	Description
InputUnflattened	InputUnflattened total object.

Source code in src/flapjax/utils/linear.py

def get_total_input(self, u: InputUnflattened) -> InputUnflattened:
    r"""
    Get the total input by adding the reference to the input perturbation.
    :param u: InputUnflattened perturbation object.
    :return: InputUnflattened total object.
    """
    return self.input_object(
        **self._get_total(
            shallow_as_dict(u), shallow_as_dict(self.reference_inputs)
        )
    )

get_total_state

get_total_state(x: StateUnflattened) -> StateUnflattened

Get the total state by adding the reference to the state perturbation.

Parameters:

Name	Type	Description	Default
x	StateUnflattened	StateUnflattened perturbation object.	required

Returns:

Type	Description
StateUnflattened	StateUnflattened total object.

Source code in src/flapjax/utils/linear.py

def get_total_state(self, x: StateUnflattened) -> StateUnflattened:
    r"""
    Get the total state by adding the reference to the state perturbation.
    :param x: StateUnflattened perturbation object.
    :return: StateUnflattened total object.
    """
    return self.state_object(
        **self._get_total(
            shallow_as_dict(x), shallow_as_dict(self.reference_states)
        )
    )

get_total_output

get_total_output(y: OutputUnflattened) -> OutputUnflattened

Get the total output by adding the reference to the output perturbation.

Parameters:

Name	Type	Description	Default
y	OutputUnflattened	OutputUnflattened perturbation object.	required

Returns:

Type	Description
OutputUnflattened	OutputUnflattened total object.

Source code in src/flapjax/utils/linear.py

def get_total_output(self, y: OutputUnflattened) -> OutputUnflattened:
    r"""
    Get the total output by adding the reference to the output perturbation.
    :param y: OutputUnflattened perturbation object.
    :return: OutputUnflattened total object.
    """
    return self.output_object(
        **self._get_total(
            shallow_as_dict(y), shallow_as_dict(self.reference_outputs)
        )
    )

get_total_input_t

get_total_input_t(
    u_t: InputUnflattened,
) -> InputUnflattened

Get the total input time history by adding the reference to the input perturbation time history.

Parameters:

Name	Type	Description	Default
u_t	InputUnflattened	InputUnflattened perturbation time history object.	required

Returns:

Type	Description
InputUnflattened	InputUnflattened total time history object.

Source code in src/flapjax/utils/linear.py

def get_total_input_t(self, u_t: InputUnflattened) -> InputUnflattened:
    r"""
    Get the total input time history by adding the reference to the input perturbation time history.
    :param u_t: InputUnflattened perturbation time history object.
    :return: InputUnflattened total time history object.
    """
    return self.input_object(
        **self._get_total(
            input_=shallow_as_dict(u_t),
            reference=shallow_as_dict(self.reference_inputs),
            add_t=True,
        )
    )

get_total_state_t

get_total_state_t(
    x_t: StateUnflattened,
) -> StateUnflattened

Get the total state time history by adding the reference to the state perturbation time history.

Parameters:

Name	Type	Description	Default
x_t	StateUnflattened	StateUnflattened perturbation time history object.	required

Returns:

Type	Description
StateUnflattened	StateUnflattened total time history object.

Source code in src/flapjax/utils/linear.py

def get_total_state_t(self, x_t: StateUnflattened) -> StateUnflattened:
    r"""
    Get the total state time history by adding the reference to the state perturbation time history.
    :param x_t: StateUnflattened perturbation time history object.
    :return: StateUnflattened total time history object.
    """
    return self.state_object(
        **self._get_total(
            shallow_as_dict(x_t),
            shallow_as_dict(self.reference_states),
            add_t=True,
        )
    )

get_total_output_t

get_total_output_t(
    y_t: OutputUnflattened,
) -> OutputUnflattened

Get the total output time history by adding the reference to the output perturbation time history.

Parameters:

Name	Type	Description	Default
y_t	OutputUnflattened	OutputUnflattened perturbation time history object.	required

Returns:

Type	Description
OutputUnflattened	OutputUnflattened total time history object.

Source code in src/flapjax/utils/linear.py

def get_total_output_t(self, y_t: OutputUnflattened) -> OutputUnflattened:
    r"""
    Get the total output time history by adding the reference to the output perturbation time history.
    :param y_t: OutputUnflattened perturbation time history object.
    :return: OutputUnflattened total time history object.
    """
    return self.output_object(
        **self._get_total(
            shallow_as_dict(y_t),
            shallow_as_dict(self.reference_outputs),
            add_t=True,
        )
    )

get_zero_input

get_zero_input() -> InputUnflattened

Get a zero input unflattened object.

Returns:

Type	Description
InputUnflattened	InputUnflattened object with zero arrays.

Source code in src/flapjax/utils/linear.py

def get_zero_input(self) -> InputUnflattened:
    r"""
    Get a zero input unflattened object.
    :return: InputUnflattened object with zero arrays.
    """
    return self.input_object(**self._get_zero(shallow_as_dict(self.input_slices)))

get_zero_state

get_zero_state() -> StateUnflattened

Get a zero state unflattened object.

Returns:

Type	Description
StateUnflattened	StateUnflattened object with zero arrays.

Source code in src/flapjax/utils/linear.py

def get_zero_state(self) -> StateUnflattened:
    r"""
    Get a zero state unflattened object.
    :return: StateUnflattened object with zero arrays.
    """
    return self.state_object(**self._get_zero(shallow_as_dict(self.state_slices)))

get_zero_output

get_zero_output() -> OutputUnflattened

Get a zero output unflattened object.

Returns:

Type	Description
OutputUnflattened	OutputUnflattened object with zero arrays.

Source code in src/flapjax/utils/linear.py

def get_zero_output(self) -> OutputUnflattened:
    r"""
    Get a zero output unflattened object.
    :return: OutputUnflattened object with zero arrays.
    """
    return self.output_object(**self._get_zero(shallow_as_dict(self.output_slices)))

nodal_to_modal

nodal_to_modal(q_nodal: Array) -> Array

Convert a nodal property to a modal property.

Parameters:

Name	Type	Description	Default
q_nodal	Array	Nodal property, (n_free_dof, ).	required

Returns:

Type	Description
Array	Modal property, (n_modes, ).

Source code in src/flapjax/structure/linear/linear_beam.py

def nodal_to_modal(self, q_nodal: Array) -> Array:
    r"""
    Convert a nodal property to a modal property.
    :param q_nodal: Nodal property, ``(n_free_dof, )``.
    :return: Modal property, ``(n_modes, )``.
    """

    return self.mode_shapes @ q_nodal

modal_to_nodal

modal_to_nodal(q_modal: Array) -> Array

Convert a modal property to a nodal property.

Parameters:

Name	Type	Description	Default
q_modal	Array	Mode property, (n_modes, ).	required

Returns:

Type	Description
Array	Nodal property, (n_free_dofs, ).

Source code in src/flapjax/structure/linear/linear_beam.py

def modal_to_nodal(self, q_modal: Array) -> Array:
    r"""
    Convert a modal property to a nodal property.
    :param q_modal: Mode property, ``(n_modes, )``.
    :return: Nodal property, ``(n_free_dofs, )``.
    """

    return q_modal @ self.mode_shapes

linearise_continuous

linearise_continuous() -> LinearSystem

Form a system of linear equations about a reference state. The system is of the form: :math:\dot{\mathbf{x}} = \mathbf{A~x + B~u}, \mathbf{y} = \mathbf{C~x + D~u}.

The state, input and output layouts each depend on their respective modal_* flag:

• States: nodal-free-dof :math:[q, \dot q] of length 2 * n_free_dof when n_modes is None, or modal :math:[q_m, \dot q_m] of length 2 * n_modes when modal_states. Modal mass and stiffness are the projected :math:\Phi M \Phi^T / :math:\Phi K \Phi^T with Phi = mode_shapes of shape [n_modes, n_free_dof] (rows are mode shapes).
• Inputs: a single global-frame external force f_ext. Non-modal inputs are per-node [n_nodes, 6]. Modal inputs are direct modal forces of length n_modes, requiring the user to have already applied the modal projection.
• Outputs: :math:[q, \dot q] in either nodal-free-dof or modal form to match modal_outputs.

Returns:

Type	Description
LinearSystem	Linearised continuous-time system.

Source code in src/flapjax/structure/linear/linear_beam.py

def linearise_continuous(self) -> LinearSystem:
    r"""
    Form a system of linear equations about a reference state. The system is of the form:
    :math:`\dot{\mathbf{x}} = \mathbf{A~x + B~u}, \mathbf{y} = \mathbf{C~x + D~u}`.

    The state, input and output layouts each depend on their respective ``modal_*`` flag:

    * States: nodal-free-dof :math:`[q, \dot q]` of length ``2 * n_free_dof`` when ``n_modes is None``, or modal
      :math:`[q_m, \dot q_m]` of length ``2 * n_modes`` when ``modal_states``. Modal mass and stiffness are the
      projected :math:`\Phi M \Phi^T` / :math:`\Phi K \Phi^T` with ``Phi = mode_shapes`` of shape
      ``[n_modes, n_free_dof]`` (rows are mode shapes).
    * Inputs: a single global-frame external force ``f_ext``. Non-modal inputs are per-node ``[n_nodes, 6]``. Modal
      inputs are direct modal forces of length ``n_modes``, requiring the user to have already applied the modal
      projection.
    * Outputs: :math:`[q, \dot q]` in either nodal-free-dof or modal form to match ``modal_outputs``.

    :return: Linearised continuous-time system.
    """

    # single LU factorisation of M reused for every M^{-1} product below
    m_lu = lu_factor(self.m)  # [q_state_size, q_state_size]
    q_state_size = self.n_modes if self.modal_states else self.n_free_dof

    # dynamics matrix: [q_dot; q_ddot] = A [q; q_dot]
    # includes Raleigh damping contribution if enabled
    m_inv_c = (
        lu_solve(m_lu, self.c)
        if self.c is not None
        else jnp.zeros((q_state_size, q_state_size))
    )
    a = jnp.block(
        [
            [jnp.zeros((q_state_size, q_state_size)), jnp.eye(q_state_size)],
            [-lu_solve(m_lu, self.k), -m_inv_c],
        ]
    )

    # input matrix: maps a single global-frame external force to state derivative
    if self.modal_inputs:
        assert self.modal_states, "modal_inputs requires modal_states"
        # user supplies modal forces directly
        b_bottom = lu_solve(m_lu, jnp.eye(q_state_size))  # (n_modes, n_modes)
    else:
        p_free = jnp.eye(self.n_nodes * 6)[
            jnp.array(self.free_dofs), :
        ]  # (n_free_dof, n_nodes * 6)

        if self.modal_states:
            nodal_to_state = self.nodal_to_modal(p_free)  # (n_modes, n_nodes * 6)
        else:
            nodal_to_state = p_free  # (n_free_dof, n_nodes * 6)
        b_bottom = lu_solve(
            m_lu, nodal_to_state
        )  # (n_modes | n_free_dof, n_nodes * 6)

    b = jnp.concatenate(
        [jnp.zeros((q_state_size, b_bottom.shape[1])), b_bottom], axis=0
    )

    if self.modal_states and not self.modal_outputs:
        assert self.mode_shapes is not None
        state_to_output = (
            self.mode_shapes.T
        )  # projects force to modes, (n_free_dof, n_modes)
    else:
        state_to_output = jnp.eye(
            q_state_size
        )  # direct projection, (n_free_dof, n_free_dof) or (n_modes, n_modes)

    c = jnp.block(
        [
            [state_to_output, jnp.zeros_like(state_to_output)],
            [jnp.zeros_like(state_to_output), state_to_output],
        ]
    )  # pass state (q, q_dot) to output (q, q_dot)
    d = jnp.zeros((c.shape[0], b.shape[1]))  # no feedthrough

    return LinearSystem(
        a=a, b=b, c=c, d=d, dt=self.dt, continuous_time=True, removed_u_np1=False
    )

run

run(
    u: StructureInputUnflattened,
    x0: StructureStateUnflattened | None = None,
) -> StructureLinearResult

Run the linear system.

Parameters:

Name	Type	Description	Default
u	StructureInputUnflattened	Total input over time (reference + pertubation).	required
x0	StructureStateUnflattened | None	Initial state perturbations, defaults to zero state.	None

Returns:

Type	Description
StructureLinearResult	Linear system results.

Source code in src/flapjax/structure/linear/linear_beam.py

def run(
    self,
    u: StructureInputUnflattened,
    x0: StructureStateUnflattened | None = None,
) -> StructureLinearResult:
    r"""
    Run the linear system.
    :param u: Total input over time (reference + pertubation).
    :param x0: Initial state perturbations, defaults to zero state.
    :return: Linear system results.
    """

    assert (
        isinstance(self.reference, StructureCase) and not self.reference.is_dynamic
    ), "Reference structure must be a static Structure"

    # has to be specified in the input object, as the linear system does not know how many time steps to run for in
    # the case where there is no external forcing
    n_tstep = u.n_tstep

    q_state_size = self.n_modes if self.modal_states else self.n_free_dof

    # initial states
    if x0 is None:
        q0_ = None
        q0_dot_ = None
    else:
        q0_ = x0.q
        q0_dot_ = x0.q_dot

    def _project_initial(q0_arr: Array | None) -> Array:
        if q0_arr is None:
            return jnp.zeros((q_state_size,))

        match q0_arr.shape:
            case (self.n_nodes, 6):
                q0_nodal = q0_arr.ravel()[self.free_dofs]  # remove prescribed dofs
                if self.modal_states:
                    q0__ = self.nodal_to_modal(q0_nodal)
                else:
                    q0__ = q0_nodal
            case (self.n_modes,):
                if not self.modal_states:
                    raise NotImplementedError(
                        "Inputs cannot be modal for a nodal system"
                    )
                q0__ = q0_arr
            case _:
                raise ValueError("Invalid input for initial states")
        return q0__

    q0 = _project_initial(q0_)  # (n_free_dof | n_modes, )
    q0_dot = _project_initial(q0_dot_)  # (n_free_dof | n_modes, )

    if q0.shape != q0_dot.shape:
        raise ValueError("q0 and q0_dot must both be modal or nodal")

    ref_f_ext = self._reference_inputs["f_ext"]
    assert ref_f_ext is None or isinstance(
        ref_f_ext, Array
    )  # type narrowing as it cannot be an ArrayList

    if self.modal_inputs:
        # user-provided input is a modal force time history of shape (n_tstep, n_modes)
        if u.f_ext is None:
            delta_f_ext_t = (
                jnp.zeros((n_tstep, self.n_modes))
                if ref_f_ext is None
                else -jnp.broadcast_to(ref_f_ext[None, :], (n_tstep, self.n_modes))
            )
        else:
            check_arr_shape(u.f_ext, (n_tstep, self.n_modes), name="f_ext_t")
            delta_f_ext_t = (
                u.f_ext if ref_f_ext is None else u.f_ext - ref_f_ext[None, :]
            )
    else:
        # user-provided input is a nodal global-frame force time history of shape (n_tstep, n_nodes, 6)
        if u.f_ext is None:
            delta_f_ext_t = (
                jnp.zeros((n_tstep, self.n_nodes, 6))
                if ref_f_ext is None
                else -jnp.broadcast_to(
                    ref_f_ext[None, :, :], (n_tstep, self.n_nodes, 6)
                )
            )
        else:
            check_arr_shape(u.f_ext, (n_tstep, self.n_nodes, 6), name="f_ext_t")
            delta_f_ext_t = (
                u.f_ext if ref_f_ext is None else u.f_ext - ref_f_ext[None, :, :]
            )

    delta_u = StructureInputUnflattened(n_tstep=n_tstep, f_ext=delta_f_ext_t)
    delta_u_vec = self._pack_input_vector_t(delta_u)

    # run linear system
    x_t, _ = self.sys.run(
        u=delta_u_vec,
        x0=jnp.concatenate((q0, q0_dot)),
    )

    # extract perturbations in displacements and velocities, reconstructing nodal form if requested
    delta_q_state = x_t[:, :q_state_size]
    delta_q_dot_state = x_t[:, q_state_size:]

    if self.modal_outputs:
        delta_q_t = delta_q_state
        delta_q_dot_t = delta_q_dot_state
        hg_t = None  # do not reconstruct coordinates
    else:
        if self.modal_states:
            assert self.mode_shapes is not None
            delta_q_t = self.modal_to_nodal(delta_q_state)  # (n_tstep, n_free_dof)
            delta_q_dot_t = self.modal_to_nodal(delta_q_dot_state)
        else:
            delta_q_t = delta_q_state
            delta_q_dot_t = delta_q_dot_state

        # add in zeros for dofs not solved for to allow for reconstructing the full configuration
        delta_q_t_full = (
            jnp.zeros((n_tstep, self.n_nodes * 6))
            .at[:, self.free_dofs]
            .set(delta_q_t)
            .reshape(n_tstep, self.n_nodes, 6)
        )

        delta_hg_t = vmap(vmap(exp_se3, 0, 0), 1, 1)(
            delta_q_t_full
        )  # (n_tstep, n_nodes, 4, 4)

        hg_t = jnp.einsum("ijk,hikl->hijl", self.reference.hg, delta_hg_t)

    return StructureLinearResult(
        reference=self.reference,
        f_ext=u.f_ext,
        delta_q=delta_q_t,
        delta_q_dot=delta_q_dot_t,
        hg=hg_t,
        t=jnp.arange(n_tstep) * self.dt,
    )

beam

BaseBeamStructure

BaseBeamStructure(
    num_nodes: int,
    connectivity: Array,
    y_vector: Array,
    k_cs_index: Array | None = None,
    m_cs_index: Array | None = None,
    m_lumped_index: Array | None = None,
    gravity: Array | Sequence[float] | None = None,
    thrust_nodes: dict[str, int] | None = None,
    thrust_direction: dict[str, Array] | None = None,
    optional_jacobians: OptionalJacobians | None = None,
    relaxation_factor: float = 1.0,
    spectral_radius: float = 0.9,
    alpha_m: float = 0.0,
    beta_k: float = 0.0,
    struct_convergence_settings: ConvergenceSettings = DEFAULT_STRUCT_CONVERGENCE_SETTINGS,
    constraints: dict[str, SoftConstraint | HardConstraint]
    | None = None,
)

Class to represent nonlinear beam structure model

Initialise BaseBeamStructure class with all non-design parameters.

Parameters:

Name	Type	Description	Default
num_nodes	int	Number of nodes in the structure.	required
connectivity	Array	Connectivity array, `(n_elem, 2)``.	required
y_vector	Array	Vector defining the y direction for each element, (n_elem, 3).	required
k_cs_index	Array | None	Array defining the index from the library of k_cs to use for each element, (n_elem, ). If None, all elements will use the first entry in the k_cs library.	None
m_cs_index	Array | None	Array defining the index from the library of m_cs to use for each element, (n_elem, ). If None, all elements will use the first entry in the m_cs library.	None
m_lumped_index	Array | None	Node index for nodes which are to have a lumped mass attached. The order is the same as that for the lumped mass data (n_lumped_mass, ).	None
gravity	Array | Sequence[float] | None	Gravity vector in global reference frame, or None for no gravity_vec, (3, ).	None
thrust_nodes	dict[str, int] | None	Dictionary of thrust node names and their corresponding node indices, {keys, int}.	None
thrust_direction	dict[str, Array] | None	Dictionary of thrust node names and their corresponding thrust direction vectors, {keys, (3, )}.	None
optional_jacobians	OptionalJacobians | None	Define which Jacobians contributions are to be used for solution.	None
relaxation_factor	float	Relaxation factor which reduces the displacement update at each iteration. A value of 1 is no relaxation, and a value of 0 is no update.	1.0
spectral_radius	float	Spectral radius for structural time integrator, where a value of 0 is highly damped and a value of 1 is undamped.	0.9
alpha_m	float	Mass-proportional Rayleigh damping coefficient.	0.0
beta_k	float	Stiffness-proportional Rayleigh damping coefficient.	0.0
struct_convergence_settings	ConvergenceSettings	Structure convergence settings.	DEFAULT_STRUCT_CONVERGENCE_SETTINGS
constraints	dict[str, SoftConstraint | HardConstraint] | None	Named dict {name: constraint}, or None, which add	None

Source code in src/flapjax/structure/beam.py

def __init__(
    self,
    num_nodes: int,
    connectivity: Array,
    y_vector: Array,
    k_cs_index: Array | None = None,
    m_cs_index: Array | None = None,
    m_lumped_index: Array | None = None,
    gravity: Array | Sequence[float] | None = None,
    thrust_nodes: dict[str, int] | None = None,
    thrust_direction: dict[str, Array] | None = None,
    optional_jacobians: OptionalJacobians | None = None,
    relaxation_factor: float = 1.0,
    spectral_radius: float = 0.9,
    alpha_m: float = 0.0,
    beta_k: float = 0.0,
    struct_convergence_settings: ConvergenceSettings = DEFAULT_STRUCT_CONVERGENCE_SETTINGS,
    constraints: (dict[str, SoftConstraint | HardConstraint] | None) = None,
) -> None:
    r"""
    Initialise BaseBeamStructure class with all non-design parameters.
    :param num_nodes: Number of nodes in the structure.
    :param connectivity: Connectivity array, `(n_elem, 2)``.
    :param y_vector: Vector defining the y direction for each element, ``(n_elem, 3)``.
    :param k_cs_index: Array defining the index from the library of k_cs to use for each element, ``(n_elem, )``.
    If ``None``, all elements will use the first entry in the k_cs library.
    :param m_cs_index: Array defining the index from the library of m_cs to use for each element, ``(n_elem, )``.
    If ``None``, all elements will use the first entry in the m_cs library.
    :param m_lumped_index: Node index for nodes which are to have a lumped mass attached. The order is the same as
    that for the lumped mass data ``(n_lumped_mass, )``.
    :param gravity: Gravity vector in global reference frame, or None for no gravity_vec, ``(3, )``.
    :param thrust_nodes: Dictionary of thrust node names and their corresponding node indices, {keys, int}.
    :param thrust_direction: Dictionary of thrust node names and their corresponding thrust direction vectors,
    ``{keys, (3, )}``.
    :param optional_jacobians: Define which Jacobians contributions are to be used for solution.
    :param relaxation_factor: Relaxation factor which reduces the displacement update at each iteration. A value of
    1 is no relaxation, and a value of 0 is no update.
    :param spectral_radius: Spectral radius for structural time integrator, where a value of 0 is highly damped and
    a value of 1 is undamped.
    :param alpha_m: Mass-proportional Rayleigh damping coefficient.
    :param beta_k: Stiffness-proportional Rayleigh damping coefficient.
    :param struct_convergence_settings: Structure convergence settings.
    :param constraints: Named dict ``{name: constraint}``, or None, which add
    """

    check_type(num_nodes, int)

    if constraints is None:
        named_constraints: dict[str, SoftConstraint | HardConstraint] = {}
    elif isinstance(constraints, dict):
        named_constraints = constraints
    else:
        raise ValueError("Invalid constraint input")

    hard_constraints = tuple(
        c for c in named_constraints.values() if isinstance(c, HardConstraint)
    )

    check_arr_shape(connectivity, (None, 2), "connectivity")
    check_arr_dtype(connectivity, int, "connectivity")
    _check_connectivity(connectivity, num_nodes)

    auto_node_sources: list[int] = []
    conn_list: list[list[int]] = connectivity.tolist()

    # add extra nodes and alter connectivity to account for constraints with the auto-generate node behaviour
    for con in hard_constraints:
        if con.node_j is None and not con.is_grounded:
            node_j_new = num_nodes + len(auto_node_sources)
            con.node_j = node_j_new
            auto_node_sources.append(con.node_i)
            conn_list = _split_connectivity(conn_list, con.node_i, node_j_new)
    self._auto_node_sources: tuple[int, ...] = tuple(auto_node_sources)

    num_nodes += len(auto_node_sources)
    connectivity = jnp.array(conn_list, dtype=int)
    self.n_nodes: int = num_nodes
    self.n_dof: int = num_nodes * 6

    self.connectivity: tuple[tuple[int, int], ...] = nested_list_to_tuple(
        connectivity.tolist()
    )  # (n_elem, 2)
    self.n_elem_per_node: tuple[int] = tuple(
        _n_elem_per_node(connectivity=connectivity, n_nodes=num_nodes).tolist()
    )  # (n_nodes, )
    self.n_elem: int = connectivity.shape[0]

    self.dof_per_elem: tuple[tuple[float]] = nested_list_to_tuple(
        jnp.zeros((self.n_elem, 12), dtype=int)
        .at[:, :6]
        .set(6 * self.connectivity_arr[:, [0]] + jnp.arange(6)[None, :])
        .at[:, 6:]
        .set(6 * self.connectivity_arr[:, [1]] + jnp.arange(6)[None, :])
        .tolist()
    )

    # allow for a single y_vector to be broadcast to all elements
    if y_vector.shape == (3,):
        y_vector = y_vector[None, :]
    if y_vector.shape == (1, 3):
        y_vector = jnp.broadcast_to(y_vector, (self.n_elem, 3))

    # y vectors in reference unoriented configuration, and placeholder for oriented equivalent.
    check_arr_shape(y_vector, (self.n_elem, 3), "y_vector")
    self.y_vector_reference: tuple[tuple[tuple[float]]] = nested_list_to_tuple(
        y_vector.tolist()
    )
    self.y_vector: Array = jnp.zeros_like(jnp.array(y_vector))

    # initialise design variables with default values
    self.x0_reference: Array = jnp.zeros((num_nodes, 3))  # unoriented
    self.x0: Array = jnp.zeros((num_nodes, 3))  # oriented

    self.m_cs = None
    self.k_cs = None
    self.m_lumped = None
    self.use_lumped_mass: bool = m_lumped_index is not None

    # initialise auxiliary arrays
    self.o0: Array = jnp.zeros((self.n_elem, 3, 3))
    self.l0: Array = jnp.zeros(self.n_elem)
    self.d0: Array = jnp.zeros((self.n_elem, 6))

    # initialise undeformed algebra and group
    self.hg0_reference: Array = jnp.zeros((self.n_nodes, 4, 4))  # unoriented
    self.hg0: Array = jnp.zeros((self.n_nodes, 4, 4))  # oriented

    # grads inverse action for the reference rotations
    self.ad_inv_o0: Array = jnp.zeros((self.n_elem, 6, 6))

    # gravity settings
    if not isinstance(gravity, jnp.ndarray) and gravity is not None:
        gravity = jnp.array(gravity)
    self.use_gravity: bool = gravity is not None and bool(jnp.any(gravity))
    if self.use_gravity:
        assert gravity is not None
        check_arr_shape(gravity, (3,), "gravity")
        self.gravity_vec: tuple[float, float, float] = tuple(gravity.tolist())
    else:
        self.gravity_vec = (0.0, 0.0, 0.0)

    # indexing
    if k_cs_index is None:
        k_cs_index_ = jnp.zeros(self.n_elem, dtype=int)
    else:
        check_arr_shape(k_cs_index, (self.n_elem,), "k_cs_index")
        check_arr_dtype(k_cs_index, int, "k_cs_index")
        k_cs_index_ = k_cs_index
    self.k_cs_index: tuple[int] = tuple(k_cs_index_.tolist())

    if m_cs_index is None:
        m_cs_index_ = jnp.zeros(self.n_elem, dtype=int)
    else:
        check_arr_shape(m_cs_index, (self.n_elem,), "m_cs_index")
        check_arr_dtype(m_cs_index, int, "m_cs_index")
        m_cs_index_ = m_cs_index
    self.m_cs_index: tuple[int] = tuple(m_cs_index_.tolist())

    self.m_lumped_index: tuple[int] | None = None
    if m_lumped_index is not None:
        check_arr_dtype(m_lumped_index, int, "m_lumped_index")
        if m_lumped_index.ndim not in (0, 1):
            raise ValueError("m_lumped_index.ndim must be 0 or 1.")
        self.m_lumped_index = tuple(jnp.atleast_1d(m_lumped_index).tolist())

    # add thrust
    self.thrust_nodes: tuple[tuple[str, int], ...] = ()
    self.thrust_direction: tuple[tuple[str, tuple[float, float, float]], ...] = ()
    if thrust_nodes is not None and thrust_direction is not None:
        if thrust_nodes.keys() != thrust_direction.keys():
            raise ValueError(
                f"Mismatch in keys of thrust_nodes ({thrust_nodes.keys()}) and thrust_direction ({thrust_direction.keys()}))."
            )

        for k, v in thrust_direction.items():
            check_arr_shape(v, (3,), f"thrust_direction[{k}]")

        self.thrust_nodes = tuple([(k, v) for k, v in thrust_nodes.items()])
        self.thrust_direction = tuple(
            [
                (k, nested_list_to_tuple((v / jnp.linalg.norm(v)).tolist()))
                for k, v in thrust_direction.items()
            ]
        )  # make unit vectors
    elif thrust_nodes is not None or thrust_direction is not None:
        warn(
            "One of thrust_nodes or thrust_direction has not been passed. Running with no thrust nodes."
        )

    # set the reference thrust to be zero, which can be overwritten later
    self.thrust_reference: dict[str, Array] = {
        k: jnp.atleast_1d(1) for k in [k_ for k_, v in self.thrust_nodes]
    }

    # set the reference orientation, which can be overwritten later.
    self.orientation_euler: Array = jnp.zeros(3)
    self.orientation: Array = jnp.eye(3)

    self.optional_jacobians: OptionalJacobians = (
        optional_jacobians
        if optional_jacobians is not None
        else OptionalJacobians()
    )
    self.struct_convergence_settings: ConvergenceSettings = (
        struct_convergence_settings
    )
    self.relaxation_factor: float = relaxation_factor
    self.spectral_radius: float = spectral_radius
    self.alpha_m: float = float(alpha_m)
    self.beta_k: float = float(beta_k)

    self.time_integrator = None

    self.constraints: dict[str, SoftConstraint | HardConstraint] = named_constraints

n_multibody_constraints property

n_multibody_constraints: int

Total number of scalar Lagrange-multiplier constraints.

n_holonomic_constraints property

n_holonomic_constraints: int

Number of scalar holonomic (position-level) Lagrange-multiplier constraints.

n_nonholonomic_constraints property

n_nonholonomic_constraints: int

Number of scalar non-holonomic (velocity-level) Lagrange-multiplier constraints.

set_design_variables

set_design_variables(
    coords: Array,
    k_cs: Array,
    m_cs: Array | None,
    m_lumped: Array | None = None,
    orientation_euler: Array | None = None,
    thrust_reference: dict[str, Array | float]
    | None = None,
    *,
    remove_checks: bool = False,
) -> None

Set design variables and compute initial configuration dependent quantities.

Parameters:

Name	Type	Description	Default
coords	Array	Node coordinates in the reference configuration, (n_nodes, 3).	required
k_cs	Array	Cross-section stiffness matrices, (n_entry, 6, 6) or (6, 6).	required
m_cs	Array | None	Cross-section mass matrices, (n_entry, 6, 6) or (6, 6).	required
m_lumped	Array | None	Lumped mass matrices at nodes, (n_entry, 6, 6).	None
orientation_euler	Array | None	Euler angles in radians which to rotate the reference configuration by, (3, ). This rotation is performed about the origin, and will default to the identity is no Array is passed. These are rotated in z-y-x order.	None
thrust_reference	dict[str, Array | float] | None	Reference thrust magnitude, {keys, (1, )}.	None
remove_checks	bool	Flag to ignore input checks, used when function is JIT compiled.	False

Source code in src/flapjax/structure/beam.py

def set_design_variables(
    self,
    coords: Array,
    k_cs: Array,
    m_cs: Array | None,
    m_lumped: Array | None = None,
    orientation_euler: Array | None = None,
    thrust_reference: dict[str, Array | float] | None = None,
    *,
    remove_checks: bool = False,
) -> None:
    r"""
    Set design variables and compute initial configuration dependent quantities.
    :param coords: Node coordinates in the reference configuration, ``(n_nodes, 3)``.
    :param k_cs: Cross-section stiffness matrices, ``(n_entry, 6, 6)`` or ``(6, 6)``.
    :param m_cs: Cross-section mass matrices, ``(n_entry, 6, 6)`` or ``(6, 6)``.
    :param m_lumped: Lumped mass matrices at nodes, ``(n_entry, 6, 6)``.
    :param orientation_euler: Euler angles in radians which to rotate the reference configuration by, ``(3, )``. This rotation
    is performed about the origin, and will default to the identity is no Array is passed. These are rotated in
    z-y-x order.
    :param thrust_reference: Reference thrust magnitude, ``{keys, (1, )}``.
    :param remove_checks: Flag to ignore input checks, used when function is JIT compiled.
    """

    # orientation
    if orientation_euler is not None:
        check_arr_shape(orientation_euler, (3,), "orientation_euler")
        self.orientation_euler = orientation_euler
        self.orientation = Rotation.from_euler(
            seq="zyx", angles=orientation_euler
        ).as_matrix()

    # rotate the y vectors
    self.y_vector = jnp.einsum(
        "jk,ik->ij",
        self.orientation,
        self.y_vector_reference_arr,
    )

    # coordinates — auto-extend for nodes created by multibody constraints
    if self._auto_node_sources and coords.shape[0] == self.n_nodes - len(
        self._auto_node_sources
    ):
        coords = jnp.concatenate(
            [coords, coords[jnp.array(self._auto_node_sources)]],
            axis=0,
        )
    check_arr_shape(coords, (self.n_nodes, 3), "coords")
    self.x0_reference = coords
    self.x0 = jnp.einsum("jk,ik->ij", self.orientation, coords)

    # populate arrays
    if k_cs.ndim == 2:
        k_cs = k_cs[None, ...]
    check_arr_shape(k_cs, (None, 6, 6), "k_cs")

    if (
        not remove_checks
        and k_cs.shape[0] != jnp.unique_values(jnp.array(self.k_cs_index)).size
    ):
        warn(
            "Redundant values in k_cs which are not used for solution due to no corresponding entry in k_cs_index."
        )

    self.k_cs = k_cs
    if m_cs is None:
        if not remove_checks and self.use_gravity and m_lumped is None:
            warn(
                "No mass matrices provided, but gravity is enabled. Assuming zero mass.",
            )
        m_cs_ = jnp.zeros((6, 6))
    else:
        m_cs_ = m_cs

    if m_cs_.ndim == 2:
        m_cs_ = m_cs_[None, ...]

    check_arr_shape(m_cs_, (None, 6, 6), "m_cs")

    if (
        not remove_checks
        and m_cs_.shape[0] != jnp.unique_values(jnp.array(self.m_cs_index)).size
        and m_cs is not None
    ):
        warn(
            "Redundant values in m_cs which are not used for solution due to no corresponding entry in "
            "m_cs_index."
        )

    self.m_cs = m_cs_

    # thrust
    if thrust_reference is not None:
        self.thrust_reference = {
            k: jnp.atleast_1d(v) for k, v in thrust_reference.items()
        }

        for k, v in self.thrust_reference.items():
            check_arr_shape(v, (1,), f"thrust_reference[{k}]")

    if m_lumped is not None:
        if not remove_checks:
            check_arr_shape(m_lumped, (None, 6, 6), "m_lumped")

            if self.m_lumped_index is None:
                raise ValueError("m_lumped_index has not been set")

            if m_lumped.shape[0] != len(self.m_lumped_index):
                raise ValueError(
                    "Number of entries in m_lumped does not match number of indices in m_lumped_index."
                )

        self.m_lumped = m_lumped

    # obtain initial orientation and length
    x_elem = jnp.take(
        self.x0_reference, self.connectivity_arr, axis=0
    )  # (n_elem, 2, 3)
    dx = x_elem[:, 1, :] - x_elem[:, 0, :]  # (n_elem, 3)

    # ensure out-of-plane vector and beam vector are not collinear
    if not remove_checks and jnp.any(
        jnp.linalg.norm(jnp.cross(dx, self.y_vector_reference_arr, 1, 1), axis=-1)
        < 1e-6
    ):
        raise ValueError(
            "y_vector is collinear with beam element direction for at least one element. "
            "Please provide a different y_vector."
        )

    self.l0 = jnp.linalg.norm(dx, axis=-1)  # (n_elem,)
    self.d0 = self.d0.at[:, 0].set(self.l0)

    dx_unit = dx / self.l0[:, None]  # unit vector in beam direction, (n_elem, 3)
    dz = jnp.cross(
        dx_unit, self.y_vector_reference_arr, axis=-1
    )  # vector in plane(n_elem, 3)
    dz_unit = dz / jnp.linalg.norm(dz, axis=-1)[:, None]  # (n_elem, 3)

    dy_unit = jnp.cross(dz_unit, dx_unit)

    self.o0 = self.o0.at[..., 0].set(dx_unit)
    self.o0 = self.o0.at[..., 1].set(dy_unit)
    self.o0 = self.o0.at[..., 2].set(dz_unit)

    self.ad_inv_o0 = vmap(rmat_to_ha_hat)(jnp.transpose(self.o0, (0, 2, 1)))

    # set unoriented initial coordinates
    self.hg0_reference = jnp.broadcast_to(
        jnp.eye(4)[None, ...], (self.n_nodes, 4, 4)
    )  # (n_nodes, 4, 4)
    self.hg0_reference = self.hg0_reference.at[:, :3, 3].set(self.x0_reference)

    # set oriented initial coordinates
    self.hg0 = self.hg0.at[:, :3, :3].set(
        jnp.broadcast_to(self.orientation[None, ...], (self.n_nodes, 3, 3))
    )  # (n_nodes, 4, 4)
    self.hg0 = self.hg0.at[:, :3, 3].set(self.x0)
    self.hg0 = self.hg0.at[:, 3, 3].set(1.0)

    # add reference frames to the nodal constraints
    for con in self.nodal_constraints:
        con.resolve_hg_ref(self.hg0)
    for con in self.multibody_constraints:
        con.resolve_hg_ref(self.hg0)

get_design_variables

get_design_variables(
    struct_case: StructureCase,
    thrust_t: dict[str, Array],
    grads_to_compute: StructureGradsToCompute | None,
) -> StructureDesignVariables

Obtain the design variables for the structural problem. As the external forcing is defined for each solve, the chosen forcing is required as input.

Parameters:

Name	Type	Description	Default
struct_case	StructureCase	Structural case	required
thrust_t	dict[str, Array]	Thrust time history, {keys, (n_tstep,)}.	required
grads_to_compute	StructureGradsToCompute | None	Data structure which describes which design variables should be obtained. If none, all variables are obtained.	required

Returns:

Type	Description
StructureDesignVariables	StructureDesignVariables dataclass containing design variables

Source code in src/flapjax/structure/beam.py

def get_design_variables(
    self,
    struct_case: StructureCase,
    thrust_t: dict[str, Array],
    grads_to_compute: StructureGradsToCompute | None,
) -> StructureDesignVariables:
    r"""
    Obtain the design variables for the structural problem. As the external forcing is defined for each solve, the
    chosen forcing is required as input.
    :param struct_case: Structural case
    :param thrust_t: Thrust time history, {keys, ``(n_tstep,)``}.
    :param grads_to_compute: Data structure which describes which design variables should be obtained. If none, all
    variables are obtained.
    :return: StructureDesignVariables dataclass containing design variables
    """

    # struct_case.f_ext_dead is stored in local frame: f_local = R^T @ f_global,
    # so recover f_global = R @ f_local
    hg = struct_case.hg
    if hg.ndim == 4:  # batched case: (n_tstep, n_nodes, 4, 4)
        rmat = hg[:, :, :3, :3]
    else:  # snapshot: (n_nodes, 4, 4)
        rmat = hg[:, :3, :3]
    f_ext_dead_global = (
        transform_nodal_vect(struct_case.f_ext_dead, rmat)
        if struct_case.f_ext_dead is not None
        else None
    )
    if isinstance(grads_to_compute, StructureGradsToCompute):
        return StructureDesignVariables(
            x0=self.x0 if grads_to_compute.x0 else None,
            orientation_euler=self.orientation_euler
            if grads_to_compute.orientation_euler
            else None,
            m_cs=self.m_cs if grads_to_compute.m_cs else None,
            k_cs=self.k_cs if grads_to_compute.k_cs else None,
            m_lumped=self._m_lumped if grads_to_compute.m_lumped else None,
            f_ext_dead=f_ext_dead_global if grads_to_compute.f_ext_dead else None,
            f_ext_follower=struct_case.f_ext_follower
            if grads_to_compute.f_ext_follower
            else None,
            thrust_t=thrust_t if grads_to_compute.thrust_t else None,
            f_shape=(),
        )
    else:
        return StructureDesignVariables(
            x0=self.x0,
            orientation_euler=self.orientation_euler,
            m_cs=self.m_cs,
            k_cs=self.k_cs,
            m_lumped=self._m_lumped,
            f_ext_dead=f_ext_dead_global,
            f_ext_follower=struct_case.f_ext_follower,
            thrust_t=thrust_t,
            f_shape=(),
        )

reference_configuration

reference_configuration(
    prescribed_dofs: Sequence[int]
    | Array
    | slice
    | int = (),
    use_f_ext_follower: bool = True,
    use_f_ext_dead: bool = True,
    use_f_aero: bool = True,
    use_f_grav: bool = True,
) -> StructureCase

Get the reference configuration of the structure.

Parameters:

Name	Type	Description	Default
prescribed_dofs	Sequence[int] | Array | slice | int	Prescribed degrees of freedom, which are not solved for. Defaults to no prescribed DoFs.	()
use_f_ext_follower	bool	Whether to include follower forces in the reference configuration.	True
use_f_ext_dead	bool	Whether to include dead forces in the reference configuration.	True
use_f_aero	bool	Whether to include aerodynamic forces in the reference configuration.	True
use_f_grav	bool	Whether to include gravitational forces in the reference configuration.	True

Returns:

Type	Description
StructureCase	Structure dataclass containing reference configuration.

Source code in src/flapjax/structure/beam.py

def reference_configuration(
    self,
    prescribed_dofs: Sequence[int] | Array | slice | int = (),
    use_f_ext_follower: bool = True,
    use_f_ext_dead: bool = True,
    use_f_aero: bool = True,
    use_f_grav: bool = True,
) -> StructureCase:
    r"""
    Get the reference configuration of the structure.
    :param prescribed_dofs: Prescribed degrees of freedom, which are not solved for. Defaults to no prescribed DoFs.
    :param use_f_ext_follower: Whether to include follower forces in the reference configuration.
    :param use_f_ext_dead: Whether to include dead forces in the reference configuration.
    :param use_f_aero: Whether to include aerodynamic forces in the reference configuration.
    :param use_f_grav: Whether to include gravitational forces in the reference configuration.
    :return: Structure dataclass containing reference configuration.
    """
    prescribed_dofs = self.make_prescribed_dofs_tuple(prescribed_dofs)
    return StructureCase(
        hg=self.hg0,
        conn=self.connectivity,
        o0=self.o0,
        d=self.d0,
        eps=jnp.zeros((self.n_elem, 6)),
        varphi=jnp.zeros((self.n_nodes, 6)),
        f_ext_follower=jnp.zeros((self.n_nodes, 6)) if use_f_ext_follower else None,
        f_ext_dead=jnp.zeros((self.n_nodes, 6)) if use_f_ext_dead else None,
        f_ext_aero=jnp.zeros((self.n_nodes, 6)) if use_f_aero else None,
        f_grav=jnp.zeros((self.n_nodes, 6)) if use_f_grav else None,
        f_int=jnp.zeros((self.n_nodes, 6)),
        f_elem=jnp.zeros((self.n_elem, 6)),
        f_res=jnp.zeros((self.n_nodes, 6)),
        thrust=self.thrust_reference,
        thrust_direction=self.thrust_direction,
        thrust_nodes=self.thrust_nodes,
        local=True,
        prescribed_dofs=prescribed_dofs,
        t=jnp.zeros(1),
    )

compute_varphi_from_hg

compute_varphi_from_hg(hg: Array) -> Array

Calculate the twist vector from the reference configuration to hg

Parameters:

Name	Type	Description	Default
hg	Array	Deformed coordinates, (n_nodes, 4, 4)	required

Returns:

Type	Description
Array	Vector of twists, (n_nodes, 6)

Source code in src/flapjax/structure/beam.py

def compute_varphi_from_hg(self, hg: Array) -> Array:
    r"""
    Calculate the twist vector from the reference configuration to hg
    :param hg: Deformed coordinates, ``(n_nodes, 4, 4)``
    :return: Vector of twists, ``(n_nodes, 6)``
    """
    return vmap(hg_to_d, (0, 0), 0)(self.hg0, hg)

assemble_matrix_from_entries

assemble_matrix_from_entries(entries: Array) -> Array

Assemble global matrix from element entries

Parameters:

Name	Type	Description	Default
entries	Array	Array of element matrix entries, (n_elem, 12, 12)	required

Returns:

Type	Description
Array	System global matrix, (n_dof, n_dof)

Source code in src/flapjax/structure/beam.py

def assemble_matrix_from_entries(self, entries: Array) -> Array:
    r"""
    Assemble global matrix from element entries
    :param entries: Array of element matrix entries, ``(n_elem, 12, 12)``
    :return: System global matrix, ``(n_dof, n_dof)``
    """

    row_idx = jnp.broadcast_to(
        self.dof_per_elem_arr[:, :, None], (self.n_elem, 12, 12)
    )
    col_idx = jnp.broadcast_to(
        self.dof_per_elem_arr[:, None, :], (self.n_elem, 12, 12)
    )
    return (
        jnp.zeros((self.n_dof, self.n_dof))
        .at[row_idx.ravel(), col_idx.ravel()]
        .add(entries.ravel())
    )

assemble_vector_from_entries

assemble_vector_from_entries(entries: Array) -> Array

Assemble global vector from element entries

Parameters:

Name	Type	Description	Default
entries	Array	Array of element vector entries, (n_elem, 12)	required

Returns:

Type	Description
Array	System global vector, (n_dof, )

Source code in src/flapjax/structure/beam.py

def assemble_vector_from_entries(self, entries: Array) -> Array:
    r"""
    Assemble global vector from element entries
    :param entries: Array of element vector entries, ``(n_elem, 12)``
    :return: System global vector, ``(n_dof, )``
    """

    vect = jnp.zeros(self.n_dof)
    vect = vect.at[self.dof_per_elem_arr[:, :6]].add(entries[:, :6])
    return vect.at[self.dof_per_elem_arr[:, 6:]].add(entries[:, 6:])

add_lumped_contributions_to_arr

add_lumped_contributions_to_arr(
    arr: Array, lumped_arr: Array
) -> Array

Add lumped contributions to an array

Parameters:

Name	Type	Description	Default
arr	Array	Full array, (6*n_node, 6*n_node)	required
lumped_arr	Array	Lumped contributions, (n_lump, 6, 6)	required

Returns:

Type	Description
Array	In-place updated array, (6*n_node, 6*n_node)

Source code in src/flapjax/structure/beam.py

def add_lumped_contributions_to_arr(self, arr: Array, lumped_arr: Array) -> Array:
    r"""
    Add lumped contributions to an array
    :param arr: Full array, ``(6*n_node, 6*n_node)``
    :param lumped_arr: Lumped contributions, ``(n_lump, 6, 6)``
    :return: In-place updated array, ``(6*n_node, 6*n_node)``
    """

    assert self.m_lumped_index is not None

    def add_block(carry, x):
        node_idx, block = x
        dofs = node_idx * 6 + jnp.arange(6)
        return carry.at[jnp.ix_(dofs, dofs)].add(block), None

    arr, _ = jax.lax.scan(
        add_block, arr, (jnp.array(self.m_lumped_index), lumped_arr)
    )
    return arr

add_lumped_contributions_to_vec

add_lumped_contributions_to_vec(
    vec: Array, lumped_vec: Array
) -> Array

Add lumped contributions to an array

Parameters:

Name	Type	Description	Default
vec	Array	Full vector, (6*n_node, )	required
lumped_vec	Array	Lumped contributions, (n_lump, 6)	required

Returns:

Type	Description
Array	In-place updated vector, (6*n_node, ).

Source code in src/flapjax/structure/beam.py

def add_lumped_contributions_to_vec(self, vec: Array, lumped_vec: Array) -> Array:
    r"""
    Add lumped contributions to an array
    :param vec: Full vector, ``(6*n_node, )``
    :param lumped_vec: Lumped contributions, ``(n_lump, 6)``
    :return: In-place updated vector, ``(6*n_node, )``.
    """

    assert self.m_lumped_index is not None

    idx = (
        jnp.array(self.m_lumped_index)[:, None] * 6 + jnp.arange(6)[None, :]
    ).ravel()  # (n_lump * 6,)

    return vec.at[idx].add(lumped_vec)

make_k_t

make_k_t(d: Array, p_d: Array, eps: Array) -> Array

Assemble tangent stiffness matrix as a function of the element relative configuration vectors

Parameters:

Name	Type	Description	Default
d	Array	Element relative configuration, (n_elem, 6).	required
p_d	Array	P(d) operator, (n_elem, 6, 12).	required
eps	Array	Element strains, (n_elem, 6).	required

Returns:

Type	Description
Array	Elementwise stiffness matrix entries, (n_elem, 12, 12).

Source code in src/flapjax/structure/beam.py

def make_k_t(
    self,
    d: Array,
    p_d: Array,
    eps: Array,
) -> Array:
    r"""
    Assemble tangent stiffness matrix as a function of the element relative configuration vectors
    :param d: Element relative configuration, ``(n_elem, 6)``.
    :param p_d: P(d) operator, ``(n_elem, 6, 12)``.
    :param eps: Element strains, ``(n_elem, 6)``.
    :return: Elementwise stiffness matrix entries, ``(n_elem, 12, 12)``.
    """
    # compute stiffness matrix entries
    return vmap(
        partial(
            _k_t_entry,
            include_geometric=self.optional_jacobians.d_f_int_d_p_d,
        ),
        (0, 0, 0, 0, 0, 0),
        0,
    )(
        d,
        p_d,
        self.l0,
        eps,
        self.k_cs[self.k_cs_index, ...],
        self.ad_inv_o0,
    )  # (n_elem, 12, 12)

make_k_t_full

make_k_t_full(
    d: Array,
    p_d: Array,
    eps: Array,
    f_ext_dead: Array | None,
    rmat: Array,
    m_t: Array | None,
) -> Array

Compute the full tangent stiffness matrix, with contributions from stiffness, dead forces and gravity.

Parameters:

Name	Type	Description	Default
d	Array	Element relative configuration, (n_elem, 6).	required
p_d	Array	P(d) operator, (n_elem, 6, 12).	required
eps	Array	Strain vectors, (n_elem, 6).	required
f_ext_dead	Array | None	External dead forces in global reference, (n_node, 6).	required
rmat	Array	Nodal rotation matrices, (n_node, 3, 3).	required
m_t	Array | None	Disassembled system mass matrix, (n_elem, 12, 12).	required

Returns:

Type	Description
Array	Tangent stiffness matrix with all contributions, (n_dof, n_dof).

Source code in src/flapjax/structure/beam.py

def make_k_t_full(
    self,
    d: Array,
    p_d: Array,
    eps: Array,
    f_ext_dead: Array | None,
    rmat: Array,
    m_t: Array | None,
) -> Array:
    r"""
    Compute the full tangent stiffness matrix, with contributions from stiffness, dead forces and gravity.
    :param d: Element relative configuration, ``(n_elem, 6)``.
    :param p_d: P(d) operator, ``(n_elem, 6, 12)``.
    :param eps: Strain vectors, ``(n_elem, 6)``.
    :param f_ext_dead: External dead forces in global reference, ``(n_node, 6)``.
    :param rmat: Nodal rotation matrices, ``(n_node, 3, 3)``.
    :param m_t: Disassembled system mass matrix, ``(n_elem, 12, 12)``.
    :return: Tangent stiffness matrix with all contributions, ``(n_dof, n_dof)``.
    """

    k_t = self.assemble_matrix_from_entries(self.make_k_t(d, p_d, eps))
    if f_ext_dead is not None and self.optional_jacobians.d_f_ext_dead_d_n:
        k_t += block_diag(*self._make_k_t_dead(rmat, f_ext_dead))

    if self.use_gravity and self.optional_jacobians.d_f_grav_d_n:
        if m_t is None:
            raise ValueError("m_t needs to be provided")
        k_t += self.assemble_matrix_from_entries(
            self._make_k_t_grav(d, p_d, rmat, m_t)
        )
        if self.use_lumped_mass:
            k_t_lumped = self._make_k_t_grav_lumped(rmat)
            k_t = self.add_lumped_contributions_to_arr(
                arr=k_t, lumped_arr=k_t_lumped
            )
    return k_t

make_m_t

make_m_t(
    d: Array,
    int_order: Literal[3, 4, 5] = BASE_LOBATTO_ORDER,
) -> Array

Assemble tangent mass matrix as a function of the element relative configuration vectors. This does not include the lumped mass contribution.

Parameters:

Name	Type	Description	Default
d	Array	Element relative configuration, (n_elem, 6)	required
int_order	Literal[3, 4, 5]	Integration order for mass matrix computation	BASE_LOBATTO_ORDER

Returns:

Type	Description
Array	Elementwise mass matrix, (n_elem, 12, 12)

Source code in src/flapjax/structure/beam.py

def make_m_t(
    self, d: Array, int_order: Literal[3, 4, 5] = BASE_LOBATTO_ORDER
) -> Array:
    r"""
    Assemble tangent mass matrix as a function of the element relative configuration vectors. This does not include
    the lumped mass contribution.
    :param d: Element relative configuration, ``(n_elem, 6)``
    :param int_order: Integration order for mass matrix computation
    :return: Elementwise mass matrix, ``(n_elem, 12, 12)``
    """
    return vmap(partial(_integrate_m_l, int_order=int_order), (0, 0, 0, 0), 0)(
        self.m_cs[self.m_cs_index, ...], d, self.ad_inv_o0, self.l0
    )

make_nodal_m_k

make_nodal_m_k(
    case: StructureCase,
    int_order: Literal[3, 4, 5] = BASE_LOBATTO_ORDER,
) -> tuple[Array, Array]

Create the global mass and stiffness matrices for a given static structure case. These can be used for modal analysis or other purposes. These matrices are the Jacobians of the local forcing residual with respect to global perturbations in acceleration and displacement, respectively.

Parameters:

Name	Type	Description	Default
case	StructureCase	Static structure case for which to compute the global mass and stiffness matrices.	required
int_order	Literal[3, 4, 5]	Integration order for mass matrix computation.	BASE_LOBATTO_ORDER

Returns:

Type	Description
tuple[Array, Array]	Global mass and stiffness matrices, (n_free_dof, n_free_dof).

Source code in src/flapjax/structure/beam.py

def make_nodal_m_k(
    self,
    case: StructureCase,
    int_order: Literal[3, 4, 5] = BASE_LOBATTO_ORDER,
) -> tuple[Array, Array]:
    r"""
    Create the global mass and stiffness matrices for a given static structure case. These can be used for modal
    analysis or other purposes. These matrices are the Jacobians of the local forcing residual with respect to
    global perturbations in acceleration and displacement, respectively.
    :param case: Static structure case for which to compute the global mass and stiffness matrices.
    :param int_order: Integration order for mass matrix computation.
    :return: Global mass and stiffness matrices, ``(n_free_dof, n_free_dof)``.
    """
    # extract variables from case
    d = case.d
    eps = self.make_eps(d=d)
    p_d = self.make_p_d(d=d)
    t_varphi = vmap(t_se3)(case.varphi)  # (n_node, 6, 6)
    rmat = case.hg[:, :3, :3]  # (n_node, 3, 3)

    # get dead external forcing as this has a stiffness contribution
    f_ext_dead_local: Array | None
    if case.f_ext_dead is not None and case.f_ext_aero is not None:
        f_ext_dead_local = case.f_ext_dead + case.f_ext_aero
    else:
        f_ext_dead_local = (
            case.f_ext_dead if case.f_ext_dead is not None else case.f_ext_aero
        )

    # convert to global frame, as it required for creating the stiffness matrix
    f_ext_dead: Array | None = (
        transform_nodal_vect(f_ext_dead_local, rmat)
        if f_ext_dead_local is not None
        else None
    )
    free_dofs = jnp.array(
        get_solve_dofs(n_dof=self.n_dof, prescribed_dofs=case.prescribed_dofs)
    )

    def transform_mat_to_global(mat: Array) -> Array:
        # function to rotate a forcing Jacobian matrix from the local frame to the global frame.
        mat_reshaped = mat.reshape(self.n_nodes, 6, self.n_dof)
        m_lin = jnp.einsum("nij,njk->nik", rmat, mat_reshaped[:, :3, :])
        m_rot = jnp.einsum("nij,njk->nik", rmat, mat_reshaped[:, 3:, :])
        return jnp.concatenate((m_lin, m_rot), axis=1).reshape(
            self.n_dof, self.n_dof
        )

    # mass
    m_t = self.assemble_matrix_from_entries(
        self.make_m_t(d=d, int_order=int_order)
    )  # (n_dof, n_dof)
    if self.use_lumped_mass:
        m_t = self.add_lumped_contributions_to_arr(
            arr=m_t, lumped_arr=self.m_lumped
        )

    m_modal_full = transform_mat_to_global(
        mat=jnp.einsum("ijk,jkl->ijl", m_t.reshape(self.n_dof, -1, 6), t_varphi)
    )

    m_modal = m_modal_full.reshape(self.n_dof, self.n_dof)[
        jnp.ix_(free_dofs, free_dofs)
    ]

    # stiffness
    k_t = self.make_k_t_full(
        d=case.d, p_d=p_d, eps=eps, f_ext_dead=f_ext_dead, rmat=rmat, m_t=m_t
    )

    k_modal_full = transform_mat_to_global(
        mat=jnp.einsum("ijk,jkl->ijl", k_t.reshape(self.n_dof, -1, 6), t_varphi)
    )

    k_modal = k_modal_full.reshape(self.n_dof, self.n_dof)[
        jnp.ix_(free_dofs, free_dofs)
    ]

    return m_modal, k_modal

modal

modal(
    case: StructureCase,
    remove_complex_conjugate: bool = True,
    int_order: Literal[3, 4, 5] = BASE_LOBATTO_ORDER,
    n_modes: int = 20,
    freq_range: tuple[float | Array, float | Array] = (
        0.0,
        jnp.inf,
    ),
    damp_range: tuple[float | Array, float | Array] = (
        -jnp.inf,
        jnp.inf,
    ),
    vtu_directory: str | PathLike = "./modal",
    n_plot_vtu: int | None = None,
    aero: UVLM | None = None,
    n_phase: int = 8,
    n_interp: int = 0,
    max_disp: float = 0.2,
    max_ang: float = 0.2,
) -> tuple[Array, Array, Array]

Perform modal analysis on the structure.

Parameters:

Name	Type	Description	Default
case	StructureCase	The static structure case for which to perform modal analysis.	required
remove_complex_conjugate	bool	If true, keep only one mode from each complex conjugate pair.	True
int_order	Literal[3, 4, 5]	Integration order for mass matrix computation.	BASE_LOBATTO_ORDER
n_modes	int	Number of modes to preserve.	20
freq_range	tuple[float | Array, float | Array]	Frequency range for filtering out modes.	(0.0, inf)
damp_range	tuple[float | Array, float | Array]	Damping range for filtering out modes.	(-inf, inf)
vtu_directory	str | PathLike	Directory to for saving the mode shapes to vtu files.	'./modal'
n_plot_vtu	int | None	Number of modes to plot to vtu files. Will default to "./modal".	None
aero	UVLM | None	UVLM aerodynamic model. If passed, the vtu files will include the aerodynamic grid. If not, they will just be the beam structure.	None
n_phase	int	Number of phases to use when plotting the modes to vtu files.	8
n_interp	int	Number of times to interpolate between beam nodes for vtu plotting.	0
max_disp	float	Maximum displacement of structure for plotted modes, used for scaling.	0.2
max_ang	float	Maximum angle of structure for plotted modes in radians, used for scaling.	0.2

Returns:

Type	Description
tuple[Array, Array, Array]	Tuple of natural frequencies (n_free_dof), damping ratios (n_free_dof), and mode shapes with no normalisation (n_modes, n_free_dof).

Source code in src/flapjax/structure/beam.py

def modal(
    self,
    case: StructureCase,
    remove_complex_conjugate: bool = True,
    int_order: Literal[3, 4, 5] = BASE_LOBATTO_ORDER,
    n_modes: int = 20,
    freq_range: tuple[float | Array, float | Array] = (0.0, jnp.inf),
    damp_range: tuple[float | Array, float | Array] = (-jnp.inf, jnp.inf),
    vtu_directory: str | os.PathLike = "./modal",
    n_plot_vtu: int | None = None,
    aero: UVLM | None = None,
    n_phase: int = 8,
    n_interp: int = 0,
    max_disp: float = 0.2,
    max_ang: float = 0.2,
) -> tuple[Array, Array, Array]:
    r"""
    Perform modal analysis on the structure.
    :param case: The static structure case for which to perform modal analysis.
    :param remove_complex_conjugate: If true, keep only one mode from each complex conjugate pair.
    :param int_order: Integration order for mass matrix computation.
    :param n_modes: Number of modes to preserve.
    :param freq_range: Frequency range for filtering out modes.
    :param damp_range: Damping range for filtering out modes.
    :param vtu_directory: Directory to for saving the mode shapes to vtu files.
    :param n_plot_vtu: Number of modes to plot to vtu files. Will default to "./modal".
    :param aero: UVLM aerodynamic model. If passed, the vtu files will include the aerodynamic grid. If not, they
    will just be the beam structure.
    :param n_phase: Number of phases to use when plotting the modes to vtu files.
    :param n_interp: Number of times to interpolate between beam nodes for vtu plotting.
    :param max_disp: Maximum displacement of structure for plotted modes, used for scaling.
    :param max_ang: Maximum angle of structure for plotted modes in radians, used for scaling.
    :return: Tuple of natural frequencies (n_free_dof), damping ratios (n_free_dof), and mode shapes with no
     normalisation ``(n_modes, n_free_dof)``.
    """
    freqs, damping, modes, *_ = self.base_modal(
        case=case,
        freq_range=freq_range,
        damp_range=damp_range,
        int_order=int_order,
        n_modes=n_modes,
        remove_complex_conjugate=remove_complex_conjugate,
    )

    if n_plot_vtu is not None:
        q_full = (
            jnp.zeros((n_plot_vtu, self.n_nodes * 6))
            .at[:, case.free_dofs]
            .set(modes[:n_plot_vtu, :])
        )

        for _i_mode in range(n_plot_vtu):
            plot_modes_vtu(
                reference=case,
                directory=vtu_directory,
                q_full=q_full.reshape(n_plot_vtu, self.n_nodes, 6),
                freqs=freqs,
                dampings=damping,
                gamma_b_full=None,
                gamma_w_full=None,
                zeta_w_full=None,
                uvlm=aero,
                n_interp=n_interp,
                n_phase=n_phase,
                max_disp=max_disp,
                max_ang=max_ang,
                max_gamma=1e6,
            )

    return freqs, damping, modes

linearise

linearise(
    reference: StructureCase,
    dt: float,
    n_modes: int | None = None,
    modal_inputs: bool = False,
    modal_outputs: bool = False,
    prescribed_dofs: Sequence[int]
    | Array
    | slice
    | int
    | None = None,
) -> LinearBeam

Linearise the beam about a given static structure case. This creates a LinearBeam object which can be used for linear dynamic analysis.

Parameters:

Name	Type	Description	Default
reference	StructureCase	Static structure case about which to linearise the beam.	required
dt	float	Time step size, used for conversions between continuous and discrete time.	required
n_modes	int | None	If not None, the linearised system uses modal state coordinates truncated to this many modes.	None
modal_inputs	bool	If True, external forcing inputs are provided as modal forces (requires n_modes).	False
modal_outputs	bool	If True, outputs are exposed as modal coordinates (requires n_modes).	False
prescribed_dofs	Sequence[int] | Array | slice | int | None	If provided, overrides the prescribed DOFs from the reference case.	None

Returns:

Type	Description
LinearBeam	Continuous-time linearised beam object.

Source code in src/flapjax/structure/beam.py

def linearise(
    self,
    reference: StructureCase,
    dt: float,
    n_modes: int | None = None,
    modal_inputs: bool = False,
    modal_outputs: bool = False,
    prescribed_dofs: Sequence[int] | Array | slice | int | None = None,
) -> LinearBeam:
    r"""
    Linearise the beam about a given static structure case. This creates a LinearBeam object which can be used for
    linear dynamic analysis.
    :param reference: Static structure case about which to linearise the beam.
    :param dt: Time step size, used for conversions between continuous and discrete time.
    :param n_modes: If not None, the linearised system uses modal state coordinates truncated to this many modes.
    :param modal_inputs: If True, external forcing inputs are provided as modal forces (requires n_modes).
    :param modal_outputs: If True, outputs are exposed as modal coordinates (requires n_modes).
    :param prescribed_dofs: If provided, overrides the prescribed DOFs from the reference case.
    :return: Continuous-time linearised beam object.
    """
    return LinearBeam(
        beam=self,
        reference=reference,
        dt=dt,
        n_modes=n_modes,
        modal_inputs=modal_inputs,
        modal_outputs=modal_outputs,
        prescribed_dofs=prescribed_dofs,
    )

apply_nodal_constraint_tangent

apply_nodal_constraint_tangent(
    mat: Array,
    hg: Array,
    i_ts: int,
    gamma_prime: float | Array | None,
) -> Array

Add nodal constraint contributions to a system matrix.

Parameters:

Name	Type	Description	Default
mat	Array	System matrix to update, (n_dof, n_dof).	required
hg	Array	SE(3) coordiantes, (n_nodes, 4, 4).	required
i_ts	int	Time-step index (0 for static solves).	required
gamma_prime	float | Array | None	Time-integrator gamma_prime for damping scaling, or None to skip damping.	required

Returns:

Type	Description
Array	Updated system matrix.

Source code in src/flapjax/structure/beam.py

def apply_nodal_constraint_tangent(
    self,
    mat: Array,
    hg: Array,
    i_ts: int,
    gamma_prime: float | Array | None,
) -> Array:
    r"""
    Add nodal constraint contributions to a system matrix.
    :param mat: System matrix to update, ``(n_dof, n_dof)``.
    :param hg: SE(3) coordiantes, ``(n_nodes, 4, 4)``.
    :param i_ts: Time-step index (0 for static solves).
    :param gamma_prime: Time-integrator gamma_prime for damping scaling, or ``None`` to skip damping.
    :return: Updated system matrix.
    """
    for con in self.nodal_constraints:
        node = con.node_index
        hg_i = hg[node]
        dofs = node * 6 + jnp.arange(6)
        mat = mat.at[jnp.ix_(dofs, dofs)].add(con.k_tangent(hg_i, i_ts))
        if gamma_prime is not None:
            mat = mat.at[jnp.ix_(dofs, dofs)].add(
                gamma_prime * con.c_tangent(hg_i, i_ts)
            )

    # hard constraint tangent contributions (e.g. hinge spring-damper)
    for con in self.multibody_constraints:
        if con.has_f_res:
            dofs_i = con.node_i * 6 + jnp.arange(6)
            dofs_j = con.node_j * 6 + jnp.arange(6)
            dofs_ij = jnp.concatenate([dofs_i, dofs_j])
            k_12 = con.k_tangent(hg[con.node_i], hg[con.node_j])
            mat = mat.at[jnp.ix_(dofs_ij, dofs_ij)].add(k_12)
            if gamma_prime is not None:
                # add damping terms
                mat = mat.at[jnp.ix_(dofs_ij, dofs_ij)].add(
                    gamma_prime * con.c_tangent(hg[con.node_i], hg[con.node_j])
                )

    return mat

postprocess_constraints

postprocess_constraints(
    hg: Array,
) -> dict[str, dict[str, Array]]

Postprocess all constraints to extract derived quantities (e.g. hinge angles).

Parameters:

Name	Type	Description	Default
hg	Array	Nodal SE(3) frames, (n_nodes, 4, 4) or (n_tstep, n_nodes, 4, 4).	required

Returns:

Type	Description
dict[str, dict[str, Array]]	Nested dict {constraint_name: {quantity_name: Array}}.

Source code in src/flapjax/structure/beam.py

def postprocess_constraints(self, hg: Array) -> dict[str, dict[str, Array]]:
    r"""
    Postprocess all constraints to extract derived quantities (e.g. hinge angles).
    :param hg: Nodal SE(3) frames, ``(n_nodes, 4, 4)`` or ``(n_tstep, n_nodes, 4, 4)``.
    :return: Nested dict ``{constraint_name: {quantity_name: Array}}``.
    """
    data: dict[str, dict[str, Array]] = {}
    for name, con in self.constraints.items():
        pp = con.postprocess(hg)
        if pp:
            data[name] = pp
    return data

solve_constrained

solve_constrained(
    sys_mat_solve: Array,
    f_res_solve: Array,
    hg_eval: Array,
    solve_dofs: Array,
    hg_base: Array | None = None,
    phi: Array | None = None,
    v: Array | None = None,
    gamma_prime: float | Array | None = None,
) -> tuple[Array, Array, Array]

Solve the augmented system with Lagrange multipliers, supporting both holonomic and non-holonomic constraints.

Parameters:

Name	Type	Description	Default
sys_mat_solve	Array	System matrix at solve DOFs, (n_solve, n_solve).	required
f_res_solve	Array	Force residual at solve DOFs, (n_solve,).	required
hg_eval	Array	SE(3) frames for constraint evaluation, (n_nodes, 4, 4).	required
solve_dofs	Array	Free DOF indices, (n_solve,).	required
hg_base	Array | None	Base frames for Jacobian computation (defaults to hg_eval).	None
phi	Array | None	Accumulated configuration increment, (n_nodes, 6).	None
v	Array | None	Current nodal velocities for non-holonomic constraints, (n_nodes, 6).	None
gamma_prime	float | Array | None	Newmark parameter for non-holonomic constraints.	None

Returns:

Type	Description
tuple[Array, Array, Array]	(delta_phi, lagrange_multipliers, constraint_violation).

Source code in src/flapjax/structure/beam.py

def solve_constrained(
    self,
    sys_mat_solve: Array,
    f_res_solve: Array,
    hg_eval: Array,
    solve_dofs: Array,
    hg_base: Array | None = None,
    phi: Array | None = None,
    v: Array | None = None,
    gamma_prime: float | Array | None = None,
) -> tuple[Array, Array, Array]:
    r"""
    Solve the augmented system with Lagrange multipliers, supporting both holonomic and non-holonomic constraints.
    :param sys_mat_solve: System matrix at solve DOFs, ``(n_solve, n_solve)``.
    :param f_res_solve: Force residual at solve DOFs, ``(n_solve,)``.
    :param hg_eval: SE(3) frames for constraint evaluation, ``(n_nodes, 4, 4)``.
    :param solve_dofs: Free DOF indices, ``(n_solve,)``.
    :param hg_base: Base frames for Jacobian computation (defaults to ``hg_eval``).
    :param phi: Accumulated configuration increment, ``(n_nodes, 6)``.
    :param v: Current nodal velocities for non-holonomic constraints, ``(n_nodes, 6)``.
    :param gamma_prime: Newmark parameter for non-holonomic constraints.
    :return: ``(delta_phi, lagrange_multipliers, constraint_violation)``.
    """
    hg_jac = hg_base if hg_base is not None else hg_eval

    n_s = sys_mat_solve.shape[0]
    n_h = self.n_holonomic_constraints
    n_nh = self.n_nonholonomic_constraints
    n_c = n_h + n_nh

    aug = jnp.zeros((n_s + n_c, n_s + n_c))
    aug = aug.at[:n_s, :n_s].set(sys_mat_solve)

    rhs_parts: list[Array] = [f_res_solve]
    violation_parts: list[Array] = []

    if n_h > 0:
        viol_h = self._compute_holonomic_violation(hg_eval)
        jac_h = self._compute_holonomic_jacobian(hg_jac, solve_dofs, phi)
        aug = aug.at[:n_s, n_s : n_s + n_h].set(-jac_h.T)
        aug = aug.at[n_s : n_s + n_h, :n_s].set(jac_h)
        rhs_parts.append(-viol_h)
        violation_parts.append(viol_h)

    if n_nh > 0:
        assert v is not None
        vel_viol_nh = self._compute_nonholonomic_vel_violation(hg_eval, v)
        a_vel_solve = self._compute_nonholonomic_a_vel(hg_eval, v, solve_dofs)
        a_phi_solve = self._compute_nonholonomic_a_phi(hg_jac, v, solve_dofs, phi)
        aug = aug.at[:n_s, n_s + n_h : n_s + n_c].set(-a_vel_solve.T)
        aug = aug.at[n_s + n_h : n_s + n_c, :n_s].set(
            a_phi_solve + gamma_prime * a_vel_solve
        )
        rhs_parts.append(-vel_viol_nh)
        violation_parts.append(vel_viol_nh)

    rhs = jnp.concatenate(rhs_parts)
    sol = jnp.linalg.solve(aug, rhs)

    return sol[:n_s], sol[n_s:], jnp.concatenate(violation_parts)

compute_centre_of_mass

compute_centre_of_mass(hg: Array) -> Array

Compute the centre of mass for an arbitrary system.

Parameters:

Name	Type	Description	Default
hg	Array	Node SE(3) coordinates, (n_node, 4, 4) or (n_tstep, n_node, 4, 4).	required

Returns:

Type	Description
Array	Centre of mass, (3) or (n_tstep, 3).

Source code in src/flapjax/structure/beam.py

def compute_centre_of_mass(self, hg: Array) -> Array:
    r"""
    Compute the centre of mass for an arbitrary system.
    :param hg: Node SE(3) coordinates, ``(n_node, 4, 4)`` or ``(n_tstep, n_node, 4, 4)``.
    :return: Centre of mass, (3) or ``(n_tstep, 3)``.
    """

    def inner_func(hg_: Array) -> Array:
        d = self.make_d(hg=hg_)
        m = self.assemble_matrix_from_entries(self.make_m_t(d=d))  # (n_dof, n_dof)
        if self.use_lumped_mass:
            m = self.add_lumped_contributions_to_arr(
                arr=m, lumped_arr=self.m_lumped
            )
        m_lin = m[::6, ::6]
        return jnp.einsum("ij,jk->k", m_lin, hg_[:, :3, 3]) / m_lin.sum()  # (3, )

    if hg.ndim == 3:
        return inner_func(hg)  # single timestep, (3, ).
    elif hg.ndim == 4:
        return vmap(inner_func, 0, 0)(hg)  # multiple timesteps, (n_tstep, 3)
    else:
        raise ValueError("hg.ndim must be 3 or 4")

make_f_elem

make_f_elem(eps: Array) -> Array

Compute the forces within the elements as :math:\mathbf{f}_{elem} = \mathcal{K}_{cs} \epsilon.

Parameters:

Name	Type	Description	Default
eps	Array	Element strain vectors, (n_elem, 6).	required

Returns:

Type	Description
Array	Element forces, (n_elem, 6).

Source code in src/flapjax/structure/beam.py

def make_f_elem(self, eps: Array) -> Array:
    r"""
    Compute the forces within the elements as :math:`\mathbf{f}_{elem} = \mathcal{K}_{cs} \epsilon`.
    :param eps: Element strain vectors, ``(n_elem, 6)``.
    :return: Element forces, ``(n_elem, 6)``.
    """
    return jnp.einsum("ijk,ik->ij", self.k_cs[self.k_cs_index, ...], eps)

make_f_int

make_f_int(p_d: Array, eps: Array) -> Array

Assemble global internal force vector as a function of the element relative configuration vectors.

Parameters:

Name	Type	Description	Default
p_d	Array	P(d) operator, (n_elem, 6, 12).	required
eps	Array	Element strain vectors, (n_elem, 6).	required

Returns:

Type	Description
Array	Internal forces, (n_elem, 12).

Source code in src/flapjax/structure/beam.py

def make_f_int(self, p_d: Array, eps: Array) -> Array:
    r"""
    Assemble global internal force vector as a function of the element relative configuration vectors.
    :param p_d: P(d) operator, ``(n_elem, 6, 12)``.
    :param eps: Element strain vectors, ``(n_elem, 6)``.
    :return: Internal forces, ``(n_elem, 12)``.
    """

    return -jnp.einsum("ikj,ikl,il->ij", p_d, self.k_cs[self.k_cs_index, ...], eps)

make_f_dead_ext staticmethod

make_f_dead_ext(f_ext: Array, rmat: Array) -> Array

Compute the global external dead force vector.

Parameters:

Name	Type	Description	Default
f_ext	Array	External forces array of dead forces in global reference, (n_node, 6)	required
rmat	Array	Deformation rotation matrices, (n_node, 3, 3)	required

Returns:

Type	Description
Array	External forces, (n_node, 6)

Source code in src/flapjax/structure/beam.py

@staticmethod
def make_f_dead_ext(f_ext: Array, rmat: Array) -> Array:
    r"""
    Compute the global external dead force vector.
    :param f_ext: External forces array of dead forces in global reference, ``(n_node, 6)``
    :param rmat: Deformation rotation matrices, ``(n_node, 3, 3)``
    :return: External forces, ``(n_node, 6)``
    """

    return transform_nodal_vect(f_ext, jnp.swapaxes(rmat, -1, -2))

add_thrust_force

add_thrust_force(
    force: Array, thrust: dict[str, Array]
) -> Array

Add thrust acting at nodes onto full system forcing.

Parameters:

Name	Type	Description	Default
force	Array	Input forcing, (n_node, 6).	required
thrust	dict[str, Array]	Input thrust at the current step, {key: ()}.	required

Returns:

Type	Description
Array	Updated forcing, (n_node, 6).

Source code in src/flapjax/structure/beam.py

def add_thrust_force(self, force: Array, thrust: dict[str, Array]) -> Array:
    r"""
    Add thrust acting at nodes onto full system forcing.
    :param force: Input forcing, ``(n_node, 6)``.
    :param thrust: Input thrust at the current step, ``{key: ()}``.
    :return: Updated forcing, ``(n_node, 6)``.
    """

    for k, v in thrust.items():
        node = dict(self.thrust_nodes)[k]
        direction = jnp.array(dict(self.thrust_direction)[k])
        force = force.at[node, :3].add(v * direction)
    return force

make_eps

make_eps(d: Array) -> Array

Compute the element strain vectors as a function of the element relative configuration vectors. Formulation from Geometrically exact beam finite element formulated on the special Euclidean group SE(3), by Sonneville et al., 2013, Eq 64.

Parameters:

Name	Type	Description	Default
d	Array	Element relative configuration, (n_elem, 6)	required

Returns:

Type	Description
Array	Element strain vectors, (n_elem, 6)

Source code in src/flapjax/structure/beam.py

def make_eps(self, d: Array) -> Array:
    r"""
    Compute the element strain vectors as a function of the element relative configuration vectors. Formulation from
    Geometrically exact beam finite element formulated on the special Euclidean group SE(3), by Sonneville et al.,
    2013, Eq 64.
    :param d: Element relative configuration, ``(n_elem, 6)``
    :return: Element strain vectors, ``(n_elem, 6)``
    """

    return (d - self.d0) / self.l0[:, None]

make_p_d

make_p_d(d: Array) -> Array

Compute the P(d) operator as a function of the element relative configuration vectors.

Parameters:

Name	Type	Description	Default
d	Array	Relative configuration vectors, (n_elem, 6)	required

Returns:

Type	Description
Array	P(d) operator, (n_elem, 6, 12)

Source code in src/flapjax/structure/beam.py

def make_p_d(self, d: Array) -> Array:
    r"""
    Compute the P(d) operator as a function of the element relative configuration vectors.
    :param d: Relative configuration vectors, ``(n_elem, 6)``
    :return: P(d) operator, ``(n_elem, 6, 12)``
    """
    return vmap(p, (0, 0), 0)(d, self.ad_inv_o0)  # [n_elem, 6, 12]

make_d

make_d(hg: Array) -> Array

Compute the element relative configuration vectors from the nodal homogeneous transformation matrices

Parameters:

Name	Type	Description	Default
hg	Array	Nodal homogeneous transformation matrices, (n_nodes, 4, 4)	required

Returns:

Type	Description
Array	Element relative configuration vectors, (n_elem, 6)

Source code in src/flapjax/structure/beam.py

def make_d(self, hg: Array) -> Array:
    r"""
    Compute the element relative configuration vectors from the nodal homogeneous transformation matrices
    :param hg: Nodal homogeneous transformation matrices, ``(n_nodes, 4, 4)``
    :return: Element relative configuration vectors, ``(n_elem, 6)``
    """

    base_hg = jnp.zeros((self.n_elem, 4, 4))
    base_hg = base_hg.at[:, :3, :3].set(self.o0)
    base_hg = base_hg.at[:, 3, 3].set(1.0)

    haha0 = jnp.einsum(
        "ijk,ikl->ijl", hg[self.connectivity_arr[:, 0], :, :], base_hg
    )  # (n_elem, 4, 4)
    haha1 = jnp.einsum(
        "ijk,ikl->ijl", hg[self.connectivity_arr[:, 1], :, :], base_hg
    )  # (n_elem, 4, 4)

    return vmap(hg_to_d, (0, 0), 0)(haha0, haha1)  # (n_elem, 6)

make_hg_dot staticmethod

make_hg_dot(hg: Array, v: Array) -> Array

Obtain the time derivative of the nodal coordinates.

Parameters:

Name	Type	Description	Default
hg	Array	Node coordinates, (n_node, 4, 4).	required
v	Array	Node local velocities, (n_node, 6)	required

Returns:

Type	Description
Array	Coordinate time derivative, (n_node, 4, 4)

Source code in src/flapjax/structure/beam.py

@staticmethod
def make_hg_dot(hg: Array, v: Array) -> Array:
    r"""
    Obtain the time derivative of the nodal coordinates.
    :param hg: Node coordinates, ``(n_node, 4, 4)``.
    :param v: Node local velocities, ``(n_node, 6)``
    :return: Coordinate time derivative, ``(n_node, 4, 4)``
    """
    return jnp.einsum(
        "ijk,ikl->ijl", hg, vmap(ha_to_ha_tilde, 0, 0)(v)
    )  # (n_nodes, 4, 4)

resolve_forces

resolve_forces(
    hg: Array,
    dynamic: Literal[True],
    f_ext_follower: Array | None,
    f_ext_dead: Array | None,
    f_ext_aero: Array | None,
    thrust: dict[str, Array],
    v: Array,
    v_dot: Array,
    approx_gradients: bool = False,
) -> tuple[
    Array,
    Array,
    Array | None,
    Array | None,
    Array | None,
    Array,
    Array,
    Array,
    Array,
]

resolve_forces(
    hg: Array,
    dynamic: Literal[False],
    f_ext_follower: Array | None,
    f_ext_dead: Array | None,
    f_ext_aero: Array | None,
    thrust: dict[str, Array],
    v: None,
    v_dot: None,
    approx_gradients: bool = False,
) -> tuple[
    Array,
    Array,
    Array | None,
    Array | None,
    Array | None,
    Array,
    None,
    None,
    Array,
]

resolve_forces(
    hg: Array,
    dynamic: bool,
    f_ext_follower: Array | None,
    f_ext_dead: Array | None,
    f_ext_aero: Array | None,
    thrust: dict[str, Array],
    v: Array | None,
    v_dot: Array | None,
    approx_gradients: bool = False,
) -> tuple[
    Array,
    Array,
    Array | None,
    Array | None,
    Array | None,
    Array,
    Array | None,
    Array | None,
    Array,
]

Obtain all components of the force from a final solution.

Parameters:

Name	Type	Description	Default
hg	Array	Nodal homogeneous transformation matrices, (n_nodes, 4, 4).	required
dynamic	bool	Whether to compute dynamic forces.	required
f_ext_follower	Array | None	External follower forces in local reference, (n_node, 6).	required
f_ext_dead	Array | None	External dead forces in global reference, (n_node, 6).	required
f_ext_aero	Array | None	External aero forces in global reference, (n_node, 6).	required
thrust	dict[str, Array]	Thrust forces at current step, {keys, ()}.	required
v	Array | None	Nodal velocities in global frame, (n_node, 6).	required
v_dot	Array | None	Nodal accelerations in global frame, (n_node, 6).	required
approx_gradients	bool	Whether to stop computing gradients of the inertial and gyroscopic forces with respect to the node coordinates, as these are small but nonzero values in practice.	False

Returns:

Type	Description
tuple[Array, Array, Array | None, Array | None, Array | None, Array, Array | None, Array | None, Array]	Configuration vectors, strain vectors, Dead external forces, aero external forces, gravitational forces, internal forces, gyroscopic forces, inertial forces and residual forces.

Source code in src/flapjax/structure/beam.py

def resolve_forces(
    self,
    hg: Array,
    dynamic: bool,
    f_ext_follower: Array | None,
    f_ext_dead: Array | None,
    f_ext_aero: Array | None,
    thrust: dict[str, Array],
    v: Array | None,
    v_dot: Array | None,
    approx_gradients: bool = False,
) -> tuple[
    Array,
    Array,
    Array | None,
    Array | None,
    Array | None,
    Array,
    Array | None,
    Array | None,
    Array,
]:
    r"""
    Obtain all components of the force from a final solution.
    :param hg: Nodal homogeneous transformation matrices, ``(n_nodes, 4, 4)``.
    :param dynamic: Whether to compute dynamic forces.
    :param f_ext_follower: External follower forces in local reference, ``(n_node, 6)``.
    :param f_ext_dead: External dead forces in global reference, ``(n_node, 6)``.
    :param f_ext_aero: External aero forces in global reference, ``(n_node, 6)``.
    :param thrust: Thrust forces at current step, {keys, ``()``}.
    :param v: Nodal velocities in global frame, ``(n_node, 6)``.
    :param v_dot: Nodal accelerations in global frame, ``(n_node, 6)``.
    :param approx_gradients: Whether to stop computing gradients of the inertial and gyroscopic forces with respect to
    the node coordinates, as these are small but nonzero values in practice.
    :return: Configuration vectors, strain vectors, Dead external forces, aero external forces, gravitational forces, internal forces,
    gyroscopic forces, inertial forces and residual forces.
    """

    def prop_grad(x: Array) -> Array:
        return jax.lax.stop_gradient(x) if approx_gradients else x

    d = self.make_d(hg)
    eps = self.make_eps(d)
    p_d = self.make_p_d(d)

    if dynamic or self.use_gravity:
        m_t = self.make_m_t(prop_grad(d))
    else:
        m_t = None

    if dynamic:
        assert v is not None

        d_dot = self._make_d_dot(p_d, v)
        c_l = self._make_c_t(prop_grad(d), prop_grad(d_dot), v)[0]
        c_l_lumped = self._make_c_t_lumped(v)[0] if self.use_lumped_mass else None
    else:
        d_dot, c_l, c_l_lumped = None, None, None

    this_f_res = self.add_thrust_force(
        force=jnp.zeros((self.n_nodes, 6)), thrust=thrust
    )

    if f_ext_dead is not None:
        this_f_ext_dead = self.make_f_dead_ext(f_ext_dead, hg[:, :3, :3])
        this_f_res += this_f_ext_dead
    else:
        this_f_ext_dead = None

    if f_ext_aero is not None:
        this_f_ext_aero = self.make_f_dead_ext(f_ext_aero, hg[:, :3, :3])
        this_f_res += this_f_ext_aero
    else:
        this_f_ext_aero = None

    if self.use_gravity:
        assert m_t is not None
        this_f_grav = self.assemble_vector_from_entries(
            self._make_f_grav(m_t, hg[:, :3, :3])
        ).reshape(-1, 6)
        if self.use_lumped_mass:
            f_grav_lumped = self._make_f_grav_lumped(hg[:, :3, :3])
            this_f_grav = self.add_lumped_contributions_to_vec(
                vec=this_f_grav.ravel(), lumped_vec=f_grav_lumped.ravel()
            ).reshape(-1, 6)
        this_f_res += this_f_grav
    else:
        this_f_grav = None

    this_f_int = self.assemble_vector_from_entries(
        self.make_f_int(p_d, eps)
    ).reshape(-1, 6)
    this_f_res += this_f_int

    if dynamic:
        assert (
            m_t is not None
            and c_l is not None
            and v is not None
            and v_dot is not None
        )
        this_f_iner, this_f_gyr = self._make_f_iner_gyr(m_t, c_l, v, v_dot)
        this_f_iner = self.assemble_vector_from_entries(this_f_iner).reshape(-1, 6)
        this_f_gyr = self.assemble_vector_from_entries(this_f_gyr).reshape(-1, 6)

        if self.use_lumped_mass:
            assert c_l_lumped is not None
            f_iner_lumped, f_gyr_lumped = self._make_f_iner_gyr_lumped(
                c_l_lumped, v, v_dot
            )
            this_f_iner = self.add_lumped_contributions_to_vec(
                this_f_iner.ravel(), (f_iner_lumped + f_gyr_lumped).ravel()
            ).reshape(-1, 6)
        this_f_res += this_f_iner
    else:
        this_f_iner = None
        this_f_gyr = None

    if f_ext_follower is not None:
        this_f_res += f_ext_follower

    return (
        d,
        eps,
        this_f_ext_dead,
        this_f_ext_aero,
        this_f_grav,
        this_f_int,
        this_f_gyr,
        this_f_iner,
        this_f_res,
    )

make_f_res

make_f_res(
    solve_dofs: Array | None,
    p_d: Array,
    eps: Array,
    hg: Array,
    f_ext_follower_n: Array | None,
    f_ext_dead_n: Array | None,
    thrust_n: dict[str, Array],
    dynamic: Literal[True],
    m_t: Array,
    c_l: Array,
    c_l_lumped: Array | None,
    v: Array,
    v_dot: Array,
    i_ts: int = 0,
    k_t_assembled: Array | None = None,
) -> tuple[Array, Array]

make_f_res(
    solve_dofs: Array | None,
    p_d: Array,
    eps: Array,
    hg: Array,
    f_ext_follower_n: Array | None,
    f_ext_dead_n: Array | None,
    thrust_n: dict[str, Array],
    dynamic: Literal[False],
    m_t: Array | None,
    c_l: None,
    c_l_lumped: None,
    v: None,
    v_dot: None,
    i_ts: int = 0,
    k_t_assembled: Array | None = None,
) -> tuple[Array, Array]

make_f_res(
    solve_dofs: Array | None,
    p_d: Array,
    eps: Array,
    hg: Array,
    f_ext_follower_n: Array | None,
    f_ext_dead_n: Array | None,
    thrust_n: dict[str, Array],
    dynamic: bool,
    m_t,
    c_l,
    c_l_lumped,
    v,
    v_dot,
    i_ts: int = 0,
    k_t_assembled: Array | None = None,
) -> tuple[Array, Array]

Compute the residual force vector for a given configuration and external forces, used in the nonlinear solve. This is the force imbalance that the nonlinear solver will seek to drive to zero. Additionally, returns an "absolute sum" of all forces, used for relative convergence checks.

Parameters:

Name	Type	Description	Default
solve_dofs	Array | None	Optional array of degrees of freedom to solve for (n_solve_dofs, ).	required
p_d	Array	P(d) operator, (n_elem, 6, 12).	required
eps	Array	Element strain vectors, (n_elem, 6).	required
hg	Array	Nodal homogeneous transformation matrices, (n_nodes, 4, 4).	required
f_ext_follower_n	Array | None	Nodal follower forces, (n_nodes, 6).	required
f_ext_dead_n	Array | None	Nodal dead forces, (n_nodes, 6).	required
thrust_n	dict[str, Array]	Thrust magnitude, {key: ()}.	required
dynamic	bool	Flag for whether to compute dynamic entries.	required
m_t		Disassembled system mass matrix, (n_elem, 12, 12).	required
c_l		Dissembled system gyroscopic matrix, (n_elem, 12, 12).	required
c_l_lumped		Lumped gyroscopic matrix, (n_nodes, 6, 6).	required
v		Nodal velocities, (n_nodes, 6).	required
v_dot		Nodal accelerations, (n_node, 6).	required
i_ts	int	Time-step index (0 for static solves).	0
k_t_assembled	Array | None	Assembled global tangent stiffness matrix (n_dof, n_dof), required for Rayleigh damping.	None

Returns:

Type	Description
tuple[Array, Array]	Residual force vector, (n_dof, ), absolute sum of forces, (n_dof, ).

Source code in src/flapjax/structure/beam.py

def make_f_res(
    self,
    solve_dofs: Array | None,
    p_d: Array,
    eps: Array,
    hg: Array,
    f_ext_follower_n: Array | None,
    f_ext_dead_n: Array | None,
    thrust_n: dict[str, Array],
    dynamic: bool,
    m_t,
    c_l,
    c_l_lumped,
    v,
    v_dot,
    i_ts: int = 0,
    k_t_assembled: Array | None = None,
) -> tuple[Array, Array]:
    r"""
    Compute the residual force vector for a given configuration and external forces, used in the nonlinear solve.
    This is the force imbalance that the nonlinear solver will seek to drive to zero. Additionally, returns an
    "absolute sum" of all forces, used for relative convergence checks.
    :param solve_dofs: Optional array of degrees of freedom to solve for ``(n_solve_dofs, )``.
    :param p_d: P(d) operator, ``(n_elem, 6, 12)``.
    :param eps: Element strain vectors, ``(n_elem, 6)``.
    :param hg: Nodal homogeneous transformation matrices, ``(n_nodes, 4, 4)``.
    :param f_ext_follower_n: Nodal follower forces, ``(n_nodes, 6)``.
    :param f_ext_dead_n: Nodal dead forces, ``(n_nodes, 6)``.
    :param thrust_n: Thrust magnitude, ``{key: ()}``.
    :param dynamic: Flag for whether to compute dynamic entries.
    :param m_t: Disassembled system mass matrix, ``(n_elem, 12, 12)``.
    :param c_l: Dissembled system gyroscopic matrix, ``(n_elem, 12, 12)``.
    :param c_l_lumped: Lumped gyroscopic matrix, ``(n_nodes, 6, 6)``.
    :param v: Nodal velocities, ``(n_nodes, 6)``.
    :param v_dot: Nodal accelerations, ``(n_node, 6)``.
    :param i_ts: Time-step index (0 for static solves).
    :param k_t_assembled: Assembled global tangent stiffness matrix ``(n_dof, n_dof)``, required for Rayleigh
    damping.
    :return: Residual force vector, ``(n_dof, )``, absolute sum of forces, ``(n_dof, )``.
    """

    f_res = self.make_f_int(p_d, eps)  # (n_elem, 12)
    f_abs_sum = jnp.abs(f_res)

    if self.use_gravity:
        f_grav = self._make_f_grav(m_t, hg[:, :3, :3])
        f_res += f_grav
        f_abs_sum += jnp.abs(f_grav)

    if dynamic:
        f_iner, f_gyr = self._make_f_iner_gyr(m_t, c_l, v, v_dot)
        f_res += f_iner + f_gyr
        f_abs_sum += jnp.abs(f_iner + f_gyr)

    f_res_vect = self.assemble_vector_from_entries(f_res)
    f_abs_sum_vect = self.assemble_vector_from_entries(f_abs_sum)

    # add external forcing contributions
    if f_ext_follower_n is not None:
        f_res_vect += f_ext_follower_n.reshape(self.n_dof).ravel()
        f_abs_sum_vect += jnp.abs(f_ext_follower_n.reshape(self.n_dof).ravel())
    if f_ext_dead_n is not None:
        f_dead = self.make_f_dead_ext(f_ext_dead_n, hg[:, :3, :3]).ravel()
        f_res_vect += f_dead
        f_abs_sum_vect += jnp.abs(f_dead)

    f_thrust = self.add_thrust_force(
        force=jnp.zeros((self.n_nodes, 6)), thrust=thrust_n
    ).ravel()
    f_res_vect += f_thrust
    f_abs_sum_vect += jnp.abs(f_thrust)

    if self.use_lumped_mass:
        if dynamic:
            f_iner_lumped, f_gyr_lumped = self._make_f_iner_gyr_lumped(
                c_l_lumped, v, v_dot
            )
            f_iner_gyr_lumped = (f_iner_lumped + f_gyr_lumped).ravel()
            f_res_vect = self.add_lumped_contributions_to_vec(
                f_res_vect, f_iner_gyr_lumped
            )
            f_abs_sum_vect = self.add_lumped_contributions_to_vec(
                f_abs_sum_vect, jnp.abs(f_iner_gyr_lumped)
            )
        if self.use_gravity:
            f_grav_lumped = self._make_f_grav_lumped(hg[:, :3, :3]).ravel()
            f_res_vect = self.add_lumped_contributions_to_vec(
                vec=f_res_vect, lumped_vec=f_grav_lumped
            )
            f_abs_sum_vect = self.add_lumped_contributions_to_vec(
                vec=f_abs_sum_vect, lumped_vec=f_grav_lumped
            )

    # nodal constraint contributions
    for con in self.nodal_constraints:
        node = con.node_index
        v_node = v[node] if dynamic else jnp.zeros(6)
        f_constraint = con.f_res(hg[node], v_node, i_ts)
        dofs = node * 6 + jnp.arange(6)
        f_res_vect = f_res_vect.at[dofs].add(f_constraint)
        f_abs_sum_vect = f_abs_sum_vect.at[dofs].add(jnp.abs(f_constraint))

    # hard constraint force contributions (e.g. hinge spring-damper)
    for con in self.multibody_constraints:
        if con.has_f_res:
            v_i = v[con.node_i] if dynamic else jnp.zeros(6)
            v_j = v[con.node_j] if dynamic else jnp.zeros(6)
            f_i, f_j = con.f_res(hg[con.node_i], hg[con.node_j], v_i, v_j)
            dofs_i = con.node_i * 6 + jnp.arange(6)
            dofs_j = con.node_j * 6 + jnp.arange(6)
            f_res_vect = f_res_vect.at[dofs_i].add(f_i)
            f_res_vect = f_res_vect.at[dofs_j].add(f_j)
            f_abs_sum_vect = f_abs_sum_vect.at[dofs_i].add(jnp.abs(f_i))
            f_abs_sum_vect = f_abs_sum_vect.at[dofs_j].add(jnp.abs(f_j))

    # Rayleigh structural damping
    if dynamic and (self.alpha_m != 0.0 or self.beta_k != 0.0):
        if self.beta_k != 0.0 and k_t_assembled is None:
            raise ValueError(
                "k_t_assembled must be provided when beta_k != 0 for dynamic residual."
            )
        f_damp = self._make_f_rayleigh_damp(
            m_t=m_t,
            k_t_assembled=k_t_assembled
            if k_t_assembled is not None
            else jnp.zeros((self.n_dof, self.n_dof)),
            v=v,
        )
        f_res_vect += f_damp
        f_abs_sum_vect += jnp.abs(f_damp)

    if solve_dofs is not None:
        return f_res_vect[solve_dofs], f_abs_sum_vect[
            solve_dofs
        ]  # (n_solve_dof, ), (n_solve_dof, )
    else:
        return f_res_vect, f_abs_sum_vect  # (n_dof, ), (n_dof, )

update_hg staticmethod

update_hg(hg: Array, phi: Array) -> Array

Update the nodal homogeneous transformation matrices with the configuration increments.

Parameters:

Name	Type	Description	Default
hg	Array	Existing nodal homogeneous transformation matrices, (n_nodes, 4, 4)	required
phi	Array	Perturbation to the configuration vector, (n_nodes, 6)	required

Returns:

Type	Description
Array	Updated nodal homogeneous transformation matrices, (n_nodes, 4, 4)

Source code in src/flapjax/structure/beam.py

@staticmethod
def update_hg(hg: Array, phi: Array) -> Array:
    r"""
    Update the nodal homogeneous transformation matrices with the configuration increments.
    :param hg: Existing nodal homogeneous transformation matrices, ``(n_nodes, 4, 4)``
    :param phi: Perturbation to the configuration vector, ``(n_nodes, 6)``
    :return: Updated nodal homogeneous transformation matrices, ``(n_nodes, 4, 4)``
    """
    return jnp.einsum(
        "ijk,ikl->ijl",
        hg,
        vmap(exp_se3, 0, 0)(phi.reshape(-1, 6)),
    )

static_solve

static_solve(
    prescribed_dofs: Sequence[int] | Array | slice | int,
    f_ext_follower: Array | None = None,
    f_ext_dead: Array | None = None,
    f_ext_aero: Array | None = None,
    load_steps: int = 1,
    *,
    print_header: bool = True,
    postprocess_constraints: bool = True,
) -> StructureCase

Perform static solve of the structure under external loads.

Parameters:

Name	Type	Description	Default
f_ext_follower	Array | None	External forces array of follower forces (n_node, 6).	None
f_ext_dead	Array | None	External forces array of dead loads (n_node, 6).	None
f_ext_aero	Array | None	External forces array of aerodynamic loads (n_node, 6).	None
prescribed_dofs	Sequence[int] | Array | slice | int	Index of degrees of freedom which are prescribed (not solved for).	required
load_steps	int	Number of load steps to apply the external loads over.	1
print_header	bool	If False, suppress the "Static Solve" table header and trailing line.	True
postprocess_constraints	bool	If True, apply constraint postprocessing to the final solution.	True

Returns:

Type	Description
StructureCase	StructureCase object containing results of the static analysis.

Source code in src/flapjax/structure/beam.py

def static_solve(
    self,
    prescribed_dofs: Sequence[int] | Array | slice | int,
    f_ext_follower: Array | None = None,
    f_ext_dead: Array | None = None,
    f_ext_aero: Array | None = None,
    load_steps: int = 1,
    *,
    print_header: bool = True,
    postprocess_constraints: bool = True,
) -> StructureCase:
    r"""
    Perform static solve of the structure under external loads.
    :param f_ext_follower: External forces array of follower forces ``(n_node, 6)``.
    :param f_ext_dead: External forces array of dead loads ``(n_node, 6)``.
    :param f_ext_aero: External forces array of aerodynamic loads ``(n_node, 6)``.
    :param prescribed_dofs: Index of degrees of freedom which are prescribed (not solved for).
    :param load_steps: Number of load steps to apply the external loads over.
    :param print_header: If False, suppress the "Static Solve" table header and trailing line.
    :param postprocess_constraints: If True, apply constraint postprocessing to the final solution.
    :return: StructureCase object containing results of the static analysis.
    """

    if load_steps < 1:
        raise ValueError("load_steps must be at least 1")

    # check inputs
    if f_ext_follower is not None:
        check_arr_shape(f_ext_follower, (self.n_nodes, 6), "f_ext_follower")
    if f_ext_dead is not None:
        check_arr_shape(f_ext_dead, (self.n_nodes, 6), "f_ext_dead")

    if not (0.0 < self.relaxation_factor <= 1.0):
        raise ValueError("struct_relaxation_factor must be in the range (0, 1]")

    # degrees of freedom to solve for
    prescribed_dofs_: tuple[int, ...] = self.make_prescribed_dofs_tuple(
        prescribed_dofs
    )
    solve_dofs: Array = jnp.array(
        get_solve_dofs(n_dof=self.n_dof, prescribed_dofs=prescribed_dofs_)
    )

    # process external forces for load stepping
    load_step_weight: Array = jnp.linspace(0.0, 1.0, load_steps + 1)[
        1:
    ]  # (load_steps, )

    f_ext_follower_steps = self._make_load_steps_f(
        f_ext_follower, load_step_weight, apply_alpha_weighting=False
    )
    f_ext_dead_steps = self._make_load_steps_f(
        f_ext_dead, load_step_weight, apply_alpha_weighting=False
    )
    f_ext_aero_steps = self._make_load_steps_f(
        f_ext_aero, load_step_weight, apply_alpha_weighting=False
    )

    def _update(
        i_load_step: int,
        converge_status: ConvergenceStatus,
        hg_n: Array,
    ) -> tuple[int, ConvergenceStatus, Array]:
        # base parameters
        d_n = self.make_d(hg_n)  # (n_elem, 6)
        p_d_n = self.make_p_d(d_n)  # (n_elem, 6, 12)
        eps_n = self.make_eps(d_n)  # (n_elem, 6)
        m_t = self.make_m_t(d_n) if self.use_gravity else None  # (n_elem, 12, 12)

        # get total dead forces for this load step, (n_node, 6)
        total_f_ext_dead_step = self.make_f_ext_dead_tot(
            f_ext_dead_steps, f_ext_aero_steps, i_load_step
        )

        # assemble tangent stiffness matrix, (n_dof, n_dof)
        k_t_full_n = self.make_k_t_full(
            d=d_n,
            p_d=p_d_n,
            eps=eps_n,
            f_ext_dead=total_f_ext_dead_step,
            rmat=hg_n[:, :3, :3],
            m_t=m_t,
        )
        # apply nodal constraint contributions
        k_t_full_n = self.apply_nodal_constraint_tangent(
            mat=k_t_full_n, hg=hg_n, i_ts=0, gamma_prime=None
        )
        k_t_solve_n = k_t_full_n[jnp.ix_(solve_dofs, solve_dofs)]

        # compute residual forces, (n_solve_dofs, )
        f_res_solve_n, f_abs_sum_n = self.make_f_res(
            solve_dofs=solve_dofs,
            p_d=p_d_n,
            eps=eps_n,
            hg=hg_n,
            f_ext_follower_n=f_ext_follower_steps[i_load_step, ...]
            if f_ext_follower_steps is not None
            else None,
            f_ext_dead_n=total_f_ext_dead_step,
            thrust_n=self.thrust_reference,  # use reference thrust in static case
            dynamic=False,
            m_t=m_t,
            c_l=None,
            c_l_lumped=None,
            v=None,
            v_dot=None,
        )

        # solve for configuration increment, (n_solve_dofs, )
        if self.n_holonomic_constraints:
            d_varphi_np1, _, _ = self.solve_constrained(
                sys_mat_solve=k_t_solve_n,
                f_res_solve=f_res_solve_n,
                hg_eval=hg_n,
                solve_dofs=solve_dofs,
            )
            d_varphi_np1 *= self.relaxation_factor
        else:
            d_varphi_np1 = (
                jnp.linalg.solve(k_t_solve_n, f_res_solve_n)
                * self.relaxation_factor
            )

        # update configuration, (n_nodes, 4, 4)
        hg_np1_full = self.update_hg(
            hg_n, jnp.zeros(self.n_dof).at[solve_dofs].set(d_varphi_np1)
        )

        # algebra between undeformed and deformed shape, used to check relative convergence, (n_solve_dofs, )
        # this is relatively expensive to compute
        if self.struct_convergence_settings.rel_disp_tol is not None:
            h_full = vmap(hg_to_d, (0, 0), 0)(self.hg0, hg_np1_full).ravel()[
                solve_dofs
            ]
        else:
            h_full = None

        # update convergence status
        converge_status.update(
            delta_disp=d_varphi_np1,
            total_disp=h_full,
            delta_force=f_res_solve_n,
            total_force=f_abs_sum_n,
        )

        if map_verbosity_level(get_verbosity()) >= map_verbosity_level("verbose"):
            converge_status.print_struct_message(
                i_ts=None, t=None, i_load_step=i_load_step
            )

        return i_load_step, converge_status, hg_np1_full

    def convergence_loop(
        i_load_step: int,
        hg_init: Array,
    ) -> Array:
        r"""
        Convergence loop
        :param i_load_step: Index of load step.
        :param hg_init: Initial coordinates, ``(n_nodes, 4, 4)``.
        :return: Converged coordinates, ``(n_nodes, 4, 4)``.
        """
        _, convergence_status, hg_solve = eqxi.while_loop(
            lambda args_: ~args_[1].get_status(),
            lambda args_: _update(*args_),
            (
                i_load_step,
                ConvergenceStatus(
                    self.struct_convergence_settings,
                ),
                hg_init,
            ),
            max_steps=self.struct_convergence_settings.max_n_iter,
            kind="bounded",
        )

        if map_verbosity_level(get_verbosity()) >= map_verbosity_level("normal"):
            convergence_status.print_struct_message(
                i_ts=None, t=None, i_load_step=i_load_step
            )

        return hg_solve

    if print_header and map_verbosity_level(get_verbosity()) >= map_verbosity_level(
        "normal"
    ):
        ConvergenceStatus.print_header(dynamic=False)

    # solve for each load step
    hg = jax.lax.fori_loop(
        0,
        load_steps,
        lambda *args: convergence_loop(*args),
        self.hg0,
    )

    if print_header and map_verbosity_level(get_verbosity()) >= map_verbosity_level(
        "normal"
    ):
        ConvergenceStatus.print_line(dynamic=False)

    # postprocess final results
    d, eps, f_ext_dead_local, f_ext_aero_local, f_grav, f_int, _, _, f_res = (
        self.resolve_forces(
            hg=hg,
            dynamic=False,
            f_ext_dead=f_ext_dead,
            f_ext_follower=f_ext_follower,
            f_ext_aero=f_ext_aero,
            thrust=self.thrust_reference,
            v=None,
            v_dot=None,
        )
    )
    varphi = self.compute_varphi_from_hg(hg)
    f_elem = self.make_f_elem(eps=eps)  # compute loads in each element

    result = StructureCase(
        hg=hg,
        conn=self.connectivity,
        o0=self.o0,
        d=d,
        eps=eps,
        varphi=varphi,
        f_int=f_int,
        f_elem=f_elem,
        f_ext_follower=f_ext_follower,
        f_ext_dead=f_ext_dead_local,
        f_ext_aero=f_ext_aero_local,
        f_grav=f_grav,
        f_res=f_res,
        thrust=self.thrust_reference,
        thrust_nodes=self.thrust_nodes,
        thrust_direction=self.thrust_direction,
        prescribed_dofs=prescribed_dofs_,
        t=jnp.zeros(1),
    )
    if postprocess_constraints:
        result.constraint_data = self.postprocess_constraints(hg)
    return result

base_dynamic_solve

base_dynamic_solve(
    struct_case: StructureCase,
    struct_convergence_status: ConvergenceStatus,
    t: Array,
    solve_dofs: tuple[int, ...],
    load_steps: int,
    f_ext_dead: Array | None,
    f_ext_follower: Array | None,
    thrust_t: dict[str, Array],
    aero_obj: None,
    aero_case: None,
    fsi_convergence_status: None,
    cs_ang_t: None,
    cs_vel_t: None,
) -> StructureCase

base_dynamic_solve(
    struct_case: StructureCase,
    struct_convergence_status: ConvergenceStatus,
    t: Array,
    solve_dofs: tuple[int, ...],
    load_steps: int,
    f_ext_dead: Array | None,
    f_ext_follower: Array | None,
    thrust_t: dict[str, Array],
    aero_obj: DynamicAeroSolver,
    aero_case: AeroCase,
    fsi_convergence_status: ConvergenceStatus,
    cs_ang_t: dict[str, Array],
    cs_vel_t: dict[str, Array],
) -> AeroelasticCase

base_dynamic_solve(
    struct_case: StructureCase,
    struct_convergence_status: ConvergenceStatus,
    t: Array,
    solve_dofs: tuple[int, ...],
    load_steps: int,
    f_ext_dead: Array | None,
    f_ext_follower: Array | None,
    thrust_t: dict[str, Array],
    aero_obj: DynamicAeroSolver | None,
    aero_case: AeroCase | None,
    fsi_convergence_status: ConvergenceStatus | None,
    cs_ang_t: dict[str, Array] | None,
    cs_vel_t: dict[str, Array] | None,
) -> StructureCase | AeroelasticCase

Generic dynamic solver. Both the structural dynamic solve, and aeroelastic dynamic solve, are formed as wrappers of this

Source code in src/flapjax/structure/beam.py

def base_dynamic_solve(
    self,
    struct_case: StructureCase,
    struct_convergence_status: ConvergenceStatus,
    t: Array,
    solve_dofs: tuple[int, ...],
    load_steps: int,
    f_ext_dead: Array | None,
    f_ext_follower: Array | None,
    thrust_t: dict[str, Array],
    aero_obj: DynamicAeroSolver | None,
    aero_case: AeroCase | None,
    fsi_convergence_status: ConvergenceStatus | None,
    cs_ang_t: dict[str, Array] | None,
    cs_vel_t: dict[str, Array] | None,
) -> StructureCase | AeroelasticCase:
    r"""
    Generic dynamic solver. Both the structural dynamic solve, and aeroelastic dynamic solve, are formed as wrappers
    of this
    """

    if not (0.0 < self.relaxation_factor <= 1.0):
        raise ValueError("Relaxation factor must be in range (0, 1]")

    n_tstep = len(t)

    include_aero: bool = aero_obj is not None

    # process external forces for load stepping
    load_step_weight: Array = jnp.linspace(0.0, 1.0, load_steps + 1)[
        1:
    ]  # (load_steps, )
    f_ext_follower_alpha_steps = self._make_load_steps_f(
        f_ext_follower, load_step_weight, apply_alpha_weighting=True
    )
    f_ext_dead_alpha_steps = self._make_load_steps_f(
        f_ext_dead, load_step_weight, apply_alpha_weighting=True
    )

    solve_dofs_arr: Array = jnp.array(solve_dofs)
    prescribed_dofs_arr: Array = jnp.array(
        sorted(set(range(self.n_dof)) - set(solve_dofs)), dtype=int
    )

    def _update(
        i_load_step: int,
        i_ts: int,
        struct_convergence_status_: ConvergenceStatus,
        hg_n: Array,
        phi_alpha: Array,
        q_alpha: StructureMinimalStates,
        f_ext_aero_alpha_steps: Array | None,
        thrust_alpha: dict[str, Array],
    ) -> tuple[
        int,
        int,
        ConvergenceStatus,
        Array,
        Array,
        StructureMinimalStates,
        Array | None,
        dict[str, Array],
    ]:
        r"""
        Solution update for a single iteration of the nonlinear solver at a given time step and load step.
        :param i_load_step: Load step index.
        :param i_ts: Time step index.
        :param struct_convergence_status_: ConvergenceStatus object for the current iteration, used to track
        convergence and print messages.
        :param hg_n: Transformation matrices at iteration varphi, ``(n_nodes, 4, 4)``.
        :param phi_alpha: Timestep increment to the alpha step, ``(n_nodes, 6)``.
        :param f_ext_aero_alpha_steps: Load steps for the external aerodynamic forcing, ``(n_steps, n_nodes, 6)``.
        :param thrust_alpha: Thrust magnitude at the alpha step, ``{keys: ()}``.
        :return: Load and time step indices, updated ConvergenceStatus object, updated transformation matrices,
        configuration, velocities and accelerations for iteration n+1.
        """

        hg_update = self.update_hg(hg_n, phi_alpha)  # (n_node, 4, 4)

        # base parameters
        d_n = self.make_d(hg_update)  # (n_elem, 6)
        p_d_n = self.make_p_d(d_n)  # (n_elem, 6, 12)
        eps_n = self.make_eps(d_n)  # (n_elem, 6)
        d_dot_n = self._make_d_dot(p_d_n, q_alpha.v)  # (n_elem, 6)
        t_n = vmap(t_se3, 0, 0)(phi_alpha)  # (n_node, 6, 6)

        # tangent matrices
        m_t = self.make_m_t(d_n)  # (n_elem, 12, 12)
        c_l, c_t = self._make_c_t(
            d_n, d_dot_n, q_alpha.v
        )  # (n_elem, 12, 12), (n_elem, 12, 12)

        total_f_ext_dead = self.make_f_ext_dead_tot(
            f_ext_dead=f_ext_dead_alpha_steps[:, i_ts, :, :]
            if f_ext_dead_alpha_steps is not None
            else None,
            f_ext_aero=f_ext_aero_alpha_steps,
            i_load_step=i_load_step,
        )  # (n_node, 6)

        k_t = self.make_k_t_full(
            d_n,
            p_d_n,
            eps_n,
            total_f_ext_dead,
            hg_update[:, :3, :3],
            m_t,
        )  # (n_dof, n_dof)

        # add lumped mass contributions if applicable
        if self.use_lumped_mass:
            c_l_lumped, c_t_lumped = self._make_c_t_lumped(
                q_alpha.v
            )  # (n_node, 6, 6), (n_node, 6, 6)
        else:
            c_l_lumped, c_t_lumped = None, None

        # residual forces, (n_solve_dofs, )
        f_res_n_solve, f_abs_sum_n = self.make_f_res(
            solve_dofs=solve_dofs_arr,
            p_d=p_d_n,
            eps=eps_n,
            hg=hg_update,
            f_ext_follower_n=f_ext_follower_alpha_steps[i_load_step, i_ts, ...]
            if f_ext_follower_alpha_steps is not None
            else None,
            f_ext_dead_n=total_f_ext_dead,
            thrust_n=thrust_alpha,
            dynamic=True,
            m_t=m_t,
            c_l=c_l,
            c_l_lumped=c_l_lumped,
            v=q_alpha.v,
            v_dot=q_alpha.v_dot,
            i_ts=i_ts,
            k_t_assembled=k_t,
        )

        # system matrix, (n_dof, n_dof)
        sys_mat_full = self._make_sys_matrix(
            m_t=m_t,
            c_t=c_t,
            c_t_lumped=c_t_lumped,
            k_t=k_t,
            t_n=t_n,
            ti=self.time_integrator,
        )
        # add nodal constraint contributions
        sys_mat_full = self.apply_nodal_constraint_tangent(
            mat=sys_mat_full,
            hg=hg_update,
            i_ts=i_ts,
            gamma_prime=self.time_integrator.gamma_prime,
        )
        sys_mat = sys_mat_full[jnp.ix_(solve_dofs_arr, solve_dofs_arr)]

        # solve for configuration increment, (n_solve_dofs, )
        if self.multibody_constraints:
            d_n_np1, _, _ = self.solve_constrained(
                sys_mat_solve=sys_mat,
                f_res_solve=f_res_n_solve,
                hg_eval=hg_update,
                solve_dofs=solve_dofs_arr,
                hg_base=hg_n,
                phi=phi_alpha,
                v=q_alpha.v,
                gamma_prime=self.time_integrator.gamma_prime,
            )
            d_n_np1 *= self.relaxation_factor
        else:
            d_n_np1 = (
                jnp.linalg.solve(sys_mat, f_res_n_solve) * self.relaxation_factor
            )
        phi_np1 = phi_alpha.ravel().at[solve_dofs_arr].add(d_n_np1).reshape(-1, 6)

        # update configuration, velocities and accelerations
        v_np1 = (
            q_alpha.v.ravel()
            .at[solve_dofs_arr]
            .add(self.time_integrator.gamma_prime * d_n_np1)
            .reshape(-1, 6)
        )
        v_dot_np1 = (
            q_alpha.v_dot.ravel()
            .at[solve_dofs_arr]
            .add(self.time_integrator.beta_prime * d_n_np1)
            .reshape(-1, 6)
        )

        # update convergence status
        struct_convergence_status_.update(
            delta_disp=d_n_np1,
            total_disp=phi_np1,
            delta_force=f_res_n_solve,
            total_force=f_abs_sum_n,
        )

        if map_verbosity_level(get_verbosity()) >= map_verbosity_level("verbose"):
            struct_convergence_status_.print_struct_message(
                i_ts=i_ts, t=t[i_ts], i_load_step=i_load_step
            )

        q_alpha_update = StructureMinimalStates(
            varphi=None, v=v_np1, v_dot=v_dot_np1, a=q_alpha.a
        )

        return (
            i_load_step,
            i_ts,
            struct_convergence_status_,
            hg_n,
            phi_np1,
            q_alpha_update,
            f_ext_aero_alpha_steps,
            thrust_alpha,
        )

    @overload
    def time_step_loop(
        i_ts: int,
        struct_sol: StructureCase,
        struct_convergence_status_: ConvergenceStatus,
        aero_sol: None,
        fsi_convergence_status_: None,
        thrust_t_: dict[str, Array],
        cs_ang_t_: None,
        cs_vel_t_: None,
    ) -> tuple[
        StructureCase,
        ConvergenceStatus,
        None,
        None,
        dict[str, Array],
        None,
        None,
    ]: ...

    @overload
    def time_step_loop(
        i_ts: int,
        struct_sol: StructureCase,
        struct_convergence_status_: ConvergenceStatus,
        aero_sol: AeroCase,
        fsi_convergence_status_: ConvergenceStatus,
        thrust_t_: dict[str, Array],
        cs_ang_t_: dict[str, Array],
        cs_vel_t_: dict[str, Array],
    ) -> tuple[
        StructureCase,
        ConvergenceStatus,
        AeroCase,
        ConvergenceStatus,
        dict[str, Array],
        dict[str, Array],
        dict[str, Array],
    ]: ...

    def time_step_loop(
        i_ts: int,
        struct_sol: StructureCase,
        struct_convergence_status_: ConvergenceStatus,
        aero_sol: AeroCase | None,
        fsi_convergence_status_: ConvergenceStatus | None,
        thrust_t_: dict[str, Array],
        cs_ang_t_: dict[str, Array] | None,
        cs_vel_t_: dict[str, Array] | None,
    ) -> tuple[
        StructureCase,
        ConvergenceStatus,
        AeroCase | None,
        ConvergenceStatus | None,
        dict[str, Array],
        dict[str, Array] | None,
        dict[str, Array] | None,
    ]:
        r"""
        Performs analysis on a single time step, including load stepping
        :param i_ts: Index of time step to solve
        :param struct_sol: Solution object, with results up to time step i_ts-1.
        :param struct_convergence_status_: Convergence status object.
        :param aero_sol: Aero solution object, with results up to time step i_ts-1, if aero is included.
        :param fsi_convergence_status_: Convergence status object.
        :param thrust_t_: Thrust magnitude time history, ``{name: (n_tstep, )}``.
        :param cs_ang_t_: Control surface angle time history, ``{name: (n_tstep, )}``.
        :param cs_vel_t_: Control surface velocity time history, ``{name: (n_tstep, )}``.
        :return: Solution object with results up to time step i_ts.
        """

        # predictor step
        q_nm1 = struct_sol.get_minimal_states(i_ts - 1)
        phi_init, q_init = self.time_integrator.predict_q(q_nm1)
        phi_alpha_init, q_alpha_init = self.time_integrator.compute_q_alpha(
            q_nm1=q_nm1,
            q_n=q_init,
            phi_n=phi_init,
        )

        # prescribed DOFs should not be influenced by the time integration
        phi_alpha_init = (
            phi_alpha_init.ravel().at[prescribed_dofs_arr].set(0.0).reshape(-1, 6)
        )
        q_alpha_init.v = (
            q_alpha_init.v.ravel()
            .at[prescribed_dofs_arr]
            .set(q_nm1.v.ravel()[prescribed_dofs_arr])
            .reshape(-1, 6)
        )
        q_alpha_init.v_dot = (
            q_alpha_init.v_dot.ravel()
            .at[prescribed_dofs_arr]
            .set(q_nm1.v_dot.ravel()[prescribed_dofs_arr])
            .reshape(-1, 6)
        )
        q_alpha_init.a = (
            q_alpha_init.a.ravel()
            .at[prescribed_dofs_arr]
            .set(q_nm1.a.ravel()[prescribed_dofs_arr])
            .reshape(-1, 6)
        )

        q_alpha_init.varphi = None  # this value is not used during the loop

        # thrust force
        thrust_alpha: dict[str, Array] = {
            k: self.time_integrator.compute_f_alpha(f_nm1=v[i_ts - 1], f_n=v[i_ts])
            for k, v in thrust_t_.items()
        }
        thrust_n: dict[str, Array] = {k: v[i_ts] for k, v in thrust_t_.items()}

        if include_aero:
            assert (
                aero_sol is not None
                and fsi_convergence_status_ is not None
                and struct_sol.f_ext_aero is not None
                and cs_ang_t_ is not None
                and cs_vel_t_ is not None
            )

            fsi_convergence_status_.reset_status()

            # f_ext_aero is stored in local frame, so we convert back to global
            # so that both operands of the alpha blend are in the same (global) frame.
            f_aero_nm1 = jnp.concatenate(
                [
                    jnp.einsum(
                        "ijk,ik->ij",
                        struct_sol.hg[i_ts - 1, :, :3, :3],
                        struct_sol.f_ext_aero[i_ts - 1, :, :3],
                    ),
                    jnp.einsum(
                        "ijk,ik->ij",
                        struct_sol.hg[i_ts - 1, :, :3, :3],
                        struct_sol.f_ext_aero[i_ts - 1, :, 3:],
                    ),
                ],
                axis=-1,
            )

            # get control surface angles and velocities
            cs_ang_nm1 = {k: v[i_ts - 1] for k, v in cs_ang_t_.items()}
            cs_ang_n = {k: v[i_ts] for k, v in cs_ang_t_.items()}
            cs_vel_n = {k: v[i_ts] for k, v in cs_vel_t_.items()}
            assert fsi_convergence_status is not None
            (
                _,
                struct_sol,
                aero_sol,
                struct_convergence_status_,
                fsi_convergence_status_,
                phi_alpha,
                q_alpha,
                *_,
            ) = eqxi.while_loop(
                lambda args_: ~cast(ConvergenceStatus, args_[4]).get_status(),
                lambda args_: fsi_convergence_loop(*args_),
                (
                    i_ts,
                    struct_sol,
                    aero_sol,
                    struct_convergence_status_,
                    fsi_convergence_status_,
                    phi_alpha_init,
                    q_alpha_init,
                    f_aero_nm1,  # this value is for the previous timesteps force, and is propagated unaltered
                    f_aero_nm1,  # first guess for forcing at alpha is to use value from i_ts=n-1
                    thrust_alpha,
                    cs_ang_n,
                    cs_ang_nm1,
                    cs_vel_n,
                ),
                max_steps=fsi_convergence_status.convergence_settings.max_n_iter,
                kind="bounded",
            )

        else:
            # solve pure structural problem
            _, struct_convergence_status_, _, phi_alpha, q_alpha, *_ = (
                load_step_loop(
                    i_ts=i_ts,
                    struct_convergence_status_=struct_convergence_status_,
                    hg_alpha=struct_sol.hg[i_ts - 1, ...],
                    phi_alpha=phi_alpha_init,
                    q_alpha=q_alpha_init,
                    f_ext_aero_steps=None,
                    thrust_alpha=thrust_alpha,
                )
            )

        # print message where we only require one message per timestep
        if map_verbosity_level(get_verbosity()) == map_verbosity_level("normal"):
            struct_convergence_status_.print_struct_message(
                i_ts=i_ts, t=struct_sol.t[i_ts], i_load_step=load_steps - 1
            )
            if include_aero and fsi_convergence_status_ is not None:
                fsi_convergence_status_.print_fsi_message(
                    i_ts=i_ts, t=struct_sol.t[i_ts]
                )

        # postprocess results for time step and store in solution object
        q_n, phi_n = self.time_integrator.compute_q_n_from_q_alpha(
            q_alpha=q_alpha,
            q_nm1=struct_sol.get_minimal_states(i_ts - 1),
            phi_alpha=phi_alpha,
        )

        # update pseudo-acceleration
        q_n.a = self.time_integrator.compute_a_n(
            a_nm1=struct_sol.a[i_ts - 1, ...],
            v_dot_nm1=struct_sol.v_dot[i_ts - 1, ...],
            v_dot_n=q_n.v_dot,
        )

        # final node coordinates
        hg_n = self.update_hg(struct_sol.hg[i_ts - 1, ...], phi_n)

        if include_aero:
            if (
                aero_sol is None
                or aero_obj is None
                or fsi_convergence_status_ is None
            ):
                raise ValueError("Missing aero arguments")

            f_ext_aero = aero_sol.project_forcing_to_beam(
                i_ts=i_ts,
                rmat=hg_n[:, :3, :3],
                x0_aero=aero_obj.zeta_b0,
                include_unsteady=aero_obj.include_unsteady_force,
            )

        else:
            f_ext_aero = None

        (
            d,
            eps,
            f_ext_dead_local,
            f_ext_aero_local,
            f_grav,
            f_int,
            f_gyr,
            f_iner,
            f_res,
        ) = self.resolve_forces(
            hg=hg_n,
            dynamic=True,
            f_ext_dead=f_ext_dead[i_ts, ...] if f_ext_dead is not None else None,
            f_ext_follower=f_ext_follower[i_ts, ...]
            if f_ext_follower is not None
            else None,
            thrust=thrust_n,
            f_ext_aero=f_ext_aero,
            v=q_n.v,
            v_dot=q_n.v_dot,
        )
        struct_sol.d = struct_sol.d.at[i_ts, ...].set(d)
        struct_sol.eps = struct_sol.eps.at[i_ts, ...].set(eps)
        struct_sol.v = struct_sol.v.at[i_ts, ...].set(q_n.v)
        struct_sol.v_dot = struct_sol.v_dot.at[i_ts, ...].set(q_n.v_dot)
        struct_sol.a = struct_sol.a.at[i_ts, ...].set(q_n.a)
        struct_sol.hg = struct_sol.hg.at[i_ts, ...].set(hg_n)
        struct_sol.varphi = struct_sol.varphi.at[i_ts, ...].set(
            vmap(hg_to_d, (0, 0), 0)(self.hg0, hg_n)
        )

        if f_ext_follower is not None and struct_sol.f_ext_follower is not None:
            struct_sol.f_ext_follower = struct_sol.f_ext_follower.at[i_ts, ...].set(
                f_ext_follower[i_ts, ...]
            )
        if f_ext_dead is not None and struct_sol.f_ext_dead is not None:
            struct_sol.f_ext_dead = struct_sol.f_ext_dead.at[i_ts, ...].set(
                f_ext_dead_local
            )

        if f_ext_aero is not None and struct_sol.f_ext_aero is not None:
            struct_sol.f_ext_aero = struct_sol.f_ext_aero.at[i_ts, ...].set(
                f_ext_aero_local
            )

        if self.use_gravity:
            if struct_sol.f_grav is None:
                raise ValueError("struct_sol.f_grav is None")
            struct_sol.f_grav = struct_sol.f_grav.at[i_ts, ...].set(f_grav)
        struct_sol.f_int = struct_sol.f_int.at[i_ts, ...].set(f_int)
        struct_sol.f_elem = struct_sol.f_elem.at[i_ts, ...].set(
            self.make_f_elem(eps=eps)
        )
        struct_sol.f_iner_gyr = struct_sol.f_iner_gyr.at[i_ts, ...].set(
            f_iner + f_gyr
        )
        struct_sol.f_res = struct_sol.f_res.at[i_ts, ...].set(f_res)

        if include_aero and aero_sol is not None:
            assert cs_ang_t_ is not None and cs_vel_t_ is not None
            cs_ang_n = {k: v[i_ts] for k, v in cs_ang_t_.items()}
            cs_vel_n = {k: v[i_ts] for k, v in cs_vel_t_.items()}
            aero_sol.cs_ang = {
                k: v.at[i_ts].set(cs_ang_n[k]) for k, v in aero_sol.cs_ang.items()
            }
            aero_sol.cs_vel = {
                k: v.at[i_ts].set(cs_vel_n[k]) for k, v in aero_sol.cs_vel.items()
            }

        return (
            struct_sol,
            struct_convergence_status_,
            aero_sol,
            fsi_convergence_status_,
            thrust_t_,
            cs_ang_t_,
            cs_vel_t_,
        )

    def fsi_convergence_loop(
        i_ts: int,
        struct_sol: StructureCase,
        aero_sol: AeroCase,
        struct_convergence_status_: ConvergenceStatus,
        fsi_convergence_status_: ConvergenceStatus,
        phi_alpha_init: Array,
        q_alpha_init: StructureMinimalStates,
        f_aero_nm1: Array,
        f_aero_alpha_prev: Array,
        thrust_alpha: dict[str, Array],
        cs_ang_n: dict[str, Array],
        cs_ang_nm1: dict[str, Array],
        cs_vel_n: dict[str, Array],
    ) -> tuple[
        int,
        StructureCase,
        AeroCase,
        ConvergenceStatus,
        ConvergenceStatus,
        Array,
        StructureMinimalStates,
        Array,
        Array,
        dict[str, Array],
        dict[str, Array],
        dict[str, Array],
        dict[str, Array],
    ]:
        # obtain coordinates at timestep (not alpha)
        phi_n = self.time_integrator.compute_phi_from_phi_alpha(
            phi_alpha=phi_alpha_init
        )
        v_n = self.time_integrator.compute_v_from_v_alpha(
            v_alpha=q_alpha_init.v, v_nm1=struct_sol.v[i_ts - 1, ...]
        )

        hg_n = self.update_hg(hg=struct_sol.hg[i_ts - 1, ...], phi=phi_n)
        hg_dot = self.make_hg_dot(hg=hg_n, v=v_n)

        if aero_obj is None or struct_sol.f_ext_aero is None:
            raise ValueError("Missing aero parameters")

        # evaluate aerodynamic forcing on beam
        aero_sol = aero_obj.case_solve(
            case=aero_sol,
            i_ts=i_ts,
            hg_n=hg_n,
            hg_nm1=struct_sol.hg[i_ts - 1, ...],
            hg_dot_n=hg_dot,
            static=False,
            horseshoe=False,
            cs_ang_n=cs_ang_n,
            cs_ang_nm1=cs_ang_nm1,
            cs_vel_n=cs_vel_n,
        )

        f_aero_n = aero_sol.project_forcing_to_beam(
            i_ts=i_ts,
            rmat=hg_n[:, :3, :3],
            x0_aero=aero_obj.zeta_b0,
            include_unsteady=aero_obj.include_unsteady_force,
        )

        # aerodynamic force at alpha point, subsequently divided into load steps
        f_aero_alpha = self.time_integrator.compute_f_alpha(
            f_nm1=f_aero_nm1, f_n=f_aero_n
        )

        f_aero_alpha_steps = self._make_load_steps_f(
            f=f_aero_alpha, weighting=load_step_weight, apply_alpha_weighting=False
        )

        # reset convergence status
        struct_convergence_status_.reset_status()

        # solve structural problem for given aero load
        _, struct_convergence_status_, _, phi_alpha, q_alpha, *_ = load_step_loop(
            i_ts,
            struct_convergence_status_,
            struct_sol.hg[i_ts - 1, ...],
            phi_alpha_init,
            q_alpha_init,
            f_aero_alpha_steps,
            thrust_alpha,
        )

        # update the FSI convergence object
        # note that for convenience we use the alpha properties
        fsi_convergence_status_.update(
            delta_disp=(phi_alpha_init - phi_alpha).ravel()[solve_dofs_arr],
            total_disp=phi_alpha.ravel()[solve_dofs_arr],
            delta_force=(f_aero_alpha - f_aero_alpha_prev).ravel()[solve_dofs_arr],
            total_force=f_aero_alpha.ravel()[solve_dofs_arr],
        )

        if map_verbosity_level(get_verbosity()) >= map_verbosity_level("verbose"):
            fsi_convergence_status_.print_fsi_message(i_ts=i_ts, t=t[i_ts])

        return (
            i_ts,
            struct_sol,
            aero_sol,
            struct_convergence_status_,
            fsi_convergence_status_,
            phi_alpha,
            q_alpha,
            f_aero_nm1,
            f_aero_alpha,
            thrust_alpha,
            cs_ang_n,
            cs_ang_nm1,
            cs_vel_n,
        )

    def struct_convergence_loop(
        i_load_step: int,
        i_ts: int,
        struct_convergence_status_: ConvergenceStatus,
        hg_alpha: Array,
        phi_alpha: Array,
        q_alpha: StructureMinimalStates,
        f_ext_aero_steps: Array | None,
        thrust_alpha: dict[str, Array],
    ) -> tuple[
        int,
        ConvergenceStatus,
        Array,
        Array,
        StructureMinimalStates,
        Array | None,
        dict[str, Array],
    ]:
        r"""
        Convergence loop within each load step of a time step.
        :param i_load_step: Load step index.
        :param i_ts: Time step index.
        :param struct_convergence_status_: ConvergenceStatus object to update with convergence information during load
        stepping.
        :param hg_alpha: Node transformations at the beginning of the load step, ``(n_nodes, 4, 4)``.
        :param phi_alpha: Node configuration increments in algebra space, ``(n_nodes, 6)``.
        :param q_alpha: Minimal states at intermediate alpha step.
        :param f_ext_aero_steps: Optional aerodynamic forcing alpha load steps ``(n_steps, n_nodes, 6)``.
        :param thrust_alpha: Thrust at the alpha step, ``{key: ()}``.
        :return: Time step index, convergence status, and updated configuration, velocities, accelerations, and
        optional aerodynamic forcing.
        """

        struct_convergence_status_.reset_status()

        _, _, struct_convergence_status_, hg_solve, phi_alpha, q_alpha, _, _ = (
            eqxi.while_loop(
                lambda args_: ~args_[2].get_status(),
                lambda args_: _update(*args_),
                (
                    i_load_step,
                    i_ts,
                    struct_convergence_status_,
                    hg_alpha,
                    phi_alpha,
                    q_alpha,
                    f_ext_aero_steps,
                    thrust_alpha,
                ),
                max_steps=self.struct_convergence_settings.max_n_iter,
                kind="bounded",
            )
        )

        if map_verbosity_level(get_verbosity()) >= map_verbosity_level("verbose"):
            struct_convergence_status_.print_struct_message(
                i_ts=i_ts, t=t[i_ts], i_load_step=i_load_step
            )

        return (
            i_ts,
            struct_convergence_status_,
            hg_solve,
            phi_alpha,
            q_alpha,
            f_ext_aero_steps,
            thrust_alpha,
        )

    def load_step_loop(
        i_ts: int,
        struct_convergence_status_: ConvergenceStatus,
        hg_alpha: Array,
        phi_alpha: Array,
        q_alpha: StructureMinimalStates,
        f_ext_aero_steps: Array | None,
        thrust_alpha: dict[str, Array],
    ) -> tuple[
        int,
        ConvergenceStatus,
        Array,
        Array,
        StructureMinimalStates,
        Array | None,
    ]:
        r"""
        Performs load stepping iterations for a given time step. Load stepping is not performed for thrust.
        :param i_ts: Timestep index for which to perform load stepping.
        :param struct_convergence_status_: ConvergenceStatus object to update with load stepping convergence information.
        :param hg_alpha: SE(3) nodal transformation matrices at the beginning of the load step, ``(n_nodes, 4, 4)``.
        :param phi_alpha: Nodal updates to the configuration in the algebra space, ``(n_nodes, 6)``.
        :param q_alpha: Minimal states at intermediate alpha step.
        :param f_ext_aero_steps: Optional aerodynamic forcing alpha load steps ``(n_steps, n_nodes, 6)``.
        :param thrust_alpha: Thrust at the alpha step, ``{key: ()}``.
        :return: Time step index, updated ConvergenceStatus object, and updated configuration, velocities and accelerations after load stepping
        """
        return jax.lax.fori_loop(
            0,
            load_steps,
            lambda i_load_step, args: struct_convergence_loop(i_load_step, *args),
            (
                i_ts,
                struct_convergence_status_,
                hg_alpha,
                phi_alpha,
                q_alpha,
                f_ext_aero_steps,
                thrust_alpha,
            ),
        )

    def time_step_loop_checked(
        i_ts: int,
        struct_sol: StructureCase,
        struct_convergence_status_: ConvergenceStatus,
        aero_sol: AeroCase | None,
        fsi_convergence_status_: ConvergenceStatus | None,
        thrust_t_: dict[str, Array],
        cs_ang_t_: dict[str, Array] | None,
        cs_vel_t_: dict[str, Array] | None,
        diverged: Array,
    ) -> tuple[
        StructureCase,
        ConvergenceStatus,
        AeroCase | None,
        ConvergenceStatus | None,
        dict[str, Array],
        dict[str, Array] | None,
        dict[str, Array] | None,
        Array,
    ]:
        r"""
        Wraps ``time_step_loop`` with a check for solution divergence. Once a NaN is detected, this becomes a no-op
        for all remaining time steps. The corresponding time history entries are left at their initialised value
        (zero).
        """

        if include_aero:
            assert aero_sol is not None
            assert fsi_convergence_status_ is not None
            assert cs_ang_t_ is not None
            assert cs_vel_t_ is not None
            aero_sol_ok: AeroCase = aero_sol
            fsi_convergence_status_ok: ConvergenceStatus = fsi_convergence_status_
            cs_ang_t_ok: dict[str, Array] = cs_ang_t_
            cs_vel_t_ok: dict[str, Array] = cs_vel_t_
            false_branch = lambda: time_step_loop(
                i_ts,
                struct_sol,
                struct_convergence_status_,
                aero_sol_ok,
                fsi_convergence_status_ok,
                thrust_t_,
                cs_ang_t_ok,
                cs_vel_t_ok,
            )
        else:
            assert aero_sol is None
            assert fsi_convergence_status_ is None
            assert cs_ang_t_ is None
            assert cs_vel_t_ is None
            aero_sol_none: None = aero_sol
            fsi_convergence_status_none: None = fsi_convergence_status_
            cs_ang_t_none: None = cs_ang_t_
            cs_vel_t_none: None = cs_vel_t_
            false_branch = lambda: time_step_loop(
                i_ts,
                struct_sol,
                struct_convergence_status_,
                aero_sol_none,
                fsi_convergence_status_none,
                thrust_t_,
                cs_ang_t_none,
                cs_vel_t_none,
            )

        (
            struct_sol,
            struct_convergence_status_,
            aero_sol,
            fsi_convergence_status_,
            thrust_t_,
            cs_ang_t_,
            cs_vel_t_,
        ) = jax.lax.cond(
            diverged,
            lambda: (
                struct_sol,
                struct_convergence_status_,
                aero_sol,
                fsi_convergence_status_,
                thrust_t_,
                cs_ang_t_,
                cs_vel_t_,
            ),
            false_branch,
        )

        has_nan = struct_convergence_status_.has_nan
        if include_aero:
            assert fsi_convergence_status_ is not None
            has_nan = has_nan | fsi_convergence_status_.has_nan
        new_diverged = diverged | has_nan

        jax.lax.cond(
            new_diverged & ~diverged,
            lambda: warn(
                "NaN detected in dynamic solve at time step {i_ts} (t={t_val:.03e}) - skipping remaining time steps.",
                i_ts=i_ts,
                t_val=t[i_ts],
            ),
            lambda: None,
        )

        return (
            struct_sol,
            struct_convergence_status_,
            aero_sol,
            fsi_convergence_status_,
            thrust_t_,
            cs_ang_t_,
            cs_vel_t_,
            new_diverged,
        )

    struct_case, _, aero_case, *_ = jax.lax.fori_loop(
        1,
        n_tstep,
        lambda i_ts, args: time_step_loop_checked(i_ts, *args),
        (
            struct_case,
            struct_convergence_status,
            aero_case,
            fsi_convergence_status,
            thrust_t,
            cs_ang_t,
            cs_vel_t,
            jnp.zeros((), dtype=bool),
        ),
    )

    struct_case.constraint_data = self.postprocess_constraints(struct_case.hg)

    if include_aero:
        if aero_case is None:
            raise ValueError("aero_case cannot be None")

        from flapjax.coupled.data_structures import (
            AeroelasticCase,
        )  # import here to prevent circular references

        return AeroelasticCase(structure=struct_case, aero=aero_case)
    else:
        return struct_case

dynamic_solve

dynamic_solve(
    init_state: StructureCase | None,
    n_tstep: int,
    dt: Array | float,
    prescribed_dofs: Sequence[int]
    | Array
    | slice
    | int
    | None = None,
    f_ext_follower: Array | None = None,
    f_ext_dead: Array | None = None,
    f_ext_aero: Array | None = None,
    thrust_t: dict[str, Array] | None = None,
    load_steps: int = 1,
) -> StructureCase

Perform dynamic solve of the structure under external loads

Parameters:

Name	Type	Description	Default
init_state	StructureCase | None	Initial state of the structure, either static or a dynamic snapshot. If None, the reference configuration is used with zero velocities.	required
prescribed_dofs	Sequence[int] | Array | slice | int | None	Degrees of freedom which are prescribed (not solved for). If None, inherit from the initial state.	None
n_tstep	int	Number of time steps to simulate.	required
dt	Array | float	Time step length.	required
f_ext_follower	Array | None	Following external forces array, (n_tstep, n_node, 6), (n_node, 6) or None for zero external follower forces.	None
f_ext_dead	Array | None	Dead external forces array, (n_tstep, n_node, 6), (n_node, 6) or None for zero external dead forces.	None
f_ext_aero	Array | None	Aerodynamic external forces array, (n_tstep, n_node, 6), (n_node, 6) or None for zero external aerodynamic forces.	None
thrust_t	dict[str, Array] | None	Thrust time history, {key: (n_tstep, )}. If none, this will use the reference value.	None
load_steps	int	Number of load steps to apply the external loads over.	1

Returns:

Type	Description
StructureCase	Structure dataclass containing results of the dynamic analysis.

Source code in src/flapjax/structure/beam.py

def dynamic_solve(
    self,
    init_state: StructureCase | None,
    n_tstep: int,
    dt: Array | float,
    prescribed_dofs: Sequence[int] | Array | slice | int | None = None,
    f_ext_follower: Array | None = None,
    f_ext_dead: Array | None = None,
    f_ext_aero: Array | None = None,
    thrust_t: dict[str, Array] | None = None,
    load_steps: int = 1,
) -> StructureCase:
    r"""
    Perform dynamic solve of the structure under external loads
    :param init_state: Initial state of the structure, either static or a
    dynamic snapshot. If None, the reference configuration is used with zero
    velocities.
    :param prescribed_dofs: Degrees of freedom which are prescribed (not solved for). If None, inherit
    from the initial state.
    :param n_tstep: Number of time steps to simulate.
    :param dt: Time step length.
    :param f_ext_follower: Following external forces array, ``(n_tstep, n_node, 6)``, ``(n_node, 6)`` or None for zero external follower forces.
    :param f_ext_dead: Dead external forces array, ``(n_tstep, n_node, 6)``, ``(n_node, 6)`` or None for zero external dead forces.
    :param f_ext_aero: Aerodynamic external forces array, ``(n_tstep, n_node, 6)``, ``(n_node, 6)`` or None for zero external aerodynamic forces.
    :param thrust_t: Thrust time history, ``{key: (n_tstep, )}``. If none, this will use the reference value.
    :param load_steps: Number of load steps to apply the external loads over.
    :return: Structure dataclass containing results of the dynamic analysis.
    """

    if load_steps <= 0:
        raise ValueError("load_steps must be a positive integer")

    # set thrust if not provided
    thrust_t_: dict[str, Array] = (
        thrust_t
        if thrust_t is not None
        else {k: jnp.full(n_tstep, v) for k, v in self.thrust_reference.items()}
    )

    if prescribed_dofs is None:
        # inherit prescribed dofs from initial state
        if init_state is None:
            raise ValueError("prescribed_dofs cannot be None if init_state is None")
        prescribed_dofs = init_state.prescribed_dofs

    # degrees of freedom to solve for
    prescribed_dofs_arr = self.make_prescribed_dofs_tuple(prescribed_dofs)
    solve_dofs = get_solve_dofs(
        n_dof=self.n_dof, prescribed_dofs=prescribed_dofs_arr
    )

    # check and process external forces
    def check_force(arr: Array | None, name: str) -> Array | None:
        if arr is None:
            return None
        match arr.ndim:
            case 2:
                out_ = jnp.broadcast_to(arr[None, ...], (n_tstep, self.n_nodes, 6))
            case 3:
                out_ = arr
            case _:
                raise ValueError(
                    f"{name} must have shape [n_node, 6] or [n_tstep, n_node, 6]"
                )
        check_arr_shape(out_, (n_tstep, self.n_nodes, 6), name)
        return out_

    f_ext_dead = check_force(f_ext_dead, "f_ext_dead")  # (n_tstep, n_node, 6)
    f_ext_follower = check_force(
        f_ext_follower, "f_ext_follower"
    )  # (n_tstep, n_node, 6)
    f_ext_aero = check_force(f_ext_aero, "f_ext_aero")  # (n_tstep, n_node, 6)

    # time integration parameters
    self.time_integrator = TimeIntegrator(
        spectral_radius=self.spectral_radius, dt=jnp.array(dt)
    )

    def evaluate_initial_equilibrium(
        init_state__: StructureCase,
    ) -> StructureCase:
        r"""
        Evaluates the forces for a given initial state to check whether it is in equilibrium. If not, a warning is
        raised with the maximum residual force. This is important to ensure that the time integration starts from a
        consistent state.
        :param init_state__: Structure containing the initial state to evaluate.
        :return: Structure with the forces evaluated for the initial state.
        """
        d, eps, f_ext_dead_, f_ext_aero_, f_grav, f_int, f_gyr, f_iner, f_res = (
            self.resolve_forces(
                hg=init_state__.hg,
                dynamic=True,
                f_ext_dead=init_state__.f_ext_dead,
                f_ext_aero=init_state__.f_ext_aero,
                thrust=init_state__.thrust,
                f_ext_follower=init_state__.f_ext_follower,
                v=init_state__.v,
                v_dot=init_state__.v_dot,
            )
        )

        max_res = jnp.max(jnp.abs(f_res))
        jax_print(
            "Initial state maximum residual force: {max_res:.3e}",
            max_res=max_res,
            verbose_level="normal",
        )

        f_elem = self.make_f_elem(eps=eps)

        return StructureCase(
            hg=init_state__.hg,
            conn=self.connectivity,
            o0=self.o0,
            d=d,
            eps=eps,
            varphi=init_state__.varphi,
            v=init_state__.v,
            v_dot=init_state__.v_dot,
            a=init_state__.v_dot,  # initial pseudo-acceleration set equal to initial acceleration
            f_ext_follower=init_state__.f_ext_follower,
            f_ext_dead=f_ext_dead_,
            f_ext_aero=f_ext_aero_,
            f_grav=f_grav,
            f_int=f_int,
            f_elem=f_elem,
            f_iner_gyr=f_iner + f_gyr,  # type: ignore
            f_res=f_res,
            thrust=init_state__.thrust,
            thrust_nodes=self.thrust_nodes,
            thrust_direction=self.thrust_direction,
            t=init_state__.t,
            i_ts=init_state__.i_ts,
            prescribed_dofs=prescribed_dofs_arr,
        )

    # time steps
    t = jnp.arange(n_tstep) * dt
    if init_state is not None and init_state.is_dynamic:
        if init_state.is_batched:
            t += init_state.t[0]
        else:
            t += init_state.t

    # set up initial state
    if init_state is None:
        init_state_: StructureCase = self.reference_configuration(
            use_f_aero=f_ext_aero is not None,
            use_f_ext_dead=f_ext_dead is not None,
            use_f_ext_follower=f_ext_follower is not None,
            prescribed_dofs=tuple(prescribed_dofs_arr),
        ).to_dynamic(t=None)
    elif not init_state.is_dynamic:
        init_state_ = init_state.to_dynamic(t=None)
    elif not init_state.is_batched:
        init_state_ = init_state
    else:
        raise TypeError(
            "dynamic_solve init_state cannot be a batched Structure; pass a "
            "snapshot or static state"
        )

    # check if initial state satisfies equilibrium
    init_state_eval = evaluate_initial_equilibrium(init_state_)
    dynamic_struct = StructureCase.initialise(
        initial_snapshot=init_state_eval,
        t=t,
        use_f_ext_follower=f_ext_follower is not None,
        use_f_ext_dead=f_ext_dead is not None,
        use_f_ext_aero=False,
    )
    converge_status = ConvergenceStatus(
        convergence_settings=self.struct_convergence_settings
    )

    ConvergenceStatus.print_header(dynamic=True)

    out = self.base_dynamic_solve(
        struct_case=dynamic_struct,
        struct_convergence_status=converge_status,
        t=t,
        solve_dofs=solve_dofs,
        load_steps=load_steps,
        f_ext_dead=f_ext_dead,
        f_ext_follower=f_ext_follower,
        aero_obj=None,
        aero_case=None,
        fsi_convergence_status=None,
        thrust_t=thrust_t_,
        cs_ang_t=None,
        cs_vel_t=None,
    )

    ConvergenceStatus.print_line(dynamic=True)
    return out

constraints

SoftConstraint

Bases: ABC

Base class for a single-node structural soft constraint (applied without Lagrange multipliers).

f_res abstractmethod

f_res(hg: Array, v: Array, i_ts: int) -> Array

Compute the forcing residual contribution from the constraint

Parameters:

Name	Type	Description	Default
hg	Array	Node SE(3) coordinate, (4, 4).	required
v	Array	Node local velocity, (6, ).	required
i_ts	int	Time-step index.	required

Returns:

Type	Description
Array	Nodal force, (6, ).

Source code in src/flapjax/structure/constraints.py

@abstractmethod
def f_res(self, hg: Array, v: Array, i_ts: int) -> Array:
    r"""
    Compute the forcing residual contribution from the constraint
    :param hg: Node SE(3) coordinate, ``(4, 4)``.
    :param v: Node local velocity, ``(6, )``.
    :param i_ts: Time-step index.
    :return: Nodal force, ``(6, )``.
    """

k_tangent

k_tangent(hg: Array, i_ts: int) -> Array

Effective stiffness contribution :math:-\partial \mathbf{f_{res}}/\partial \boldsymbol{\varphi}.

Parameters:

Name	Type	Description	Default
hg	Array	Node SE(3) coordinate, (4, 4).	required
i_ts	int	Time-step index.	required

Returns:

Type	Description
Array	Nodal stiffness contribution, (6, 6).

Source code in src/flapjax/structure/constraints.py

def k_tangent(self, hg: Array, i_ts: int) -> Array:
    r"""
    Effective stiffness contribution :math:`-\partial \mathbf{f_{res}}/\partial \boldsymbol{\varphi}`.
    :param hg: Node SE(3) coordinate, ``(4, 4)``.
    :param i_ts: Time-step index.
    :return: Nodal stiffness contribution, ``(6, 6)``.
    """
    return jnp.zeros((6, 6))

c_tangent

c_tangent(hg: Array, i_ts: int) -> Array

Effective damping contribution :math:-\partial \mathbf{f_{res}}/\partial \mathbf{v}.

Parameters:

Name	Type	Description	Default
hg	Array	Node SE(3) coordinate, (4, 4).	required
i_ts	int	Time-step index.	required

Returns:

Type	Description
Array	Nodal damping contribution, (6, 6). Default implementation returns zero damping.

Source code in src/flapjax/structure/constraints.py

def c_tangent(self, hg: Array, i_ts: int) -> Array:
    r"""
    Effective damping contribution :math:`-\partial \mathbf{f_{res}}/\partial \mathbf{v}`.
    :param hg: Node SE(3) coordinate, ``(4, 4)``.
    :param i_ts: Time-step index.
    :return: Nodal damping contribution, ``(6, 6)``. Default implementation returns zero damping.
    """
    return jnp.zeros((6, 6))

resolve_hg_ref

resolve_hg_ref(hg0: Array) -> None

Called once hg0 is populated so subclasses can default their reference frame to the node's initial pose.

Parameters:

Name	Type	Description	Default
hg0	Array	Full nodal initial-frame array, (n_nodes, 4, 4).	required

Source code in src/flapjax/structure/constraints.py

def resolve_hg_ref(self, hg0: Array) -> None:
    r"""
    Called once ``hg0`` is populated so subclasses can default their reference frame to the node's initial pose.
    :param hg0: Full nodal initial-frame array, ``(n_nodes, 4, 4)``.
    """

postprocess

postprocess(hg: Array) -> dict[str, Array]

Extract derived quantities from the converged solution that are given in the returned structure object.

Parameters:

Name	Type	Description	Default
hg	Array	Nodal SE(3) frames, (n_nodes, 4, 4) or (n_tstep, n_nodes, 4, 4).	required

Returns:

Type	Description
dict[str, Array]	Dict of named result arrays.

Source code in src/flapjax/structure/constraints.py

def postprocess(self, hg: Array) -> dict[str, Array]:
    r"""
    Extract derived quantities from the converged solution that are given in the returned structure object.
    :param hg: Nodal SE(3) frames, ``(n_nodes, 4, 4)`` or ``(n_tstep, n_nodes, 4, 4)``.
    :return: Dict of named result arrays.
    """

SpringDamper

SpringDamper(
    node_index: int,
    k: Array,
    hg_ref: Array | None = None,
    c: Array | None = None,
)

Bases: SoftConstraint

6-DOF spring-damper attached to a fixed reference frame.

Parameters:

Name	Type	Description	Default
node_index	int	Index of the node to attach to.	required
k	Array	Stiffness matrix, (6, 6).	required
hg_ref	Array | None	Reference SE(3) frame, (4, 4), or None to default to the node's initial pose.	None
c	Array | None	Damping matrix, (6, 6) or None for no damping.	None

Source code in src/flapjax/structure/constraints.py

def __init__(
    self,
    node_index: int,
    k: Array,
    hg_ref: Array | None = None,
    c: Array | None = None,
) -> None:
    r"""
    :param node_index: Index of the node to attach to.
    :param k: Stiffness matrix, ``(6, 6)``.
    :param hg_ref: Reference SE(3) frame, ``(4, 4)``, or ``None`` to default
    to the node's initial pose.
    :param c: Damping matrix, ``(6, 6)`` or ``None`` for no damping.
    """
    self.node_index: int = int(node_index)
    self.k: Array = k
    # sentinel: hg_ref is always a (4, 4) Array so the pytree structure is stable.
    # The static flag records whether the beam should overwrite it with hg0[node].
    self._hg_ref_from_hg0: bool = hg_ref is None
    self.hg_ref: Array = jnp.eye(4) if hg_ref is None else hg_ref
    self.c: Array = jnp.zeros((6, 6)) if c is None else c

postprocess

postprocess(hg: Array) -> dict[str, Array]

Extract derived quantities from the converged solution that are given in the returned structure object.

Parameters:

Name	Type	Description	Default
hg	Array	Nodal SE(3) frames, (n_nodes, 4, 4) or (n_tstep, n_nodes, 4, 4).	required

Returns:

Type	Description
dict[str, Array]	Dict of named result arrays.

Source code in src/flapjax/structure/constraints.py

def postprocess(self, hg: Array) -> dict[str, Array]:
    r"""
    Extract derived quantities from the converged solution that are given in the returned structure object.
    :param hg: Nodal SE(3) frames, ``(n_nodes, 4, 4)`` or ``(n_tstep, n_nodes, 4, 4)``.
    :return: Dict of named result arrays.
    """

PrescribedMotion

PrescribedMotion(
    node_index: int,
    k: Array,
    hg_ref_t: Array,
    c: Array | None = None,
)

Bases: SoftConstraint

Prescribe the trajectory of a node along a reference SE(3) history. The node is driven through a 6-DOF spring-damper to the reference trajectory.

Parameters:

Name	Type	Description	Default
node_index	int	Index of the node to drive.	required
k	Array	Tracking stiffness, (6, 6).	required
hg_ref_t	Array	Reference SE(3) trajectory, (n_tstep, 4, 4).	required
c	Array | None	Tracking damping, (6, 6) or None for no damping.	None

Source code in src/flapjax/structure/constraints.py

def __init__(
    self,
    node_index: int,
    k: Array,
    hg_ref_t: Array,
    c: Array | None = None,
) -> None:
    r"""
    :param node_index: Index of the node to drive.
    :param k: Tracking stiffness, ``(6, 6)``.
    :param hg_ref_t: Reference SE(3) trajectory, ``(n_tstep, 4, 4)``.
    :param c: Tracking damping, ``(6, 6)`` or ``None`` for no damping.
    """
    self.node_index: int = int(node_index)
    self.k: Array = k
    self.hg_ref_t: Array = hg_ref_t
    self.c: Array = jnp.zeros((6, 6)) if c is None else c

resolve_hg_ref

resolve_hg_ref(hg0: Array) -> None

Called once hg0 is populated so subclasses can default their reference frame to the node's initial pose.

Parameters:

Name	Type	Description	Default
hg0	Array	Full nodal initial-frame array, (n_nodes, 4, 4).	required

Source code in src/flapjax/structure/constraints.py

def resolve_hg_ref(self, hg0: Array) -> None:
    r"""
    Called once ``hg0`` is populated so subclasses can default their reference frame to the node's initial pose.
    :param hg0: Full nodal initial-frame array, ``(n_nodes, 4, 4)``.
    """

postprocess

postprocess(hg: Array) -> dict[str, Array]

Extract derived quantities from the converged solution that are given in the returned structure object.

Parameters:

Name	Type	Description	Default
hg	Array	Nodal SE(3) frames, (n_nodes, 4, 4) or (n_tstep, n_nodes, 4, 4).	required

Returns:

Type	Description
dict[str, Array]	Dict of named result arrays.

Source code in src/flapjax/structure/constraints.py

def postprocess(self, hg: Array) -> dict[str, Array]:
    r"""
    Extract derived quantities from the converged solution that are given in the returned structure object.
    :param hg: Nodal SE(3) frames, ``(n_nodes, 4, 4)`` or ``(n_tstep, n_nodes, 4, 4)``.
    :return: Dict of named result arrays.
    """

HardConstraint

Bases: ABC

Base class for a constraint applied through Lagrange multipliers.

Subclasses are either holonomic or non-holonomic

Grounded constraints constrain a single node relative to a fixed reference frame stored on the object. They have no node_j and must provide a reference coordinate hg_ref.

hg_ref property

hg_ref: Array

Fixed reference SE(3) frame for grounded constraints, (4, 4).

n_constraints abstractmethod property

n_constraints: int

Number of scalar constraints.

has_f_res property

has_f_res: bool

Whether this constraint contributes internal forces, for example if it includes a srping or damper.

violation

violation(hg_i: Array, hg_j: Array) -> Array

Compute the position-level constraint violation vector. Must be overridden by holonomic constraints.

Parameters:

Name	Type	Description	Default
hg_i	Array	SE(3) frame of node i, (4, 4).	required
hg_j	Array	SE(3) frame of node j, (4, 4).	required

Returns:

Type	Description
Array	Constraint violation, (n_constraints,).

Source code in src/flapjax/structure/constraints.py

def violation(self, hg_i: Array, hg_j: Array) -> Array:
    r"""
    Compute the position-level constraint violation vector. Must be overridden
    by holonomic constraints.
    :param hg_i: SE(3) frame of node i, ``(4, 4)``.
    :param hg_j: SE(3) frame of node j, ``(4, 4)``.
    :return: Constraint violation, ``(n_constraints,)``.
    """
    raise NotImplementedError(
        f"{type(self).__name__} is non-holonomic and has no position-level constraint"
    )

vel_violation

vel_violation(
    hg_i: Array, hg_j: Array, v_i: Array, v_j: Array
) -> Array

Compute the velocity-level constraint violation vector. Must be overridden by non-holonomic constraints.

Parameters:

Name	Type	Description	Default
hg_i	Array	SE(3) frame of node i, (4, 4).	required
hg_j	Array	SE(3) frame of node j, (4, 4).	required
v_i	Array	Local velocity of node i, (6, ).	required
v_j	Array	Local velocity of node j, (6, ).	required

Returns:

Type	Description
Array	Velocity constraint violation, (n_constraints,).

Source code in src/flapjax/structure/constraints.py

def vel_violation(self, hg_i: Array, hg_j: Array, v_i: Array, v_j: Array) -> Array:
    r"""
    Compute the velocity-level constraint violation vector. Must be overridden
    by non-holonomic constraints.
    :param hg_i: SE(3) frame of node i, ``(4, 4)``.
    :param hg_j: SE(3) frame of node j, ``(4, 4)``.
    :param v_i: Local velocity of node i, ``(6, )``.
    :param v_j: Local velocity of node j, ``(6, )``.
    :return: Velocity constraint violation, ``(n_constraints,)``.
    """
    raise NotImplementedError(
        f"{type(self).__name__} is holonomic and has no velocity-level constraint"
    )

jacobian_local

jacobian_local(
    hg_base_i: Array,
    hg_base_j: Array,
    phi_i: Array | None = None,
    phi_j: Array | None = None,
) -> tuple[Array, Array]

Compute constraint Jacobian blocks via forward-mode AD.

Parameters:

Name	Type	Description	Default
hg_base_i	Array	Base SE(3) frame of node i, (4, 4).	required
hg_base_j	Array	Base SE(3) frame of node j, (4, 4).	required
phi_i	Array | None	Accumulated configuration increment for node i, (6, ).	None
phi_j	Array | None	Accumulated configuration increment for node j, (6, ).	None

Returns:

Type	Description
tuple[Array, Array]	Jacobian blocks (jac_i, jac_j) each of shape (n_constraints, 6).

Source code in src/flapjax/structure/constraints.py

def jacobian_local(
    self,
    hg_base_i: Array,
    hg_base_j: Array,
    phi_i: Array | None = None,
    phi_j: Array | None = None,
) -> tuple[Array, Array]:
    r"""
    Compute constraint Jacobian blocks via forward-mode AD.
    :param hg_base_i: Base SE(3) frame of node i, ``(4, 4)``.
    :param hg_base_j: Base SE(3) frame of node j, ``(4, 4)``.
    :param phi_i: Accumulated configuration increment for node i, ``(6, )``.
    :param phi_j: Accumulated configuration increment for node j, ``(6, )``.
    :return: Jacobian blocks ``(jac_i, jac_j)`` each of shape ``(n_constraints, 6)``.
    """
    _phi_i = jnp.zeros(6) if phi_i is None else phi_i
    _phi_j = jnp.zeros(6) if phi_j is None else phi_j

    def violation_delta_i(delta_i: Array) -> Array:
        return self.violation(
            hg_base_i @ exp_se3(_phi_i + delta_i),
            hg_base_j @ exp_se3(_phi_j),
        )

    def violation_delta_j(delta_j: Array) -> Array:
        return self.violation(
            hg_base_i @ exp_se3(_phi_i),
            hg_base_j @ exp_se3(_phi_j + delta_j),
        )

    jac_i = jax.jacfwd(violation_delta_i)(jnp.zeros(6))
    jac_j = jax.jacfwd(violation_delta_j)(jnp.zeros(6))
    return jac_i, jac_j

a_vel_local

a_vel_local(
    hg_i: Array, hg_j: Array, v_i: Array, v_j: Array
) -> tuple[Array, Array]

Velocity Jacobian blocks :math:\partial g_{vel}/\partial v_i and :math:\partial g_{vel}/\partial v_j via forward-mode AD.

Parameters:

Name	Type	Description	Default
hg_i	Array	SE(3) frame of node i, (4, 4).	required
hg_j	Array	SE(3) frame of node j, (4, 4).	required
v_i	Array	Local velocity of node i, (6, ).	required
v_j	Array	Local velocity of node j, (6, ).	required

Returns:

Type	Description
tuple[Array, Array]	Jacobian blocks (av_i, av_j) each (n_constraints, 6).

Source code in src/flapjax/structure/constraints.py

def a_vel_local(
    self,
    hg_i: Array,
    hg_j: Array,
    v_i: Array,
    v_j: Array,
) -> tuple[Array, Array]:
    r"""
    Velocity Jacobian blocks :math:`\partial g_{vel}/\partial v_i` and
    :math:`\partial g_{vel}/\partial v_j` via forward-mode AD.
    :param hg_i: SE(3) frame of node i, ``(4, 4)``.
    :param hg_j: SE(3) frame of node j, ``(4, 4)``.
    :param v_i: Local velocity of node i, ``(6, )``.
    :param v_j: Local velocity of node j, ``(6, )``.
    :return: Jacobian blocks ``(av_i, av_j)`` each ``(n_constraints, 6)``.
    """

    def vel_violation_vi(vi: Array) -> Array:
        return self.vel_violation(hg_i, hg_j, vi, v_j)

    def vel_violation_vj(vj: Array) -> Array:
        return self.vel_violation(hg_i, hg_j, v_i, vj)

    return jax.jacfwd(vel_violation_vi)(v_i), jax.jacfwd(vel_violation_vj)(v_j)

a_phi_local

a_phi_local(
    hg_base_i: Array,
    hg_base_j: Array,
    v_i: Array,
    v_j: Array,
    phi_i: Array | None = None,
    phi_j: Array | None = None,
) -> tuple[Array, Array]

Configuration Jacobian blocks of :meth:vel_violation with respect to SE(3) configuration increments, via forward-mode AD.

Parameters:

Name	Type	Description	Default
hg_base_i	Array	Base SE(3) frame of node i, (4, 4).	required
hg_base_j	Array	Base SE(3) frame of node j, (4, 4).	required
v_i	Array	Local velocity of node i, (6, ).	required
v_j	Array	Local velocity of node j, (6, ).	required
phi_i	Array | None	Accumulated configuration increment for node i, (6, ).	None
phi_j	Array | None	Accumulated configuration increment for node j, (6, ).	None

Returns:

Type	Description
tuple[Array, Array]	Jacobian blocks (ap_i, ap_j) each (n_constraints, 6).

Source code in src/flapjax/structure/constraints.py

def a_phi_local(
    self,
    hg_base_i: Array,
    hg_base_j: Array,
    v_i: Array,
    v_j: Array,
    phi_i: Array | None = None,
    phi_j: Array | None = None,
) -> tuple[Array, Array]:
    r"""
    Configuration Jacobian blocks of :meth:`vel_violation` with respect to SE(3)
    configuration increments, via forward-mode AD.
    :param hg_base_i: Base SE(3) frame of node i, ``(4, 4)``.
    :param hg_base_j: Base SE(3) frame of node j, ``(4, 4)``.
    :param v_i: Local velocity of node i, ``(6, )``.
    :param v_j: Local velocity of node j, ``(6, )``.
    :param phi_i: Accumulated configuration increment for node i, ``(6, )``.
    :param phi_j: Accumulated configuration increment for node j, ``(6, )``.
    :return: Jacobian blocks ``(ap_i, ap_j)`` each ``(n_constraints, 6)``.
    """
    _phi_i = jnp.zeros(6) if phi_i is None else phi_i
    _phi_j = jnp.zeros(6) if phi_j is None else phi_j

    def vel_violation_delta_i(delta_i: Array) -> Array:
        return self.vel_violation(
            hg_base_i @ exp_se3(_phi_i + delta_i),
            hg_base_j @ exp_se3(_phi_j),
            v_i,
            v_j,
        )

    def vel_violation_delta_j(delta_j: Array) -> Array:
        return self.vel_violation(
            hg_base_i @ exp_se3(_phi_i),
            hg_base_j @ exp_se3(_phi_j + delta_j),
            v_i,
            v_j,
        )

    return jax.jacfwd(vel_violation_delta_i)(jnp.zeros(6)), jax.jacfwd(
        vel_violation_delta_j
    )(jnp.zeros(6))

f_res

f_res(
    hg_i: Array, hg_j: Array, v_i: Array, v_j: Array
) -> tuple[Array, Array]

Internal force contributions at nodes i and j.

Parameters:

Name	Type	Description	Default
hg_i	Array	SE(3) frame of node i, (4, 4).	required
hg_j	Array	SE(3) frame of node j, (4, 4).	required
v_i	Array	Local velocity of node i, (6, ).	required
v_j	Array	Local velocity of node j, (6, ).	required

Returns:

Type	Description
tuple[Array, Array]	(f_i, f_j) each (6, ).

Source code in src/flapjax/structure/constraints.py

def f_res(
    self,
    hg_i: Array,
    hg_j: Array,
    v_i: Array,
    v_j: Array,
) -> tuple[Array, Array]:
    r"""
    Internal force contributions at nodes i and j.
    :param hg_i: SE(3) frame of node i, ``(4, 4)``.
    :param hg_j: SE(3) frame of node j, ``(4, 4)``.
    :param v_i: Local velocity of node i, ``(6, )``.
    :param v_j: Local velocity of node j, ``(6, )``.
    :return: ``(f_i, f_j)`` each ``(6, )``.
    """
    return jnp.zeros(6), jnp.zeros(6)

k_tangent

k_tangent(hg_i: Array, hg_j: Array) -> Array

Tangent stiffness contribution from constraint, (12, 12).

Parameters:

Name	Type	Description	Default
hg_i	Array	SE(3) frame of node i, (4, 4).	required
hg_j	Array	SE(3) frame of node j, (4, 4).	required

Returns:

Type	Description
Array	Tangent stiffness, (12, 12).

Source code in src/flapjax/structure/constraints.py

def k_tangent(self, hg_i: Array, hg_j: Array) -> Array:
    r"""
    Tangent stiffness contribution from constraint, ``(12, 12)``.
    :param hg_i: SE(3) frame of node i, ``(4, 4)``.
    :param hg_j: SE(3) frame of node j, ``(4, 4)``.
    :return: Tangent stiffness, ``(12, 12)``.
    """
    return jnp.zeros((12, 12))

c_tangent

c_tangent(hg_i: Array, hg_j: Array) -> Array

Tangent damping contribution from constraint, (12, 12).

Parameters:

Name	Type	Description	Default
hg_i	Array	SE(3) frame of node i, (4, 4).	required
hg_j	Array	SE(3) frame of node j, (4, 4).	required

Returns:

Type	Description
Array	Tangent damping, (12, 12).

Source code in src/flapjax/structure/constraints.py

def c_tangent(self, hg_i: Array, hg_j: Array) -> Array:
    r"""
    Tangent damping contribution from constraint, ``(12, 12)``.
    :param hg_i: SE(3) frame of node i, ``(4, 4)``.
    :param hg_j: SE(3) frame of node j, ``(4, 4)``.
    :return: Tangent damping, ``(12, 12)``.
    """
    return jnp.zeros((12, 12))

resolve_hg_ref

resolve_hg_ref(hg0: Array) -> None

Called once hg0 is populated so subclasses can default their reference frame to the initial relative pose.

Parameters:

Name	Type	Description	Default
hg0	Array	Full nodal initial-frame array, (n_nodes, 4, 4).	required

Source code in src/flapjax/structure/constraints.py

def resolve_hg_ref(self, hg0: Array) -> None:
    r"""
    Called once ``hg0`` is populated so subclasses can default their reference frame to the initial relative pose.
    :param hg0: Full nodal initial-frame array, ``(n_nodes, 4, 4)``.
    """

postprocess

postprocess(hg: Array) -> dict[str, Array]

Extract derived quantities from the converged solution.

Parameters:

Name	Type	Description	Default
hg	Array	Nodal SE(3) frames, (n_nodes, 4, 4) or (n_tstep, n_nodes, 4, 4).	required

Returns:

Type	Description
dict[str, Array]	Dict of named result arrays.

Source code in src/flapjax/structure/constraints.py

def postprocess(self, hg: Array) -> dict[str, Array]:
    r"""
    Extract derived quantities from the converged solution.
    :param hg: Nodal SE(3) frames, ``(n_nodes, 4, 4)`` or ``(n_tstep, n_nodes, 4, 4)``.
    :return: Dict of named result arrays.
    """
    return {}

MultibodyHinge

MultibodyHinge(
    node_i: int,
    node_j: int | None = None,
    *,
    axis: Array,
    hg_rel_ref: Array | None = None,
    spring_stiffness: float = 0.0,
    damping: float = 0.0,
    prescribed_angle: Array | float | None = None,
)

Bases: HardConstraint

Hinge joint between two nodes. Constrains the relative configuration to allow only rotation about the specified axis (5 scalar constraints: 3 translation + 2 perpendicular rotation).

Optionally includes a linear rotational spring and/or damper about the hinge axis.

Parameters:

Name	Type	Description	Default
node_i	int	Index of first node.	required
node_j	int | None	Index of second node, or None to have the beam structure automatically create a co-located node.	None
axis	Array	Hinge axis direction (3, ), given in the frame of hg_rel_ref.	required
hg_rel_ref	Array | None	Reference relative SE(3) pose of the hinge (4, 4), or None to default to the initial relative pose.	None
spring_stiffness	float	Scalar rotational spring stiffness about the hinge axis (N·m/rad).	0.0
damping	float	Scalar rotational damping about the hinge axis (N·m·s/rad).	0.0
prescribed_angle	Array | float | None	If not None, adds a 6th holonomic constraint pinning the hinge rotation to this angle. This makes the joint fully rigid, and allows for a determined static problem.	None

Source code in src/flapjax/structure/constraints.py

def __init__(
    self,
    node_i: int,
    node_j: int | None = None,
    *,
    axis: Array,
    hg_rel_ref: Array | None = None,
    spring_stiffness: float = 0.0,
    damping: float = 0.0,
    prescribed_angle: Array | float | None = None,
) -> None:
    r"""
    :param node_i: Index of first node.
    :param node_j: Index of second node, or ``None`` to have the beam structure automatically create a co-located
    node.
    :param axis: Hinge axis direction ``(3, )``, given in the frame of ``hg_rel_ref``.
    :param hg_rel_ref: Reference relative SE(3) pose of the hinge ``(4, 4)``, or ``None`` to default to the initial
    relative pose.
    :param spring_stiffness: Scalar rotational spring stiffness about the hinge axis (N·m/rad).
    :param damping: Scalar rotational damping about the hinge axis (N·m·s/rad).
    :param prescribed_angle: If not ``None``, adds a 6th holonomic constraint pinning the hinge rotation to this
    angle. This makes the joint fully rigid, and allows for a determined static problem.
    """
    self.node_i: int = int(node_i)
    self.node_j: int | None = None if node_j is None else int(node_j)

    self.projection, self.axis = _hinge_projection(axis)

    self._hg_ref_from_hg0: bool = hg_rel_ref is None
    self.hg_rel_ref: Array = jnp.eye(4) if hg_rel_ref is None else hg_rel_ref

    self.spring_stiffness: float = float(spring_stiffness)
    self.damping: float = float(damping)

    self._prescribed: bool = prescribed_angle is not None
    self.prescribed_angle: Array = jnp.asarray(
        0.0 if prescribed_angle is None else prescribed_angle
    )

hg_ref property

hg_ref: Array

Fixed reference SE(3) frame for grounded constraints, (4, 4).

hinge_angle

hinge_angle(hg_i: Array, hg_j: Array) -> Array

Extract rotation angle about the hinge axis, relative to the reference configuration.

Parameters:

Name	Type	Description	Default
hg_i	Array	SE(3) frame of node i, (4, 4).	required
hg_j	Array	SE(3) frame of node j, (4, 4).	required

Returns:

Type	Description
Array	Hinge angle (rad), scalar.

Source code in src/flapjax/structure/constraints.py

def hinge_angle(self, hg_i: Array, hg_j: Array) -> Array:
    r"""
    Extract rotation angle about the hinge axis, relative to the reference configuration.
    :param hg_i: SE(3) frame of node i, ``(4, 4)``.
    :param hg_j: SE(3) frame of node j, ``(4, 4)``.
    :return: Hinge angle (rad), scalar.
    """
    hg_rel = hg_inv(hg_i) @ hg_j
    d_error = hg_to_d(self.hg_rel_ref, hg_rel)
    return self.axis @ d_error[3:]

f_res

f_res(
    hg_i: Array, hg_j: Array, v_i: Array, v_j: Array
) -> tuple[Array, Array]

Spring-damper force contributions at nodes i and j.

Parameters:

Name	Type	Description	Default
hg_i	Array	SE(3) frame of node i, (4, 4).	required
hg_j	Array	SE(3) frame of node j, (4, 4).	required
v_i	Array	Local velocity of node i, (6, ).	required
v_j	Array	Local velocity of node j, (6, ).	required

Returns:

Type	Description
tuple[Array, Array]	Forces (f_i, f_j), each (6, ).

Source code in src/flapjax/structure/constraints.py

def f_res(
    self,
    hg_i: Array,
    hg_j: Array,
    v_i: Array,
    v_j: Array,
) -> tuple[Array, Array]:
    r"""
    Spring-damper force contributions at nodes i and j.
    :param hg_i: SE(3) frame of node i, ``(4, 4)``.
    :param hg_j: SE(3) frame of node j, ``(4, 4)``.
    :param v_i: Local velocity of node i, ``(6, )``.
    :param v_j: Local velocity of node j, ``(6, )``.
    :return: Forces ``(f_i, f_j)``, each ``(6, )``.
    """

    def _theta(d_ij: Array) -> Array:
        return self.hinge_angle(hg_i @ exp_se3(d_ij[:6]), hg_j @ exp_se3(d_ij[6:]))

    z12 = jnp.zeros(12)
    theta = _theta(z12)
    dtheta = jax.grad(_theta)(z12)
    dtheta_di, dtheta_dj = dtheta[:6], dtheta[6:]

    theta_dot = dtheta_di @ v_i + dtheta_dj @ v_j
    f_scalar = -self.spring_stiffness * theta - self.damping * theta_dot

    return f_scalar * dtheta_di, f_scalar * dtheta_dj

k_tangent

k_tangent(hg_i: Array, hg_j: Array) -> Array

Tangent stiffness from the hinge spring, (12, 12).

Source code in src/flapjax/structure/constraints.py

def k_tangent(self, hg_i: Array, hg_j: Array) -> Array:
    r"""
    Tangent stiffness from the hinge spring, ``(12, 12)``.
    """

    def _potential(d_ij: Array) -> Array:
        theta = self.hinge_angle(hg_i @ exp_se3(d_ij[:6]), hg_j @ exp_se3(d_ij[6:]))
        return 0.5 * self.spring_stiffness * theta**2

    return jax.jacfwd(jax.grad(_potential))(jnp.zeros(12))

c_tangent

c_tangent(hg_i: Array, hg_j: Array) -> Array

Tangent damping from the hinge damper, (12, 12).

Source code in src/flapjax/structure/constraints.py

def c_tangent(self, hg_i: Array, hg_j: Array) -> Array:
    r"""
    Tangent damping from the hinge damper, ``(12, 12)``.
    """

    def _theta(d_ij: Array) -> Array:
        return self.hinge_angle(hg_i @ exp_se3(d_ij[:6]), hg_j @ exp_se3(d_ij[6:]))

    dtheta = jax.grad(_theta)(jnp.zeros(12))
    return self.damping * jnp.outer(dtheta, dtheta)

vel_violation

vel_violation(
    hg_i: Array, hg_j: Array, v_i: Array, v_j: Array
) -> Array

Compute the velocity-level constraint violation vector. Must be overridden by non-holonomic constraints.

Parameters:

Name	Type	Description	Default
hg_i	Array	SE(3) frame of node i, (4, 4).	required
hg_j	Array	SE(3) frame of node j, (4, 4).	required
v_i	Array	Local velocity of node i, (6, ).	required
v_j	Array	Local velocity of node j, (6, ).	required

Returns:

Type	Description
Array	Velocity constraint violation, (n_constraints,).

Source code in src/flapjax/structure/constraints.py

def vel_violation(self, hg_i: Array, hg_j: Array, v_i: Array, v_j: Array) -> Array:
    r"""
    Compute the velocity-level constraint violation vector. Must be overridden
    by non-holonomic constraints.
    :param hg_i: SE(3) frame of node i, ``(4, 4)``.
    :param hg_j: SE(3) frame of node j, ``(4, 4)``.
    :param v_i: Local velocity of node i, ``(6, )``.
    :param v_j: Local velocity of node j, ``(6, )``.
    :return: Velocity constraint violation, ``(n_constraints,)``.
    """
    raise NotImplementedError(
        f"{type(self).__name__} is holonomic and has no velocity-level constraint"
    )

jacobian_local

jacobian_local(
    hg_base_i: Array,
    hg_base_j: Array,
    phi_i: Array | None = None,
    phi_j: Array | None = None,
) -> tuple[Array, Array]

Compute constraint Jacobian blocks via forward-mode AD.

Parameters:

Name	Type	Description	Default
hg_base_i	Array	Base SE(3) frame of node i, (4, 4).	required
hg_base_j	Array	Base SE(3) frame of node j, (4, 4).	required
phi_i	Array | None	Accumulated configuration increment for node i, (6, ).	None
phi_j	Array | None	Accumulated configuration increment for node j, (6, ).	None

Returns:

Type	Description
tuple[Array, Array]	Jacobian blocks (jac_i, jac_j) each of shape (n_constraints, 6).

Source code in src/flapjax/structure/constraints.py

def jacobian_local(
    self,
    hg_base_i: Array,
    hg_base_j: Array,
    phi_i: Array | None = None,
    phi_j: Array | None = None,
) -> tuple[Array, Array]:
    r"""
    Compute constraint Jacobian blocks via forward-mode AD.
    :param hg_base_i: Base SE(3) frame of node i, ``(4, 4)``.
    :param hg_base_j: Base SE(3) frame of node j, ``(4, 4)``.
    :param phi_i: Accumulated configuration increment for node i, ``(6, )``.
    :param phi_j: Accumulated configuration increment for node j, ``(6, )``.
    :return: Jacobian blocks ``(jac_i, jac_j)`` each of shape ``(n_constraints, 6)``.
    """
    _phi_i = jnp.zeros(6) if phi_i is None else phi_i
    _phi_j = jnp.zeros(6) if phi_j is None else phi_j

    def violation_delta_i(delta_i: Array) -> Array:
        return self.violation(
            hg_base_i @ exp_se3(_phi_i + delta_i),
            hg_base_j @ exp_se3(_phi_j),
        )

    def violation_delta_j(delta_j: Array) -> Array:
        return self.violation(
            hg_base_i @ exp_se3(_phi_i),
            hg_base_j @ exp_se3(_phi_j + delta_j),
        )

    jac_i = jax.jacfwd(violation_delta_i)(jnp.zeros(6))
    jac_j = jax.jacfwd(violation_delta_j)(jnp.zeros(6))
    return jac_i, jac_j

a_vel_local

a_vel_local(
    hg_i: Array, hg_j: Array, v_i: Array, v_j: Array
) -> tuple[Array, Array]

Velocity Jacobian blocks :math:\partial g_{vel}/\partial v_i and :math:\partial g_{vel}/\partial v_j via forward-mode AD.

Parameters:

Name	Type	Description	Default
hg_i	Array	SE(3) frame of node i, (4, 4).	required
hg_j	Array	SE(3) frame of node j, (4, 4).	required
v_i	Array	Local velocity of node i, (6, ).	required
v_j	Array	Local velocity of node j, (6, ).	required

Returns:

Type	Description
tuple[Array, Array]	Jacobian blocks (av_i, av_j) each (n_constraints, 6).

Source code in src/flapjax/structure/constraints.py

def a_vel_local(
    self,
    hg_i: Array,
    hg_j: Array,
    v_i: Array,
    v_j: Array,
) -> tuple[Array, Array]:
    r"""
    Velocity Jacobian blocks :math:`\partial g_{vel}/\partial v_i` and
    :math:`\partial g_{vel}/\partial v_j` via forward-mode AD.
    :param hg_i: SE(3) frame of node i, ``(4, 4)``.
    :param hg_j: SE(3) frame of node j, ``(4, 4)``.
    :param v_i: Local velocity of node i, ``(6, )``.
    :param v_j: Local velocity of node j, ``(6, )``.
    :return: Jacobian blocks ``(av_i, av_j)`` each ``(n_constraints, 6)``.
    """

    def vel_violation_vi(vi: Array) -> Array:
        return self.vel_violation(hg_i, hg_j, vi, v_j)

    def vel_violation_vj(vj: Array) -> Array:
        return self.vel_violation(hg_i, hg_j, v_i, vj)

    return jax.jacfwd(vel_violation_vi)(v_i), jax.jacfwd(vel_violation_vj)(v_j)

a_phi_local

a_phi_local(
    hg_base_i: Array,
    hg_base_j: Array,
    v_i: Array,
    v_j: Array,
    phi_i: Array | None = None,
    phi_j: Array | None = None,
) -> tuple[Array, Array]

Configuration Jacobian blocks of :meth:vel_violation with respect to SE(3) configuration increments, via forward-mode AD.

Parameters:

Name	Type	Description	Default
hg_base_i	Array	Base SE(3) frame of node i, (4, 4).	required
hg_base_j	Array	Base SE(3) frame of node j, (4, 4).	required
v_i	Array	Local velocity of node i, (6, ).	required
v_j	Array	Local velocity of node j, (6, ).	required
phi_i	Array | None	Accumulated configuration increment for node i, (6, ).	None
phi_j	Array | None	Accumulated configuration increment for node j, (6, ).	None

Returns:

Type	Description
tuple[Array, Array]	Jacobian blocks (ap_i, ap_j) each (n_constraints, 6).

Source code in src/flapjax/structure/constraints.py

def a_phi_local(
    self,
    hg_base_i: Array,
    hg_base_j: Array,
    v_i: Array,
    v_j: Array,
    phi_i: Array | None = None,
    phi_j: Array | None = None,
) -> tuple[Array, Array]:
    r"""
    Configuration Jacobian blocks of :meth:`vel_violation` with respect to SE(3)
    configuration increments, via forward-mode AD.
    :param hg_base_i: Base SE(3) frame of node i, ``(4, 4)``.
    :param hg_base_j: Base SE(3) frame of node j, ``(4, 4)``.
    :param v_i: Local velocity of node i, ``(6, )``.
    :param v_j: Local velocity of node j, ``(6, )``.
    :param phi_i: Accumulated configuration increment for node i, ``(6, )``.
    :param phi_j: Accumulated configuration increment for node j, ``(6, )``.
    :return: Jacobian blocks ``(ap_i, ap_j)`` each ``(n_constraints, 6)``.
    """
    _phi_i = jnp.zeros(6) if phi_i is None else phi_i
    _phi_j = jnp.zeros(6) if phi_j is None else phi_j

    def vel_violation_delta_i(delta_i: Array) -> Array:
        return self.vel_violation(
            hg_base_i @ exp_se3(_phi_i + delta_i),
            hg_base_j @ exp_se3(_phi_j),
            v_i,
            v_j,
        )

    def vel_violation_delta_j(delta_j: Array) -> Array:
        return self.vel_violation(
            hg_base_i @ exp_se3(_phi_i),
            hg_base_j @ exp_se3(_phi_j + delta_j),
            v_i,
            v_j,
        )

    return jax.jacfwd(vel_violation_delta_i)(jnp.zeros(6)), jax.jacfwd(
        vel_violation_delta_j
    )(jnp.zeros(6))

GroundedHinge

GroundedHinge(
    node_i: int, *, axis: Array, hg_ref: Array | None = None
)

Bases: HardConstraint

Hinge joint pinning a single node to a fixed point in space.

Parameters:

Name	Type	Description	Default
node_i	int	Index of the constrained node.	required
axis	Array	Hinge axis direction (3, ), given in the global frame.	required
hg_ref	Array | None	Reference SE(3) frame (4, 4) the node is pinned to, or None to default to the node's initial pose.	None

Source code in src/flapjax/structure/constraints.py

def __init__(
    self,
    node_i: int,
    *,
    axis: Array,
    hg_ref: Array | None = None,
) -> None:
    r"""
    :param node_i: Index of the constrained node.
    :param axis: Hinge axis direction ``(3, )``, given in the global frame.
    :param hg_ref: Reference SE(3) frame ``(4, 4)`` the node is pinned to,
        or ``None`` to default to the node's initial pose.
    """
    self.node_i: int = int(node_i)
    self.node_j: None = None

    self.projection, self.axis = _hinge_projection(axis)

    self._hg_ref_from_hg0: bool = hg_ref is None
    self._hg_ref: Array = jnp.eye(4) if hg_ref is None else hg_ref

has_f_res property

has_f_res: bool

Whether this constraint contributes internal forces, for example if it includes a srping or damper.

hinge_angle

hinge_angle(hg_i: Array) -> Array

Scalar rotation angle about the hinge axis, relative to the reference frame.

Parameters:

Name	Type	Description	Default
hg_i	Array	SE(3) frame of the constrained node, (4, 4).	required

Returns:

Type	Description
Array	Hinge angle (rad), scalar.

Source code in src/flapjax/structure/constraints.py

def hinge_angle(self, hg_i: Array) -> Array:
    r"""
    Scalar rotation angle about the hinge axis, relative to the reference frame.
    :param hg_i: SE(3) frame of the constrained node, ``(4, 4)``.
    :return: Hinge angle (rad), scalar.
    """
    d_error = hg_to_d(self._hg_ref, hg_i)
    return self.axis @ d_error[3:]

vel_violation

vel_violation(
    hg_i: Array, hg_j: Array, v_i: Array, v_j: Array
) -> Array

Compute the velocity-level constraint violation vector. Must be overridden by non-holonomic constraints.

Parameters:

Name	Type	Description	Default
hg_i	Array	SE(3) frame of node i, (4, 4).	required
hg_j	Array	SE(3) frame of node j, (4, 4).	required
v_i	Array	Local velocity of node i, (6, ).	required
v_j	Array	Local velocity of node j, (6, ).	required

Returns:

Type	Description
Array	Velocity constraint violation, (n_constraints,).

Source code in src/flapjax/structure/constraints.py

def vel_violation(self, hg_i: Array, hg_j: Array, v_i: Array, v_j: Array) -> Array:
    r"""
    Compute the velocity-level constraint violation vector. Must be overridden
    by non-holonomic constraints.
    :param hg_i: SE(3) frame of node i, ``(4, 4)``.
    :param hg_j: SE(3) frame of node j, ``(4, 4)``.
    :param v_i: Local velocity of node i, ``(6, )``.
    :param v_j: Local velocity of node j, ``(6, )``.
    :return: Velocity constraint violation, ``(n_constraints,)``.
    """
    raise NotImplementedError(
        f"{type(self).__name__} is holonomic and has no velocity-level constraint"
    )

a_vel_local

a_vel_local(
    hg_i: Array, hg_j: Array, v_i: Array, v_j: Array
) -> tuple[Array, Array]

Velocity Jacobian blocks :math:\partial g_{vel}/\partial v_i and :math:\partial g_{vel}/\partial v_j via forward-mode AD.

Parameters:

Name	Type	Description	Default
hg_i	Array	SE(3) frame of node i, (4, 4).	required
hg_j	Array	SE(3) frame of node j, (4, 4).	required
v_i	Array	Local velocity of node i, (6, ).	required
v_j	Array	Local velocity of node j, (6, ).	required

Returns:

Type	Description
tuple[Array, Array]	Jacobian blocks (av_i, av_j) each (n_constraints, 6).

Source code in src/flapjax/structure/constraints.py

def a_vel_local(
    self,
    hg_i: Array,
    hg_j: Array,
    v_i: Array,
    v_j: Array,
) -> tuple[Array, Array]:
    r"""
    Velocity Jacobian blocks :math:`\partial g_{vel}/\partial v_i` and
    :math:`\partial g_{vel}/\partial v_j` via forward-mode AD.
    :param hg_i: SE(3) frame of node i, ``(4, 4)``.
    :param hg_j: SE(3) frame of node j, ``(4, 4)``.
    :param v_i: Local velocity of node i, ``(6, )``.
    :param v_j: Local velocity of node j, ``(6, )``.
    :return: Jacobian blocks ``(av_i, av_j)`` each ``(n_constraints, 6)``.
    """

    def vel_violation_vi(vi: Array) -> Array:
        return self.vel_violation(hg_i, hg_j, vi, v_j)

    def vel_violation_vj(vj: Array) -> Array:
        return self.vel_violation(hg_i, hg_j, v_i, vj)

    return jax.jacfwd(vel_violation_vi)(v_i), jax.jacfwd(vel_violation_vj)(v_j)

a_phi_local

a_phi_local(
    hg_base_i: Array,
    hg_base_j: Array,
    v_i: Array,
    v_j: Array,
    phi_i: Array | None = None,
    phi_j: Array | None = None,
) -> tuple[Array, Array]

Configuration Jacobian blocks of :meth:vel_violation with respect to SE(3) configuration increments, via forward-mode AD.

Parameters:

Name	Type	Description	Default
hg_base_i	Array	Base SE(3) frame of node i, (4, 4).	required
hg_base_j	Array	Base SE(3) frame of node j, (4, 4).	required
v_i	Array	Local velocity of node i, (6, ).	required
v_j	Array	Local velocity of node j, (6, ).	required
phi_i	Array | None	Accumulated configuration increment for node i, (6, ).	None
phi_j	Array | None	Accumulated configuration increment for node j, (6, ).	None

Returns:

Type	Description
tuple[Array, Array]	Jacobian blocks (ap_i, ap_j) each (n_constraints, 6).

Source code in src/flapjax/structure/constraints.py

def a_phi_local(
    self,
    hg_base_i: Array,
    hg_base_j: Array,
    v_i: Array,
    v_j: Array,
    phi_i: Array | None = None,
    phi_j: Array | None = None,
) -> tuple[Array, Array]:
    r"""
    Configuration Jacobian blocks of :meth:`vel_violation` with respect to SE(3)
    configuration increments, via forward-mode AD.
    :param hg_base_i: Base SE(3) frame of node i, ``(4, 4)``.
    :param hg_base_j: Base SE(3) frame of node j, ``(4, 4)``.
    :param v_i: Local velocity of node i, ``(6, )``.
    :param v_j: Local velocity of node j, ``(6, )``.
    :param phi_i: Accumulated configuration increment for node i, ``(6, )``.
    :param phi_j: Accumulated configuration increment for node j, ``(6, )``.
    :return: Jacobian blocks ``(ap_i, ap_j)`` each ``(n_constraints, 6)``.
    """
    _phi_i = jnp.zeros(6) if phi_i is None else phi_i
    _phi_j = jnp.zeros(6) if phi_j is None else phi_j

    def vel_violation_delta_i(delta_i: Array) -> Array:
        return self.vel_violation(
            hg_base_i @ exp_se3(_phi_i + delta_i),
            hg_base_j @ exp_se3(_phi_j),
            v_i,
            v_j,
        )

    def vel_violation_delta_j(delta_j: Array) -> Array:
        return self.vel_violation(
            hg_base_i @ exp_se3(_phi_i),
            hg_base_j @ exp_se3(_phi_j + delta_j),
            v_i,
            v_j,
        )

    return jax.jacfwd(vel_violation_delta_i)(jnp.zeros(6)), jax.jacfwd(
        vel_violation_delta_j
    )(jnp.zeros(6))

f_res

f_res(
    hg_i: Array, hg_j: Array, v_i: Array, v_j: Array
) -> tuple[Array, Array]

Internal force contributions at nodes i and j.

Parameters:

Name	Type	Description	Default
hg_i	Array	SE(3) frame of node i, (4, 4).	required
hg_j	Array	SE(3) frame of node j, (4, 4).	required
v_i	Array	Local velocity of node i, (6, ).	required
v_j	Array	Local velocity of node j, (6, ).	required

Returns:

Type	Description
tuple[Array, Array]	(f_i, f_j) each (6, ).

Source code in src/flapjax/structure/constraints.py

def f_res(
    self,
    hg_i: Array,
    hg_j: Array,
    v_i: Array,
    v_j: Array,
) -> tuple[Array, Array]:
    r"""
    Internal force contributions at nodes i and j.
    :param hg_i: SE(3) frame of node i, ``(4, 4)``.
    :param hg_j: SE(3) frame of node j, ``(4, 4)``.
    :param v_i: Local velocity of node i, ``(6, )``.
    :param v_j: Local velocity of node j, ``(6, )``.
    :return: ``(f_i, f_j)`` each ``(6, )``.
    """
    return jnp.zeros(6), jnp.zeros(6)

k_tangent

k_tangent(hg_i: Array, hg_j: Array) -> Array

Tangent stiffness contribution from constraint, (12, 12).

Parameters:

Name	Type	Description	Default
hg_i	Array	SE(3) frame of node i, (4, 4).	required
hg_j	Array	SE(3) frame of node j, (4, 4).	required

Returns:

Type	Description
Array	Tangent stiffness, (12, 12).

Source code in src/flapjax/structure/constraints.py

def k_tangent(self, hg_i: Array, hg_j: Array) -> Array:
    r"""
    Tangent stiffness contribution from constraint, ``(12, 12)``.
    :param hg_i: SE(3) frame of node i, ``(4, 4)``.
    :param hg_j: SE(3) frame of node j, ``(4, 4)``.
    :return: Tangent stiffness, ``(12, 12)``.
    """
    return jnp.zeros((12, 12))

c_tangent

c_tangent(hg_i: Array, hg_j: Array) -> Array

Tangent damping contribution from constraint, (12, 12).

Parameters:

Name	Type	Description	Default
hg_i	Array	SE(3) frame of node i, (4, 4).	required
hg_j	Array	SE(3) frame of node j, (4, 4).	required

Returns:

Type	Description
Array	Tangent damping, (12, 12).

Source code in src/flapjax/structure/constraints.py

def c_tangent(self, hg_i: Array, hg_j: Array) -> Array:
    r"""
    Tangent damping contribution from constraint, ``(12, 12)``.
    :param hg_i: SE(3) frame of node i, ``(4, 4)``.
    :param hg_j: SE(3) frame of node j, ``(4, 4)``.
    :return: Tangent damping, ``(12, 12)``.
    """
    return jnp.zeros((12, 12))

data_structures

StructureCase

StructureCase(
    hg: Array,
    conn: tuple[tuple[int, int], ...],
    o0: Array,
    d: Array,
    eps: Array,
    varphi: Array,
    f_ext_follower: Array | None,
    f_ext_dead: Array | None,
    f_ext_aero: Array | None,
    f_grav: Array | None,
    f_int: Array,
    f_elem: Array,
    f_res: Array,
    t: Array,
    thrust: dict[str, Array],
    thrust_nodes: tuple[tuple[str, int], ...],
    thrust_direction: tuple[
        tuple[str, tuple[float, float, float]], ...
    ],
    prescribed_dofs: tuple[int, ...] | Array,
    v: Array | None = None,
    v_dot: Array | None = None,
    a: Array | None = None,
    f_iner_gyr: Array | None = None,
    i_ts: int | None = None,
    local: bool = True,
    constraint_data: dict[str, dict[str, Array]]
    | None = None,
)

Object to hold the full state and forces of a structure analysis.

A single instance may represent any of three flavours:

• Static: dynamic-only fields (v, v_dot, a, f_iner_gyr) are not set — their public properties return an all-zeros array of appropriate shape. Array shapes are (n_nodes, ...).
• Dynamic snapshot: all dynamic fields populated for a single timestep. Array shapes are (n_nodes, ...).
• Dynamic trajectory: all dynamic fields populated for multiple timesteps. t is a (n_tstep,) array and i_ts is None. Array shapes are (n_tstep, n_nodes, ...).

Use :attr:is_dynamic and :attr:is_batched to distinguish at runtime.

Source code in src/flapjax/structure/data_structures.py

def __init__(
    self,
    hg: Array,
    conn: tuple[tuple[int, int], ...],
    o0: Array,
    d: Array,
    eps: Array,
    varphi: Array,
    f_ext_follower: Array | None,
    f_ext_dead: Array | None,
    f_ext_aero: Array | None,
    f_grav: Array | None,
    f_int: Array,
    f_elem: Array,
    f_res: Array,
    t: Array,
    thrust: dict[str, Array],
    thrust_nodes: tuple[tuple[str, int], ...],
    thrust_direction: tuple[tuple[str, tuple[float, float, float]], ...],
    prescribed_dofs: tuple[int, ...] | Array,
    v: Array | None = None,
    v_dot: Array | None = None,
    a: Array | None = None,
    f_iner_gyr: Array | None = None,
    i_ts: int | None = None,
    local: bool = True,
    constraint_data: dict[str, dict[str, Array]] | None = None,
):
    self.hg: Array = hg
    self.conn: tuple[tuple[int, int], ...] = conn
    self.o0: Array = o0
    self.d: Array = d
    self.eps: Array = eps
    self.varphi: Array = varphi
    self._v: Array | None = v
    self._v_dot: Array | None = v_dot
    self._a: Array | None = a
    self.f_ext_follower: Array | None = f_ext_follower
    self.f_ext_dead: Array | None = f_ext_dead
    self.f_ext_aero: Array | None = f_ext_aero
    self.f_grav: Array | None = f_grav
    self.f_int: Array = f_int
    self.f_elem: Array = f_elem
    self._f_iner_gyr: Array | None = f_iner_gyr
    self.f_res: Array = f_res
    self.thrust: dict[str, Array] = thrust
    self.thrust_nodes: tuple[tuple[str, int], ...] = thrust_nodes
    self.thrust_direction: tuple[tuple[str, tuple[float, float, float]], ...] = (
        thrust_direction
    )
    self.t: Array = t
    self.i_ts: int | None = i_ts
    self.prescribed_dofs: tuple[int, ...] = input_dof_index_to_tuple(
        prescribed_dofs
    )
    self.free_dofs: tuple[int, ...] = get_solve_dofs(
        n_dof=varphi.shape[-2] * 6, prescribed_dofs=self.prescribed_dofs
    )
    self.local: bool = local
    self.constraint_data: dict[str, dict[str, Array]] = constraint_data if constraint_data is not None else {}

to_dynamic

to_dynamic() -> StructureCase

to_dynamic(t: None) -> StructureCase

to_dynamic(t: Array) -> StructureCase

to_dynamic(t: Array | None = None) -> StructureCase

Convert static structure results to a dynamic snapshot (t=None) or a batched trajectory (t provided), zeroing velocity/acceleration fields. Calling on a Structure that is already dynamic returns self.

Source code in src/flapjax/structure/data_structures.py

def to_dynamic(self, t: Array | None = None) -> StructureCase:
    """Convert static structure results to a dynamic snapshot (``t=None``) or
    a batched trajectory (``t`` provided), zeroing velocity/acceleration
    fields. Calling on a Structure that is already dynamic returns ``self``.
    """
    if self.is_dynamic:
        return self

    dyn_snapshot = StructureCase(
        hg=self.hg,
        conn=self.conn,
        o0=self.o0,
        d=self.d,
        eps=self.eps,
        varphi=self.varphi,
        v=self.v,
        v_dot=self.v_dot,
        a=self.a,
        f_ext_follower=self.f_ext_follower,
        f_ext_dead=self.f_ext_dead,
        f_ext_aero=self.f_ext_aero,
        f_grav=self.f_grav,
        f_int=self.f_int,
        f_elem=self.f_elem,
        f_iner_gyr=self.f_iner_gyr,
        f_res=self.f_res,
        thrust=self.thrust,
        thrust_nodes=self.thrust_nodes,
        thrust_direction=self.thrust_direction,
        t=jnp.array(0.0),
        i_ts=-1,
        prescribed_dofs=self.prescribed_dofs,
        constraint_data=self.constraint_data,
    )

    if t is None:
        return dyn_snapshot
    return StructureCase.initialise(
        initial_snapshot=dyn_snapshot,
        t=t,
        use_f_ext_aero=self.f_ext_aero is not None,
        use_f_ext_follower=self.f_ext_follower is not None,
        use_f_ext_dead=self.f_ext_dead is not None,
    )

to_static

to_static() -> StructureCase

Return a static Structure, dropping velocity/acceleration fields. If already static, returns self. For a batched trajectory, raises: use self[i_ts].to_static() to extract a single time step first.

Source code in src/flapjax/structure/data_structures.py

def to_static(self) -> StructureCase:
    """Return a static Structure, dropping velocity/acceleration
    fields. If already static, returns ``self``. For a batched trajectory,
    raises: use ``self[i_ts].to_static()`` to extract a single time step first.
    """
    if not self.is_dynamic:
        return self
    if self.is_batched:
        raise ValueError(
            "to_static() on a batched Structure is ambiguous; index a "
            "single time step first (e.g. `structure[i_ts].to_static()`)."
        )
    return StructureCase(
        hg=self.hg,
        conn=self.conn,
        o0=self.o0,
        d=self.d,
        eps=self.eps,
        varphi=self.varphi,
        f_ext_follower=self.f_ext_follower,
        f_ext_dead=self.f_ext_dead,
        f_ext_aero=self.f_ext_aero,
        f_grav=self.f_grav,
        f_int=self.f_int,
        f_elem=self.f_elem,
        f_res=self.f_res,
        t=self.t,
        thrust=self.thrust,
        thrust_nodes=self.thrust_nodes,
        thrust_direction=self.thrust_direction,
        prescribed_dofs=self.prescribed_dofs,
        constraint_data=self.constraint_data,
    )

initialise classmethod

initialise(
    initial_snapshot: StructureCase,
    t: Array,
    use_f_ext_follower: bool,
    use_f_ext_dead: bool,
    use_f_ext_aero: bool,
) -> StructureCase

Initialise a batched dynamic Structure from a single dynamic snapshot.

Parameters:

Name	Type	Description	Default
initial_snapshot	StructureCase	Snapshot at initial time step (must be dynamic, i.e. is_dynamic == True and not batched).	required
t	Array	Time step array, (n_tstep, )	required
use_f_ext_follower	bool	Whether to include follower force array	required
use_f_ext_dead	bool	Whether to include dead force array	required
use_f_ext_aero	bool	Whether to include aero force array	required

Returns:

Type	Description
StructureCase	Batched Structure with arrays initialised to zero except for the first time step.

Source code in src/flapjax/structure/data_structures.py

@classmethod
def initialise(
    cls,
    initial_snapshot: StructureCase,
    t: Array,
    use_f_ext_follower: bool,
    use_f_ext_dead: bool,
    use_f_ext_aero: bool,
) -> StructureCase:
    r"""
    Initialise a batched dynamic Structure from a single dynamic snapshot.
    :param initial_snapshot: Snapshot at initial time step (must be dynamic,
    i.e. ``is_dynamic == True`` and not batched).
    :param t: Time step array, ``(n_tstep, )``
    :param use_f_ext_follower: Whether to include follower force array
    :param use_f_ext_dead: Whether to include dead force array
    :param use_f_ext_aero: Whether to include aero force array
    :return: Batched Structure with arrays initialised to zero except for
    the first time step.
    """
    if not initial_snapshot.is_dynamic:
        raise ValueError(
            "initial_snapshot must be dynamic; call to_dynamic() first"
        )
    if initial_snapshot.is_batched:
        raise ValueError("initial_snapshot must be a single snapshot, not batched")

    n_node = initial_snapshot.hg.shape[0]
    n_elem = initial_snapshot.d.shape[0]
    n_tstep = t.shape[0]

    hg = jnp.zeros((n_tstep, n_node, 4, 4)).at[0, ...].set(initial_snapshot.hg)
    d = jnp.zeros((n_tstep, n_elem, 6)).at[0, ...].set(initial_snapshot.d)
    eps = jnp.zeros((n_tstep, n_elem, 6)).at[0, ...].set(initial_snapshot.eps)
    varphi = jnp.zeros((n_tstep, n_node, 6)).at[0, ...].set(initial_snapshot.varphi)
    v = jnp.zeros((n_tstep, n_node, 6)).at[0, ...].set(initial_snapshot.v)
    v_dot = jnp.zeros((n_tstep, n_node, 6)).at[0, ...].set(initial_snapshot.v_dot)
    a = jnp.zeros((n_tstep, n_node, 6)).at[0, ...].set(initial_snapshot.a)

    if use_f_ext_follower:
        f_ext_follower = jnp.zeros((n_tstep, n_node, 6))
        if initial_snapshot.f_ext_follower is not None:
            f_ext_follower = f_ext_follower.at[0, ...].set(
                initial_snapshot.f_ext_follower
            )
    else:
        f_ext_follower = None

    if use_f_ext_dead:
        f_ext_dead = jnp.zeros((n_tstep, n_node, 6))
        if initial_snapshot.f_ext_dead is not None:
            f_ext_dead = f_ext_dead.at[0, ...].set(initial_snapshot.f_ext_dead)
    else:
        f_ext_dead = None

    if use_f_ext_aero:
        f_ext_aero = jnp.zeros((n_tstep, n_node, 6))
        if initial_snapshot.f_ext_aero is not None:
            f_ext_aero = f_ext_aero.at[0, ...].set(initial_snapshot.f_ext_aero)
    else:
        f_ext_aero = None

    f_grav = (
        jnp.zeros((n_tstep, n_node, 6)).at[0, ...].set(initial_snapshot.f_grav)
        if initial_snapshot.f_grav is not None
        else None
    )
    f_int = jnp.zeros((n_tstep, n_node, 6)).at[0, ...].set(initial_snapshot.f_int)
    f_elem = jnp.zeros((n_tstep, n_elem, 6)).at[0, ...].set(initial_snapshot.f_elem)
    f_iner_gyr = (
        jnp.zeros((n_tstep, n_node, 6)).at[0, ...].set(initial_snapshot.f_iner_gyr)
    )
    f_res = jnp.zeros((n_tstep, n_node, 6)).at[0, ...].set(initial_snapshot.f_res)

    thrust = {k: jnp.full(n_tstep, v) for k, v in initial_snapshot.thrust.items()}
    return cls(
        hg=hg,
        conn=initial_snapshot.conn,
        o0=initial_snapshot.o0,
        d=d,
        eps=eps,
        varphi=varphi,
        v=v,
        v_dot=v_dot,
        a=a,
        f_ext_follower=f_ext_follower,
        f_ext_dead=f_ext_dead,
        f_ext_aero=f_ext_aero,
        f_grav=f_grav,
        f_int=f_int,
        f_elem=f_elem,
        f_iner_gyr=f_iner_gyr,
        f_res=f_res,
        thrust=thrust,
        thrust_nodes=initial_snapshot.thrust_nodes,
        thrust_direction=initial_snapshot.thrust_direction,
        t=t,
        prescribed_dofs=initial_snapshot.prescribed_dofs,
    )

to_global

to_global() -> None

Convert local structure results to global frame.

Source code in src/flapjax/structure/data_structures.py

def to_global(self) -> None:
    """Convert local structure results to global frame."""
    if not self.local:
        warn("Results already in global frame, skipping conversion.")
        return
    self.local = False
    self._transform(rmat=self.rmat)

to_local

to_local() -> None

Convert global structure results to local frame.

Source code in src/flapjax/structure/data_structures.py

def to_local(self) -> None:
    """Convert global structure results to local frame."""
    if self.local:
        warn("Results already in local frame, skipping conversion.")
        return
    self.local = True
    if self.is_batched:
        rmat_t = jnp.transpose(self.rmat, (0, 1, 3, 2))
    else:
        rmat_t = jnp.transpose(self.rmat, (0, 2, 1))
    self._transform(rmat=rmat_t)

plot

plot(directory: PathLike | str, n_interp: int = 0) -> Path

plot(
    directory: PathLike | str,
    n_interp: int = 0,
    *,
    index: slice
    | Sequence[int]
    | int
    | Array
    | None = None,
) -> Path

plot(
    directory: PathLike | str,
    n_interp: int = 0,
    *,
    index: slice
    | Sequence[int]
    | int
    | Array
    | None = None,
) -> Path

Plot beam results to VTK/VTU files in the specified directory. For a batched Structure, a PVD is written alongside per-timestep VTUs.

Parameters:

Name	Type	Description	Default
directory	PathLike | str	Path to write files to.	required
n_interp	int	Number of interpolation points to add between each element for smoother visualisation.	0
index	slice | Sequence[int] | int | Array | None	For batched Structures only, time step indices to plot.	None

Source code in src/flapjax/structure/data_structures.py

def plot(
    self,
    directory: os.PathLike | str,
    n_interp: int = 0,
    *,
    index: slice | Sequence[int] | int | Array | None = None,
) -> Path:
    r"""
    Plot beam results to VTK/VTU files in the specified directory. For a
    batched Structure, a PVD is written alongside per-timestep VTUs.
    :param directory: Path to write files to.
    :param n_interp: Number of interpolation points to add between each element for smoother visualisation.
    :param index: For batched Structures only, time step indices to plot.
    """
    if self.is_batched:
        index_ = index_to_arr(index=index, n_entries=self.n_tstep)
        directory_path = Path(directory).resolve()
        directory_path.mkdir(parents=True, exist_ok=True)

        paths = [self[i_ts]._plot_single(directory, n_interp) for i_ts in index_]

        assert self.t is not None
        return write_pvd(directory, "beam_dynamic_ts", paths, list(self.t[index_]))

    if index is not None:
        raise ValueError("`index` is only used for batched Structure")

    if not self.is_dynamic:
        return self.to_dynamic()._plot_single(directory, n_interp)
    return self._plot_single(directory, n_interp)

gradients

beam

BeamStructure

BeamStructure(
    num_nodes: int,
    connectivity: Array,
    y_vector: Array,
    k_cs_index: Array | None = None,
    m_cs_index: Array | None = None,
    m_lumped_index: Array | None = None,
    gravity: Array | Sequence[float] | None = None,
    thrust_nodes: dict[str, int] | None = None,
    thrust_direction: dict[str, Array] | None = None,
    optional_jacobians: OptionalJacobians | None = None,
    relaxation_factor: float = 1.0,
    spectral_radius: float = 0.9,
    alpha_m: float = 0.0,
    beta_k: float = 0.0,
    struct_convergence_settings: ConvergenceSettings = DEFAULT_STRUCT_CONVERGENCE_SETTINGS,
    constraints: dict[str, SoftConstraint | HardConstraint]
    | None = None,
)

Bases: BaseBeamStructure

Initialise BaseBeamStructure class with all non-design parameters.

Parameters:

Name	Type	Description	Default
num_nodes	int	Number of nodes in the structure.	required
connectivity	Array	Connectivity array, `(n_elem, 2)``.	required
y_vector	Array	Vector defining the y direction for each element, (n_elem, 3).	required
k_cs_index	Array | None	Array defining the index from the library of k_cs to use for each element, (n_elem, ). If None, all elements will use the first entry in the k_cs library.	None
m_cs_index	Array | None	Array defining the index from the library of m_cs to use for each element, (n_elem, ). If None, all elements will use the first entry in the m_cs library.	None
m_lumped_index	Array | None	Node index for nodes which are to have a lumped mass attached. The order is the same as that for the lumped mass data (n_lumped_mass, ).	None
gravity	Array | Sequence[float] | None	Gravity vector in global reference frame, or None for no gravity_vec, (3, ).	None
thrust_nodes	dict[str, int] | None	Dictionary of thrust node names and their corresponding node indices, {keys, int}.	None
thrust_direction	dict[str, Array] | None	Dictionary of thrust node names and their corresponding thrust direction vectors, {keys, (3, )}.	None
optional_jacobians	OptionalJacobians | None	Define which Jacobians contributions are to be used for solution.	None
relaxation_factor	float	Relaxation factor which reduces the displacement update at each iteration. A value of 1 is no relaxation, and a value of 0 is no update.	1.0
spectral_radius	float	Spectral radius for structural time integrator, where a value of 0 is highly damped and a value of 1 is undamped.	0.9
alpha_m	float	Mass-proportional Rayleigh damping coefficient.	0.0
beta_k	float	Stiffness-proportional Rayleigh damping coefficient.	0.0
struct_convergence_settings	ConvergenceSettings	Structure convergence settings.	DEFAULT_STRUCT_CONVERGENCE_SETTINGS
constraints	dict[str, SoftConstraint | HardConstraint] | None	Named dict {name: constraint}, or None, which add	None

Source code in src/flapjax/structure/beam.py

def __init__(
    self,
    num_nodes: int,
    connectivity: Array,
    y_vector: Array,
    k_cs_index: Array | None = None,
    m_cs_index: Array | None = None,
    m_lumped_index: Array | None = None,
    gravity: Array | Sequence[float] | None = None,
    thrust_nodes: dict[str, int] | None = None,
    thrust_direction: dict[str, Array] | None = None,
    optional_jacobians: OptionalJacobians | None = None,
    relaxation_factor: float = 1.0,
    spectral_radius: float = 0.9,
    alpha_m: float = 0.0,
    beta_k: float = 0.0,
    struct_convergence_settings: ConvergenceSettings = DEFAULT_STRUCT_CONVERGENCE_SETTINGS,
    constraints: (dict[str, SoftConstraint | HardConstraint] | None) = None,
) -> None:
    r"""
    Initialise BaseBeamStructure class with all non-design parameters.
    :param num_nodes: Number of nodes in the structure.
    :param connectivity: Connectivity array, `(n_elem, 2)``.
    :param y_vector: Vector defining the y direction for each element, ``(n_elem, 3)``.
    :param k_cs_index: Array defining the index from the library of k_cs to use for each element, ``(n_elem, )``.
    If ``None``, all elements will use the first entry in the k_cs library.
    :param m_cs_index: Array defining the index from the library of m_cs to use for each element, ``(n_elem, )``.
    If ``None``, all elements will use the first entry in the m_cs library.
    :param m_lumped_index: Node index for nodes which are to have a lumped mass attached. The order is the same as
    that for the lumped mass data ``(n_lumped_mass, )``.
    :param gravity: Gravity vector in global reference frame, or None for no gravity_vec, ``(3, )``.
    :param thrust_nodes: Dictionary of thrust node names and their corresponding node indices, {keys, int}.
    :param thrust_direction: Dictionary of thrust node names and their corresponding thrust direction vectors,
    ``{keys, (3, )}``.
    :param optional_jacobians: Define which Jacobians contributions are to be used for solution.
    :param relaxation_factor: Relaxation factor which reduces the displacement update at each iteration. A value of
    1 is no relaxation, and a value of 0 is no update.
    :param spectral_radius: Spectral radius for structural time integrator, where a value of 0 is highly damped and
    a value of 1 is undamped.
    :param alpha_m: Mass-proportional Rayleigh damping coefficient.
    :param beta_k: Stiffness-proportional Rayleigh damping coefficient.
    :param struct_convergence_settings: Structure convergence settings.
    :param constraints: Named dict ``{name: constraint}``, or None, which add
    """

    check_type(num_nodes, int)

    if constraints is None:
        named_constraints: dict[str, SoftConstraint | HardConstraint] = {}
    elif isinstance(constraints, dict):
        named_constraints = constraints
    else:
        raise ValueError("Invalid constraint input")

    hard_constraints = tuple(
        c for c in named_constraints.values() if isinstance(c, HardConstraint)
    )

    check_arr_shape(connectivity, (None, 2), "connectivity")
    check_arr_dtype(connectivity, int, "connectivity")
    _check_connectivity(connectivity, num_nodes)

    auto_node_sources: list[int] = []
    conn_list: list[list[int]] = connectivity.tolist()

    # add extra nodes and alter connectivity to account for constraints with the auto-generate node behaviour
    for con in hard_constraints:
        if con.node_j is None and not con.is_grounded:
            node_j_new = num_nodes + len(auto_node_sources)
            con.node_j = node_j_new
            auto_node_sources.append(con.node_i)
            conn_list = _split_connectivity(conn_list, con.node_i, node_j_new)
    self._auto_node_sources: tuple[int, ...] = tuple(auto_node_sources)

    num_nodes += len(auto_node_sources)
    connectivity = jnp.array(conn_list, dtype=int)
    self.n_nodes: int = num_nodes
    self.n_dof: int = num_nodes * 6

    self.connectivity: tuple[tuple[int, int], ...] = nested_list_to_tuple(
        connectivity.tolist()
    )  # (n_elem, 2)
    self.n_elem_per_node: tuple[int] = tuple(
        _n_elem_per_node(connectivity=connectivity, n_nodes=num_nodes).tolist()
    )  # (n_nodes, )
    self.n_elem: int = connectivity.shape[0]

    self.dof_per_elem: tuple[tuple[float]] = nested_list_to_tuple(
        jnp.zeros((self.n_elem, 12), dtype=int)
        .at[:, :6]
        .set(6 * self.connectivity_arr[:, [0]] + jnp.arange(6)[None, :])
        .at[:, 6:]
        .set(6 * self.connectivity_arr[:, [1]] + jnp.arange(6)[None, :])
        .tolist()
    )

    # allow for a single y_vector to be broadcast to all elements
    if y_vector.shape == (3,):
        y_vector = y_vector[None, :]
    if y_vector.shape == (1, 3):
        y_vector = jnp.broadcast_to(y_vector, (self.n_elem, 3))

    # y vectors in reference unoriented configuration, and placeholder for oriented equivalent.
    check_arr_shape(y_vector, (self.n_elem, 3), "y_vector")
    self.y_vector_reference: tuple[tuple[tuple[float]]] = nested_list_to_tuple(
        y_vector.tolist()
    )
    self.y_vector: Array = jnp.zeros_like(jnp.array(y_vector))

    # initialise design variables with default values
    self.x0_reference: Array = jnp.zeros((num_nodes, 3))  # unoriented
    self.x0: Array = jnp.zeros((num_nodes, 3))  # oriented

    self.m_cs = None
    self.k_cs = None
    self.m_lumped = None
    self.use_lumped_mass: bool = m_lumped_index is not None

    # initialise auxiliary arrays
    self.o0: Array = jnp.zeros((self.n_elem, 3, 3))
    self.l0: Array = jnp.zeros(self.n_elem)
    self.d0: Array = jnp.zeros((self.n_elem, 6))

    # initialise undeformed algebra and group
    self.hg0_reference: Array = jnp.zeros((self.n_nodes, 4, 4))  # unoriented
    self.hg0: Array = jnp.zeros((self.n_nodes, 4, 4))  # oriented

    # grads inverse action for the reference rotations
    self.ad_inv_o0: Array = jnp.zeros((self.n_elem, 6, 6))

    # gravity settings
    if not isinstance(gravity, jnp.ndarray) and gravity is not None:
        gravity = jnp.array(gravity)
    self.use_gravity: bool = gravity is not None and bool(jnp.any(gravity))
    if self.use_gravity:
        assert gravity is not None
        check_arr_shape(gravity, (3,), "gravity")
        self.gravity_vec: tuple[float, float, float] = tuple(gravity.tolist())
    else:
        self.gravity_vec = (0.0, 0.0, 0.0)

    # indexing
    if k_cs_index is None:
        k_cs_index_ = jnp.zeros(self.n_elem, dtype=int)
    else:
        check_arr_shape(k_cs_index, (self.n_elem,), "k_cs_index")
        check_arr_dtype(k_cs_index, int, "k_cs_index")
        k_cs_index_ = k_cs_index
    self.k_cs_index: tuple[int] = tuple(k_cs_index_.tolist())

    if m_cs_index is None:
        m_cs_index_ = jnp.zeros(self.n_elem, dtype=int)
    else:
        check_arr_shape(m_cs_index, (self.n_elem,), "m_cs_index")
        check_arr_dtype(m_cs_index, int, "m_cs_index")
        m_cs_index_ = m_cs_index
    self.m_cs_index: tuple[int] = tuple(m_cs_index_.tolist())

    self.m_lumped_index: tuple[int] | None = None
    if m_lumped_index is not None:
        check_arr_dtype(m_lumped_index, int, "m_lumped_index")
        if m_lumped_index.ndim not in (0, 1):
            raise ValueError("m_lumped_index.ndim must be 0 or 1.")
        self.m_lumped_index = tuple(jnp.atleast_1d(m_lumped_index).tolist())

    # add thrust
    self.thrust_nodes: tuple[tuple[str, int], ...] = ()
    self.thrust_direction: tuple[tuple[str, tuple[float, float, float]], ...] = ()
    if thrust_nodes is not None and thrust_direction is not None:
        if thrust_nodes.keys() != thrust_direction.keys():
            raise ValueError(
                f"Mismatch in keys of thrust_nodes ({thrust_nodes.keys()}) and thrust_direction ({thrust_direction.keys()}))."
            )

        for k, v in thrust_direction.items():
            check_arr_shape(v, (3,), f"thrust_direction[{k}]")

        self.thrust_nodes = tuple([(k, v) for k, v in thrust_nodes.items()])
        self.thrust_direction = tuple(
            [
                (k, nested_list_to_tuple((v / jnp.linalg.norm(v)).tolist()))
                for k, v in thrust_direction.items()
            ]
        )  # make unit vectors
    elif thrust_nodes is not None or thrust_direction is not None:
        warn(
            "One of thrust_nodes or thrust_direction has not been passed. Running with no thrust nodes."
        )

    # set the reference thrust to be zero, which can be overwritten later
    self.thrust_reference: dict[str, Array] = {
        k: jnp.atleast_1d(1) for k in [k_ for k_, v in self.thrust_nodes]
    }

    # set the reference orientation, which can be overwritten later.
    self.orientation_euler: Array = jnp.zeros(3)
    self.orientation: Array = jnp.eye(3)

    self.optional_jacobians: OptionalJacobians = (
        optional_jacobians
        if optional_jacobians is not None
        else OptionalJacobians()
    )
    self.struct_convergence_settings: ConvergenceSettings = (
        struct_convergence_settings
    )
    self.relaxation_factor: float = relaxation_factor
    self.spectral_radius: float = spectral_radius
    self.alpha_m: float = float(alpha_m)
    self.beta_k: float = float(beta_k)

    self.time_integrator = None

    self.constraints: dict[str, SoftConstraint | HardConstraint] = named_constraints

n_multibody_constraints property

n_multibody_constraints: int

Total number of scalar Lagrange-multiplier constraints.

n_holonomic_constraints property

n_holonomic_constraints: int

Number of scalar holonomic (position-level) Lagrange-multiplier constraints.

n_nonholonomic_constraints property

n_nonholonomic_constraints: int

Number of scalar non-holonomic (velocity-level) Lagrange-multiplier constraints.

case_from_dv

case_from_dv(dv: StructureDesignVariables) -> BeamStructure

Obtain a structural object as a function of design variables, allowing it to have defined gradients w.r.t. design variables.

Parameters:

Name	Type	Description	Default
dv	StructureDesignVariables	Design variables.	required

Returns:

Type	Description
BeamStructure	Beam structure object with the same functionality as self.

Source code in src/flapjax/structure/gradients/beam.py

def case_from_dv(self, dv: StructureDesignVariables) -> BeamStructure:
    r"""
    Obtain a structural object as a function of design variables, allowing it to have defined gradients w.r.t. design variables.
    :param dv: Design variables.
    :return: Beam structure object with the same functionality as self.
    """
    inner_case = pytree_clone(self)
    inner_case.set_design_variables(
        coords=dv_or(dv.x0, self.x0),
        k_cs=dv_or(dv.k_cs, self.k_cs),
        m_cs=dv_or(dv.m_cs, self.m_cs),
        m_lumped=dv_or(dv.m_lumped, self._m_lumped),
        remove_checks=True,
    )

    return inner_case

minimal_states_to_full_states

minimal_states_to_full_states(
    i_ts: int,
    q: StructureMinimalStates,
    dv: StructureDesignVariables,
    dv_full: StructureDesignVariables,
) -> StructureFullStates

Obtain the full set of states from the minimal states and the design variables.

Parameters:

Name	Type	Description	Default
i_ts	int	Index of the time step.	required
q	StructureMinimalStates	Minimal dynamic structure states.	required
dv	StructureDesignVariables	Design variables, where entries for gradients which aren't needed are set to None.	required
dv_full	StructureDesignVariables	Design variables, without omissions. These values are fallen back to when an entry in dv is none, to give an equivalent with zero gradient.	required

Returns:

Type	Description
StructureFullStates	Full set of structural states used inside objective function.

Source code in src/flapjax/structure/gradients/beam.py

def minimal_states_to_full_states(
    self,
    i_ts: int,
    q: StructureMinimalStates,
    dv: StructureDesignVariables,
    dv_full: StructureDesignVariables,
) -> StructureFullStates:
    r"""
    Obtain the full set of states from the minimal states and the design variables.
    :param i_ts: Index of the time step.
    :param q: Minimal dynamic structure states.
    :param dv: Design variables, where entries for gradients which aren't needed are set to None.
    :param dv_full: Design variables, without omissions. These values are fallen back to when an entry in ``dv`` is
    none, to give an equivalent with zero gradient.
    :return: Full set of structural states used inside objective function.
    """
    struct = self.case_from_dv(dv)
    hg = struct.compute_hg_from_varphi(q.varphi)
    d = struct.make_d(hg=hg)
    p_d = struct.make_p_d(d=d)
    eps = struct.make_eps(d=d)
    f_elem = struct.make_f_elem(eps=eps)

    f_ext_dead_i = (
        dv.f_ext_dead[i_ts, ...]
        if dv.f_ext_dead is not None
        else dv_full.f_ext_dead[i_ts, ...]
        if dv_full.f_ext_dead is not None
        else None
    )
    m_t = struct.make_m_t(d=d)

    # k_t_assembled is only required when Rayleigh damping is active. Skip the extra assembly when unused.
    k_t_assembled = (
        struct.make_k_t_full(
            d=d,
            p_d=p_d,
            eps=eps,
            f_ext_dead=f_ext_dead_i,
            rmat=hg[:, :3, :3],
            m_t=m_t,
        )
        if struct.beta_k != 0.0
        else None
    )

    assert dv_full.thrust_t is not None
    f_res = struct.make_f_res(
        solve_dofs=None,
        p_d=p_d,
        eps=eps,
        hg=hg,
        f_ext_follower_n=dv.f_ext_follower[i_ts, ...]
        if dv.f_ext_follower is not None
        else dv_full.f_ext_follower[i_ts, ...]
        if dv_full.f_ext_follower is not None
        else None,
        f_ext_dead_n=f_ext_dead_i,
        thrust_n={k: (v[i_ts] if v.ndim > 0 else v) for k, v in dv.thrust_t.items()}
        if dv.thrust_t is not None
        else {
            k: (v[i_ts] if v.ndim > 0 else v) for k, v in dv_full.thrust_t.items()
        },
        dynamic=True,
        m_t=m_t,
        c_l=self._make_c_t(d=d, d_dot=self._make_d_dot(p_d=p_d, v=q.v), v=q.v)[0],
        c_l_lumped=self._make_c_t_lumped(v=q.v)[0]
        if self.use_lumped_mass
        else None,
        v=q.v,
        v_dot=q.v_dot,
        i_ts=i_ts,
        k_t_assembled=k_t_assembled,
    )[0]
    return StructureFullStates(
        v=q.v,
        v_dot=q.v_dot,
        eps=eps,
        varphi=q.varphi,
        hg=hg,
        f_elem=f_elem,
        f_res=f_res,
    )

static_adjoint

static_adjoint(
    structure: StructureCase,
    objective: StructureObjectiveFunction,
    optional_jacobians: OptionalJacobians
    | None = OPTIONAL_JACOBIANS_DEFAULT,
    ad_mode: ADMode = "reverse",
) -> tuple[StructureDesignVariables, Array]

Computes the static grads of the structure, which is used to compute gradients of the loss with respect to the structure's parameters.

Parameters:

Name	Type	Description	Default
structure	StructureCase	StructureCase containing the current state of the structure.	required
objective	StructureObjectiveFunction	Objective function that takes the structure and design variables and returns an array	required
optional_jacobians	OptionalJacobians | None	OptionalJacobians object specifying which Jacobians to compute.	OPTIONAL_JACOBIANS_DEFAULT
ad_mode	ADMode	Flag on which to use of the forward or reverse adjoint.	'reverse'

Returns:

Type	Description
tuple[StructureDesignVariables, Array]	Gradient of objective function output with respect to design variables, and adjoint states.

Source code in src/flapjax/structure/gradients/beam.py

def static_adjoint(
    self,
    structure: StructureCase,
    objective: StructureObjectiveFunction,
    optional_jacobians: OptionalJacobians | None = OPTIONAL_JACOBIANS_DEFAULT,
    ad_mode: ADMode = "reverse",
) -> tuple[StructureDesignVariables, Array]:
    r"""
    Computes the static grads of the structure, which is used to compute gradients of the loss with respect to
    the structure's parameters.
    :param structure: StructureCase containing the current state of the structure.
    :param objective: Objective function that takes the structure and design variables and returns an array
    :param optional_jacobians: OptionalJacobians object specifying which Jacobians to compute.
    :param ad_mode: Flag on which to use of the forward or reverse adjoint.
    :return: Gradient of objective function output with respect to design variables, and adjoint states.
    """

    solve_dofs = jnp.array(
        get_solve_dofs(n_dof=self.n_dof, prescribed_dofs=structure.prescribed_dofs)
    )

    if optional_jacobians is not None:
        self.optional_jacobians = optional_jacobians

    # Recover original global dead force: structure.f_ext_dead is stored in local frame as
    # f_local = R^T @ f_global, so f_global = R @ f_local
    rmat = structure.hg[:, :3, :3]
    f_ext_dead_global = (
        transform_nodal_vect(structure.f_ext_dead, rmat)
        if structure.f_ext_dead is not None
        else None
    )

    # make design variables for current state of structure
    dv = StructureDesignVariables(
        x0=self.x0,
        orientation_euler=self.orientation_euler,
        k_cs=self.k_cs,
        m_cs=self._m_cs,
        m_lumped=self.m_lumped if self.use_lumped_mass else None,
        f_ext_follower=structure.f_ext_follower,
        f_ext_dead=f_ext_dead_global,
        thrust_t=structure.thrust,
        f_shape=(),
    )

    struct_states = structure.get_full_states()

    # find shape of objective function output without evaluating function
    f_properties = jax.eval_shape(lambda: objective(struct_states, dv, None))
    f_shape = f_properties.shape
    n_f = f_properties.size
    n_x = dv.n_x
    n_u = len(solve_dofs)
    n_u_full = self.n_dof

    # gradient of objective w.r.t. minimal states
    p_f_p_n, p_f_p_x = jax.jacrev(
        lambda varphi_, dv_: objective(
            self._structural_states_res_from_dv_varphi(
                dv=dv_, varphi=varphi_, thrust=structure.thrust
            ),
            dv_,
            None,
        ),
        argnums=(0, 1),
        allow_int=True,
    )(structure.varphi, dv)

    p_f_p_n = p_f_p_n.reshape(n_f, n_u_full)[:, solve_dofs]  # (n_f, n_u)
    p_f_p_x = p_f_p_x.ravel_jacobian(n_f, n_x)  # (n_f, n_x)

    # gradient of residual w.r.t. design variables and minimal states
    p_res_p_x, p_res_p_varphi = (jax.jacfwd if n_u > n_x else jax.jacrev)(
        lambda dv_, varphi_: (
            self._structural_states_res_from_dv_varphi(
                dv=dv_, varphi=varphi_, thrust=structure.thrust
            ).f_res
        ),
        argnums=(0, 1),
        allow_int=True,
    )(dv, structure.varphi)

    p_res_p_x = p_res_p_x.ravel_jacobian(n_u_full, n_x)[solve_dofs, :]  # (n_u, n_x)
    p_res_p_varphi = p_res_p_varphi.reshape(n_u_full, n_u_full)[
        jnp.ix_(solve_dofs, solve_dofs)
    ]  # (n_u, n_u)

    if ad_mode == "forward":
        # forward mode
        adj = jnp.linalg.solve(p_res_p_varphi, p_res_p_x)  # (n_u, n_x)
        rhs = p_f_p_n @ adj  # (n_f, n_x)
    elif ad_mode == "reverse":
        # reverse mode
        adj = jnp.linalg.solve(p_res_p_varphi.T, p_f_p_n.T).T  # (n_f, n_u)
        rhs = adj @ p_res_p_x  # (n_f, n_x)
    else:
        raise ValueError("AD mode must be either 'forward' or 'reverse'")

    return StructureDesignVariables(
        **dv.from_adjoint(f_shape, p_f_p_x - rhs), f_shape=f_shape
    ), adj

timestep_residual

timestep_residual(
    i_ts: int | Array,
    q_nm1: StructureMinimalStates,
    q_n: StructureMinimalStates,
    dv_: StructureDesignVariables,
    thrust_t: dict[str, Array],
    solve_dofs: tuple[int, ...],
    approx_grads: bool,
) -> Array

Routine to compute the full residual for the structural dynamic problem.

Parameters:

Name	Type	Description	Default
i_ts	int | Array	Time step index.	required
q_nm1	StructureMinimalStates	Previous minimal state.	required
q_n	StructureMinimalStates	Current minimal state.	required
dv_	StructureDesignVariables	Design variables.	required
thrust_t	dict[str, Array]	Thrust time history, {key, (n_tstep, )}.	required
solve_dofs	tuple[int, ...]	Solve degrees of freedom.	required
approx_grads	bool	If true, block gradients from some parts of the solution.	required

Returns:

Type	Description
Array	Residual vector, (4 * n_solve_dof,).

Source code in src/flapjax/structure/gradients/beam.py

def timestep_residual(
    self,
    i_ts: int | Array,
    q_nm1: StructureMinimalStates,
    q_n: StructureMinimalStates,
    dv_: StructureDesignVariables,
    thrust_t: dict[str, Array],
    solve_dofs: tuple[int, ...],
    approx_grads: bool,
) -> Array:
    r"""
    Routine to compute the full residual for the structural dynamic problem.
    :param i_ts: Time step index.
    :param q_nm1: Previous minimal state.
    :param q_n: Current minimal state.
    :param dv_: Design variables.
    :param thrust_t: Thrust time history, ``{key, (n_tstep, )}``.
    :param solve_dofs: Solve degrees of freedom.
    :param approx_grads: If true, block gradients from some parts of the solution.
    :return: Residual vector, ``(4 * n_solve_dof,)``.
    """
    return jnp.stack(
        (
            self.varphi_res_func(
                varphi_nm1=q_nm1.varphi,
                varphi_n=q_n.varphi,
                v_nm1=q_nm1.v,
                a_nm1=q_nm1.a,
                a_n=q_n.a,
                solve_dofs=solve_dofs,
            ),
            self.v_res_func(
                v_nm1=q_nm1.v,
                v_n=q_n.v,
                a_nm1=q_nm1.a,
                a_n=q_n.a,
                solve_dofs=solve_dofs,
            ),
            self.v_dot_res_func(
                i_ts=i_ts,
                varphi_nm1=q_nm1.varphi,
                varphi_n=q_n.varphi,
                v_nm1=q_nm1.v,
                v_n=q_n.v,
                v_dot_nm1=q_nm1.v_dot,
                v_dot_n=q_n.v_dot,
                approx_grads=approx_grads,
                f_aero_nm1=q_nm1.f_ext_aero,
                f_aero_n=q_n.f_ext_aero,
                thrust_t=thrust_t,
                dv=dv_,
                solve_dofs=solve_dofs,
            ),
            self.a_res_func(
                v_dot_nm1=q_nm1.v_dot,
                v_dot_n=q_n.v_dot,
                a_nm1=q_nm1.a,
                a_n=q_n.a,
                solve_dofs=solve_dofs,
            ),
        ),
        axis=0,
    ).ravel()  # [4*n_free_dof]

timestep_residual_jacobians

timestep_residual_jacobians(
    i_ts: int | Array,
    q_nm1: StructureMinimalStates,
    q_n: StructureMinimalStates,
    f_ext_aero_nm1: Array | None,
    f_ext_aero_n: Array | None,
    dv: StructureDesignVariables,
    thrust_t: dict[str, Array],
    solve_dofs: tuple[int, ...],
    approx_grads: bool,
    n_profile_loops: int | None,
    jac_options: dict[
        str, dict[str, Callable[..., Any] | None]
    ],
    mode: ADMode = "reverse",
) -> tuple[
    Array,
    Array,
    StructureDesignVariables,
    Array | None,
    Array | None,
    dict[str, dict[str, float]] | None,
    dict[str, dict[str, float]] | None,
]

Obtain the Jacobians of the structural residual with respect to the current states and previous states.

Parameters:

Name	Type	Description	Default
i_ts	int | Array	Time step index.	required
q_nm1	StructureMinimalStates	Previous minimal states.	required
q_n	StructureMinimalStates	Current minimal states.	required
f_ext_aero_nm1	Array | None	Optional aerodynamic forcing for previous time step, (n_nodes, 6).	required
f_ext_aero_n	Array | None	Optional aerodynamic forcing for current time step, (n_nodes, 6).	required
dv	StructureDesignVariables	Design variables.	required
thrust_t	dict[str, Array]	Thrust time history, {key, (n_tstep, )}.	required
solve_dofs	tuple[int, ...]	Index of degrees of freedom to solve for.	required
approx_grads	bool	If True, remove some gradient terms which are generally small.	required
n_profile_loops	int | None	Number of profile loops to run for timing function. If None, no profiling is done.	required
jac_options	dict[str, dict[str, Callable[..., Any] | None]]	Input which passes functions which can be used to approximate the Jacobians. If entries are None, AD is used.	required
mode	ADMode	AD mode used for Jacobian construction.	'reverse'

Returns:

Type	Description
tuple[Array, Array, StructureDesignVariables, Array | None, Array | None, dict[str, dict[str, float]] | None, dict[str, dict[str, float]] | None]	Jacobians with respect to previous state and current state, gradients with respect to design variables, previous, and current aerodynamic forces respectively, and profiling times for compilation and run time.

Source code in src/flapjax/structure/gradients/beam.py

def timestep_residual_jacobians(
    self,
    i_ts: int | Array,
    q_nm1: StructureMinimalStates,
    q_n: StructureMinimalStates,
    f_ext_aero_nm1: Array | None,
    f_ext_aero_n: Array | None,
    dv: StructureDesignVariables,
    thrust_t: dict[str, Array],
    solve_dofs: tuple[int, ...],
    approx_grads: bool,
    n_profile_loops: int | None,
    jac_options: dict[str, dict[str, Callable[..., Any] | None]],
    mode: ADMode = "reverse",
) -> tuple[
    Array,
    Array,
    StructureDesignVariables,
    Array | None,
    Array | None,
    dict[str, dict[str, float]] | None,
    dict[str, dict[str, float]] | None,
]:
    r"""
    Obtain the Jacobians of the structural residual with respect to the current states and previous states.
    :param i_ts: Time step index.
    :param q_nm1: Previous minimal states.
    :param q_n: Current minimal states.
    :param f_ext_aero_nm1: Optional aerodynamic forcing for previous time step, ``(n_nodes, 6)``.
    :param f_ext_aero_n: Optional aerodynamic forcing for current time step, ``(n_nodes, 6)``.
    :param dv: Design variables.
    :param thrust_t: Thrust time history, ``{key, (n_tstep, )}``.
    :param solve_dofs: Index of degrees of freedom to solve for.
    :param approx_grads: If True, remove some gradient terms which are generally small.
    :param n_profile_loops: Number of profile loops to run for timing function. If None, no profiling is done.
    :param jac_options: Input which passes functions which can be used to approximate the Jacobians. If entries are
    None, AD is used.
    :param mode: AD mode used for Jacobian construction.
    :return: Jacobians with respect to previous state and current state, gradients with respect to design variables,
    previous, and current aerodynamic forces respectively, and profiling times for compilation and run time.
    """

    compile_time: dict[str, dict[str, float]] = {}
    run_time: dict[str, dict[str, float]] = {}

    compute_f_aero_grads = f_ext_aero_nm1 is not None and f_ext_aero_n is not None
    if compute_f_aero_grads:
        jac_options["v_dot"].update({"f_aero_nm1": None, "f_aero_n": None})

    # varphi
    d_varphi, compile_time["varphi"], run_time["varphi"] = jacrev_custom(
        func=self.varphi_res_func,
        jac_options=jac_options["varphi"],
        n_profile_loops=n_profile_loops,
        func_name="varphi",
        static_argnames=("solve_dofs",),
        mode=mode,
    )(
        varphi_nm1=q_nm1.varphi.ravel(),
        varphi_n=q_n.varphi.ravel(),
        v_nm1=q_nm1.v.ravel(),
        a_nm1=q_nm1.a.ravel(),
        a_n=q_n.a.ravel(),
        solve_dofs=solve_dofs,
    )

    # velocity
    d_v, compile_time["v"], run_time["v"] = jacrev_custom(
        func=self.v_res_func,
        jac_options=jac_options["v"],
        n_profile_loops=n_profile_loops,
        func_name="v",
        static_argnames=("solve_dofs",),
        mode=mode,
    )(
        v_nm1=q_nm1.v.ravel(),
        v_n=q_n.v.ravel(),
        a_nm1=q_nm1.a.ravel(),
        a_n=q_n.a.ravel(),
        solve_dofs=solve_dofs,
    )

    # acceleration
    d_v_dot, compile_time["v_dot"], run_time["v_dot"] = jacrev_custom(
        func=self.v_dot_res_func,
        jac_options=jac_options["v_dot"],
        n_profile_loops=n_profile_loops,
        func_name="v_dot",
        static_argnames=("solve_dofs", "approx_grads"),
        mode=mode,
    )(
        i_ts=i_ts,
        varphi_nm1=q_nm1.varphi.ravel(),
        varphi_n=q_n.varphi.ravel(),
        v_nm1=q_nm1.v.ravel(),
        v_n=q_n.v.ravel(),
        v_dot_nm1=q_nm1.v_dot.ravel(),
        v_dot_n=q_n.v_dot.ravel(),
        dv=dv,
        f_aero_nm1=f_ext_aero_nm1.ravel() if compute_f_aero_grads else None,  # type: ignore
        f_aero_n=f_ext_aero_n.ravel() if compute_f_aero_grads else None,  # type: ignore
        thrust_t=thrust_t,
        solve_dofs=solve_dofs,
        approx_grads=approx_grads,
    )

    if not compute_f_aero_grads:
        # no Jacobians for aero case, but include a None to keep the keys consistent
        d_v_dot.update({"f_aero_nm1": None, "f_aero_n": None})

    # pseudo-acceleration
    d_a, compile_time["a"], run_time["a"] = jacrev_custom(
        func=self.a_res_func,
        jac_options=jac_options["a"],
        n_profile_loops=n_profile_loops,
        func_name="a",
        static_argnames=("solve_dofs",),
        mode=mode,
    )(
        v_dot_nm1=q_nm1.v_dot.ravel(),
        v_dot_n=q_n.v_dot.ravel(),
        a_nm1=q_nm1.a.ravel(),
        a_n=q_n.a.ravel(),
        solve_dofs=solve_dofs,
    )

    struct_sizes = (
        len(solve_dofs),
        len(solve_dofs),
        len(solve_dofs),
        len(solve_dofs),
    )

    nm1_keys = ("varphi_nm1", "v_nm1", "v_dot_nm1", "a_nm1")
    p_r_n_p_q_nm1 = construct_named_block_jacobian(
        entries=tuple(
            [
                {k: v[:, solve_dofs] for k, v in jacs.items() if k in nm1_keys}
                for jacs in (d_varphi, d_v, d_v_dot, d_a)
            ]
        ),
        keys=nm1_keys,
        widths=struct_sizes,
        heights=struct_sizes,
    )

    n_keys = ("varphi_n", "v_n", "v_dot_n", "a_n")
    p_r_n_p_q_n = construct_named_block_jacobian(
        entries=tuple(
            [
                {k: v[:, solve_dofs] for k, v in jacs.items() if k in n_keys}
                for jacs in (d_varphi, d_v, d_v_dot, d_a)
            ]
        ),
        keys=n_keys,
        widths=struct_sizes,
        heights=struct_sizes,
    )

    return (
        p_r_n_p_q_nm1,
        p_r_n_p_q_n,
        d_v_dot["dv"],
        d_v_dot["f_aero_nm1"],
        d_v_dot["f_aero_n"],
        compile_time if n_profile_loops is not None else None,
        run_time if n_profile_loops is not None else None,
    )

j_from_q_x

j_from_q_x(
    q_n_mat: Array,
    dv: StructureDesignVariables,
    dv_full: StructureDesignVariables,
    objective: StructureObjectiveFunction,
    i_ts: int,
) -> Array

Obtain the objective as a function of the minimal states and design variables.

Parameters:

Name	Type	Description	Default
q_n_mat	Array	Matrix representation of the minimal states.	required
dv	StructureDesignVariables	Design variables, with unwanted entries replaced with None.	required
dv_full	StructureDesignVariables	Design variables which are defined for all entries.	required
objective	StructureObjectiveFunction	Objective function.	required
i_ts	int	Time step index.	required

Returns:

Type	Description
Array	Objective value.

Source code in src/flapjax/structure/gradients/beam.py

def j_from_q_x(
    self,
    q_n_mat: Array,
    dv: StructureDesignVariables,
    dv_full: StructureDesignVariables,
    objective: StructureObjectiveFunction,
    i_ts: int,
) -> Array:
    r"""
    Obtain the objective as a function of the minimal states and design variables.
    :param q_n_mat: Matrix representation of the minimal states.
    :param dv: Design variables, with unwanted entries replaced with None.
    :param dv_full: Design variables which are defined for all entries.
    :param objective: Objective function.
    :param i_ts: Time step index.
    :return: Objective value.
    """
    full_states = self.minimal_states_to_full_states(
        i_ts=i_ts,
        q=StructureMinimalStates.from_mat(q_n_mat),
        dv=dv,
        dv_full=dv_full,
    )
    return jnp.atleast_1d(objective(full_states, dv, i_ts))

p_j

p_j(
    objective: StructureObjectiveFunction,
    i_ts: int,
    dv: StructureDesignVariables,
    dv_full: StructureDesignVariables,
    q_n: StructureMinimalStates,
) -> tuple[Array, StructureDesignVariables]

Obtains Jacobians of the objective function.

Parameters:

Name	Type	Description	Default
objective	StructureObjectiveFunction	Objective function.	required
i_ts	int	Time step index.	required
dv	StructureDesignVariables	Design variables.	required
dv_full	StructureDesignVariables	Design variables which are defined for all entries.	required
q_n	StructureMinimalStates	Current minimal states.	required

Returns:

Type	Description
tuple[Array, StructureDesignVariables]	Jacobian with respect to minimal states and design variables.

Source code in src/flapjax/structure/gradients/beam.py

@jax.jit(static_argnums=(0, 1, 3, 4))
def p_j(
    self,
    objective: StructureObjectiveFunction,
    i_ts: int,
    dv: StructureDesignVariables,
    dv_full: StructureDesignVariables,
    q_n: StructureMinimalStates,
) -> tuple[Array, StructureDesignVariables]:
    r"""
    Obtains Jacobians of the objective function.
    :param objective: Objective function.
    :param i_ts: Time step index.
    :param dv: Design variables.
    :param dv_full: Design variables which are defined for all entries.
    :param q_n: Current minimal states.
    :return: Jacobian with respect to minimal states and design variables.
    """

    def _j(q_n_mat: Array, dv_: StructureDesignVariables) -> Array:
        return self.j_from_q_x(
            q_n_mat=q_n_mat, dv=dv_, dv_full=dv_full, objective=objective, i_ts=i_ts
        )

    p_j_n_p_q_n, p_j_n_p_x = jax.jacrev(_j, argnums=(0, 1), allow_int=True)(
        q_n.to_mat(), dv
    )

    return cast(Array, p_j_n_p_q_n), cast(StructureDesignVariables, p_j_n_p_x)

adjoint_time_loop

adjoint_time_loop(
    rev_i_ts: int,
    d_j_d_x_: StructureDesignVariables,
    adj_: Array,
    p_r_np1_p_q_n: Array | None,
    adj_t_p_r_np1_p_q_n: Array | None,
    q_n: StructureMinimalStates,
    structure: StructureCase,
    objective: StructureObjectiveFunction,
    dv: StructureDesignVariables,
    dv_full: StructureDesignVariables,
    thrust_t: dict[str, Array],
    solve_dofs: tuple[int, ...],
    approx_grads: bool,
    save_adjoint: bool,
    matrix_free: bool,
    n_j: int,
    jac_options: dict[
        str, dict[str, Callable[..., Any] | None]
    ],
    i_ts_end: int | None = None,
) -> tuple[
    StructureDesignVariables,
    Array,
    Array,
    StructureMinimalStates,
]

Function to obtain the grads states at timestep varphi, which is dependent on the grads at timestep varphi+1.

Parameters:

Name	Type	Description	Default
rev_i_ts	int	Reversed timestep index. JAX loop does not allow for reverse indexing, and so this is. explicitly reversed within the function body to obtain i_ts.	required
d_j_d_x_	StructureDesignVariables	Design gradient to accumulate.	required
adj_	Array	Full grads matrix which is updated inplace, (n_tstep, *j_shape, 5*n_dof).	required
p_r_np1_p_q_n	Array | None	Gradient of future step with respect to current state, used when computing the full Jacobian (5*n_dof, 5*n_dof).	required
adj_t_p_r_np1_p_q_n	Array | None	VJP of the future adjoint step and the Jacobian of the future residual with respect to the current state, (n_adj_dof, ).	required
q_n	StructureMinimalStates	Current minimal states.	required
structure	StructureCase	Dynamic structure solution.	required
objective	StructureObjectiveFunction	Objective function.	required
dv	StructureDesignVariables	Structure design variables.	required
dv_full	StructureDesignVariables	Structure design variables which are defined for all entries.	required
thrust_t	dict[str, Array]	Thrust time history, {key, (n_tstep, )}.	required
solve_dofs	tuple[int, ...]	Tuple of dof index to solve.	required
approx_grads	bool	Whether to approximate the gradient or not.	required
save_adjoint	bool	Whether to save the full adjoint time history.	required
matrix_free	bool	If False, solve the system using the residual Jacobian-vector product using GMRES.	required
n_j	int	Number of objective function outputs.	required
jac_options	dict[str, dict[str, Callable[..., Any] | None]]	Input which passes functions which can be used to approximate the Jacobians. If entries are None, AD is used.	required
i_ts_end	int | None	Largest time step index for which the adjoint is computed. Defaults to structure.n_tstep - 1 when None	None

Returns:

Type	Description
tuple[StructureDesignVariables, Array, Array, StructureMinimalStates]	Updated grads matrix, gradient of current step with respect to previous state and current state.

Source code in src/flapjax/structure/gradients/beam.py

def adjoint_time_loop(
    self,
    rev_i_ts: int,
    d_j_d_x_: StructureDesignVariables,
    adj_: Array,
    p_r_np1_p_q_n: Array | None,
    adj_t_p_r_np1_p_q_n: Array | None,
    q_n: StructureMinimalStates,
    structure: StructureCase,
    objective: StructureObjectiveFunction,
    dv: StructureDesignVariables,
    dv_full: StructureDesignVariables,
    thrust_t: dict[str, Array],
    solve_dofs: tuple[int, ...],
    approx_grads: bool,
    save_adjoint: bool,
    matrix_free: bool,
    n_j: int,
    jac_options: dict[str, dict[str, Callable[..., Any] | None]],
    i_ts_end: int | None = None,
) -> tuple[StructureDesignVariables, Array, Array, StructureMinimalStates]:
    r"""
    Function to obtain the grads states at timestep varphi, which is dependent on the grads at timestep varphi+1.
    :param rev_i_ts: Reversed timestep index. JAX loop does not allow for reverse indexing, and so this is.
    explicitly reversed within the function body to obtain i_ts.
    :param d_j_d_x_: Design gradient to accumulate.
    :param adj_: Full grads matrix which is updated inplace, ``(n_tstep, *j_shape, 5*n_dof)``.
    :param p_r_np1_p_q_n: Gradient of future step with respect to current state, used when computing the full
    Jacobian ``(5*n_dof, 5*n_dof)``.
    :param adj_t_p_r_np1_p_q_n: VJP of the future adjoint step and the Jacobian of the future residual with respect
    to the current state, ``(n_adj_dof, )``.
    :param q_n: Current minimal states.
    :param structure: Dynamic structure solution.
    :param objective: Objective function.
    :param dv: Structure design variables.
    :param dv_full: Structure design variables which are defined for all entries.
    :param thrust_t: Thrust time history, ``{key, (n_tstep, )}``.
    :param solve_dofs: Tuple of dof index to solve.
    :param approx_grads: Whether to approximate the gradient or not.
    :param save_adjoint: Whether to save the full adjoint time history.
    :param matrix_free: If False, solve the system using the residual Jacobian-vector product using GMRES.
    :param n_j: Number of objective function outputs.
    :param jac_options: Input which passes functions which can be used to approximate the Jacobians. If entries are
    None, AD is used.
    :param i_ts_end: Largest time step index for which the adjoint is computed. Defaults to
    ``structure.n_tstep - 1`` when ``None``
    :return: Updated grads matrix, gradient of current step with respect to previous state and current state.
    """

    i_ts_end_ = structure.n_tstep - 1 if i_ts_end is None else i_ts_end
    i_ts = i_ts_end_ - rev_i_ts  # index for timestep n, which decrements

    i_ts_nm1 = jnp.maximum(i_ts - 1, 0)  # index for timestep varphi-1

    solve_idx = jnp.array(solve_dofs)

    # find minimal states for timestep varphi-1
    q_nm1 = structure.get_minimal_states(i_ts_nm1)

    # Objective sensitivities
    p_j_n_p_q_n, p_j_n_p_x = self.p_j(
        objective=objective, i_ts=i_ts, dv=dv, dv_full=dv_full, q_n=q_n
    )

    if matrix_free:

        def _residual_states(
            q_n_: StructureMinimalStates, q_nm1_: StructureMinimalStates
        ):
            return self.timestep_residual(
                i_ts=i_ts,
                q_nm1=q_nm1_,
                q_n=q_n_,
                dv_=dv,
                thrust_t=thrust_t,
                solve_dofs=solve_dofs,
                approx_grads=approx_grads,
            )

        # Linearise the timestep residual around (q_n, q_nm1). This single VJP returns:
        # p_r_n_dot_v(v)[0] = (p_r_n/p_q_n).T @ v, p_r_n_dot_v(v)[1] = (p_r_n/p_q_nm1).T @ v
        _, p_r_n_dot_v = jax.vjp(_residual_states, q_n, q_nm1)

        def _cot_to_solve_vec(cot: StructureMinimalStates) -> Array:
            # collapse cotangent to [n_adj_dof]
            mat = cot.to_mat()  # [4, n_nodes, 6]
            return mat.reshape(mat.shape[0], -1)[:, solve_idx].ravel()

        def matvec_qn_t(v: Array) -> Array:
            # function to compute (p_r_n/p_q_n).T @ v for some vector v
            return _cot_to_solve_vec(p_r_n_dot_v(v)[0])

        # sensitivity of objective to degrees of freedom, (n_j, n_adj_dof)
        p_j_solve = (
            p_j_n_p_q_n.reshape(n_j, 4, -1, 6)
            .reshape(n_j, 4, -1)[..., solve_idx]
            .reshape(n_j, -1)
        )
        assert adj_t_p_r_np1_p_q_n is not None, (
            "The adjoint-Jacobian product has not been passed"
        )
        b_rhs = -(p_j_solve + adj_t_p_r_np1_p_q_n)  # (n_j, n_adj_dof)

        # solve for the adjoint vector at timestep n, batched along the size of the objective.
        def _solve_row(b_row: Array) -> Array:
            # noinspection PyTypeChecker
            x, _ = jax.scipy.sparse.linalg.gmres(
                matvec_qn_t,
                b_row,
                tol=1e-10,
                atol=1e-10,
                maxiter=50,
                solve_method="batched",
            )
            return x

        adj_n = jax.vmap(_solve_row)(b_rhs)  # (n_j, n_adj_dof)

        # Design gradient accumulation via a separate VJP to obtain adj.T @ p_r_v_dot_n_p_dv.
        def _residual_dv(dv_: StructureDesignVariables) -> Array:
            return self.timestep_residual(
                i_ts=i_ts,
                q_nm1=q_nm1,
                q_n=q_n,
                dv_=dv_,
                thrust_t=thrust_t,
                solve_dofs=solve_dofs,
                approx_grads=approx_grads,
            )

        _, pull_dv = jax.vjp(_residual_dv, dv)
        dv_grads = jax.vmap(pull_dv)(adj_n)[0]

        # accumulate with seperate statements as there is no __add__ member
        d_j_d_x_ += dv_grads
        d_j_d_x_ += p_j_n_p_x

        # compute adj_n @ p_r_n/p_q_nm1 for next iteration
        def _coupling_row(a: Array) -> Array:
            _, cot_qnm1 = p_r_n_dot_v(a)
            return _cot_to_solve_vec(cot_qnm1)

        adj_t_p_r_n_p_q_nm1 = jax.vmap(_coupling_row)(adj_n)  # (n_j, n_adj_dof)

        p_r_n_p_q_nm1: Array | None = None  # unused
    else:
        # find gradients of residual function (state Jacobians only)
        p_r_n_p_q_nm1, p_r_n_p_q_n, p_r_v_dot_n_p_dv, *_ = (
            self.timestep_residual_jacobians(
                i_ts=i_ts,
                q_n=q_n,
                q_nm1=q_nm1,
                dv=dv,
                solve_dofs=solve_dofs,
                approx_grads=approx_grads,
                f_ext_aero_nm1=None,
                f_ext_aero_n=None,
                thrust_t=thrust_t,
                n_profile_loops=None,
                jac_options=jac_options,
            )
        )

        # solve for adjoint at current timestep
        prev_adjoint = adj_[i_ts + 1, ...] if save_adjoint else adj_
        b: Array = -(p_j_n_p_q_n.reshape(n_j, -1) + prev_adjoint @ p_r_np1_p_q_n).T
        adj_n = jnp.linalg.solve(p_r_n_p_q_n.T, b).T

        # accumulate design derivative
        d_j_d_x_ += p_r_v_dot_n_p_dv.premultiply_adj(
            adj_n[:, solve_idx + 2 * len(solve_dofs)]
        )

        # add on direct contribution from objective
        d_j_d_x_ += p_j_n_p_x

        adj_t_p_r_n_p_q_nm1 = None  # unused

    # update adjoint vector time history if requested
    if save_adjoint:
        adj_ = adj_.at[i_ts, ...].set(adj_n)

    # print to console
    jax_print(
        "Adjoint step: {i_ts}",
        i_ts=i_ts,
        verbose_level="normal",
    )

    if matrix_free:
        assert adj_t_p_r_n_p_q_nm1 is not None
        return d_j_d_x_, adj_ if save_adjoint else adj_n, adj_t_p_r_n_p_q_nm1, q_nm1
    else:
        assert p_r_n_p_q_nm1 is not None
        return d_j_d_x_, adj_ if save_adjoint else adj_n, p_r_n_p_q_nm1, q_nm1

construct_approximate_jacobians

construct_approximate_jacobians(
    sol: StructureCase,
    jacobian_approximations: StructureJacobianApproximations,
) -> dict[str, dict[str, Callable[..., Any] | None]]

Compute approximations for Jacobians which are specified in the jacobian_approximations data structure.

Parameters:

Name	Type	Description	Default
sol	StructureCase	Solution for which approximations will be created for the initial time step.	required
jacobian_approximations	StructureJacobianApproximations	Data structure which defines which approximations to create.	required

Returns:

Type	Description
dict[str, dict[str, Callable[..., Any] | None]]	Dictionary of approximations.

Source code in src/flapjax/structure/gradients/beam.py

def construct_approximate_jacobians(
    self,
    sol: StructureCase,
    jacobian_approximations: StructureJacobianApproximations,
) -> dict[str, dict[str, Callable[..., Any] | None]]:
    r"""
    Compute approximations for Jacobians which are specified in the jacobian_approximations data structure.
    :param sol: Solution for which approximations will be created for the initial time step.
    :param jacobian_approximations: Data structure which defines which approximations to create.
    :return: Dictionary of approximations.
    """
    q_nm1 = sol.get_minimal_states(0)
    q_n = sol.get_minimal_states(1)
    dv = self.get_design_variables(
        struct_case=sol, thrust_t=sol.thrust, grads_to_compute=None
    )
    solve_dofs = tuple(
        int(i)
        for i in get_solve_dofs(
            n_dof=self.n_dof, prescribed_dofs=sol.prescribed_dofs
        )
    )
    if sol.f_ext_aero is not None:
        f_aero_nm1 = sol.f_ext_aero[0, ...].ravel()
        f_aero_n = sol.f_ext_aero[1, ...].ravel()
    else:
        f_aero_nm1 = None
        f_aero_n = None

    res_args: dict[
        str, tuple[Callable[..., Array], dict[str, Any], Sequence[str]]
    ] = {
        "varphi": (
            self.varphi_res_func,
            {
                "varphi_nm1": q_nm1.varphi.ravel(),
                "varphi_n": q_n.varphi.ravel(),
                "v_nm1": q_nm1.v.ravel(),
                "a_nm1": q_nm1.a.ravel(),
                "a_n": q_n.a.ravel(),
                "solve_dofs": solve_dofs,
            },
            [f.name for f in fields(VarphiApprox)],
        ),
        "v": (
            self.v_res_func,
            {
                "v_nm1": q_nm1.v.ravel(),
                "v_n": q_n.v.ravel(),
                "a_nm1": q_nm1.a.ravel(),
                "a_n": q_n.a.ravel(),
                "solve_dofs": solve_dofs,
            },
            [f.name for f in fields(VApprox)],
        ),
        "v_dot": (
            self.v_dot_res_func,
            {
                "i_ts": 1,
                "varphi_nm1": q_nm1.varphi.ravel(),
                "varphi_n": q_n.varphi.ravel(),
                "v_nm1": q_nm1.v.ravel(),
                "v_n": q_n.v.ravel(),
                "v_dot_nm1": q_nm1.v_dot.ravel(),
                "v_dot_n": q_n.v_dot.ravel(),
                "dv": dv,
                "f_aero_nm1": f_aero_nm1,
                "f_aero_n": f_aero_n,
                "thrust_t": sol.thrust,
                "solve_dofs": solve_dofs,
                "approx_grads": True,
            },
            [f.name for f in fields(VDotApprox)],
        ),
        "a": (
            self.a_res_func,
            {
                "v_dot_nm1": q_nm1.v_dot.ravel(),
                "v_dot_n": q_n.v_dot.ravel(),
                "a_nm1": q_nm1.a.ravel(),
                "a_n": q_n.a.ravel(),
                "solve_dofs": solve_dofs,
            },
            [f.name for f in fields(AApprox)],
        ),
    }

    return construct_approximation(
        res_args=res_args, jacobian_approximations=jacobian_approximations
    )

dynamic_adjoint

dynamic_adjoint(
    structure: StructureCase,
    objective: StructureObjectiveFunction,
    matrix_free: bool = False,
    jacobian_approximations: StructureJacobianApproximations = JACOBIAN_APPROXIMATIONS_DEFAULT,
    p_q0_p_x: StructureDesignVariables | None = None,
    save_adjoint: bool = False,
    approx_grads: bool = True,
    grads_to_compute: StructureGradsToCompute = GRADS_TO_COMPUTE_DEFAULT,
    i_ts_adjoint_range: tuple[int | None, int | None] = (
        None,
        None,
    ),
) -> tuple[StructureDesignVariables, Array | None]

Dynamic structure grads problem. This computes the gradient of the objective of the dynamic response with respect to design variables. The objective has structure :math:J = \sum_{i=1}^N \left(j(\mathbf{x}, \mathbf{y}_i)\right) where :math:\mathbf{x} are the design variables and :math:\mathbf{y} are the structural states at each timestep, which depend on the design variables through the dynamic structure equations. The gradient is computed by first solving a backward pass to obtain the grads states, and then using these to compute the gradient w.r.t. design variables in a forward pass.

Parameters:

Name	Type	Description	Default
structure	StructureCase	Dynamic structure solution object.	required
objective	StructureObjectiveFunction	Objective function :math:j(\mathbf{x}, \mathbf{y}_i).	required
matrix_free	bool	Whether to use matrix-free methods for solving the linear systems. Default is False, as structural problems generally do not benefit from this solve.	False
jacobian_approximations	StructureJacobianApproximations	Data structure which specifies Jacobian approximations to use for each part of the problem. The value can either be None for no approximation, constant for the assumption that the Jacobian does not vary with any variables. Alternatively, it can be tuple pairs with first entry being dense_linear or lazy_linear, with the second entry being a sequence of argument names for which to obtain the Hessian.	JACOBIAN_APPROXIMATIONS_DEFAULT
p_q0_p_x	StructureDesignVariables | None	Optional Jacobian used to describe the sensitivities of the initial structural degrees of freedom to the design variables.	None
save_adjoint	bool	Whether to save the full adjoint vectors.	False
approx_grads	bool	If true, some gradient contributions which are assumed to be near-zero are removed to decrease computational cost.	True
grads_to_compute	StructureGradsToCompute	Design variables with which to compute design gradients for.	GRADS_TO_COMPUTE_DEFAULT
i_ts_adjoint_range	tuple[int | None, int | None]	Optional (start, end) window of time step indices over which the objective contributes to the gradient. Either entry may be None to leave that side untruncated. Defining a start time step that is nonzero will skip the initial state adjoint contribution.	(None, None)

Returns:

Type	Description
tuple[StructureDesignVariables, Array | None]	Objective gradient :math:\frac{dJ}{d\mathbf{x}} and adjoint states

Source code in src/flapjax/structure/gradients/beam.py

def dynamic_adjoint(
    self,
    structure: StructureCase,
    objective: StructureObjectiveFunction,
    matrix_free: bool = False,
    jacobian_approximations: StructureJacobianApproximations = JACOBIAN_APPROXIMATIONS_DEFAULT,
    p_q0_p_x: StructureDesignVariables | None = None,
    save_adjoint: bool = False,
    approx_grads: bool = True,
    grads_to_compute: StructureGradsToCompute = GRADS_TO_COMPUTE_DEFAULT,
    i_ts_adjoint_range: tuple[int | None, int | None] = (None, None),
) -> tuple[StructureDesignVariables, Array | None]:
    r"""
    Dynamic structure grads problem. This computes the gradient of the objective of the dynamic response with
    respect to design variables. The objective has structure
    :math:`J = \sum_{i=1}^N \left(j(\mathbf{x}, \mathbf{y}_i)\right)` where :math:`\mathbf{x}` are the design variables
    and :math:`\mathbf{y}` are the structural states at each timestep, which depend on the design variables through
    the dynamic structure equations. The gradient is computed by first solving a backward pass to obtain the grads
    states, and then using these to compute the gradient w.r.t. design variables in a forward pass.
    :param structure: Dynamic structure solution object.
    :param objective: Objective function :math:`j(\mathbf{x}, \mathbf{y}_i)`.
    :param matrix_free: Whether to use matrix-free methods for solving the linear systems. Default is False, as
    structural problems generally do not benefit from this solve.
    :param jacobian_approximations: Data structure which specifies Jacobian approximations to use for each part of
    the problem. The value can either be None for no approximation, `constant` for the assumption that the Jacobian
    does not vary with any variables. Alternatively, it can be tuple pairs with first entry being `dense_linear` or
    `lazy_linear`, with the second entry being a sequence of argument names for which to obtain the Hessian.
    :param p_q0_p_x: Optional Jacobian used to describe the sensitivities of the initial structural degrees of
    freedom to the design variables.
    :param save_adjoint: Whether to save the full adjoint vectors.
    :param approx_grads: If true, some gradient contributions which are assumed to be near-zero are removed to
    decrease computational cost.
    :param grads_to_compute: Design variables with which to compute design gradients for.
    :param i_ts_adjoint_range: Optional ``(start, end)`` window of time step indices over which the objective
    contributes to the gradient. Either entry may be ``None`` to leave that side untruncated. Defining a start
    time step that is nonzero will  skip the initial state adjoint contribution.
    :return: Objective gradient :math:`\frac{dJ}{d\mathbf{x}}` and adjoint states
    """

    dv = self.get_design_variables(
        struct_case=structure,
        thrust_t=structure.thrust,
        grads_to_compute=grads_to_compute,
    )

    dv_full = self.get_design_variables(
        struct_case=structure, thrust_t=structure.thrust, grads_to_compute=None
    )

    struct_states_init = structure.get_full_states(i_ts=0)

    j_properties = jax.eval_shape(
        lambda: jnp.atleast_1d(objective(struct_states_init, dv, None))
    )
    j_shape = j_properties.shape
    n_j = j_properties.size

    # assemble
    solve_dofs: tuple[int, ...] = tuple(
        int(i)
        for i in get_solve_dofs(
            n_dof=self.n_dof, prescribed_dofs=structure.prescribed_dofs
        )
    )

    dv_grad_init = StructureDesignVariables(
        x0=jnp.zeros((*j_shape, *self.x0.shape)) if dv.x0 is not None else None,
        orientation_euler=jnp.zeros((*j_shape, 3))
        if dv.orientation_euler is not None
        else None,
        k_cs=jnp.zeros((*j_shape, *self.k_cs.shape))
        if dv.k_cs is not None
        else None,
        m_cs=jnp.zeros((*j_shape, *self.m_cs.shape))
        if dv.m_cs is not None
        else None,
        m_lumped=jnp.zeros((*j_shape, *self.m_lumped.shape))
        if self.use_lumped_mass and dv.m_lumped is not None
        else None,
        f_ext_dead=jnp.zeros((*j_shape, *structure.f_ext_dead.shape))
        if structure.f_ext_dead is not None and dv.f_ext_dead is not None
        else None,
        f_ext_follower=jnp.zeros((*j_shape, *structure.f_ext_follower.shape))
        if structure.f_ext_follower is not None and dv.f_ext_follower is not None
        else None,
        thrust_t={
            k: jnp.zeros((*j_shape, *v.shape)) for k, v in structure.thrust.items()
        }
        if dv.thrust_t is not None
        else None,
        f_shape=(),
    )

    n_adj_dof = 4 * (
        self.n_dof - len(structure.prescribed_dofs)
    )  # number of grads degrees of freedom

    # compute Jacobian approximations, if requested
    jac_options = self.construct_approximate_jacobians(
        sol=structure, jacobian_approximations=jacobian_approximations
    )

    # check adjoint window
    i_ts_start_adj, i_ts_end_adj = i_ts_adjoint_range
    i_ts_start_adj_: int = 1 if i_ts_start_adj is None else i_ts_start_adj
    i_ts_end_adj_: int = (
        structure.n_tstep - 1 if i_ts_end_adj is None else i_ts_end_adj
    )
    if i_ts_start_adj_ < 1:
        raise ValueError(
            f"i_ts_adjoint_range start must be >= 1, got {i_ts_start_adj_}"
        )
    if i_ts_end_adj_ > structure.n_tstep - 1:
        raise ValueError(
            f"i_ts_adjoint_range end must be <= n_tstep - 1 = "
            f"{structure.n_tstep - 1}, got {i_ts_end_adj_}"
        )
    if i_ts_end_adj_ < i_ts_start_adj_:
        raise ValueError(
            f"i_ts_adjoint_range end ({i_ts_end_adj_}) must be >= start "
            f"({i_ts_start_adj_})"
        )
    n_adj_iters: int = i_ts_end_adj_ - i_ts_start_adj_ + 1

    # wrap in a local JIT so structure/aero_dv become closure constants
    @jax.jit
    def adjoint_step(
        rev_i_ts_: int,
        d_j_d_x_: StructureDesignVariables,
        adj_: Array,
        coupling_arr: Array,
        q_n: StructureMinimalStates,
    ) -> tuple[StructureDesignVariables, Array, Array, StructureMinimalStates]:
        return self.adjoint_time_loop(
            rev_i_ts=rev_i_ts_,
            d_j_d_x_=d_j_d_x_,
            adj_=adj_,
            p_r_np1_p_q_n=None if matrix_free else coupling_arr,
            adj_t_p_r_np1_p_q_n=coupling_arr if matrix_free else None,
            q_n=q_n,
            structure=structure,
            objective=objective,
            dv=dv,
            dv_full=dv_full,
            thrust_t=structure.thrust,
            solve_dofs=solve_dofs,
            approx_grads=approx_grads,
            save_adjoint=save_adjoint,
            matrix_free=matrix_free,
            n_j=n_j,
            jac_options=jac_options,
            i_ts_end=i_ts_end_adj_,
        )

    # coupling array is either a Jacobian or a VJP depending on if using matrix free or not
    coupling_init = (
        jnp.zeros((n_j, n_adj_dof))
        if matrix_free
        else jnp.zeros((n_adj_dof, n_adj_dof))
    )

    # pass through time steps backwards to obtain adjoints
    # coupling0 is p_r1_p_q0 when matrix_free is False, and adj_1 @ p_r1_p_q0 when matrix_free is True
    d_j_d_x, adj, coupling0, _ = jax.lax.fori_loop(
        lower=0,
        upper=n_adj_iters,
        body_fun=lambda i_ts_, args: adjoint_step(i_ts_, *args),
        init_val=(
            dv_grad_init,
            jnp.zeros((structure.n_tstep + 1, n_j, n_adj_dof))
            if save_adjoint
            else jnp.zeros((n_j, n_adj_dof)),
            coupling_init,
            structure.get_minimal_states(i_ts_end_adj_),
        ),
    )

    # solve initial timestep adjoint, as there is no r0. Skipped when the adjoint window truncates early time steps
    if i_ts_start_adj_ <= 1:
        p_j0_p_q0, p_j0_p_x = self.p_j(
            objective=objective,
            i_ts=0,
            dv=dv,
            dv_full=dv_full,
            q_n=structure.get_minimal_states(0),
        )

        if matrix_free:
            adj0 = -p_j0_p_q0.reshape(n_j, -1) - coupling0
        else:
            adj0 = (
                -p_j0_p_q0.reshape(n_j, -1)
                - (adj[1, ...] if save_adjoint else adj) @ coupling0
            )

        # add initial direct sensitivity
        d_j_d_x += p_j0_p_x

        # include initial state sensitivity
        if p_q0_p_x is not None:
            d_j_d_x += p_q0_p_x.premultiply_adj(-adj0)
    else:
        adj0 = jnp.zeros((n_j, n_adj_dof))

    # restore original shape of j, and cut off zeros for past-end timestep
    if save_adjoint:
        adj = adj.at[0, ...].set(adj0)
        return d_j_d_x, adj.reshape(adj.shape[0], *j_shape, *adj.shape[2:])[:-1]
    else:
        return d_j_d_x, None

dynamic_adjoint_jacobian_profile

dynamic_adjoint_jacobian_profile(
    sol: StructureCase,
    approx_grads: bool,
    jacobian_approximations: StructureJacobianApproximations = JACOBIAN_APPROXIMATIONS_DEFAULT,
    grads_to_compute: StructureGradsToCompute | None = None,
    f_aero_nm1_n: tuple[Array, Array] | None = None,
    i_ts: int = 1,
    n_loop: int = 10,
    *,
    print_header: bool = True,
) -> tuple[
    dict[str, dict[str, float]], dict[str, dict[str, float]]
]

Function to time evaluation of the Jacobians used for the adjoint solution.

Parameters:

Name	Type	Description	Default
sol	StructureCase	Dynamic structural solution to extract states from.	required
approx_grads	bool	If True, neglect small gradient terms.	required
jacobian_approximations	StructureJacobianApproximations	Data structure which specifies Jacobian approximations to use for each part of the problem.	JACOBIAN_APPROXIMATIONS_DEFAULT
grads_to_compute	StructureGradsToCompute | None	StructureGradsToCompute object which describes which design gradients to compute. If None, all gradients will be computed.	None
f_aero_nm1_n	tuple[Array, Array] | None	Tuple of [f_aero_nm1, f_aero_n] which are passed from the aero problem. If None, no aerodynamic force gradients will be computed.	None
i_ts	int	Time step index where to evaluate residual Jacobians.	1
n_loop	int	Number of times to loop the Jacobian evaluation time for averaging the runtime.	10
print_header	bool	Flag used to prevent heading printer when called by the coupled profiler.	True

Returns:

Type	Description
tuple[dict[str, dict[str, float]], dict[str, dict[str, float]]]	Dictionary of {residual_name: {gradient_argument: val}} for compile time and run time respectively.

Source code in src/flapjax/structure/gradients/beam.py

def dynamic_adjoint_jacobian_profile(
    self,
    sol: StructureCase,
    approx_grads: bool,
    jacobian_approximations: StructureJacobianApproximations = JACOBIAN_APPROXIMATIONS_DEFAULT,
    grads_to_compute: StructureGradsToCompute | None = None,
    f_aero_nm1_n: tuple[Array, Array] | None = None,
    i_ts: int = 1,
    n_loop: int = 10,
    *,
    print_header: bool = True,
) -> tuple[dict[str, dict[str, float]], dict[str, dict[str, float]]]:
    r"""
    Function to time evaluation of the Jacobians used for the adjoint solution.
    :param sol: Dynamic structural solution to extract states from.
    :param approx_grads: If True, neglect small gradient terms.
    :param jacobian_approximations: Data structure which specifies Jacobian approximations to use for each part of
    the problem.
    :param grads_to_compute: StructureGradsToCompute object which describes which design gradients to compute. If
    None, all gradients will be computed.
    :param f_aero_nm1_n: Tuple of [f_aero_nm1, f_aero_n] which are passed from the aero problem. If None, no
    aerodynamic force gradients will be computed.
    :param i_ts: Time step index where to evaluate residual Jacobians.
    :param n_loop: Number of times to loop the Jacobian evaluation time for averaging the runtime.
    :param print_header: Flag used to prevent heading printer when called by the coupled profiler.
    :return: Dictionary of {residual_name: {gradient_argument: val}} for compile time and run time respectively.
    """

    if print_header:
        print_table_title(inner_width=95, title="Structure Adjoint Profile")

    # compute Jacobian approximations, if requested
    jac_options = self.construct_approximate_jacobians(
        sol=sol, jacobian_approximations=jacobian_approximations
    )

    common_kwargs = {
        "i_ts": i_ts,
        "q_nm1": sol.get_minimal_states(i_ts - 1),
        "q_n": sol.get_minimal_states(i_ts),
        "dv": self.get_design_variables(
            struct_case=sol, thrust_t=sol.thrust, grads_to_compute=grads_to_compute
        ),
        "thrust_t": sol.thrust,
        "solve_dofs": tuple(
            get_solve_dofs(n_dof=self.n_dof, prescribed_dofs=sol.prescribed_dofs)
        ),
        "approx_grads": approx_grads,
        "n_profile_loops": n_loop,
        "jac_options": jac_options,
    }

    if f_aero_nm1_n is not None:
        *_, compile_time, run_time = self.timestep_residual_jacobians(
            f_ext_aero_nm1=f_aero_nm1_n[0],
            f_ext_aero_n=f_aero_nm1_n[1],
            **common_kwargs,
        )
    else:
        *_, compile_time, run_time = self.timestep_residual_jacobians(
            f_ext_aero_nm1=None,
            f_ext_aero_n=None,
            **common_kwargs,
        )

    if print_header:
        print_table_line(inner_width=95)

    assert compile_time is not None and run_time is not None, (
        "No output timings passed"
    )

    return compile_time, run_time

set_design_variables

set_design_variables(
    coords: Array,
    k_cs: Array,
    m_cs: Array | None,
    m_lumped: Array | None = None,
    orientation_euler: Array | None = None,
    thrust_reference: dict[str, Array | float]
    | None = None,
    *,
    remove_checks: bool = False,
) -> None

Set design variables and compute initial configuration dependent quantities.

Parameters:

Name	Type	Description	Default
coords	Array	Node coordinates in the reference configuration, (n_nodes, 3).	required
k_cs	Array	Cross-section stiffness matrices, (n_entry, 6, 6) or (6, 6).	required
m_cs	Array | None	Cross-section mass matrices, (n_entry, 6, 6) or (6, 6).	required
m_lumped	Array | None	Lumped mass matrices at nodes, (n_entry, 6, 6).	None
orientation_euler	Array | None	Euler angles in radians which to rotate the reference configuration by, (3, ). This rotation is performed about the origin, and will default to the identity is no Array is passed. These are rotated in z-y-x order.	None
thrust_reference	dict[str, Array | float] | None	Reference thrust magnitude, {keys, (1, )}.	None
remove_checks	bool	Flag to ignore input checks, used when function is JIT compiled.	False

Source code in src/flapjax/structure/beam.py

def set_design_variables(
    self,
    coords: Array,
    k_cs: Array,
    m_cs: Array | None,
    m_lumped: Array | None = None,
    orientation_euler: Array | None = None,
    thrust_reference: dict[str, Array | float] | None = None,
    *,
    remove_checks: bool = False,
) -> None:
    r"""
    Set design variables and compute initial configuration dependent quantities.
    :param coords: Node coordinates in the reference configuration, ``(n_nodes, 3)``.
    :param k_cs: Cross-section stiffness matrices, ``(n_entry, 6, 6)`` or ``(6, 6)``.
    :param m_cs: Cross-section mass matrices, ``(n_entry, 6, 6)`` or ``(6, 6)``.
    :param m_lumped: Lumped mass matrices at nodes, ``(n_entry, 6, 6)``.
    :param orientation_euler: Euler angles in radians which to rotate the reference configuration by, ``(3, )``. This rotation
    is performed about the origin, and will default to the identity is no Array is passed. These are rotated in
    z-y-x order.
    :param thrust_reference: Reference thrust magnitude, ``{keys, (1, )}``.
    :param remove_checks: Flag to ignore input checks, used when function is JIT compiled.
    """

    # orientation
    if orientation_euler is not None:
        check_arr_shape(orientation_euler, (3,), "orientation_euler")
        self.orientation_euler = orientation_euler
        self.orientation = Rotation.from_euler(
            seq="zyx", angles=orientation_euler
        ).as_matrix()

    # rotate the y vectors
    self.y_vector = jnp.einsum(
        "jk,ik->ij",
        self.orientation,
        self.y_vector_reference_arr,
    )

    # coordinates — auto-extend for nodes created by multibody constraints
    if self._auto_node_sources and coords.shape[0] == self.n_nodes - len(
        self._auto_node_sources
    ):
        coords = jnp.concatenate(
            [coords, coords[jnp.array(self._auto_node_sources)]],
            axis=0,
        )
    check_arr_shape(coords, (self.n_nodes, 3), "coords")
    self.x0_reference = coords
    self.x0 = jnp.einsum("jk,ik->ij", self.orientation, coords)

    # populate arrays
    if k_cs.ndim == 2:
        k_cs = k_cs[None, ...]
    check_arr_shape(k_cs, (None, 6, 6), "k_cs")

    if (
        not remove_checks
        and k_cs.shape[0] != jnp.unique_values(jnp.array(self.k_cs_index)).size
    ):
        warn(
            "Redundant values in k_cs which are not used for solution due to no corresponding entry in k_cs_index."
        )

    self.k_cs = k_cs
    if m_cs is None:
        if not remove_checks and self.use_gravity and m_lumped is None:
            warn(
                "No mass matrices provided, but gravity is enabled. Assuming zero mass.",
            )
        m_cs_ = jnp.zeros((6, 6))
    else:
        m_cs_ = m_cs

    if m_cs_.ndim == 2:
        m_cs_ = m_cs_[None, ...]

    check_arr_shape(m_cs_, (None, 6, 6), "m_cs")

    if (
        not remove_checks
        and m_cs_.shape[0] != jnp.unique_values(jnp.array(self.m_cs_index)).size
        and m_cs is not None
    ):
        warn(
            "Redundant values in m_cs which are not used for solution due to no corresponding entry in "
            "m_cs_index."
        )

    self.m_cs = m_cs_

    # thrust
    if thrust_reference is not None:
        self.thrust_reference = {
            k: jnp.atleast_1d(v) for k, v in thrust_reference.items()
        }

        for k, v in self.thrust_reference.items():
            check_arr_shape(v, (1,), f"thrust_reference[{k}]")

    if m_lumped is not None:
        if not remove_checks:
            check_arr_shape(m_lumped, (None, 6, 6), "m_lumped")

            if self.m_lumped_index is None:
                raise ValueError("m_lumped_index has not been set")

            if m_lumped.shape[0] != len(self.m_lumped_index):
                raise ValueError(
                    "Number of entries in m_lumped does not match number of indices in m_lumped_index."
                )

        self.m_lumped = m_lumped

    # obtain initial orientation and length
    x_elem = jnp.take(
        self.x0_reference, self.connectivity_arr, axis=0
    )  # (n_elem, 2, 3)
    dx = x_elem[:, 1, :] - x_elem[:, 0, :]  # (n_elem, 3)

    # ensure out-of-plane vector and beam vector are not collinear
    if not remove_checks and jnp.any(
        jnp.linalg.norm(jnp.cross(dx, self.y_vector_reference_arr, 1, 1), axis=-1)
        < 1e-6
    ):
        raise ValueError(
            "y_vector is collinear with beam element direction for at least one element. "
            "Please provide a different y_vector."
        )

    self.l0 = jnp.linalg.norm(dx, axis=-1)  # (n_elem,)
    self.d0 = self.d0.at[:, 0].set(self.l0)

    dx_unit = dx / self.l0[:, None]  # unit vector in beam direction, (n_elem, 3)
    dz = jnp.cross(
        dx_unit, self.y_vector_reference_arr, axis=-1
    )  # vector in plane(n_elem, 3)
    dz_unit = dz / jnp.linalg.norm(dz, axis=-1)[:, None]  # (n_elem, 3)

    dy_unit = jnp.cross(dz_unit, dx_unit)

    self.o0 = self.o0.at[..., 0].set(dx_unit)
    self.o0 = self.o0.at[..., 1].set(dy_unit)
    self.o0 = self.o0.at[..., 2].set(dz_unit)

    self.ad_inv_o0 = vmap(rmat_to_ha_hat)(jnp.transpose(self.o0, (0, 2, 1)))

    # set unoriented initial coordinates
    self.hg0_reference = jnp.broadcast_to(
        jnp.eye(4)[None, ...], (self.n_nodes, 4, 4)
    )  # (n_nodes, 4, 4)
    self.hg0_reference = self.hg0_reference.at[:, :3, 3].set(self.x0_reference)

    # set oriented initial coordinates
    self.hg0 = self.hg0.at[:, :3, :3].set(
        jnp.broadcast_to(self.orientation[None, ...], (self.n_nodes, 3, 3))
    )  # (n_nodes, 4, 4)
    self.hg0 = self.hg0.at[:, :3, 3].set(self.x0)
    self.hg0 = self.hg0.at[:, 3, 3].set(1.0)

    # add reference frames to the nodal constraints
    for con in self.nodal_constraints:
        con.resolve_hg_ref(self.hg0)
    for con in self.multibody_constraints:
        con.resolve_hg_ref(self.hg0)

get_design_variables

get_design_variables(
    struct_case: StructureCase,
    thrust_t: dict[str, Array],
    grads_to_compute: StructureGradsToCompute | None,
) -> StructureDesignVariables

Obtain the design variables for the structural problem. As the external forcing is defined for each solve, the chosen forcing is required as input.

Parameters:

Name	Type	Description	Default
struct_case	StructureCase	Structural case	required
thrust_t	dict[str, Array]	Thrust time history, {keys, (n_tstep,)}.	required
grads_to_compute	StructureGradsToCompute | None	Data structure which describes which design variables should be obtained. If none, all variables are obtained.	required

Returns:

Type	Description
StructureDesignVariables	StructureDesignVariables dataclass containing design variables

Source code in src/flapjax/structure/beam.py

def get_design_variables(
    self,
    struct_case: StructureCase,
    thrust_t: dict[str, Array],
    grads_to_compute: StructureGradsToCompute | None,
) -> StructureDesignVariables:
    r"""
    Obtain the design variables for the structural problem. As the external forcing is defined for each solve, the
    chosen forcing is required as input.
    :param struct_case: Structural case
    :param thrust_t: Thrust time history, {keys, ``(n_tstep,)``}.
    :param grads_to_compute: Data structure which describes which design variables should be obtained. If none, all
    variables are obtained.
    :return: StructureDesignVariables dataclass containing design variables
    """

    # struct_case.f_ext_dead is stored in local frame: f_local = R^T @ f_global,
    # so recover f_global = R @ f_local
    hg = struct_case.hg
    if hg.ndim == 4:  # batched case: (n_tstep, n_nodes, 4, 4)
        rmat = hg[:, :, :3, :3]
    else:  # snapshot: (n_nodes, 4, 4)
        rmat = hg[:, :3, :3]
    f_ext_dead_global = (
        transform_nodal_vect(struct_case.f_ext_dead, rmat)
        if struct_case.f_ext_dead is not None
        else None
    )
    if isinstance(grads_to_compute, StructureGradsToCompute):
        return StructureDesignVariables(
            x0=self.x0 if grads_to_compute.x0 else None,
            orientation_euler=self.orientation_euler
            if grads_to_compute.orientation_euler
            else None,
            m_cs=self.m_cs if grads_to_compute.m_cs else None,
            k_cs=self.k_cs if grads_to_compute.k_cs else None,
            m_lumped=self._m_lumped if grads_to_compute.m_lumped else None,
            f_ext_dead=f_ext_dead_global if grads_to_compute.f_ext_dead else None,
            f_ext_follower=struct_case.f_ext_follower
            if grads_to_compute.f_ext_follower
            else None,
            thrust_t=thrust_t if grads_to_compute.thrust_t else None,
            f_shape=(),
        )
    else:
        return StructureDesignVariables(
            x0=self.x0,
            orientation_euler=self.orientation_euler,
            m_cs=self.m_cs,
            k_cs=self.k_cs,
            m_lumped=self._m_lumped,
            f_ext_dead=f_ext_dead_global,
            f_ext_follower=struct_case.f_ext_follower,
            thrust_t=thrust_t,
            f_shape=(),
        )

reference_configuration

reference_configuration(
    prescribed_dofs: Sequence[int]
    | Array
    | slice
    | int = (),
    use_f_ext_follower: bool = True,
    use_f_ext_dead: bool = True,
    use_f_aero: bool = True,
    use_f_grav: bool = True,
) -> StructureCase

Get the reference configuration of the structure.

Parameters:

Name	Type	Description	Default
prescribed_dofs	Sequence[int] | Array | slice | int	Prescribed degrees of freedom, which are not solved for. Defaults to no prescribed DoFs.	()
use_f_ext_follower	bool	Whether to include follower forces in the reference configuration.	True
use_f_ext_dead	bool	Whether to include dead forces in the reference configuration.	True
use_f_aero	bool	Whether to include aerodynamic forces in the reference configuration.	True
use_f_grav	bool	Whether to include gravitational forces in the reference configuration.	True

Returns:

Type	Description
StructureCase	Structure dataclass containing reference configuration.

Source code in src/flapjax/structure/beam.py

def reference_configuration(
    self,
    prescribed_dofs: Sequence[int] | Array | slice | int = (),
    use_f_ext_follower: bool = True,
    use_f_ext_dead: bool = True,
    use_f_aero: bool = True,
    use_f_grav: bool = True,
) -> StructureCase:
    r"""
    Get the reference configuration of the structure.
    :param prescribed_dofs: Prescribed degrees of freedom, which are not solved for. Defaults to no prescribed DoFs.
    :param use_f_ext_follower: Whether to include follower forces in the reference configuration.
    :param use_f_ext_dead: Whether to include dead forces in the reference configuration.
    :param use_f_aero: Whether to include aerodynamic forces in the reference configuration.
    :param use_f_grav: Whether to include gravitational forces in the reference configuration.
    :return: Structure dataclass containing reference configuration.
    """
    prescribed_dofs = self.make_prescribed_dofs_tuple(prescribed_dofs)
    return StructureCase(
        hg=self.hg0,
        conn=self.connectivity,
        o0=self.o0,
        d=self.d0,
        eps=jnp.zeros((self.n_elem, 6)),
        varphi=jnp.zeros((self.n_nodes, 6)),
        f_ext_follower=jnp.zeros((self.n_nodes, 6)) if use_f_ext_follower else None,
        f_ext_dead=jnp.zeros((self.n_nodes, 6)) if use_f_ext_dead else None,
        f_ext_aero=jnp.zeros((self.n_nodes, 6)) if use_f_aero else None,
        f_grav=jnp.zeros((self.n_nodes, 6)) if use_f_grav else None,
        f_int=jnp.zeros((self.n_nodes, 6)),
        f_elem=jnp.zeros((self.n_elem, 6)),
        f_res=jnp.zeros((self.n_nodes, 6)),
        thrust=self.thrust_reference,
        thrust_direction=self.thrust_direction,
        thrust_nodes=self.thrust_nodes,
        local=True,
        prescribed_dofs=prescribed_dofs,
        t=jnp.zeros(1),
    )

compute_varphi_from_hg

compute_varphi_from_hg(hg: Array) -> Array

Calculate the twist vector from the reference configuration to hg

Parameters:

Name	Type	Description	Default
hg	Array	Deformed coordinates, (n_nodes, 4, 4)	required

Returns:

Type	Description
Array	Vector of twists, (n_nodes, 6)

Source code in src/flapjax/structure/beam.py

def compute_varphi_from_hg(self, hg: Array) -> Array:
    r"""
    Calculate the twist vector from the reference configuration to hg
    :param hg: Deformed coordinates, ``(n_nodes, 4, 4)``
    :return: Vector of twists, ``(n_nodes, 6)``
    """
    return vmap(hg_to_d, (0, 0), 0)(self.hg0, hg)

assemble_matrix_from_entries

assemble_matrix_from_entries(entries: Array) -> Array

Assemble global matrix from element entries

Parameters:

Name	Type	Description	Default
entries	Array	Array of element matrix entries, (n_elem, 12, 12)	required

Returns:

Type	Description
Array	System global matrix, (n_dof, n_dof)

Source code in src/flapjax/structure/beam.py

def assemble_matrix_from_entries(self, entries: Array) -> Array:
    r"""
    Assemble global matrix from element entries
    :param entries: Array of element matrix entries, ``(n_elem, 12, 12)``
    :return: System global matrix, ``(n_dof, n_dof)``
    """

    row_idx = jnp.broadcast_to(
        self.dof_per_elem_arr[:, :, None], (self.n_elem, 12, 12)
    )
    col_idx = jnp.broadcast_to(
        self.dof_per_elem_arr[:, None, :], (self.n_elem, 12, 12)
    )
    return (
        jnp.zeros((self.n_dof, self.n_dof))
        .at[row_idx.ravel(), col_idx.ravel()]
        .add(entries.ravel())
    )

assemble_vector_from_entries

assemble_vector_from_entries(entries: Array) -> Array

Assemble global vector from element entries

Parameters:

Name	Type	Description	Default
entries	Array	Array of element vector entries, (n_elem, 12)	required

Returns:

Type	Description
Array	System global vector, (n_dof, )

Source code in src/flapjax/structure/beam.py

def assemble_vector_from_entries(self, entries: Array) -> Array:
    r"""
    Assemble global vector from element entries
    :param entries: Array of element vector entries, ``(n_elem, 12)``
    :return: System global vector, ``(n_dof, )``
    """

    vect = jnp.zeros(self.n_dof)
    vect = vect.at[self.dof_per_elem_arr[:, :6]].add(entries[:, :6])
    return vect.at[self.dof_per_elem_arr[:, 6:]].add(entries[:, 6:])

add_lumped_contributions_to_arr

add_lumped_contributions_to_arr(
    arr: Array, lumped_arr: Array
) -> Array

Add lumped contributions to an array

Parameters:

Name	Type	Description	Default
arr	Array	Full array, (6*n_node, 6*n_node)	required
lumped_arr	Array	Lumped contributions, (n_lump, 6, 6)	required

Returns:

Type	Description
Array	In-place updated array, (6*n_node, 6*n_node)

Source code in src/flapjax/structure/beam.py

def add_lumped_contributions_to_arr(self, arr: Array, lumped_arr: Array) -> Array:
    r"""
    Add lumped contributions to an array
    :param arr: Full array, ``(6*n_node, 6*n_node)``
    :param lumped_arr: Lumped contributions, ``(n_lump, 6, 6)``
    :return: In-place updated array, ``(6*n_node, 6*n_node)``
    """

    assert self.m_lumped_index is not None

    def add_block(carry, x):
        node_idx, block = x
        dofs = node_idx * 6 + jnp.arange(6)
        return carry.at[jnp.ix_(dofs, dofs)].add(block), None

    arr, _ = jax.lax.scan(
        add_block, arr, (jnp.array(self.m_lumped_index), lumped_arr)
    )
    return arr

add_lumped_contributions_to_vec

add_lumped_contributions_to_vec(
    vec: Array, lumped_vec: Array
) -> Array

Add lumped contributions to an array

Parameters:

Name	Type	Description	Default
vec	Array	Full vector, (6*n_node, )	required
lumped_vec	Array	Lumped contributions, (n_lump, 6)	required

Returns:

Type	Description
Array	In-place updated vector, (6*n_node, ).

Source code in src/flapjax/structure/beam.py

def add_lumped_contributions_to_vec(self, vec: Array, lumped_vec: Array) -> Array:
    r"""
    Add lumped contributions to an array
    :param vec: Full vector, ``(6*n_node, )``
    :param lumped_vec: Lumped contributions, ``(n_lump, 6)``
    :return: In-place updated vector, ``(6*n_node, )``.
    """

    assert self.m_lumped_index is not None

    idx = (
        jnp.array(self.m_lumped_index)[:, None] * 6 + jnp.arange(6)[None, :]
    ).ravel()  # (n_lump * 6,)

    return vec.at[idx].add(lumped_vec)

make_k_t

make_k_t(d: Array, p_d: Array, eps: Array) -> Array

Assemble tangent stiffness matrix as a function of the element relative configuration vectors

Parameters:

Name	Type	Description	Default
d	Array	Element relative configuration, (n_elem, 6).	required
p_d	Array	P(d) operator, (n_elem, 6, 12).	required
eps	Array	Element strains, (n_elem, 6).	required

Returns:

Type	Description
Array	Elementwise stiffness matrix entries, (n_elem, 12, 12).

Source code in src/flapjax/structure/beam.py

def make_k_t(
    self,
    d: Array,
    p_d: Array,
    eps: Array,
) -> Array:
    r"""
    Assemble tangent stiffness matrix as a function of the element relative configuration vectors
    :param d: Element relative configuration, ``(n_elem, 6)``.
    :param p_d: P(d) operator, ``(n_elem, 6, 12)``.
    :param eps: Element strains, ``(n_elem, 6)``.
    :return: Elementwise stiffness matrix entries, ``(n_elem, 12, 12)``.
    """
    # compute stiffness matrix entries
    return vmap(
        partial(
            _k_t_entry,
            include_geometric=self.optional_jacobians.d_f_int_d_p_d,
        ),
        (0, 0, 0, 0, 0, 0),
        0,
    )(
        d,
        p_d,
        self.l0,
        eps,
        self.k_cs[self.k_cs_index, ...],
        self.ad_inv_o0,
    )  # (n_elem, 12, 12)

make_k_t_full

make_k_t_full(
    d: Array,
    p_d: Array,
    eps: Array,
    f_ext_dead: Array | None,
    rmat: Array,
    m_t: Array | None,
) -> Array

Compute the full tangent stiffness matrix, with contributions from stiffness, dead forces and gravity.

Parameters:

Name	Type	Description	Default
d	Array	Element relative configuration, (n_elem, 6).	required
p_d	Array	P(d) operator, (n_elem, 6, 12).	required
eps	Array	Strain vectors, (n_elem, 6).	required
f_ext_dead	Array | None	External dead forces in global reference, (n_node, 6).	required
rmat	Array	Nodal rotation matrices, (n_node, 3, 3).	required
m_t	Array | None	Disassembled system mass matrix, (n_elem, 12, 12).	required

Returns:

Type	Description
Array	Tangent stiffness matrix with all contributions, (n_dof, n_dof).

Source code in src/flapjax/structure/beam.py

def make_k_t_full(
    self,
    d: Array,
    p_d: Array,
    eps: Array,
    f_ext_dead: Array | None,
    rmat: Array,
    m_t: Array | None,
) -> Array:
    r"""
    Compute the full tangent stiffness matrix, with contributions from stiffness, dead forces and gravity.
    :param d: Element relative configuration, ``(n_elem, 6)``.
    :param p_d: P(d) operator, ``(n_elem, 6, 12)``.
    :param eps: Strain vectors, ``(n_elem, 6)``.
    :param f_ext_dead: External dead forces in global reference, ``(n_node, 6)``.
    :param rmat: Nodal rotation matrices, ``(n_node, 3, 3)``.
    :param m_t: Disassembled system mass matrix, ``(n_elem, 12, 12)``.
    :return: Tangent stiffness matrix with all contributions, ``(n_dof, n_dof)``.
    """

    k_t = self.assemble_matrix_from_entries(self.make_k_t(d, p_d, eps))
    if f_ext_dead is not None and self.optional_jacobians.d_f_ext_dead_d_n:
        k_t += block_diag(*self._make_k_t_dead(rmat, f_ext_dead))

    if self.use_gravity and self.optional_jacobians.d_f_grav_d_n:
        if m_t is None:
            raise ValueError("m_t needs to be provided")
        k_t += self.assemble_matrix_from_entries(
            self._make_k_t_grav(d, p_d, rmat, m_t)
        )
        if self.use_lumped_mass:
            k_t_lumped = self._make_k_t_grav_lumped(rmat)
            k_t = self.add_lumped_contributions_to_arr(
                arr=k_t, lumped_arr=k_t_lumped
            )
    return k_t

make_m_t

make_m_t(
    d: Array,
    int_order: Literal[3, 4, 5] = BASE_LOBATTO_ORDER,
) -> Array

Assemble tangent mass matrix as a function of the element relative configuration vectors. This does not include the lumped mass contribution.

Parameters:

Name	Type	Description	Default
d	Array	Element relative configuration, (n_elem, 6)	required
int_order	Literal[3, 4, 5]	Integration order for mass matrix computation	BASE_LOBATTO_ORDER

Returns:

Type	Description
Array	Elementwise mass matrix, (n_elem, 12, 12)

Source code in src/flapjax/structure/beam.py

def make_m_t(
    self, d: Array, int_order: Literal[3, 4, 5] = BASE_LOBATTO_ORDER
) -> Array:
    r"""
    Assemble tangent mass matrix as a function of the element relative configuration vectors. This does not include
    the lumped mass contribution.
    :param d: Element relative configuration, ``(n_elem, 6)``
    :param int_order: Integration order for mass matrix computation
    :return: Elementwise mass matrix, ``(n_elem, 12, 12)``
    """
    return vmap(partial(_integrate_m_l, int_order=int_order), (0, 0, 0, 0), 0)(
        self.m_cs[self.m_cs_index, ...], d, self.ad_inv_o0, self.l0
    )

make_nodal_m_k

make_nodal_m_k(
    case: StructureCase,
    int_order: Literal[3, 4, 5] = BASE_LOBATTO_ORDER,
) -> tuple[Array, Array]

Create the global mass and stiffness matrices for a given static structure case. These can be used for modal analysis or other purposes. These matrices are the Jacobians of the local forcing residual with respect to global perturbations in acceleration and displacement, respectively.

Parameters:

Name	Type	Description	Default
case	StructureCase	Static structure case for which to compute the global mass and stiffness matrices.	required
int_order	Literal[3, 4, 5]	Integration order for mass matrix computation.	BASE_LOBATTO_ORDER

Returns:

Type	Description
tuple[Array, Array]	Global mass and stiffness matrices, (n_free_dof, n_free_dof).

Source code in src/flapjax/structure/beam.py

def make_nodal_m_k(
    self,
    case: StructureCase,
    int_order: Literal[3, 4, 5] = BASE_LOBATTO_ORDER,
) -> tuple[Array, Array]:
    r"""
    Create the global mass and stiffness matrices for a given static structure case. These can be used for modal
    analysis or other purposes. These matrices are the Jacobians of the local forcing residual with respect to
    global perturbations in acceleration and displacement, respectively.
    :param case: Static structure case for which to compute the global mass and stiffness matrices.
    :param int_order: Integration order for mass matrix computation.
    :return: Global mass and stiffness matrices, ``(n_free_dof, n_free_dof)``.
    """
    # extract variables from case
    d = case.d
    eps = self.make_eps(d=d)
    p_d = self.make_p_d(d=d)
    t_varphi = vmap(t_se3)(case.varphi)  # (n_node, 6, 6)
    rmat = case.hg[:, :3, :3]  # (n_node, 3, 3)

    # get dead external forcing as this has a stiffness contribution
    f_ext_dead_local: Array | None
    if case.f_ext_dead is not None and case.f_ext_aero is not None:
        f_ext_dead_local = case.f_ext_dead + case.f_ext_aero
    else:
        f_ext_dead_local = (
            case.f_ext_dead if case.f_ext_dead is not None else case.f_ext_aero
        )

    # convert to global frame, as it required for creating the stiffness matrix
    f_ext_dead: Array | None = (
        transform_nodal_vect(f_ext_dead_local, rmat)
        if f_ext_dead_local is not None
        else None
    )
    free_dofs = jnp.array(
        get_solve_dofs(n_dof=self.n_dof, prescribed_dofs=case.prescribed_dofs)
    )

    def transform_mat_to_global(mat: Array) -> Array:
        # function to rotate a forcing Jacobian matrix from the local frame to the global frame.
        mat_reshaped = mat.reshape(self.n_nodes, 6, self.n_dof)
        m_lin = jnp.einsum("nij,njk->nik", rmat, mat_reshaped[:, :3, :])
        m_rot = jnp.einsum("nij,njk->nik", rmat, mat_reshaped[:, 3:, :])
        return jnp.concatenate((m_lin, m_rot), axis=1).reshape(
            self.n_dof, self.n_dof
        )

    # mass
    m_t = self.assemble_matrix_from_entries(
        self.make_m_t(d=d, int_order=int_order)
    )  # (n_dof, n_dof)
    if self.use_lumped_mass:
        m_t = self.add_lumped_contributions_to_arr(
            arr=m_t, lumped_arr=self.m_lumped
        )

    m_modal_full = transform_mat_to_global(
        mat=jnp.einsum("ijk,jkl->ijl", m_t.reshape(self.n_dof, -1, 6), t_varphi)
    )

    m_modal = m_modal_full.reshape(self.n_dof, self.n_dof)[
        jnp.ix_(free_dofs, free_dofs)
    ]

    # stiffness
    k_t = self.make_k_t_full(
        d=case.d, p_d=p_d, eps=eps, f_ext_dead=f_ext_dead, rmat=rmat, m_t=m_t
    )

    k_modal_full = transform_mat_to_global(
        mat=jnp.einsum("ijk,jkl->ijl", k_t.reshape(self.n_dof, -1, 6), t_varphi)
    )

    k_modal = k_modal_full.reshape(self.n_dof, self.n_dof)[
        jnp.ix_(free_dofs, free_dofs)
    ]

    return m_modal, k_modal

modal

modal(
    case: StructureCase,
    remove_complex_conjugate: bool = True,
    int_order: Literal[3, 4, 5] = BASE_LOBATTO_ORDER,
    n_modes: int = 20,
    freq_range: tuple[float | Array, float | Array] = (
        0.0,
        jnp.inf,
    ),
    damp_range: tuple[float | Array, float | Array] = (
        -jnp.inf,
        jnp.inf,
    ),
    vtu_directory: str | PathLike = "./modal",
    n_plot_vtu: int | None = None,
    aero: UVLM | None = None,
    n_phase: int = 8,
    n_interp: int = 0,
    max_disp: float = 0.2,
    max_ang: float = 0.2,
) -> tuple[Array, Array, Array]

Perform modal analysis on the structure.

Parameters:

Name	Type	Description	Default
case	StructureCase	The static structure case for which to perform modal analysis.	required
remove_complex_conjugate	bool	If true, keep only one mode from each complex conjugate pair.	True
int_order	Literal[3, 4, 5]	Integration order for mass matrix computation.	BASE_LOBATTO_ORDER
n_modes	int	Number of modes to preserve.	20
freq_range	tuple[float | Array, float | Array]	Frequency range for filtering out modes.	(0.0, inf)
damp_range	tuple[float | Array, float | Array]	Damping range for filtering out modes.	(-inf, inf)
vtu_directory	str | PathLike	Directory to for saving the mode shapes to vtu files.	'./modal'
n_plot_vtu	int | None	Number of modes to plot to vtu files. Will default to "./modal".	None
aero	UVLM | None	UVLM aerodynamic model. If passed, the vtu files will include the aerodynamic grid. If not, they will just be the beam structure.	None
n_phase	int	Number of phases to use when plotting the modes to vtu files.	8
n_interp	int	Number of times to interpolate between beam nodes for vtu plotting.	0
max_disp	float	Maximum displacement of structure for plotted modes, used for scaling.	0.2
max_ang	float	Maximum angle of structure for plotted modes in radians, used for scaling.	0.2

Returns:

Type	Description
tuple[Array, Array, Array]	Tuple of natural frequencies (n_free_dof), damping ratios (n_free_dof), and mode shapes with no normalisation (n_modes, n_free_dof).

Source code in src/flapjax/structure/beam.py

def modal(
    self,
    case: StructureCase,
    remove_complex_conjugate: bool = True,
    int_order: Literal[3, 4, 5] = BASE_LOBATTO_ORDER,
    n_modes: int = 20,
    freq_range: tuple[float | Array, float | Array] = (0.0, jnp.inf),
    damp_range: tuple[float | Array, float | Array] = (-jnp.inf, jnp.inf),
    vtu_directory: str | os.PathLike = "./modal",
    n_plot_vtu: int | None = None,
    aero: UVLM | None = None,
    n_phase: int = 8,
    n_interp: int = 0,
    max_disp: float = 0.2,
    max_ang: float = 0.2,
) -> tuple[Array, Array, Array]:
    r"""
    Perform modal analysis on the structure.
    :param case: The static structure case for which to perform modal analysis.
    :param remove_complex_conjugate: If true, keep only one mode from each complex conjugate pair.
    :param int_order: Integration order for mass matrix computation.
    :param n_modes: Number of modes to preserve.
    :param freq_range: Frequency range for filtering out modes.
    :param damp_range: Damping range for filtering out modes.
    :param vtu_directory: Directory to for saving the mode shapes to vtu files.
    :param n_plot_vtu: Number of modes to plot to vtu files. Will default to "./modal".
    :param aero: UVLM aerodynamic model. If passed, the vtu files will include the aerodynamic grid. If not, they
    will just be the beam structure.
    :param n_phase: Number of phases to use when plotting the modes to vtu files.
    :param n_interp: Number of times to interpolate between beam nodes for vtu plotting.
    :param max_disp: Maximum displacement of structure for plotted modes, used for scaling.
    :param max_ang: Maximum angle of structure for plotted modes in radians, used for scaling.
    :return: Tuple of natural frequencies (n_free_dof), damping ratios (n_free_dof), and mode shapes with no
     normalisation ``(n_modes, n_free_dof)``.
    """
    freqs, damping, modes, *_ = self.base_modal(
        case=case,
        freq_range=freq_range,
        damp_range=damp_range,
        int_order=int_order,
        n_modes=n_modes,
        remove_complex_conjugate=remove_complex_conjugate,
    )

    if n_plot_vtu is not None:
        q_full = (
            jnp.zeros((n_plot_vtu, self.n_nodes * 6))
            .at[:, case.free_dofs]
            .set(modes[:n_plot_vtu, :])
        )

        for _i_mode in range(n_plot_vtu):
            plot_modes_vtu(
                reference=case,
                directory=vtu_directory,
                q_full=q_full.reshape(n_plot_vtu, self.n_nodes, 6),
                freqs=freqs,
                dampings=damping,
                gamma_b_full=None,
                gamma_w_full=None,
                zeta_w_full=None,
                uvlm=aero,
                n_interp=n_interp,
                n_phase=n_phase,
                max_disp=max_disp,
                max_ang=max_ang,
                max_gamma=1e6,
            )

    return freqs, damping, modes

linearise

linearise(
    reference: StructureCase,
    dt: float,
    n_modes: int | None = None,
    modal_inputs: bool = False,
    modal_outputs: bool = False,
    prescribed_dofs: Sequence[int]
    | Array
    | slice
    | int
    | None = None,
) -> LinearBeam

Linearise the beam about a given static structure case. This creates a LinearBeam object which can be used for linear dynamic analysis.

Parameters:

Name	Type	Description	Default
reference	StructureCase	Static structure case about which to linearise the beam.	required
dt	float	Time step size, used for conversions between continuous and discrete time.	required
n_modes	int | None	If not None, the linearised system uses modal state coordinates truncated to this many modes.	None
modal_inputs	bool	If True, external forcing inputs are provided as modal forces (requires n_modes).	False
modal_outputs	bool	If True, outputs are exposed as modal coordinates (requires n_modes).	False
prescribed_dofs	Sequence[int] | Array | slice | int | None	If provided, overrides the prescribed DOFs from the reference case.	None

Returns:

Type	Description
LinearBeam	Continuous-time linearised beam object.

Source code in src/flapjax/structure/beam.py

def linearise(
    self,
    reference: StructureCase,
    dt: float,
    n_modes: int | None = None,
    modal_inputs: bool = False,
    modal_outputs: bool = False,
    prescribed_dofs: Sequence[int] | Array | slice | int | None = None,
) -> LinearBeam:
    r"""
    Linearise the beam about a given static structure case. This creates a LinearBeam object which can be used for
    linear dynamic analysis.
    :param reference: Static structure case about which to linearise the beam.
    :param dt: Time step size, used for conversions between continuous and discrete time.
    :param n_modes: If not None, the linearised system uses modal state coordinates truncated to this many modes.
    :param modal_inputs: If True, external forcing inputs are provided as modal forces (requires n_modes).
    :param modal_outputs: If True, outputs are exposed as modal coordinates (requires n_modes).
    :param prescribed_dofs: If provided, overrides the prescribed DOFs from the reference case.
    :return: Continuous-time linearised beam object.
    """
    return LinearBeam(
        beam=self,
        reference=reference,
        dt=dt,
        n_modes=n_modes,
        modal_inputs=modal_inputs,
        modal_outputs=modal_outputs,
        prescribed_dofs=prescribed_dofs,
    )

apply_nodal_constraint_tangent

apply_nodal_constraint_tangent(
    mat: Array,
    hg: Array,
    i_ts: int,
    gamma_prime: float | Array | None,
) -> Array

Add nodal constraint contributions to a system matrix.

Parameters:

Name	Type	Description	Default
mat	Array	System matrix to update, (n_dof, n_dof).	required
hg	Array	SE(3) coordiantes, (n_nodes, 4, 4).	required
i_ts	int	Time-step index (0 for static solves).	required
gamma_prime	float | Array | None	Time-integrator gamma_prime for damping scaling, or None to skip damping.	required

Returns:

Type	Description
Array	Updated system matrix.

Source code in src/flapjax/structure/beam.py

def apply_nodal_constraint_tangent(
    self,
    mat: Array,
    hg: Array,
    i_ts: int,
    gamma_prime: float | Array | None,
) -> Array:
    r"""
    Add nodal constraint contributions to a system matrix.
    :param mat: System matrix to update, ``(n_dof, n_dof)``.
    :param hg: SE(3) coordiantes, ``(n_nodes, 4, 4)``.
    :param i_ts: Time-step index (0 for static solves).
    :param gamma_prime: Time-integrator gamma_prime for damping scaling, or ``None`` to skip damping.
    :return: Updated system matrix.
    """
    for con in self.nodal_constraints:
        node = con.node_index
        hg_i = hg[node]
        dofs = node * 6 + jnp.arange(6)
        mat = mat.at[jnp.ix_(dofs, dofs)].add(con.k_tangent(hg_i, i_ts))
        if gamma_prime is not None:
            mat = mat.at[jnp.ix_(dofs, dofs)].add(
                gamma_prime * con.c_tangent(hg_i, i_ts)
            )

    # hard constraint tangent contributions (e.g. hinge spring-damper)
    for con in self.multibody_constraints:
        if con.has_f_res:
            dofs_i = con.node_i * 6 + jnp.arange(6)
            dofs_j = con.node_j * 6 + jnp.arange(6)
            dofs_ij = jnp.concatenate([dofs_i, dofs_j])
            k_12 = con.k_tangent(hg[con.node_i], hg[con.node_j])
            mat = mat.at[jnp.ix_(dofs_ij, dofs_ij)].add(k_12)
            if gamma_prime is not None:
                # add damping terms
                mat = mat.at[jnp.ix_(dofs_ij, dofs_ij)].add(
                    gamma_prime * con.c_tangent(hg[con.node_i], hg[con.node_j])
                )

    return mat

postprocess_constraints

postprocess_constraints(
    hg: Array,
) -> dict[str, dict[str, Array]]

Postprocess all constraints to extract derived quantities (e.g. hinge angles).

Parameters:

Name	Type	Description	Default
hg	Array	Nodal SE(3) frames, (n_nodes, 4, 4) or (n_tstep, n_nodes, 4, 4).	required

Returns:

Type	Description
dict[str, dict[str, Array]]	Nested dict {constraint_name: {quantity_name: Array}}.

Source code in src/flapjax/structure/beam.py

def postprocess_constraints(self, hg: Array) -> dict[str, dict[str, Array]]:
    r"""
    Postprocess all constraints to extract derived quantities (e.g. hinge angles).
    :param hg: Nodal SE(3) frames, ``(n_nodes, 4, 4)`` or ``(n_tstep, n_nodes, 4, 4)``.
    :return: Nested dict ``{constraint_name: {quantity_name: Array}}``.
    """
    data: dict[str, dict[str, Array]] = {}
    for name, con in self.constraints.items():
        pp = con.postprocess(hg)
        if pp:
            data[name] = pp
    return data

solve_constrained

solve_constrained(
    sys_mat_solve: Array,
    f_res_solve: Array,
    hg_eval: Array,
    solve_dofs: Array,
    hg_base: Array | None = None,
    phi: Array | None = None,
    v: Array | None = None,
    gamma_prime: float | Array | None = None,
) -> tuple[Array, Array, Array]

Solve the augmented system with Lagrange multipliers, supporting both holonomic and non-holonomic constraints.

Parameters:

Name	Type	Description	Default
sys_mat_solve	Array	System matrix at solve DOFs, (n_solve, n_solve).	required
f_res_solve	Array	Force residual at solve DOFs, (n_solve,).	required
hg_eval	Array	SE(3) frames for constraint evaluation, (n_nodes, 4, 4).	required
solve_dofs	Array	Free DOF indices, (n_solve,).	required
hg_base	Array | None	Base frames for Jacobian computation (defaults to hg_eval).	None
phi	Array | None	Accumulated configuration increment, (n_nodes, 6).	None
v	Array | None	Current nodal velocities for non-holonomic constraints, (n_nodes, 6).	None
gamma_prime	float | Array | None	Newmark parameter for non-holonomic constraints.	None

Returns:

Type	Description
tuple[Array, Array, Array]	(delta_phi, lagrange_multipliers, constraint_violation).

Source code in src/flapjax/structure/beam.py

def solve_constrained(
    self,
    sys_mat_solve: Array,
    f_res_solve: Array,
    hg_eval: Array,
    solve_dofs: Array,
    hg_base: Array | None = None,
    phi: Array | None = None,
    v: Array | None = None,
    gamma_prime: float | Array | None = None,
) -> tuple[Array, Array, Array]:
    r"""
    Solve the augmented system with Lagrange multipliers, supporting both holonomic and non-holonomic constraints.
    :param sys_mat_solve: System matrix at solve DOFs, ``(n_solve, n_solve)``.
    :param f_res_solve: Force residual at solve DOFs, ``(n_solve,)``.
    :param hg_eval: SE(3) frames for constraint evaluation, ``(n_nodes, 4, 4)``.
    :param solve_dofs: Free DOF indices, ``(n_solve,)``.
    :param hg_base: Base frames for Jacobian computation (defaults to ``hg_eval``).
    :param phi: Accumulated configuration increment, ``(n_nodes, 6)``.
    :param v: Current nodal velocities for non-holonomic constraints, ``(n_nodes, 6)``.
    :param gamma_prime: Newmark parameter for non-holonomic constraints.
    :return: ``(delta_phi, lagrange_multipliers, constraint_violation)``.
    """
    hg_jac = hg_base if hg_base is not None else hg_eval

    n_s = sys_mat_solve.shape[0]
    n_h = self.n_holonomic_constraints
    n_nh = self.n_nonholonomic_constraints
    n_c = n_h + n_nh

    aug = jnp.zeros((n_s + n_c, n_s + n_c))
    aug = aug.at[:n_s, :n_s].set(sys_mat_solve)

    rhs_parts: list[Array] = [f_res_solve]
    violation_parts: list[Array] = []

    if n_h > 0:
        viol_h = self._compute_holonomic_violation(hg_eval)
        jac_h = self._compute_holonomic_jacobian(hg_jac, solve_dofs, phi)
        aug = aug.at[:n_s, n_s : n_s + n_h].set(-jac_h.T)
        aug = aug.at[n_s : n_s + n_h, :n_s].set(jac_h)
        rhs_parts.append(-viol_h)
        violation_parts.append(viol_h)

    if n_nh > 0:
        assert v is not None
        vel_viol_nh = self._compute_nonholonomic_vel_violation(hg_eval, v)
        a_vel_solve = self._compute_nonholonomic_a_vel(hg_eval, v, solve_dofs)
        a_phi_solve = self._compute_nonholonomic_a_phi(hg_jac, v, solve_dofs, phi)
        aug = aug.at[:n_s, n_s + n_h : n_s + n_c].set(-a_vel_solve.T)
        aug = aug.at[n_s + n_h : n_s + n_c, :n_s].set(
            a_phi_solve + gamma_prime * a_vel_solve
        )
        rhs_parts.append(-vel_viol_nh)
        violation_parts.append(vel_viol_nh)

    rhs = jnp.concatenate(rhs_parts)
    sol = jnp.linalg.solve(aug, rhs)

    return sol[:n_s], sol[n_s:], jnp.concatenate(violation_parts)

compute_centre_of_mass

compute_centre_of_mass(hg: Array) -> Array

Compute the centre of mass for an arbitrary system.

Parameters:

Name	Type	Description	Default
hg	Array	Node SE(3) coordinates, (n_node, 4, 4) or (n_tstep, n_node, 4, 4).	required

Returns:

Type	Description
Array	Centre of mass, (3) or (n_tstep, 3).

Source code in src/flapjax/structure/beam.py

def compute_centre_of_mass(self, hg: Array) -> Array:
    r"""
    Compute the centre of mass for an arbitrary system.
    :param hg: Node SE(3) coordinates, ``(n_node, 4, 4)`` or ``(n_tstep, n_node, 4, 4)``.
    :return: Centre of mass, (3) or ``(n_tstep, 3)``.
    """

    def inner_func(hg_: Array) -> Array:
        d = self.make_d(hg=hg_)
        m = self.assemble_matrix_from_entries(self.make_m_t(d=d))  # (n_dof, n_dof)
        if self.use_lumped_mass:
            m = self.add_lumped_contributions_to_arr(
                arr=m, lumped_arr=self.m_lumped
            )
        m_lin = m[::6, ::6]
        return jnp.einsum("ij,jk->k", m_lin, hg_[:, :3, 3]) / m_lin.sum()  # (3, )

    if hg.ndim == 3:
        return inner_func(hg)  # single timestep, (3, ).
    elif hg.ndim == 4:
        return vmap(inner_func, 0, 0)(hg)  # multiple timesteps, (n_tstep, 3)
    else:
        raise ValueError("hg.ndim must be 3 or 4")

make_f_elem

make_f_elem(eps: Array) -> Array

Compute the forces within the elements as :math:\mathbf{f}_{elem} = \mathcal{K}_{cs} \epsilon.

Parameters:

Name	Type	Description	Default
eps	Array	Element strain vectors, (n_elem, 6).	required

Returns:

Type	Description
Array	Element forces, (n_elem, 6).

Source code in src/flapjax/structure/beam.py

def make_f_elem(self, eps: Array) -> Array:
    r"""
    Compute the forces within the elements as :math:`\mathbf{f}_{elem} = \mathcal{K}_{cs} \epsilon`.
    :param eps: Element strain vectors, ``(n_elem, 6)``.
    :return: Element forces, ``(n_elem, 6)``.
    """
    return jnp.einsum("ijk,ik->ij", self.k_cs[self.k_cs_index, ...], eps)

make_f_int

make_f_int(p_d: Array, eps: Array) -> Array

Assemble global internal force vector as a function of the element relative configuration vectors.

Parameters:

Name	Type	Description	Default
p_d	Array	P(d) operator, (n_elem, 6, 12).	required
eps	Array	Element strain vectors, (n_elem, 6).	required

Returns:

Type	Description
Array	Internal forces, (n_elem, 12).

Source code in src/flapjax/structure/beam.py

def make_f_int(self, p_d: Array, eps: Array) -> Array:
    r"""
    Assemble global internal force vector as a function of the element relative configuration vectors.
    :param p_d: P(d) operator, ``(n_elem, 6, 12)``.
    :param eps: Element strain vectors, ``(n_elem, 6)``.
    :return: Internal forces, ``(n_elem, 12)``.
    """

    return -jnp.einsum("ikj,ikl,il->ij", p_d, self.k_cs[self.k_cs_index, ...], eps)

make_f_dead_ext staticmethod

make_f_dead_ext(f_ext: Array, rmat: Array) -> Array

Compute the global external dead force vector.

Parameters:

Name	Type	Description	Default
f_ext	Array	External forces array of dead forces in global reference, (n_node, 6)	required
rmat	Array	Deformation rotation matrices, (n_node, 3, 3)	required

Returns:

Type	Description
Array	External forces, (n_node, 6)

Source code in src/flapjax/structure/beam.py

@staticmethod
def make_f_dead_ext(f_ext: Array, rmat: Array) -> Array:
    r"""
    Compute the global external dead force vector.
    :param f_ext: External forces array of dead forces in global reference, ``(n_node, 6)``
    :param rmat: Deformation rotation matrices, ``(n_node, 3, 3)``
    :return: External forces, ``(n_node, 6)``
    """

    return transform_nodal_vect(f_ext, jnp.swapaxes(rmat, -1, -2))

add_thrust_force

add_thrust_force(
    force: Array, thrust: dict[str, Array]
) -> Array

Add thrust acting at nodes onto full system forcing.

Parameters:

Name	Type	Description	Default
force	Array	Input forcing, (n_node, 6).	required
thrust	dict[str, Array]	Input thrust at the current step, {key: ()}.	required

Returns:

Type	Description
Array	Updated forcing, (n_node, 6).

Source code in src/flapjax/structure/beam.py

def add_thrust_force(self, force: Array, thrust: dict[str, Array]) -> Array:
    r"""
    Add thrust acting at nodes onto full system forcing.
    :param force: Input forcing, ``(n_node, 6)``.
    :param thrust: Input thrust at the current step, ``{key: ()}``.
    :return: Updated forcing, ``(n_node, 6)``.
    """

    for k, v in thrust.items():
        node = dict(self.thrust_nodes)[k]
        direction = jnp.array(dict(self.thrust_direction)[k])
        force = force.at[node, :3].add(v * direction)
    return force

make_eps

make_eps(d: Array) -> Array

Compute the element strain vectors as a function of the element relative configuration vectors. Formulation from Geometrically exact beam finite element formulated on the special Euclidean group SE(3), by Sonneville et al., 2013, Eq 64.

Parameters:

Name	Type	Description	Default
d	Array	Element relative configuration, (n_elem, 6)	required

Returns:

Type	Description
Array	Element strain vectors, (n_elem, 6)

Source code in src/flapjax/structure/beam.py

def make_eps(self, d: Array) -> Array:
    r"""
    Compute the element strain vectors as a function of the element relative configuration vectors. Formulation from
    Geometrically exact beam finite element formulated on the special Euclidean group SE(3), by Sonneville et al.,
    2013, Eq 64.
    :param d: Element relative configuration, ``(n_elem, 6)``
    :return: Element strain vectors, ``(n_elem, 6)``
    """

    return (d - self.d0) / self.l0[:, None]

make_p_d

make_p_d(d: Array) -> Array

Compute the P(d) operator as a function of the element relative configuration vectors.

Parameters:

Name	Type	Description	Default
d	Array	Relative configuration vectors, (n_elem, 6)	required

Returns:

Type	Description
Array	P(d) operator, (n_elem, 6, 12)

Source code in src/flapjax/structure/beam.py

def make_p_d(self, d: Array) -> Array:
    r"""
    Compute the P(d) operator as a function of the element relative configuration vectors.
    :param d: Relative configuration vectors, ``(n_elem, 6)``
    :return: P(d) operator, ``(n_elem, 6, 12)``
    """
    return vmap(p, (0, 0), 0)(d, self.ad_inv_o0)  # [n_elem, 6, 12]

make_d

make_d(hg: Array) -> Array

Compute the element relative configuration vectors from the nodal homogeneous transformation matrices

Parameters:

Name	Type	Description	Default
hg	Array	Nodal homogeneous transformation matrices, (n_nodes, 4, 4)	required

Returns:

Type	Description
Array	Element relative configuration vectors, (n_elem, 6)

Source code in src/flapjax/structure/beam.py

def make_d(self, hg: Array) -> Array:
    r"""
    Compute the element relative configuration vectors from the nodal homogeneous transformation matrices
    :param hg: Nodal homogeneous transformation matrices, ``(n_nodes, 4, 4)``
    :return: Element relative configuration vectors, ``(n_elem, 6)``
    """

    base_hg = jnp.zeros((self.n_elem, 4, 4))
    base_hg = base_hg.at[:, :3, :3].set(self.o0)
    base_hg = base_hg.at[:, 3, 3].set(1.0)

    haha0 = jnp.einsum(
        "ijk,ikl->ijl", hg[self.connectivity_arr[:, 0], :, :], base_hg
    )  # (n_elem, 4, 4)
    haha1 = jnp.einsum(
        "ijk,ikl->ijl", hg[self.connectivity_arr[:, 1], :, :], base_hg
    )  # (n_elem, 4, 4)

    return vmap(hg_to_d, (0, 0), 0)(haha0, haha1)  # (n_elem, 6)

make_hg_dot staticmethod

make_hg_dot(hg: Array, v: Array) -> Array

Obtain the time derivative of the nodal coordinates.

Parameters:

Name	Type	Description	Default
hg	Array	Node coordinates, (n_node, 4, 4).	required
v	Array	Node local velocities, (n_node, 6)	required

Returns:

Type	Description
Array	Coordinate time derivative, (n_node, 4, 4)

Source code in src/flapjax/structure/beam.py

@staticmethod
def make_hg_dot(hg: Array, v: Array) -> Array:
    r"""
    Obtain the time derivative of the nodal coordinates.
    :param hg: Node coordinates, ``(n_node, 4, 4)``.
    :param v: Node local velocities, ``(n_node, 6)``
    :return: Coordinate time derivative, ``(n_node, 4, 4)``
    """
    return jnp.einsum(
        "ijk,ikl->ijl", hg, vmap(ha_to_ha_tilde, 0, 0)(v)
    )  # (n_nodes, 4, 4)

resolve_forces

resolve_forces(
    hg: Array,
    dynamic: Literal[True],
    f_ext_follower: Array | None,
    f_ext_dead: Array | None,
    f_ext_aero: Array | None,
    thrust: dict[str, Array],
    v: Array,
    v_dot: Array,
    approx_gradients: bool = False,
) -> tuple[
    Array,
    Array,
    Array | None,
    Array | None,
    Array | None,
    Array,
    Array,
    Array,
    Array,
]

resolve_forces(
    hg: Array,
    dynamic: Literal[False],
    f_ext_follower: Array | None,
    f_ext_dead: Array | None,
    f_ext_aero: Array | None,
    thrust: dict[str, Array],
    v: None,
    v_dot: None,
    approx_gradients: bool = False,
) -> tuple[
    Array,
    Array,
    Array | None,
    Array | None,
    Array | None,
    Array,
    None,
    None,
    Array,
]

resolve_forces(
    hg: Array,
    dynamic: bool,
    f_ext_follower: Array | None,
    f_ext_dead: Array | None,
    f_ext_aero: Array | None,
    thrust: dict[str, Array],
    v: Array | None,
    v_dot: Array | None,
    approx_gradients: bool = False,
) -> tuple[
    Array,
    Array,
    Array | None,
    Array | None,
    Array | None,
    Array,
    Array | None,
    Array | None,
    Array,
]

Obtain all components of the force from a final solution.

Parameters:

Name	Type	Description	Default
hg	Array	Nodal homogeneous transformation matrices, (n_nodes, 4, 4).	required
dynamic	bool	Whether to compute dynamic forces.	required
f_ext_follower	Array | None	External follower forces in local reference, (n_node, 6).	required
f_ext_dead	Array | None	External dead forces in global reference, (n_node, 6).	required
f_ext_aero	Array | None	External aero forces in global reference, (n_node, 6).	required
thrust	dict[str, Array]	Thrust forces at current step, {keys, ()}.	required
v	Array | None	Nodal velocities in global frame, (n_node, 6).	required
v_dot	Array | None	Nodal accelerations in global frame, (n_node, 6).	required
approx_gradients	bool	Whether to stop computing gradients of the inertial and gyroscopic forces with respect to the node coordinates, as these are small but nonzero values in practice.	False

Returns:

Type	Description
tuple[Array, Array, Array | None, Array | None, Array | None, Array, Array | None, Array | None, Array]	Configuration vectors, strain vectors, Dead external forces, aero external forces, gravitational forces, internal forces, gyroscopic forces, inertial forces and residual forces.

Source code in src/flapjax/structure/beam.py

def resolve_forces(
    self,
    hg: Array,
    dynamic: bool,
    f_ext_follower: Array | None,
    f_ext_dead: Array | None,
    f_ext_aero: Array | None,
    thrust: dict[str, Array],
    v: Array | None,
    v_dot: Array | None,
    approx_gradients: bool = False,
) -> tuple[
    Array,
    Array,
    Array | None,
    Array | None,
    Array | None,
    Array,
    Array | None,
    Array | None,
    Array,
]:
    r"""
    Obtain all components of the force from a final solution.
    :param hg: Nodal homogeneous transformation matrices, ``(n_nodes, 4, 4)``.
    :param dynamic: Whether to compute dynamic forces.
    :param f_ext_follower: External follower forces in local reference, ``(n_node, 6)``.
    :param f_ext_dead: External dead forces in global reference, ``(n_node, 6)``.
    :param f_ext_aero: External aero forces in global reference, ``(n_node, 6)``.
    :param thrust: Thrust forces at current step, {keys, ``()``}.
    :param v: Nodal velocities in global frame, ``(n_node, 6)``.
    :param v_dot: Nodal accelerations in global frame, ``(n_node, 6)``.
    :param approx_gradients: Whether to stop computing gradients of the inertial and gyroscopic forces with respect to
    the node coordinates, as these are small but nonzero values in practice.
    :return: Configuration vectors, strain vectors, Dead external forces, aero external forces, gravitational forces, internal forces,
    gyroscopic forces, inertial forces and residual forces.
    """

    def prop_grad(x: Array) -> Array:
        return jax.lax.stop_gradient(x) if approx_gradients else x

    d = self.make_d(hg)
    eps = self.make_eps(d)
    p_d = self.make_p_d(d)

    if dynamic or self.use_gravity:
        m_t = self.make_m_t(prop_grad(d))
    else:
        m_t = None

    if dynamic:
        assert v is not None

        d_dot = self._make_d_dot(p_d, v)
        c_l = self._make_c_t(prop_grad(d), prop_grad(d_dot), v)[0]
        c_l_lumped = self._make_c_t_lumped(v)[0] if self.use_lumped_mass else None
    else:
        d_dot, c_l, c_l_lumped = None, None, None

    this_f_res = self.add_thrust_force(
        force=jnp.zeros((self.n_nodes, 6)), thrust=thrust
    )

    if f_ext_dead is not None:
        this_f_ext_dead = self.make_f_dead_ext(f_ext_dead, hg[:, :3, :3])
        this_f_res += this_f_ext_dead
    else:
        this_f_ext_dead = None

    if f_ext_aero is not None:
        this_f_ext_aero = self.make_f_dead_ext(f_ext_aero, hg[:, :3, :3])
        this_f_res += this_f_ext_aero
    else:
        this_f_ext_aero = None

    if self.use_gravity:
        assert m_t is not None
        this_f_grav = self.assemble_vector_from_entries(
            self._make_f_grav(m_t, hg[:, :3, :3])
        ).reshape(-1, 6)
        if self.use_lumped_mass:
            f_grav_lumped = self._make_f_grav_lumped(hg[:, :3, :3])
            this_f_grav = self.add_lumped_contributions_to_vec(
                vec=this_f_grav.ravel(), lumped_vec=f_grav_lumped.ravel()
            ).reshape(-1, 6)
        this_f_res += this_f_grav
    else:
        this_f_grav = None

    this_f_int = self.assemble_vector_from_entries(
        self.make_f_int(p_d, eps)
    ).reshape(-1, 6)
    this_f_res += this_f_int

    if dynamic:
        assert (
            m_t is not None
            and c_l is not None
            and v is not None
            and v_dot is not None
        )
        this_f_iner, this_f_gyr = self._make_f_iner_gyr(m_t, c_l, v, v_dot)
        this_f_iner = self.assemble_vector_from_entries(this_f_iner).reshape(-1, 6)
        this_f_gyr = self.assemble_vector_from_entries(this_f_gyr).reshape(-1, 6)

        if self.use_lumped_mass:
            assert c_l_lumped is not None
            f_iner_lumped, f_gyr_lumped = self._make_f_iner_gyr_lumped(
                c_l_lumped, v, v_dot
            )
            this_f_iner = self.add_lumped_contributions_to_vec(
                this_f_iner.ravel(), (f_iner_lumped + f_gyr_lumped).ravel()
            ).reshape(-1, 6)
        this_f_res += this_f_iner
    else:
        this_f_iner = None
        this_f_gyr = None

    if f_ext_follower is not None:
        this_f_res += f_ext_follower

    return (
        d,
        eps,
        this_f_ext_dead,
        this_f_ext_aero,
        this_f_grav,
        this_f_int,
        this_f_gyr,
        this_f_iner,
        this_f_res,
    )

make_f_res

make_f_res(
    solve_dofs: Array | None,
    p_d: Array,
    eps: Array,
    hg: Array,
    f_ext_follower_n: Array | None,
    f_ext_dead_n: Array | None,
    thrust_n: dict[str, Array],
    dynamic: Literal[True],
    m_t: Array,
    c_l: Array,
    c_l_lumped: Array | None,
    v: Array,
    v_dot: Array,
    i_ts: int = 0,
    k_t_assembled: Array | None = None,
) -> tuple[Array, Array]

make_f_res(
    solve_dofs: Array | None,
    p_d: Array,
    eps: Array,
    hg: Array,
    f_ext_follower_n: Array | None,
    f_ext_dead_n: Array | None,
    thrust_n: dict[str, Array],
    dynamic: Literal[False],
    m_t: Array | None,
    c_l: None,
    c_l_lumped: None,
    v: None,
    v_dot: None,
    i_ts: int = 0,
    k_t_assembled: Array | None = None,
) -> tuple[Array, Array]

make_f_res(
    solve_dofs: Array | None,
    p_d: Array,
    eps: Array,
    hg: Array,
    f_ext_follower_n: Array | None,
    f_ext_dead_n: Array | None,
    thrust_n: dict[str, Array],
    dynamic: bool,
    m_t,
    c_l,
    c_l_lumped,
    v,
    v_dot,
    i_ts: int = 0,
    k_t_assembled: Array | None = None,
) -> tuple[Array, Array]

Compute the residual force vector for a given configuration and external forces, used in the nonlinear solve. This is the force imbalance that the nonlinear solver will seek to drive to zero. Additionally, returns an "absolute sum" of all forces, used for relative convergence checks.

Parameters:

Name	Type	Description	Default
solve_dofs	Array | None	Optional array of degrees of freedom to solve for (n_solve_dofs, ).	required
p_d	Array	P(d) operator, (n_elem, 6, 12).	required
eps	Array	Element strain vectors, (n_elem, 6).	required
hg	Array	Nodal homogeneous transformation matrices, (n_nodes, 4, 4).	required
f_ext_follower_n	Array | None	Nodal follower forces, (n_nodes, 6).	required
f_ext_dead_n	Array | None	Nodal dead forces, (n_nodes, 6).	required
thrust_n	dict[str, Array]	Thrust magnitude, {key: ()}.	required
dynamic	bool	Flag for whether to compute dynamic entries.	required
m_t		Disassembled system mass matrix, (n_elem, 12, 12).	required
c_l		Dissembled system gyroscopic matrix, (n_elem, 12, 12).	required
c_l_lumped		Lumped gyroscopic matrix, (n_nodes, 6, 6).	required
v		Nodal velocities, (n_nodes, 6).	required
v_dot		Nodal accelerations, (n_node, 6).	required
i_ts	int	Time-step index (0 for static solves).	0
k_t_assembled	Array | None	Assembled global tangent stiffness matrix (n_dof, n_dof), required for Rayleigh damping.	None

Returns:

Type	Description
tuple[Array, Array]	Residual force vector, (n_dof, ), absolute sum of forces, (n_dof, ).

Source code in src/flapjax/structure/beam.py

def make_f_res(
    self,
    solve_dofs: Array | None,
    p_d: Array,
    eps: Array,
    hg: Array,
    f_ext_follower_n: Array | None,
    f_ext_dead_n: Array | None,
    thrust_n: dict[str, Array],
    dynamic: bool,
    m_t,
    c_l,
    c_l_lumped,
    v,
    v_dot,
    i_ts: int = 0,
    k_t_assembled: Array | None = None,
) -> tuple[Array, Array]:
    r"""
    Compute the residual force vector for a given configuration and external forces, used in the nonlinear solve.
    This is the force imbalance that the nonlinear solver will seek to drive to zero. Additionally, returns an
    "absolute sum" of all forces, used for relative convergence checks.
    :param solve_dofs: Optional array of degrees of freedom to solve for ``(n_solve_dofs, )``.
    :param p_d: P(d) operator, ``(n_elem, 6, 12)``.
    :param eps: Element strain vectors, ``(n_elem, 6)``.
    :param hg: Nodal homogeneous transformation matrices, ``(n_nodes, 4, 4)``.
    :param f_ext_follower_n: Nodal follower forces, ``(n_nodes, 6)``.
    :param f_ext_dead_n: Nodal dead forces, ``(n_nodes, 6)``.
    :param thrust_n: Thrust magnitude, ``{key: ()}``.
    :param dynamic: Flag for whether to compute dynamic entries.
    :param m_t: Disassembled system mass matrix, ``(n_elem, 12, 12)``.
    :param c_l: Dissembled system gyroscopic matrix, ``(n_elem, 12, 12)``.
    :param c_l_lumped: Lumped gyroscopic matrix, ``(n_nodes, 6, 6)``.
    :param v: Nodal velocities, ``(n_nodes, 6)``.
    :param v_dot: Nodal accelerations, ``(n_node, 6)``.
    :param i_ts: Time-step index (0 for static solves).
    :param k_t_assembled: Assembled global tangent stiffness matrix ``(n_dof, n_dof)``, required for Rayleigh
    damping.
    :return: Residual force vector, ``(n_dof, )``, absolute sum of forces, ``(n_dof, )``.
    """

    f_res = self.make_f_int(p_d, eps)  # (n_elem, 12)
    f_abs_sum = jnp.abs(f_res)

    if self.use_gravity:
        f_grav = self._make_f_grav(m_t, hg[:, :3, :3])
        f_res += f_grav
        f_abs_sum += jnp.abs(f_grav)

    if dynamic:
        f_iner, f_gyr = self._make_f_iner_gyr(m_t, c_l, v, v_dot)
        f_res += f_iner + f_gyr
        f_abs_sum += jnp.abs(f_iner + f_gyr)

    f_res_vect = self.assemble_vector_from_entries(f_res)
    f_abs_sum_vect = self.assemble_vector_from_entries(f_abs_sum)

    # add external forcing contributions
    if f_ext_follower_n is not None:
        f_res_vect += f_ext_follower_n.reshape(self.n_dof).ravel()
        f_abs_sum_vect += jnp.abs(f_ext_follower_n.reshape(self.n_dof).ravel())
    if f_ext_dead_n is not None:
        f_dead = self.make_f_dead_ext(f_ext_dead_n, hg[:, :3, :3]).ravel()
        f_res_vect += f_dead
        f_abs_sum_vect += jnp.abs(f_dead)

    f_thrust = self.add_thrust_force(
        force=jnp.zeros((self.n_nodes, 6)), thrust=thrust_n
    ).ravel()
    f_res_vect += f_thrust
    f_abs_sum_vect += jnp.abs(f_thrust)

    if self.use_lumped_mass:
        if dynamic:
            f_iner_lumped, f_gyr_lumped = self._make_f_iner_gyr_lumped(
                c_l_lumped, v, v_dot
            )
            f_iner_gyr_lumped = (f_iner_lumped + f_gyr_lumped).ravel()
            f_res_vect = self.add_lumped_contributions_to_vec(
                f_res_vect, f_iner_gyr_lumped
            )
            f_abs_sum_vect = self.add_lumped_contributions_to_vec(
                f_abs_sum_vect, jnp.abs(f_iner_gyr_lumped)
            )
        if self.use_gravity:
            f_grav_lumped = self._make_f_grav_lumped(hg[:, :3, :3]).ravel()
            f_res_vect = self.add_lumped_contributions_to_vec(
                vec=f_res_vect, lumped_vec=f_grav_lumped
            )
            f_abs_sum_vect = self.add_lumped_contributions_to_vec(
                vec=f_abs_sum_vect, lumped_vec=f_grav_lumped
            )

    # nodal constraint contributions
    for con in self.nodal_constraints:
        node = con.node_index
        v_node = v[node] if dynamic else jnp.zeros(6)
        f_constraint = con.f_res(hg[node], v_node, i_ts)
        dofs = node * 6 + jnp.arange(6)
        f_res_vect = f_res_vect.at[dofs].add(f_constraint)
        f_abs_sum_vect = f_abs_sum_vect.at[dofs].add(jnp.abs(f_constraint))

    # hard constraint force contributions (e.g. hinge spring-damper)
    for con in self.multibody_constraints:
        if con.has_f_res:
            v_i = v[con.node_i] if dynamic else jnp.zeros(6)
            v_j = v[con.node_j] if dynamic else jnp.zeros(6)
            f_i, f_j = con.f_res(hg[con.node_i], hg[con.node_j], v_i, v_j)
            dofs_i = con.node_i * 6 + jnp.arange(6)
            dofs_j = con.node_j * 6 + jnp.arange(6)
            f_res_vect = f_res_vect.at[dofs_i].add(f_i)
            f_res_vect = f_res_vect.at[dofs_j].add(f_j)
            f_abs_sum_vect = f_abs_sum_vect.at[dofs_i].add(jnp.abs(f_i))
            f_abs_sum_vect = f_abs_sum_vect.at[dofs_j].add(jnp.abs(f_j))

    # Rayleigh structural damping
    if dynamic and (self.alpha_m != 0.0 or self.beta_k != 0.0):
        if self.beta_k != 0.0 and k_t_assembled is None:
            raise ValueError(
                "k_t_assembled must be provided when beta_k != 0 for dynamic residual."
            )
        f_damp = self._make_f_rayleigh_damp(
            m_t=m_t,
            k_t_assembled=k_t_assembled
            if k_t_assembled is not None
            else jnp.zeros((self.n_dof, self.n_dof)),
            v=v,
        )
        f_res_vect += f_damp
        f_abs_sum_vect += jnp.abs(f_damp)

    if solve_dofs is not None:
        return f_res_vect[solve_dofs], f_abs_sum_vect[
            solve_dofs
        ]  # (n_solve_dof, ), (n_solve_dof, )
    else:
        return f_res_vect, f_abs_sum_vect  # (n_dof, ), (n_dof, )

update_hg staticmethod

update_hg(hg: Array, phi: Array) -> Array

Update the nodal homogeneous transformation matrices with the configuration increments.

Parameters:

Name	Type	Description	Default
hg	Array	Existing nodal homogeneous transformation matrices, (n_nodes, 4, 4)	required
phi	Array	Perturbation to the configuration vector, (n_nodes, 6)	required

Returns:

Type	Description
Array	Updated nodal homogeneous transformation matrices, (n_nodes, 4, 4)

Source code in src/flapjax/structure/beam.py

@staticmethod
def update_hg(hg: Array, phi: Array) -> Array:
    r"""
    Update the nodal homogeneous transformation matrices with the configuration increments.
    :param hg: Existing nodal homogeneous transformation matrices, ``(n_nodes, 4, 4)``
    :param phi: Perturbation to the configuration vector, ``(n_nodes, 6)``
    :return: Updated nodal homogeneous transformation matrices, ``(n_nodes, 4, 4)``
    """
    return jnp.einsum(
        "ijk,ikl->ijl",
        hg,
        vmap(exp_se3, 0, 0)(phi.reshape(-1, 6)),
    )

static_solve

static_solve(
    prescribed_dofs: Sequence[int] | Array | slice | int,
    f_ext_follower: Array | None = None,
    f_ext_dead: Array | None = None,
    f_ext_aero: Array | None = None,
    load_steps: int = 1,
    *,
    print_header: bool = True,
    postprocess_constraints: bool = True,
) -> StructureCase

Perform static solve of the structure under external loads.

Parameters:

Name	Type	Description	Default
f_ext_follower	Array | None	External forces array of follower forces (n_node, 6).	None
f_ext_dead	Array | None	External forces array of dead loads (n_node, 6).	None
f_ext_aero	Array | None	External forces array of aerodynamic loads (n_node, 6).	None
prescribed_dofs	Sequence[int] | Array | slice | int	Index of degrees of freedom which are prescribed (not solved for).	required
load_steps	int	Number of load steps to apply the external loads over.	1
print_header	bool	If False, suppress the "Static Solve" table header and trailing line.	True
postprocess_constraints	bool	If True, apply constraint postprocessing to the final solution.	True

Returns:

Type	Description
StructureCase	StructureCase object containing results of the static analysis.

Source code in src/flapjax/structure/beam.py

def static_solve(
    self,
    prescribed_dofs: Sequence[int] | Array | slice | int,
    f_ext_follower: Array | None = None,
    f_ext_dead: Array | None = None,
    f_ext_aero: Array | None = None,
    load_steps: int = 1,
    *,
    print_header: bool = True,
    postprocess_constraints: bool = True,
) -> StructureCase:
    r"""
    Perform static solve of the structure under external loads.
    :param f_ext_follower: External forces array of follower forces ``(n_node, 6)``.
    :param f_ext_dead: External forces array of dead loads ``(n_node, 6)``.
    :param f_ext_aero: External forces array of aerodynamic loads ``(n_node, 6)``.
    :param prescribed_dofs: Index of degrees of freedom which are prescribed (not solved for).
    :param load_steps: Number of load steps to apply the external loads over.
    :param print_header: If False, suppress the "Static Solve" table header and trailing line.
    :param postprocess_constraints: If True, apply constraint postprocessing to the final solution.
    :return: StructureCase object containing results of the static analysis.
    """

    if load_steps < 1:
        raise ValueError("load_steps must be at least 1")

    # check inputs
    if f_ext_follower is not None:
        check_arr_shape(f_ext_follower, (self.n_nodes, 6), "f_ext_follower")
    if f_ext_dead is not None:
        check_arr_shape(f_ext_dead, (self.n_nodes, 6), "f_ext_dead")

    if not (0.0 < self.relaxation_factor <= 1.0):
        raise ValueError("struct_relaxation_factor must be in the range (0, 1]")

    # degrees of freedom to solve for
    prescribed_dofs_: tuple[int, ...] = self.make_prescribed_dofs_tuple(
        prescribed_dofs
    )
    solve_dofs: Array = jnp.array(
        get_solve_dofs(n_dof=self.n_dof, prescribed_dofs=prescribed_dofs_)
    )

    # process external forces for load stepping
    load_step_weight: Array = jnp.linspace(0.0, 1.0, load_steps + 1)[
        1:
    ]  # (load_steps, )

    f_ext_follower_steps = self._make_load_steps_f(
        f_ext_follower, load_step_weight, apply_alpha_weighting=False
    )
    f_ext_dead_steps = self._make_load_steps_f(
        f_ext_dead, load_step_weight, apply_alpha_weighting=False
    )
    f_ext_aero_steps = self._make_load_steps_f(
        f_ext_aero, load_step_weight, apply_alpha_weighting=False
    )

    def _update(
        i_load_step: int,
        converge_status: ConvergenceStatus,
        hg_n: Array,
    ) -> tuple[int, ConvergenceStatus, Array]:
        # base parameters
        d_n = self.make_d(hg_n)  # (n_elem, 6)
        p_d_n = self.make_p_d(d_n)  # (n_elem, 6, 12)
        eps_n = self.make_eps(d_n)  # (n_elem, 6)
        m_t = self.make_m_t(d_n) if self.use_gravity else None  # (n_elem, 12, 12)

        # get total dead forces for this load step, (n_node, 6)
        total_f_ext_dead_step = self.make_f_ext_dead_tot(
            f_ext_dead_steps, f_ext_aero_steps, i_load_step
        )

        # assemble tangent stiffness matrix, (n_dof, n_dof)
        k_t_full_n = self.make_k_t_full(
            d=d_n,
            p_d=p_d_n,
            eps=eps_n,
            f_ext_dead=total_f_ext_dead_step,
            rmat=hg_n[:, :3, :3],
            m_t=m_t,
        )
        # apply nodal constraint contributions
        k_t_full_n = self.apply_nodal_constraint_tangent(
            mat=k_t_full_n, hg=hg_n, i_ts=0, gamma_prime=None
        )
        k_t_solve_n = k_t_full_n[jnp.ix_(solve_dofs, solve_dofs)]

        # compute residual forces, (n_solve_dofs, )
        f_res_solve_n, f_abs_sum_n = self.make_f_res(
            solve_dofs=solve_dofs,
            p_d=p_d_n,
            eps=eps_n,
            hg=hg_n,
            f_ext_follower_n=f_ext_follower_steps[i_load_step, ...]
            if f_ext_follower_steps is not None
            else None,
            f_ext_dead_n=total_f_ext_dead_step,
            thrust_n=self.thrust_reference,  # use reference thrust in static case
            dynamic=False,
            m_t=m_t,
            c_l=None,
            c_l_lumped=None,
            v=None,
            v_dot=None,
        )

        # solve for configuration increment, (n_solve_dofs, )
        if self.n_holonomic_constraints:
            d_varphi_np1, _, _ = self.solve_constrained(
                sys_mat_solve=k_t_solve_n,
                f_res_solve=f_res_solve_n,
                hg_eval=hg_n,
                solve_dofs=solve_dofs,
            )
            d_varphi_np1 *= self.relaxation_factor
        else:
            d_varphi_np1 = (
                jnp.linalg.solve(k_t_solve_n, f_res_solve_n)
                * self.relaxation_factor
            )

        # update configuration, (n_nodes, 4, 4)
        hg_np1_full = self.update_hg(
            hg_n, jnp.zeros(self.n_dof).at[solve_dofs].set(d_varphi_np1)
        )

        # algebra between undeformed and deformed shape, used to check relative convergence, (n_solve_dofs, )
        # this is relatively expensive to compute
        if self.struct_convergence_settings.rel_disp_tol is not None:
            h_full = vmap(hg_to_d, (0, 0), 0)(self.hg0, hg_np1_full).ravel()[
                solve_dofs
            ]
        else:
            h_full = None

        # update convergence status
        converge_status.update(
            delta_disp=d_varphi_np1,
            total_disp=h_full,
            delta_force=f_res_solve_n,
            total_force=f_abs_sum_n,
        )

        if map_verbosity_level(get_verbosity()) >= map_verbosity_level("verbose"):
            converge_status.print_struct_message(
                i_ts=None, t=None, i_load_step=i_load_step
            )

        return i_load_step, converge_status, hg_np1_full

    def convergence_loop(
        i_load_step: int,
        hg_init: Array,
    ) -> Array:
        r"""
        Convergence loop
        :param i_load_step: Index of load step.
        :param hg_init: Initial coordinates, ``(n_nodes, 4, 4)``.
        :return: Converged coordinates, ``(n_nodes, 4, 4)``.
        """
        _, convergence_status, hg_solve = eqxi.while_loop(
            lambda args_: ~args_[1].get_status(),
            lambda args_: _update(*args_),
            (
                i_load_step,
                ConvergenceStatus(
                    self.struct_convergence_settings,
                ),
                hg_init,
            ),
            max_steps=self.struct_convergence_settings.max_n_iter,
            kind="bounded",
        )

        if map_verbosity_level(get_verbosity()) >= map_verbosity_level("normal"):
            convergence_status.print_struct_message(
                i_ts=None, t=None, i_load_step=i_load_step
            )

        return hg_solve

    if print_header and map_verbosity_level(get_verbosity()) >= map_verbosity_level(
        "normal"
    ):
        ConvergenceStatus.print_header(dynamic=False)

    # solve for each load step
    hg = jax.lax.fori_loop(
        0,
        load_steps,
        lambda *args: convergence_loop(*args),
        self.hg0,
    )

    if print_header and map_verbosity_level(get_verbosity()) >= map_verbosity_level(
        "normal"
    ):
        ConvergenceStatus.print_line(dynamic=False)

    # postprocess final results
    d, eps, f_ext_dead_local, f_ext_aero_local, f_grav, f_int, _, _, f_res = (
        self.resolve_forces(
            hg=hg,
            dynamic=False,
            f_ext_dead=f_ext_dead,
            f_ext_follower=f_ext_follower,
            f_ext_aero=f_ext_aero,
            thrust=self.thrust_reference,
            v=None,
            v_dot=None,
        )
    )
    varphi = self.compute_varphi_from_hg(hg)
    f_elem = self.make_f_elem(eps=eps)  # compute loads in each element

    result = StructureCase(
        hg=hg,
        conn=self.connectivity,
        o0=self.o0,
        d=d,
        eps=eps,
        varphi=varphi,
        f_int=f_int,
        f_elem=f_elem,
        f_ext_follower=f_ext_follower,
        f_ext_dead=f_ext_dead_local,
        f_ext_aero=f_ext_aero_local,
        f_grav=f_grav,
        f_res=f_res,
        thrust=self.thrust_reference,
        thrust_nodes=self.thrust_nodes,
        thrust_direction=self.thrust_direction,
        prescribed_dofs=prescribed_dofs_,
        t=jnp.zeros(1),
    )
    if postprocess_constraints:
        result.constraint_data = self.postprocess_constraints(hg)
    return result

base_dynamic_solve

base_dynamic_solve(
    struct_case: StructureCase,
    struct_convergence_status: ConvergenceStatus,
    t: Array,
    solve_dofs: tuple[int, ...],
    load_steps: int,
    f_ext_dead: Array | None,
    f_ext_follower: Array | None,
    thrust_t: dict[str, Array],
    aero_obj: None,
    aero_case: None,
    fsi_convergence_status: None,
    cs_ang_t: None,
    cs_vel_t: None,
) -> StructureCase

base_dynamic_solve(
    struct_case: StructureCase,
    struct_convergence_status: ConvergenceStatus,
    t: Array,
    solve_dofs: tuple[int, ...],
    load_steps: int,
    f_ext_dead: Array | None,
    f_ext_follower: Array | None,
    thrust_t: dict[str, Array],
    aero_obj: DynamicAeroSolver,
    aero_case: AeroCase,
    fsi_convergence_status: ConvergenceStatus,
    cs_ang_t: dict[str, Array],
    cs_vel_t: dict[str, Array],
) -> AeroelasticCase

base_dynamic_solve(
    struct_case: StructureCase,
    struct_convergence_status: ConvergenceStatus,
    t: Array,
    solve_dofs: tuple[int, ...],
    load_steps: int,
    f_ext_dead: Array | None,
    f_ext_follower: Array | None,
    thrust_t: dict[str, Array],
    aero_obj: DynamicAeroSolver | None,
    aero_case: AeroCase | None,
    fsi_convergence_status: ConvergenceStatus | None,
    cs_ang_t: dict[str, Array] | None,
    cs_vel_t: dict[str, Array] | None,
) -> StructureCase | AeroelasticCase

Generic dynamic solver. Both the structural dynamic solve, and aeroelastic dynamic solve, are formed as wrappers of this

Source code in src/flapjax/structure/beam.py

def base_dynamic_solve(
    self,
    struct_case: StructureCase,
    struct_convergence_status: ConvergenceStatus,
    t: Array,
    solve_dofs: tuple[int, ...],
    load_steps: int,
    f_ext_dead: Array | None,
    f_ext_follower: Array | None,
    thrust_t: dict[str, Array],
    aero_obj: DynamicAeroSolver | None,
    aero_case: AeroCase | None,
    fsi_convergence_status: ConvergenceStatus | None,
    cs_ang_t: dict[str, Array] | None,
    cs_vel_t: dict[str, Array] | None,
) -> StructureCase | AeroelasticCase:
    r"""
    Generic dynamic solver. Both the structural dynamic solve, and aeroelastic dynamic solve, are formed as wrappers
    of this
    """

    if not (0.0 < self.relaxation_factor <= 1.0):
        raise ValueError("Relaxation factor must be in range (0, 1]")

    n_tstep = len(t)

    include_aero: bool = aero_obj is not None

    # process external forces for load stepping
    load_step_weight: Array = jnp.linspace(0.0, 1.0, load_steps + 1)[
        1:
    ]  # (load_steps, )
    f_ext_follower_alpha_steps = self._make_load_steps_f(
        f_ext_follower, load_step_weight, apply_alpha_weighting=True
    )
    f_ext_dead_alpha_steps = self._make_load_steps_f(
        f_ext_dead, load_step_weight, apply_alpha_weighting=True
    )

    solve_dofs_arr: Array = jnp.array(solve_dofs)
    prescribed_dofs_arr: Array = jnp.array(
        sorted(set(range(self.n_dof)) - set(solve_dofs)), dtype=int
    )

    def _update(
        i_load_step: int,
        i_ts: int,
        struct_convergence_status_: ConvergenceStatus,
        hg_n: Array,
        phi_alpha: Array,
        q_alpha: StructureMinimalStates,
        f_ext_aero_alpha_steps: Array | None,
        thrust_alpha: dict[str, Array],
    ) -> tuple[
        int,
        int,
        ConvergenceStatus,
        Array,
        Array,
        StructureMinimalStates,
        Array | None,
        dict[str, Array],
    ]:
        r"""
        Solution update for a single iteration of the nonlinear solver at a given time step and load step.
        :param i_load_step: Load step index.
        :param i_ts: Time step index.
        :param struct_convergence_status_: ConvergenceStatus object for the current iteration, used to track
        convergence and print messages.
        :param hg_n: Transformation matrices at iteration varphi, ``(n_nodes, 4, 4)``.
        :param phi_alpha: Timestep increment to the alpha step, ``(n_nodes, 6)``.
        :param f_ext_aero_alpha_steps: Load steps for the external aerodynamic forcing, ``(n_steps, n_nodes, 6)``.
        :param thrust_alpha: Thrust magnitude at the alpha step, ``{keys: ()}``.
        :return: Load and time step indices, updated ConvergenceStatus object, updated transformation matrices,
        configuration, velocities and accelerations for iteration n+1.
        """

        hg_update = self.update_hg(hg_n, phi_alpha)  # (n_node, 4, 4)

        # base parameters
        d_n = self.make_d(hg_update)  # (n_elem, 6)
        p_d_n = self.make_p_d(d_n)  # (n_elem, 6, 12)
        eps_n = self.make_eps(d_n)  # (n_elem, 6)
        d_dot_n = self._make_d_dot(p_d_n, q_alpha.v)  # (n_elem, 6)
        t_n = vmap(t_se3, 0, 0)(phi_alpha)  # (n_node, 6, 6)

        # tangent matrices
        m_t = self.make_m_t(d_n)  # (n_elem, 12, 12)
        c_l, c_t = self._make_c_t(
            d_n, d_dot_n, q_alpha.v
        )  # (n_elem, 12, 12), (n_elem, 12, 12)

        total_f_ext_dead = self.make_f_ext_dead_tot(
            f_ext_dead=f_ext_dead_alpha_steps[:, i_ts, :, :]
            if f_ext_dead_alpha_steps is not None
            else None,
            f_ext_aero=f_ext_aero_alpha_steps,
            i_load_step=i_load_step,
        )  # (n_node, 6)

        k_t = self.make_k_t_full(
            d_n,
            p_d_n,
            eps_n,
            total_f_ext_dead,
            hg_update[:, :3, :3],
            m_t,
        )  # (n_dof, n_dof)

        # add lumped mass contributions if applicable
        if self.use_lumped_mass:
            c_l_lumped, c_t_lumped = self._make_c_t_lumped(
                q_alpha.v
            )  # (n_node, 6, 6), (n_node, 6, 6)
        else:
            c_l_lumped, c_t_lumped = None, None

        # residual forces, (n_solve_dofs, )
        f_res_n_solve, f_abs_sum_n = self.make_f_res(
            solve_dofs=solve_dofs_arr,
            p_d=p_d_n,
            eps=eps_n,
            hg=hg_update,
            f_ext_follower_n=f_ext_follower_alpha_steps[i_load_step, i_ts, ...]
            if f_ext_follower_alpha_steps is not None
            else None,
            f_ext_dead_n=total_f_ext_dead,
            thrust_n=thrust_alpha,
            dynamic=True,
            m_t=m_t,
            c_l=c_l,
            c_l_lumped=c_l_lumped,
            v=q_alpha.v,
            v_dot=q_alpha.v_dot,
            i_ts=i_ts,
            k_t_assembled=k_t,
        )

        # system matrix, (n_dof, n_dof)
        sys_mat_full = self._make_sys_matrix(
            m_t=m_t,
            c_t=c_t,
            c_t_lumped=c_t_lumped,
            k_t=k_t,
            t_n=t_n,
            ti=self.time_integrator,
        )
        # add nodal constraint contributions
        sys_mat_full = self.apply_nodal_constraint_tangent(
            mat=sys_mat_full,
            hg=hg_update,
            i_ts=i_ts,
            gamma_prime=self.time_integrator.gamma_prime,
        )
        sys_mat = sys_mat_full[jnp.ix_(solve_dofs_arr, solve_dofs_arr)]

        # solve for configuration increment, (n_solve_dofs, )
        if self.multibody_constraints:
            d_n_np1, _, _ = self.solve_constrained(
                sys_mat_solve=sys_mat,
                f_res_solve=f_res_n_solve,
                hg_eval=hg_update,
                solve_dofs=solve_dofs_arr,
                hg_base=hg_n,
                phi=phi_alpha,
                v=q_alpha.v,
                gamma_prime=self.time_integrator.gamma_prime,
            )
            d_n_np1 *= self.relaxation_factor
        else:
            d_n_np1 = (
                jnp.linalg.solve(sys_mat, f_res_n_solve) * self.relaxation_factor
            )
        phi_np1 = phi_alpha.ravel().at[solve_dofs_arr].add(d_n_np1).reshape(-1, 6)

        # update configuration, velocities and accelerations
        v_np1 = (
            q_alpha.v.ravel()
            .at[solve_dofs_arr]
            .add(self.time_integrator.gamma_prime * d_n_np1)
            .reshape(-1, 6)
        )
        v_dot_np1 = (
            q_alpha.v_dot.ravel()
            .at[solve_dofs_arr]
            .add(self.time_integrator.beta_prime * d_n_np1)
            .reshape(-1, 6)
        )

        # update convergence status
        struct_convergence_status_.update(
            delta_disp=d_n_np1,
            total_disp=phi_np1,
            delta_force=f_res_n_solve,
            total_force=f_abs_sum_n,
        )

        if map_verbosity_level(get_verbosity()) >= map_verbosity_level("verbose"):
            struct_convergence_status_.print_struct_message(
                i_ts=i_ts, t=t[i_ts], i_load_step=i_load_step
            )

        q_alpha_update = StructureMinimalStates(
            varphi=None, v=v_np1, v_dot=v_dot_np1, a=q_alpha.a
        )

        return (
            i_load_step,
            i_ts,
            struct_convergence_status_,
            hg_n,
            phi_np1,
            q_alpha_update,
            f_ext_aero_alpha_steps,
            thrust_alpha,
        )

    @overload
    def time_step_loop(
        i_ts: int,
        struct_sol: StructureCase,
        struct_convergence_status_: ConvergenceStatus,
        aero_sol: None,
        fsi_convergence_status_: None,
        thrust_t_: dict[str, Array],
        cs_ang_t_: None,
        cs_vel_t_: None,
    ) -> tuple[
        StructureCase,
        ConvergenceStatus,
        None,
        None,
        dict[str, Array],
        None,
        None,
    ]: ...

    @overload
    def time_step_loop(
        i_ts: int,
        struct_sol: StructureCase,
        struct_convergence_status_: ConvergenceStatus,
        aero_sol: AeroCase,
        fsi_convergence_status_: ConvergenceStatus,
        thrust_t_: dict[str, Array],
        cs_ang_t_: dict[str, Array],
        cs_vel_t_: dict[str, Array],
    ) -> tuple[
        StructureCase,
        ConvergenceStatus,
        AeroCase,
        ConvergenceStatus,
        dict[str, Array],
        dict[str, Array],
        dict[str, Array],
    ]: ...

    def time_step_loop(
        i_ts: int,
        struct_sol: StructureCase,
        struct_convergence_status_: ConvergenceStatus,
        aero_sol: AeroCase | None,
        fsi_convergence_status_: ConvergenceStatus | None,
        thrust_t_: dict[str, Array],
        cs_ang_t_: dict[str, Array] | None,
        cs_vel_t_: dict[str, Array] | None,
    ) -> tuple[
        StructureCase,
        ConvergenceStatus,
        AeroCase | None,
        ConvergenceStatus | None,
        dict[str, Array],
        dict[str, Array] | None,
        dict[str, Array] | None,
    ]:
        r"""
        Performs analysis on a single time step, including load stepping
        :param i_ts: Index of time step to solve
        :param struct_sol: Solution object, with results up to time step i_ts-1.
        :param struct_convergence_status_: Convergence status object.
        :param aero_sol: Aero solution object, with results up to time step i_ts-1, if aero is included.
        :param fsi_convergence_status_: Convergence status object.
        :param thrust_t_: Thrust magnitude time history, ``{name: (n_tstep, )}``.
        :param cs_ang_t_: Control surface angle time history, ``{name: (n_tstep, )}``.
        :param cs_vel_t_: Control surface velocity time history, ``{name: (n_tstep, )}``.
        :return: Solution object with results up to time step i_ts.
        """

        # predictor step
        q_nm1 = struct_sol.get_minimal_states(i_ts - 1)
        phi_init, q_init = self.time_integrator.predict_q(q_nm1)
        phi_alpha_init, q_alpha_init = self.time_integrator.compute_q_alpha(
            q_nm1=q_nm1,
            q_n=q_init,
            phi_n=phi_init,
        )

        # prescribed DOFs should not be influenced by the time integration
        phi_alpha_init = (
            phi_alpha_init.ravel().at[prescribed_dofs_arr].set(0.0).reshape(-1, 6)
        )
        q_alpha_init.v = (
            q_alpha_init.v.ravel()
            .at[prescribed_dofs_arr]
            .set(q_nm1.v.ravel()[prescribed_dofs_arr])
            .reshape(-1, 6)
        )
        q_alpha_init.v_dot = (
            q_alpha_init.v_dot.ravel()
            .at[prescribed_dofs_arr]
            .set(q_nm1.v_dot.ravel()[prescribed_dofs_arr])
            .reshape(-1, 6)
        )
        q_alpha_init.a = (
            q_alpha_init.a.ravel()
            .at[prescribed_dofs_arr]
            .set(q_nm1.a.ravel()[prescribed_dofs_arr])
            .reshape(-1, 6)
        )

        q_alpha_init.varphi = None  # this value is not used during the loop

        # thrust force
        thrust_alpha: dict[str, Array] = {
            k: self.time_integrator.compute_f_alpha(f_nm1=v[i_ts - 1], f_n=v[i_ts])
            for k, v in thrust_t_.items()
        }
        thrust_n: dict[str, Array] = {k: v[i_ts] for k, v in thrust_t_.items()}

        if include_aero:
            assert (
                aero_sol is not None
                and fsi_convergence_status_ is not None
                and struct_sol.f_ext_aero is not None
                and cs_ang_t_ is not None
                and cs_vel_t_ is not None
            )

            fsi_convergence_status_.reset_status()

            # f_ext_aero is stored in local frame, so we convert back to global
            # so that both operands of the alpha blend are in the same (global) frame.
            f_aero_nm1 = jnp.concatenate(
                [
                    jnp.einsum(
                        "ijk,ik->ij",
                        struct_sol.hg[i_ts - 1, :, :3, :3],
                        struct_sol.f_ext_aero[i_ts - 1, :, :3],
                    ),
                    jnp.einsum(
                        "ijk,ik->ij",
                        struct_sol.hg[i_ts - 1, :, :3, :3],
                        struct_sol.f_ext_aero[i_ts - 1, :, 3:],
                    ),
                ],
                axis=-1,
            )

            # get control surface angles and velocities
            cs_ang_nm1 = {k: v[i_ts - 1] for k, v in cs_ang_t_.items()}
            cs_ang_n = {k: v[i_ts] for k, v in cs_ang_t_.items()}
            cs_vel_n = {k: v[i_ts] for k, v in cs_vel_t_.items()}
            assert fsi_convergence_status is not None
            (
                _,
                struct_sol,
                aero_sol,
                struct_convergence_status_,
                fsi_convergence_status_,
                phi_alpha,
                q_alpha,
                *_,
            ) = eqxi.while_loop(
                lambda args_: ~cast(ConvergenceStatus, args_[4]).get_status(),
                lambda args_: fsi_convergence_loop(*args_),
                (
                    i_ts,
                    struct_sol,
                    aero_sol,
                    struct_convergence_status_,
                    fsi_convergence_status_,
                    phi_alpha_init,
                    q_alpha_init,
                    f_aero_nm1,  # this value is for the previous timesteps force, and is propagated unaltered
                    f_aero_nm1,  # first guess for forcing at alpha is to use value from i_ts=n-1
                    thrust_alpha,
                    cs_ang_n,
                    cs_ang_nm1,
                    cs_vel_n,
                ),
                max_steps=fsi_convergence_status.convergence_settings.max_n_iter,
                kind="bounded",
            )

        else:
            # solve pure structural problem
            _, struct_convergence_status_, _, phi_alpha, q_alpha, *_ = (
                load_step_loop(
                    i_ts=i_ts,
                    struct_convergence_status_=struct_convergence_status_,
                    hg_alpha=struct_sol.hg[i_ts - 1, ...],
                    phi_alpha=phi_alpha_init,
                    q_alpha=q_alpha_init,
                    f_ext_aero_steps=None,
                    thrust_alpha=thrust_alpha,
                )
            )

        # print message where we only require one message per timestep
        if map_verbosity_level(get_verbosity()) == map_verbosity_level("normal"):
            struct_convergence_status_.print_struct_message(
                i_ts=i_ts, t=struct_sol.t[i_ts], i_load_step=load_steps - 1
            )
            if include_aero and fsi_convergence_status_ is not None:
                fsi_convergence_status_.print_fsi_message(
                    i_ts=i_ts, t=struct_sol.t[i_ts]
                )

        # postprocess results for time step and store in solution object
        q_n, phi_n = self.time_integrator.compute_q_n_from_q_alpha(
            q_alpha=q_alpha,
            q_nm1=struct_sol.get_minimal_states(i_ts - 1),
            phi_alpha=phi_alpha,
        )

        # update pseudo-acceleration
        q_n.a = self.time_integrator.compute_a_n(
            a_nm1=struct_sol.a[i_ts - 1, ...],
            v_dot_nm1=struct_sol.v_dot[i_ts - 1, ...],
            v_dot_n=q_n.v_dot,
        )

        # final node coordinates
        hg_n = self.update_hg(struct_sol.hg[i_ts - 1, ...], phi_n)

        if include_aero:
            if (
                aero_sol is None
                or aero_obj is None
                or fsi_convergence_status_ is None
            ):
                raise ValueError("Missing aero arguments")

            f_ext_aero = aero_sol.project_forcing_to_beam(
                i_ts=i_ts,
                rmat=hg_n[:, :3, :3],
                x0_aero=aero_obj.zeta_b0,
                include_unsteady=aero_obj.include_unsteady_force,
            )

        else:
            f_ext_aero = None

        (
            d,
            eps,
            f_ext_dead_local,
            f_ext_aero_local,
            f_grav,
            f_int,
            f_gyr,
            f_iner,
            f_res,
        ) = self.resolve_forces(
            hg=hg_n,
            dynamic=True,
            f_ext_dead=f_ext_dead[i_ts, ...] if f_ext_dead is not None else None,
            f_ext_follower=f_ext_follower[i_ts, ...]
            if f_ext_follower is not None
            else None,
            thrust=thrust_n,
            f_ext_aero=f_ext_aero,
            v=q_n.v,
            v_dot=q_n.v_dot,
        )
        struct_sol.d = struct_sol.d.at[i_ts, ...].set(d)
        struct_sol.eps = struct_sol.eps.at[i_ts, ...].set(eps)
        struct_sol.v = struct_sol.v.at[i_ts, ...].set(q_n.v)
        struct_sol.v_dot = struct_sol.v_dot.at[i_ts, ...].set(q_n.v_dot)
        struct_sol.a = struct_sol.a.at[i_ts, ...].set(q_n.a)
        struct_sol.hg = struct_sol.hg.at[i_ts, ...].set(hg_n)
        struct_sol.varphi = struct_sol.varphi.at[i_ts, ...].set(
            vmap(hg_to_d, (0, 0), 0)(self.hg0, hg_n)
        )

        if f_ext_follower is not None and struct_sol.f_ext_follower is not None:
            struct_sol.f_ext_follower = struct_sol.f_ext_follower.at[i_ts, ...].set(
                f_ext_follower[i_ts, ...]
            )
        if f_ext_dead is not None and struct_sol.f_ext_dead is not None:
            struct_sol.f_ext_dead = struct_sol.f_ext_dead.at[i_ts, ...].set(
                f_ext_dead_local
            )

        if f_ext_aero is not None and struct_sol.f_ext_aero is not None:
            struct_sol.f_ext_aero = struct_sol.f_ext_aero.at[i_ts, ...].set(
                f_ext_aero_local
            )

        if self.use_gravity:
            if struct_sol.f_grav is None:
                raise ValueError("struct_sol.f_grav is None")
            struct_sol.f_grav = struct_sol.f_grav.at[i_ts, ...].set(f_grav)
        struct_sol.f_int = struct_sol.f_int.at[i_ts, ...].set(f_int)
        struct_sol.f_elem = struct_sol.f_elem.at[i_ts, ...].set(
            self.make_f_elem(eps=eps)
        )
        struct_sol.f_iner_gyr = struct_sol.f_iner_gyr.at[i_ts, ...].set(
            f_iner + f_gyr
        )
        struct_sol.f_res = struct_sol.f_res.at[i_ts, ...].set(f_res)

        if include_aero and aero_sol is not None:
            assert cs_ang_t_ is not None and cs_vel_t_ is not None
            cs_ang_n = {k: v[i_ts] for k, v in cs_ang_t_.items()}
            cs_vel_n = {k: v[i_ts] for k, v in cs_vel_t_.items()}
            aero_sol.cs_ang = {
                k: v.at[i_ts].set(cs_ang_n[k]) for k, v in aero_sol.cs_ang.items()
            }
            aero_sol.cs_vel = {
                k: v.at[i_ts].set(cs_vel_n[k]) for k, v in aero_sol.cs_vel.items()
            }

        return (
            struct_sol,
            struct_convergence_status_,
            aero_sol,
            fsi_convergence_status_,
            thrust_t_,
            cs_ang_t_,
            cs_vel_t_,
        )

    def fsi_convergence_loop(
        i_ts: int,
        struct_sol: StructureCase,
        aero_sol: AeroCase,
        struct_convergence_status_: ConvergenceStatus,
        fsi_convergence_status_: ConvergenceStatus,
        phi_alpha_init: Array,
        q_alpha_init: StructureMinimalStates,
        f_aero_nm1: Array,
        f_aero_alpha_prev: Array,
        thrust_alpha: dict[str, Array],
        cs_ang_n: dict[str, Array],
        cs_ang_nm1: dict[str, Array],
        cs_vel_n: dict[str, Array],
    ) -> tuple[
        int,
        StructureCase,
        AeroCase,
        ConvergenceStatus,
        ConvergenceStatus,
        Array,
        StructureMinimalStates,
        Array,
        Array,
        dict[str, Array],
        dict[str, Array],
        dict[str, Array],
        dict[str, Array],
    ]:
        # obtain coordinates at timestep (not alpha)
        phi_n = self.time_integrator.compute_phi_from_phi_alpha(
            phi_alpha=phi_alpha_init
        )
        v_n = self.time_integrator.compute_v_from_v_alpha(
            v_alpha=q_alpha_init.v, v_nm1=struct_sol.v[i_ts - 1, ...]
        )

        hg_n = self.update_hg(hg=struct_sol.hg[i_ts - 1, ...], phi=phi_n)
        hg_dot = self.make_hg_dot(hg=hg_n, v=v_n)

        if aero_obj is None or struct_sol.f_ext_aero is None:
            raise ValueError("Missing aero parameters")

        # evaluate aerodynamic forcing on beam
        aero_sol = aero_obj.case_solve(
            case=aero_sol,
            i_ts=i_ts,
            hg_n=hg_n,
            hg_nm1=struct_sol.hg[i_ts - 1, ...],
            hg_dot_n=hg_dot,
            static=False,
            horseshoe=False,
            cs_ang_n=cs_ang_n,
            cs_ang_nm1=cs_ang_nm1,
            cs_vel_n=cs_vel_n,
        )

        f_aero_n = aero_sol.project_forcing_to_beam(
            i_ts=i_ts,
            rmat=hg_n[:, :3, :3],
            x0_aero=aero_obj.zeta_b0,
            include_unsteady=aero_obj.include_unsteady_force,
        )

        # aerodynamic force at alpha point, subsequently divided into load steps
        f_aero_alpha = self.time_integrator.compute_f_alpha(
            f_nm1=f_aero_nm1, f_n=f_aero_n
        )

        f_aero_alpha_steps = self._make_load_steps_f(
            f=f_aero_alpha, weighting=load_step_weight, apply_alpha_weighting=False
        )

        # reset convergence status
        struct_convergence_status_.reset_status()

        # solve structural problem for given aero load
        _, struct_convergence_status_, _, phi_alpha, q_alpha, *_ = load_step_loop(
            i_ts,
            struct_convergence_status_,
            struct_sol.hg[i_ts - 1, ...],
            phi_alpha_init,
            q_alpha_init,
            f_aero_alpha_steps,
            thrust_alpha,
        )

        # update the FSI convergence object
        # note that for convenience we use the alpha properties
        fsi_convergence_status_.update(
            delta_disp=(phi_alpha_init - phi_alpha).ravel()[solve_dofs_arr],
            total_disp=phi_alpha.ravel()[solve_dofs_arr],
            delta_force=(f_aero_alpha - f_aero_alpha_prev).ravel()[solve_dofs_arr],
            total_force=f_aero_alpha.ravel()[solve_dofs_arr],
        )

        if map_verbosity_level(get_verbosity()) >= map_verbosity_level("verbose"):
            fsi_convergence_status_.print_fsi_message(i_ts=i_ts, t=t[i_ts])

        return (
            i_ts,
            struct_sol,
            aero_sol,
            struct_convergence_status_,
            fsi_convergence_status_,
            phi_alpha,
            q_alpha,
            f_aero_nm1,
            f_aero_alpha,
            thrust_alpha,
            cs_ang_n,
            cs_ang_nm1,
            cs_vel_n,
        )

    def struct_convergence_loop(
        i_load_step: int,
        i_ts: int,
        struct_convergence_status_: ConvergenceStatus,
        hg_alpha: Array,
        phi_alpha: Array,
        q_alpha: StructureMinimalStates,
        f_ext_aero_steps: Array | None,
        thrust_alpha: dict[str, Array],
    ) -> tuple[
        int,
        ConvergenceStatus,
        Array,
        Array,
        StructureMinimalStates,
        Array | None,
        dict[str, Array],
    ]:
        r"""
        Convergence loop within each load step of a time step.
        :param i_load_step: Load step index.
        :param i_ts: Time step index.
        :param struct_convergence_status_: ConvergenceStatus object to update with convergence information during load
        stepping.
        :param hg_alpha: Node transformations at the beginning of the load step, ``(n_nodes, 4, 4)``.
        :param phi_alpha: Node configuration increments in algebra space, ``(n_nodes, 6)``.
        :param q_alpha: Minimal states at intermediate alpha step.
        :param f_ext_aero_steps: Optional aerodynamic forcing alpha load steps ``(n_steps, n_nodes, 6)``.
        :param thrust_alpha: Thrust at the alpha step, ``{key: ()}``.
        :return: Time step index, convergence status, and updated configuration, velocities, accelerations, and
        optional aerodynamic forcing.
        """

        struct_convergence_status_.reset_status()

        _, _, struct_convergence_status_, hg_solve, phi_alpha, q_alpha, _, _ = (
            eqxi.while_loop(
                lambda args_: ~args_[2].get_status(),
                lambda args_: _update(*args_),
                (
                    i_load_step,
                    i_ts,
                    struct_convergence_status_,
                    hg_alpha,
                    phi_alpha,
                    q_alpha,
                    f_ext_aero_steps,
                    thrust_alpha,
                ),
                max_steps=self.struct_convergence_settings.max_n_iter,
                kind="bounded",
            )
        )

        if map_verbosity_level(get_verbosity()) >= map_verbosity_level("verbose"):
            struct_convergence_status_.print_struct_message(
                i_ts=i_ts, t=t[i_ts], i_load_step=i_load_step
            )

        return (
            i_ts,
            struct_convergence_status_,
            hg_solve,
            phi_alpha,
            q_alpha,
            f_ext_aero_steps,
            thrust_alpha,
        )

    def load_step_loop(
        i_ts: int,
        struct_convergence_status_: ConvergenceStatus,
        hg_alpha: Array,
        phi_alpha: Array,
        q_alpha: StructureMinimalStates,
        f_ext_aero_steps: Array | None,
        thrust_alpha: dict[str, Array],
    ) -> tuple[
        int,
        ConvergenceStatus,
        Array,
        Array,
        StructureMinimalStates,
        Array | None,
    ]:
        r"""
        Performs load stepping iterations for a given time step. Load stepping is not performed for thrust.
        :param i_ts: Timestep index for which to perform load stepping.
        :param struct_convergence_status_: ConvergenceStatus object to update with load stepping convergence information.
        :param hg_alpha: SE(3) nodal transformation matrices at the beginning of the load step, ``(n_nodes, 4, 4)``.
        :param phi_alpha: Nodal updates to the configuration in the algebra space, ``(n_nodes, 6)``.
        :param q_alpha: Minimal states at intermediate alpha step.
        :param f_ext_aero_steps: Optional aerodynamic forcing alpha load steps ``(n_steps, n_nodes, 6)``.
        :param thrust_alpha: Thrust at the alpha step, ``{key: ()}``.
        :return: Time step index, updated ConvergenceStatus object, and updated configuration, velocities and accelerations after load stepping
        """
        return jax.lax.fori_loop(
            0,
            load_steps,
            lambda i_load_step, args: struct_convergence_loop(i_load_step, *args),
            (
                i_ts,
                struct_convergence_status_,
                hg_alpha,
                phi_alpha,
                q_alpha,
                f_ext_aero_steps,
                thrust_alpha,
            ),
        )

    def time_step_loop_checked(
        i_ts: int,
        struct_sol: StructureCase,
        struct_convergence_status_: ConvergenceStatus,
        aero_sol: AeroCase | None,
        fsi_convergence_status_: ConvergenceStatus | None,
        thrust_t_: dict[str, Array],
        cs_ang_t_: dict[str, Array] | None,
        cs_vel_t_: dict[str, Array] | None,
        diverged: Array,
    ) -> tuple[
        StructureCase,
        ConvergenceStatus,
        AeroCase | None,
        ConvergenceStatus | None,
        dict[str, Array],
        dict[str, Array] | None,
        dict[str, Array] | None,
        Array,
    ]:
        r"""
        Wraps ``time_step_loop`` with a check for solution divergence. Once a NaN is detected, this becomes a no-op
        for all remaining time steps. The corresponding time history entries are left at their initialised value
        (zero).
        """

        if include_aero:
            assert aero_sol is not None
            assert fsi_convergence_status_ is not None
            assert cs_ang_t_ is not None
            assert cs_vel_t_ is not None
            aero_sol_ok: AeroCase = aero_sol
            fsi_convergence_status_ok: ConvergenceStatus = fsi_convergence_status_
            cs_ang_t_ok: dict[str, Array] = cs_ang_t_
            cs_vel_t_ok: dict[str, Array] = cs_vel_t_
            false_branch = lambda: time_step_loop(
                i_ts,
                struct_sol,
                struct_convergence_status_,
                aero_sol_ok,
                fsi_convergence_status_ok,
                thrust_t_,
                cs_ang_t_ok,
                cs_vel_t_ok,
            )
        else:
            assert aero_sol is None
            assert fsi_convergence_status_ is None
            assert cs_ang_t_ is None
            assert cs_vel_t_ is None
            aero_sol_none: None = aero_sol
            fsi_convergence_status_none: None = fsi_convergence_status_
            cs_ang_t_none: None = cs_ang_t_
            cs_vel_t_none: None = cs_vel_t_
            false_branch = lambda: time_step_loop(
                i_ts,
                struct_sol,
                struct_convergence_status_,
                aero_sol_none,
                fsi_convergence_status_none,
                thrust_t_,
                cs_ang_t_none,
                cs_vel_t_none,
            )

        (
            struct_sol,
            struct_convergence_status_,
            aero_sol,
            fsi_convergence_status_,
            thrust_t_,
            cs_ang_t_,
            cs_vel_t_,
        ) = jax.lax.cond(
            diverged,
            lambda: (
                struct_sol,
                struct_convergence_status_,
                aero_sol,
                fsi_convergence_status_,
                thrust_t_,
                cs_ang_t_,
                cs_vel_t_,
            ),
            false_branch,
        )

        has_nan = struct_convergence_status_.has_nan
        if include_aero:
            assert fsi_convergence_status_ is not None
            has_nan = has_nan | fsi_convergence_status_.has_nan
        new_diverged = diverged | has_nan

        jax.lax.cond(
            new_diverged & ~diverged,
            lambda: warn(
                "NaN detected in dynamic solve at time step {i_ts} (t={t_val:.03e}) - skipping remaining time steps.",
                i_ts=i_ts,
                t_val=t[i_ts],
            ),
            lambda: None,
        )

        return (
            struct_sol,
            struct_convergence_status_,
            aero_sol,
            fsi_convergence_status_,
            thrust_t_,
            cs_ang_t_,
            cs_vel_t_,
            new_diverged,
        )

    struct_case, _, aero_case, *_ = jax.lax.fori_loop(
        1,
        n_tstep,
        lambda i_ts, args: time_step_loop_checked(i_ts, *args),
        (
            struct_case,
            struct_convergence_status,
            aero_case,
            fsi_convergence_status,
            thrust_t,
            cs_ang_t,
            cs_vel_t,
            jnp.zeros((), dtype=bool),
        ),
    )

    struct_case.constraint_data = self.postprocess_constraints(struct_case.hg)

    if include_aero:
        if aero_case is None:
            raise ValueError("aero_case cannot be None")

        from flapjax.coupled.data_structures import (
            AeroelasticCase,
        )  # import here to prevent circular references

        return AeroelasticCase(structure=struct_case, aero=aero_case)
    else:
        return struct_case

dynamic_solve

dynamic_solve(
    init_state: StructureCase | None,
    n_tstep: int,
    dt: Array | float,
    prescribed_dofs: Sequence[int]
    | Array
    | slice
    | int
    | None = None,
    f_ext_follower: Array | None = None,
    f_ext_dead: Array | None = None,
    f_ext_aero: Array | None = None,
    thrust_t: dict[str, Array] | None = None,
    load_steps: int = 1,
) -> StructureCase

Perform dynamic solve of the structure under external loads

Parameters:

Name	Type	Description	Default
init_state	StructureCase | None	Initial state of the structure, either static or a dynamic snapshot. If None, the reference configuration is used with zero velocities.	required
prescribed_dofs	Sequence[int] | Array | slice | int | None	Degrees of freedom which are prescribed (not solved for). If None, inherit from the initial state.	None
n_tstep	int	Number of time steps to simulate.	required
dt	Array | float	Time step length.	required
f_ext_follower	Array | None	Following external forces array, (n_tstep, n_node, 6), (n_node, 6) or None for zero external follower forces.	None
f_ext_dead	Array | None	Dead external forces array, (n_tstep, n_node, 6), (n_node, 6) or None for zero external dead forces.	None
f_ext_aero	Array | None	Aerodynamic external forces array, (n_tstep, n_node, 6), (n_node, 6) or None for zero external aerodynamic forces.	None
thrust_t	dict[str, Array] | None	Thrust time history, {key: (n_tstep, )}. If none, this will use the reference value.	None
load_steps	int	Number of load steps to apply the external loads over.	1

Returns:

Type	Description
StructureCase	Structure dataclass containing results of the dynamic analysis.

Source code in src/flapjax/structure/beam.py

def dynamic_solve(
    self,
    init_state: StructureCase | None,
    n_tstep: int,
    dt: Array | float,
    prescribed_dofs: Sequence[int] | Array | slice | int | None = None,
    f_ext_follower: Array | None = None,
    f_ext_dead: Array | None = None,
    f_ext_aero: Array | None = None,
    thrust_t: dict[str, Array] | None = None,
    load_steps: int = 1,
) -> StructureCase:
    r"""
    Perform dynamic solve of the structure under external loads
    :param init_state: Initial state of the structure, either static or a
    dynamic snapshot. If None, the reference configuration is used with zero
    velocities.
    :param prescribed_dofs: Degrees of freedom which are prescribed (not solved for). If None, inherit
    from the initial state.
    :param n_tstep: Number of time steps to simulate.
    :param dt: Time step length.
    :param f_ext_follower: Following external forces array, ``(n_tstep, n_node, 6)``, ``(n_node, 6)`` or None for zero external follower forces.
    :param f_ext_dead: Dead external forces array, ``(n_tstep, n_node, 6)``, ``(n_node, 6)`` or None for zero external dead forces.
    :param f_ext_aero: Aerodynamic external forces array, ``(n_tstep, n_node, 6)``, ``(n_node, 6)`` or None for zero external aerodynamic forces.
    :param thrust_t: Thrust time history, ``{key: (n_tstep, )}``. If none, this will use the reference value.
    :param load_steps: Number of load steps to apply the external loads over.
    :return: Structure dataclass containing results of the dynamic analysis.
    """

    if load_steps <= 0:
        raise ValueError("load_steps must be a positive integer")

    # set thrust if not provided
    thrust_t_: dict[str, Array] = (
        thrust_t
        if thrust_t is not None
        else {k: jnp.full(n_tstep, v) for k, v in self.thrust_reference.items()}
    )

    if prescribed_dofs is None:
        # inherit prescribed dofs from initial state
        if init_state is None:
            raise ValueError("prescribed_dofs cannot be None if init_state is None")
        prescribed_dofs = init_state.prescribed_dofs

    # degrees of freedom to solve for
    prescribed_dofs_arr = self.make_prescribed_dofs_tuple(prescribed_dofs)
    solve_dofs = get_solve_dofs(
        n_dof=self.n_dof, prescribed_dofs=prescribed_dofs_arr
    )

    # check and process external forces
    def check_force(arr: Array | None, name: str) -> Array | None:
        if arr is None:
            return None
        match arr.ndim:
            case 2:
                out_ = jnp.broadcast_to(arr[None, ...], (n_tstep, self.n_nodes, 6))
            case 3:
                out_ = arr
            case _:
                raise ValueError(
                    f"{name} must have shape [n_node, 6] or [n_tstep, n_node, 6]"
                )
        check_arr_shape(out_, (n_tstep, self.n_nodes, 6), name)
        return out_

    f_ext_dead = check_force(f_ext_dead, "f_ext_dead")  # (n_tstep, n_node, 6)
    f_ext_follower = check_force(
        f_ext_follower, "f_ext_follower"
    )  # (n_tstep, n_node, 6)
    f_ext_aero = check_force(f_ext_aero, "f_ext_aero")  # (n_tstep, n_node, 6)

    # time integration parameters
    self.time_integrator = TimeIntegrator(
        spectral_radius=self.spectral_radius, dt=jnp.array(dt)
    )

    def evaluate_initial_equilibrium(
        init_state__: StructureCase,
    ) -> StructureCase:
        r"""
        Evaluates the forces for a given initial state to check whether it is in equilibrium. If not, a warning is
        raised with the maximum residual force. This is important to ensure that the time integration starts from a
        consistent state.
        :param init_state__: Structure containing the initial state to evaluate.
        :return: Structure with the forces evaluated for the initial state.
        """
        d, eps, f_ext_dead_, f_ext_aero_, f_grav, f_int, f_gyr, f_iner, f_res = (
            self.resolve_forces(
                hg=init_state__.hg,
                dynamic=True,
                f_ext_dead=init_state__.f_ext_dead,
                f_ext_aero=init_state__.f_ext_aero,
                thrust=init_state__.thrust,
                f_ext_follower=init_state__.f_ext_follower,
                v=init_state__.v,
                v_dot=init_state__.v_dot,
            )
        )

        max_res = jnp.max(jnp.abs(f_res))
        jax_print(
            "Initial state maximum residual force: {max_res:.3e}",
            max_res=max_res,
            verbose_level="normal",
        )

        f_elem = self.make_f_elem(eps=eps)

        return StructureCase(
            hg=init_state__.hg,
            conn=self.connectivity,
            o0=self.o0,
            d=d,
            eps=eps,
            varphi=init_state__.varphi,
            v=init_state__.v,
            v_dot=init_state__.v_dot,
            a=init_state__.v_dot,  # initial pseudo-acceleration set equal to initial acceleration
            f_ext_follower=init_state__.f_ext_follower,
            f_ext_dead=f_ext_dead_,
            f_ext_aero=f_ext_aero_,
            f_grav=f_grav,
            f_int=f_int,
            f_elem=f_elem,
            f_iner_gyr=f_iner + f_gyr,  # type: ignore
            f_res=f_res,
            thrust=init_state__.thrust,
            thrust_nodes=self.thrust_nodes,
            thrust_direction=self.thrust_direction,
            t=init_state__.t,
            i_ts=init_state__.i_ts,
            prescribed_dofs=prescribed_dofs_arr,
        )

    # time steps
    t = jnp.arange(n_tstep) * dt
    if init_state is not None and init_state.is_dynamic:
        if init_state.is_batched:
            t += init_state.t[0]
        else:
            t += init_state.t

    # set up initial state
    if init_state is None:
        init_state_: StructureCase = self.reference_configuration(
            use_f_aero=f_ext_aero is not None,
            use_f_ext_dead=f_ext_dead is not None,
            use_f_ext_follower=f_ext_follower is not None,
            prescribed_dofs=tuple(prescribed_dofs_arr),
        ).to_dynamic(t=None)
    elif not init_state.is_dynamic:
        init_state_ = init_state.to_dynamic(t=None)
    elif not init_state.is_batched:
        init_state_ = init_state
    else:
        raise TypeError(
            "dynamic_solve init_state cannot be a batched Structure; pass a "
            "snapshot or static state"
        )

    # check if initial state satisfies equilibrium
    init_state_eval = evaluate_initial_equilibrium(init_state_)
    dynamic_struct = StructureCase.initialise(
        initial_snapshot=init_state_eval,
        t=t,
        use_f_ext_follower=f_ext_follower is not None,
        use_f_ext_dead=f_ext_dead is not None,
        use_f_ext_aero=False,
    )
    converge_status = ConvergenceStatus(
        convergence_settings=self.struct_convergence_settings
    )

    ConvergenceStatus.print_header(dynamic=True)

    out = self.base_dynamic_solve(
        struct_case=dynamic_struct,
        struct_convergence_status=converge_status,
        t=t,
        solve_dofs=solve_dofs,
        load_steps=load_steps,
        f_ext_dead=f_ext_dead,
        f_ext_follower=f_ext_follower,
        aero_obj=None,
        aero_case=None,
        fsi_convergence_status=None,
        thrust_t=thrust_t_,
        cs_ang_t=None,
        cs_vel_t=None,
    )

    ConvergenceStatus.print_line(dynamic=True)
    return out

linear

linear_beam

LinearBeam

LinearBeam(
    beam: BaseBeamStructure,
    reference: StructureCase,
    n_modes: int | None,
    dt: float | Array,
    modal_inputs: bool = False,
    modal_outputs: bool = False,
    int_order: Literal[3, 4, 5] = BASE_LOBATTO_ORDER,
    prescribed_dofs: Sequence[int]
    | Array
    | slice
    | int
    | None = None,
)

Bases: LinearModel[StructureCase, StructureInputUnflattened, StructureStateUnflattened, StructureOutputUnflattened, StructureLinearResult]

Class to represent a linearised beam system about a reference state.

Source code in src/flapjax/structure/linear/linear_beam.py

def __init__(
    self,
    beam: BaseBeamStructure,
    reference: StructureCase,
    n_modes: int | None,
    dt: float | Array,
    modal_inputs: bool = False,
    modal_outputs: bool = False,
    int_order: Literal[3, 4, 5] = BASE_LOBATTO_ORDER,
    prescribed_dofs: Sequence[int] | Array | slice | int | None = None,
):
    if prescribed_dofs is not None:
        prescribed_dofs = beam.make_prescribed_dofs_tuple(prescribed_dofs)
        free_dofs = get_solve_dofs(
            n_dof=beam.n_dof, prescribed_dofs=prescribed_dofs
        )
    else:
        # inherit from the reference by default
        free_dofs = reference.free_dofs
    self.free_dofs: tuple = free_dofs
    self.n_free_dof: int = len(free_dofs)
    self.n_nodes: int = beam.n_nodes
    self.modal_states: bool = n_modes is not None
    self.modal_inputs: bool = modal_inputs
    self.modal_outputs: bool = modal_outputs
    self._n_modes: int | None = n_modes

    if not self.modal_states and (self.modal_inputs or self.modal_outputs):
        raise ValueError(
            "Modal projection for inputs or outputs requires the system states to be modal."
        )

    # compute m, k, mode_shapes so that the superclass has access when initialised
    if self.modal_states:
        assert n_modes is not None
        self.m, self.k, self.mode_shapes = beam.make_modal_m_k(
            case=reference,
            int_order=int_order,
            n_modes=n_modes,
        )
    else:
        self.m, self.k = beam.make_nodal_m_k(case=reference, int_order=int_order)
        self.mode_shapes = None

    # Rayleigh damping matrix
    self.alpha_m: float = beam.alpha_m
    self.beta_k: float = beam.beta_k
    if beam.alpha_m != 0.0 or beam.beta_k != 0.0:
        self.c = beam.alpha_m * self.m + beam.beta_k * self.k
    else:
        self.c = None

    super().__init__(reference=reference, dt=dt)

    self.sys: LinearSystem = self.linearise()

nodal_to_modal

nodal_to_modal(q_nodal: Array) -> Array

Convert a nodal property to a modal property.

Parameters:

Name	Type	Description	Default
q_nodal	Array	Nodal property, (n_free_dof, ).	required

Returns:

Type	Description
Array	Modal property, (n_modes, ).

Source code in src/flapjax/structure/linear/linear_beam.py

def nodal_to_modal(self, q_nodal: Array) -> Array:
    r"""
    Convert a nodal property to a modal property.
    :param q_nodal: Nodal property, ``(n_free_dof, )``.
    :return: Modal property, ``(n_modes, )``.
    """

    return self.mode_shapes @ q_nodal

modal_to_nodal

modal_to_nodal(q_modal: Array) -> Array

Convert a modal property to a nodal property.

Parameters:

Name	Type	Description	Default
q_modal	Array	Mode property, (n_modes, ).	required

Returns:

Type	Description
Array	Nodal property, (n_free_dofs, ).

Source code in src/flapjax/structure/linear/linear_beam.py

def modal_to_nodal(self, q_modal: Array) -> Array:
    r"""
    Convert a modal property to a nodal property.
    :param q_modal: Mode property, ``(n_modes, )``.
    :return: Nodal property, ``(n_free_dofs, )``.
    """

    return q_modal @ self.mode_shapes

linearise_continuous

linearise_continuous() -> LinearSystem

Form a system of linear equations about a reference state. The system is of the form: :math:\dot{\mathbf{x}} = \mathbf{A~x + B~u}, \mathbf{y} = \mathbf{C~x + D~u}.

The state, input and output layouts each depend on their respective modal_* flag:

• States: nodal-free-dof :math:[q, \dot q] of length 2 * n_free_dof when n_modes is None, or modal :math:[q_m, \dot q_m] of length 2 * n_modes when modal_states. Modal mass and stiffness are the projected :math:\Phi M \Phi^T / :math:\Phi K \Phi^T with Phi = mode_shapes of shape [n_modes, n_free_dof] (rows are mode shapes).
• Inputs: a single global-frame external force f_ext. Non-modal inputs are per-node [n_nodes, 6]. Modal inputs are direct modal forces of length n_modes, requiring the user to have already applied the modal projection.
• Outputs: :math:[q, \dot q] in either nodal-free-dof or modal form to match modal_outputs.

Returns:

Type	Description
LinearSystem	Linearised continuous-time system.

Source code in src/flapjax/structure/linear/linear_beam.py

def linearise_continuous(self) -> LinearSystem:
    r"""
    Form a system of linear equations about a reference state. The system is of the form:
    :math:`\dot{\mathbf{x}} = \mathbf{A~x + B~u}, \mathbf{y} = \mathbf{C~x + D~u}`.

    The state, input and output layouts each depend on their respective ``modal_*`` flag:

    * States: nodal-free-dof :math:`[q, \dot q]` of length ``2 * n_free_dof`` when ``n_modes is None``, or modal
      :math:`[q_m, \dot q_m]` of length ``2 * n_modes`` when ``modal_states``. Modal mass and stiffness are the
      projected :math:`\Phi M \Phi^T` / :math:`\Phi K \Phi^T` with ``Phi = mode_shapes`` of shape
      ``[n_modes, n_free_dof]`` (rows are mode shapes).
    * Inputs: a single global-frame external force ``f_ext``. Non-modal inputs are per-node ``[n_nodes, 6]``. Modal
      inputs are direct modal forces of length ``n_modes``, requiring the user to have already applied the modal
      projection.
    * Outputs: :math:`[q, \dot q]` in either nodal-free-dof or modal form to match ``modal_outputs``.

    :return: Linearised continuous-time system.
    """

    # single LU factorisation of M reused for every M^{-1} product below
    m_lu = lu_factor(self.m)  # [q_state_size, q_state_size]
    q_state_size = self.n_modes if self.modal_states else self.n_free_dof

    # dynamics matrix: [q_dot; q_ddot] = A [q; q_dot]
    # includes Raleigh damping contribution if enabled
    m_inv_c = (
        lu_solve(m_lu, self.c)
        if self.c is not None
        else jnp.zeros((q_state_size, q_state_size))
    )
    a = jnp.block(
        [
            [jnp.zeros((q_state_size, q_state_size)), jnp.eye(q_state_size)],
            [-lu_solve(m_lu, self.k), -m_inv_c],
        ]
    )

    # input matrix: maps a single global-frame external force to state derivative
    if self.modal_inputs:
        assert self.modal_states, "modal_inputs requires modal_states"
        # user supplies modal forces directly
        b_bottom = lu_solve(m_lu, jnp.eye(q_state_size))  # (n_modes, n_modes)
    else:
        p_free = jnp.eye(self.n_nodes * 6)[
            jnp.array(self.free_dofs), :
        ]  # (n_free_dof, n_nodes * 6)

        if self.modal_states:
            nodal_to_state = self.nodal_to_modal(p_free)  # (n_modes, n_nodes * 6)
        else:
            nodal_to_state = p_free  # (n_free_dof, n_nodes * 6)
        b_bottom = lu_solve(
            m_lu, nodal_to_state
        )  # (n_modes | n_free_dof, n_nodes * 6)

    b = jnp.concatenate(
        [jnp.zeros((q_state_size, b_bottom.shape[1])), b_bottom], axis=0
    )

    if self.modal_states and not self.modal_outputs:
        assert self.mode_shapes is not None
        state_to_output = (
            self.mode_shapes.T
        )  # projects force to modes, (n_free_dof, n_modes)
    else:
        state_to_output = jnp.eye(
            q_state_size
        )  # direct projection, (n_free_dof, n_free_dof) or (n_modes, n_modes)

    c = jnp.block(
        [
            [state_to_output, jnp.zeros_like(state_to_output)],
            [jnp.zeros_like(state_to_output), state_to_output],
        ]
    )  # pass state (q, q_dot) to output (q, q_dot)
    d = jnp.zeros((c.shape[0], b.shape[1]))  # no feedthrough

    return LinearSystem(
        a=a, b=b, c=c, d=d, dt=self.dt, continuous_time=True, removed_u_np1=False
    )

run

run(
    u: StructureInputUnflattened,
    x0: StructureStateUnflattened | None = None,
) -> StructureLinearResult

Run the linear system.

Parameters:

Name	Type	Description	Default
u	StructureInputUnflattened	Total input over time (reference + pertubation).	required
x0	StructureStateUnflattened | None	Initial state perturbations, defaults to zero state.	None

Returns:

Type	Description
StructureLinearResult	Linear system results.

Source code in src/flapjax/structure/linear/linear_beam.py

def run(
    self,
    u: StructureInputUnflattened,
    x0: StructureStateUnflattened | None = None,
) -> StructureLinearResult:
    r"""
    Run the linear system.
    :param u: Total input over time (reference + pertubation).
    :param x0: Initial state perturbations, defaults to zero state.
    :return: Linear system results.
    """

    assert (
        isinstance(self.reference, StructureCase) and not self.reference.is_dynamic
    ), "Reference structure must be a static Structure"

    # has to be specified in the input object, as the linear system does not know how many time steps to run for in
    # the case where there is no external forcing
    n_tstep = u.n_tstep

    q_state_size = self.n_modes if self.modal_states else self.n_free_dof

    # initial states
    if x0 is None:
        q0_ = None
        q0_dot_ = None
    else:
        q0_ = x0.q
        q0_dot_ = x0.q_dot

    def _project_initial(q0_arr: Array | None) -> Array:
        if q0_arr is None:
            return jnp.zeros((q_state_size,))

        match q0_arr.shape:
            case (self.n_nodes, 6):
                q0_nodal = q0_arr.ravel()[self.free_dofs]  # remove prescribed dofs
                if self.modal_states:
                    q0__ = self.nodal_to_modal(q0_nodal)
                else:
                    q0__ = q0_nodal
            case (self.n_modes,):
                if not self.modal_states:
                    raise NotImplementedError(
                        "Inputs cannot be modal for a nodal system"
                    )
                q0__ = q0_arr
            case _:
                raise ValueError("Invalid input for initial states")
        return q0__

    q0 = _project_initial(q0_)  # (n_free_dof | n_modes, )
    q0_dot = _project_initial(q0_dot_)  # (n_free_dof | n_modes, )

    if q0.shape != q0_dot.shape:
        raise ValueError("q0 and q0_dot must both be modal or nodal")

    ref_f_ext = self._reference_inputs["f_ext"]
    assert ref_f_ext is None or isinstance(
        ref_f_ext, Array
    )  # type narrowing as it cannot be an ArrayList

    if self.modal_inputs:
        # user-provided input is a modal force time history of shape (n_tstep, n_modes)
        if u.f_ext is None:
            delta_f_ext_t = (
                jnp.zeros((n_tstep, self.n_modes))
                if ref_f_ext is None
                else -jnp.broadcast_to(ref_f_ext[None, :], (n_tstep, self.n_modes))
            )
        else:
            check_arr_shape(u.f_ext, (n_tstep, self.n_modes), name="f_ext_t")
            delta_f_ext_t = (
                u.f_ext if ref_f_ext is None else u.f_ext - ref_f_ext[None, :]
            )
    else:
        # user-provided input is a nodal global-frame force time history of shape (n_tstep, n_nodes, 6)
        if u.f_ext is None:
            delta_f_ext_t = (
                jnp.zeros((n_tstep, self.n_nodes, 6))
                if ref_f_ext is None
                else -jnp.broadcast_to(
                    ref_f_ext[None, :, :], (n_tstep, self.n_nodes, 6)
                )
            )
        else:
            check_arr_shape(u.f_ext, (n_tstep, self.n_nodes, 6), name="f_ext_t")
            delta_f_ext_t = (
                u.f_ext if ref_f_ext is None else u.f_ext - ref_f_ext[None, :, :]
            )

    delta_u = StructureInputUnflattened(n_tstep=n_tstep, f_ext=delta_f_ext_t)
    delta_u_vec = self._pack_input_vector_t(delta_u)

    # run linear system
    x_t, _ = self.sys.run(
        u=delta_u_vec,
        x0=jnp.concatenate((q0, q0_dot)),
    )

    # extract perturbations in displacements and velocities, reconstructing nodal form if requested
    delta_q_state = x_t[:, :q_state_size]
    delta_q_dot_state = x_t[:, q_state_size:]

    if self.modal_outputs:
        delta_q_t = delta_q_state
        delta_q_dot_t = delta_q_dot_state
        hg_t = None  # do not reconstruct coordinates
    else:
        if self.modal_states:
            assert self.mode_shapes is not None
            delta_q_t = self.modal_to_nodal(delta_q_state)  # (n_tstep, n_free_dof)
            delta_q_dot_t = self.modal_to_nodal(delta_q_dot_state)
        else:
            delta_q_t = delta_q_state
            delta_q_dot_t = delta_q_dot_state

        # add in zeros for dofs not solved for to allow for reconstructing the full configuration
        delta_q_t_full = (
            jnp.zeros((n_tstep, self.n_nodes * 6))
            .at[:, self.free_dofs]
            .set(delta_q_t)
            .reshape(n_tstep, self.n_nodes, 6)
        )

        delta_hg_t = vmap(vmap(exp_se3, 0, 0), 1, 1)(
            delta_q_t_full
        )  # (n_tstep, n_nodes, 4, 4)

        hg_t = jnp.einsum("ijk,hikl->hijl", self.reference.hg, delta_hg_t)

    return StructureLinearResult(
        reference=self.reference,
        f_ext=u.f_ext,
        delta_q=delta_q_t,
        delta_q_dot=delta_q_dot_t,
        hg=hg_t,
        t=jnp.arange(n_tstep) * self.dt,
    )

unpack_state_vector

unpack_state_vector(x: Array) -> S

Unpack a state vector into its components.

Parameters:

Name	Type	Description	Default
x	Array	State vector, (n_states, )	required

Returns:

Type	Description
S	StateUnflattened object.

Source code in src/flapjax/utils/linear.py

def unpack_state_vector(self, x: Array) -> S:
    r"""
    Unpack a state vector into its components.
    :param x: State vector, ``(n_states, )``
    :return: StateUnflattened object.
    """
    return self.state_object(**self._unpack_vector(x=x, slices=self.state_slices))

unpack_output_vector

unpack_output_vector(y: Array) -> O

Unpack an output vector into its components.

Parameters:

Name	Type	Description	Default
y	Array	Output vector, (n_outputs, )	required

Returns:

Type	Description
O	OutputUnflattened object.

Source code in src/flapjax/utils/linear.py

def unpack_output_vector(self, y: Array) -> O:
    r"""
    Unpack an output vector into its components.
    :param y: Output vector, ``(n_outputs, )``
    :return: OutputUnflattened object.
    """
    return self.output_object(**self._unpack_vector(x=y, slices=self.output_slices))

pack_input_vector

pack_input_vector(u_input: InputUnflattened) -> Array

Pack an input unflattened object into a vector.

Parameters:

Name	Type	Description	Default
u_input	InputUnflattened	InputUnflattened object.	required

Returns:

Type	Description
Array	Input vector, (n_inputs, ).

Source code in src/flapjax/utils/linear.py

def pack_input_vector(self, u_input: InputUnflattened) -> Array:
    r"""
    Pack an input unflattened object into a vector.
    :param u_input: InputUnflattened object.
    :return: Input vector, ``(n_inputs, )``.
    """
    return self._pack_vector(
        slices=self.input_slices,
        vec_length=self.n_inputs,
        arrs=shallow_as_dict(u_input),
    )

pack_state_vector

pack_state_vector(x_state: StateUnflattened) -> Array

Pack a state unflattened object into a vector.

Parameters:

Name	Type	Description	Default
x_state	StateUnflattened	StateUnflattened object.	required

Returns:

Type	Description
Array	State vector, (n_states, ).

Source code in src/flapjax/utils/linear.py

def pack_state_vector(self, x_state: StateUnflattened) -> Array:
    r"""
    Pack a state unflattened object into a vector.
    :param x_state: StateUnflattened object.
    :return: State vector, ``(n_states, )``.
    """
    return self._pack_vector(
        slices=self.state_slices,
        vec_length=self.n_states,
        arrs=shallow_as_dict(x_state),
    )

pack_output_vector

pack_output_vector(y_output: OutputUnflattened) -> Array

Pack an output unflattened object into a vector.

Parameters:

Name	Type	Description	Default
y_output	OutputUnflattened	OutputUnflattened object.	required

Returns:

Type	Description
Array	Output vector, (n_outputs, ).

Source code in src/flapjax/utils/linear.py

def pack_output_vector(self, y_output: OutputUnflattened) -> Array:
    r"""
    Pack an output unflattened object into a vector.
    :param y_output: OutputUnflattened object.
    :return: Output vector, ``(n_outputs, )``.
    """
    return self._pack_vector(
        slices=self.output_slices,
        vec_length=self.n_outputs,
        arrs=shallow_as_dict(y_output),
    )

get_total_input

get_total_input(u: InputUnflattened) -> InputUnflattened

Get the total input by adding the reference to the input perturbation.

Parameters:

Name	Type	Description	Default
u	InputUnflattened	InputUnflattened perturbation object.	required

Returns:

Type	Description
InputUnflattened	InputUnflattened total object.

Source code in src/flapjax/utils/linear.py

def get_total_input(self, u: InputUnflattened) -> InputUnflattened:
    r"""
    Get the total input by adding the reference to the input perturbation.
    :param u: InputUnflattened perturbation object.
    :return: InputUnflattened total object.
    """
    return self.input_object(
        **self._get_total(
            shallow_as_dict(u), shallow_as_dict(self.reference_inputs)
        )
    )

get_total_state

get_total_state(x: StateUnflattened) -> StateUnflattened

Get the total state by adding the reference to the state perturbation.

Parameters:

Name	Type	Description	Default
x	StateUnflattened	StateUnflattened perturbation object.	required

Returns:

Type	Description
StateUnflattened	StateUnflattened total object.

Source code in src/flapjax/utils/linear.py

def get_total_state(self, x: StateUnflattened) -> StateUnflattened:
    r"""
    Get the total state by adding the reference to the state perturbation.
    :param x: StateUnflattened perturbation object.
    :return: StateUnflattened total object.
    """
    return self.state_object(
        **self._get_total(
            shallow_as_dict(x), shallow_as_dict(self.reference_states)
        )
    )

get_total_output

get_total_output(y: OutputUnflattened) -> OutputUnflattened

Get the total output by adding the reference to the output perturbation.

Parameters:

Name	Type	Description	Default
y	OutputUnflattened	OutputUnflattened perturbation object.	required

Returns:

Type	Description
OutputUnflattened	OutputUnflattened total object.

Source code in src/flapjax/utils/linear.py

def get_total_output(self, y: OutputUnflattened) -> OutputUnflattened:
    r"""
    Get the total output by adding the reference to the output perturbation.
    :param y: OutputUnflattened perturbation object.
    :return: OutputUnflattened total object.
    """
    return self.output_object(
        **self._get_total(
            shallow_as_dict(y), shallow_as_dict(self.reference_outputs)
        )
    )

get_total_input_t

get_total_input_t(
    u_t: InputUnflattened,
) -> InputUnflattened

Get the total input time history by adding the reference to the input perturbation time history.

Parameters:

Name	Type	Description	Default
u_t	InputUnflattened	InputUnflattened perturbation time history object.	required

Returns:

Type	Description
InputUnflattened	InputUnflattened total time history object.

Source code in src/flapjax/utils/linear.py

def get_total_input_t(self, u_t: InputUnflattened) -> InputUnflattened:
    r"""
    Get the total input time history by adding the reference to the input perturbation time history.
    :param u_t: InputUnflattened perturbation time history object.
    :return: InputUnflattened total time history object.
    """
    return self.input_object(
        **self._get_total(
            input_=shallow_as_dict(u_t),
            reference=shallow_as_dict(self.reference_inputs),
            add_t=True,
        )
    )

get_total_state_t

get_total_state_t(
    x_t: StateUnflattened,
) -> StateUnflattened

Get the total state time history by adding the reference to the state perturbation time history.

Parameters:

Name	Type	Description	Default
x_t	StateUnflattened	StateUnflattened perturbation time history object.	required

Returns:

Type	Description
StateUnflattened	StateUnflattened total time history object.

Source code in src/flapjax/utils/linear.py

def get_total_state_t(self, x_t: StateUnflattened) -> StateUnflattened:
    r"""
    Get the total state time history by adding the reference to the state perturbation time history.
    :param x_t: StateUnflattened perturbation time history object.
    :return: StateUnflattened total time history object.
    """
    return self.state_object(
        **self._get_total(
            shallow_as_dict(x_t),
            shallow_as_dict(self.reference_states),
            add_t=True,
        )
    )

get_total_output_t

get_total_output_t(
    y_t: OutputUnflattened,
) -> OutputUnflattened

Get the total output time history by adding the reference to the output perturbation time history.

Parameters:

Name	Type	Description	Default
y_t	OutputUnflattened	OutputUnflattened perturbation time history object.	required

Returns:

Type	Description
OutputUnflattened	OutputUnflattened total time history object.

Source code in src/flapjax/utils/linear.py

def get_total_output_t(self, y_t: OutputUnflattened) -> OutputUnflattened:
    r"""
    Get the total output time history by adding the reference to the output perturbation time history.
    :param y_t: OutputUnflattened perturbation time history object.
    :return: OutputUnflattened total time history object.
    """
    return self.output_object(
        **self._get_total(
            shallow_as_dict(y_t),
            shallow_as_dict(self.reference_outputs),
            add_t=True,
        )
    )

get_zero_input

get_zero_input() -> InputUnflattened

Get a zero input unflattened object.

Returns:

Type	Description
InputUnflattened	InputUnflattened object with zero arrays.

Source code in src/flapjax/utils/linear.py

def get_zero_input(self) -> InputUnflattened:
    r"""
    Get a zero input unflattened object.
    :return: InputUnflattened object with zero arrays.
    """
    return self.input_object(**self._get_zero(shallow_as_dict(self.input_slices)))

get_zero_state

get_zero_state() -> StateUnflattened

Get a zero state unflattened object.

Returns:

Type	Description
StateUnflattened	StateUnflattened object with zero arrays.

Source code in src/flapjax/utils/linear.py

def get_zero_state(self) -> StateUnflattened:
    r"""
    Get a zero state unflattened object.
    :return: StateUnflattened object with zero arrays.
    """
    return self.state_object(**self._get_zero(shallow_as_dict(self.state_slices)))

get_zero_output

get_zero_output() -> OutputUnflattened

Get a zero output unflattened object.

Returns:

Type	Description
OutputUnflattened	OutputUnflattened object with zero arrays.

Source code in src/flapjax/utils/linear.py

def get_zero_output(self) -> OutputUnflattened:
    r"""
    Get a zero output unflattened object.
    :return: OutputUnflattened object with zero arrays.
    """
    return self.output_object(**self._get_zero(shallow_as_dict(self.output_slices)))

time_integration

TimeIntegrator

TimeIntegrator(spectral_radius: float, dt: Array)

Container for time integration parameters.

Source code in src/flapjax/structure/time_integration.py

def __init__(
    self,
    spectral_radius: float,
    dt: Array,
):
    if 1.0 < spectral_radius < 0.0:
        warn(
            "Spectral radius should be between 0.0 and 1.0 to guarantee stability."
        )
    self.spectral_radius: float = spectral_radius
    self.dt: Array = dt
    self.spectral_radius: float = spectral_radius
    self.alpha_m: float = (2.0 * spectral_radius - 1.0) / (spectral_radius + 1.0)
    self.alpha_f: float = spectral_radius / (spectral_radius + 1.0)
    self.gamma: float = (3.0 - spectral_radius) / (2.0 + 2.0 * spectral_radius)
    self.beta: float = 1.0 / ((spectral_radius + 1.0) ** 2)
    self.gamma_prime: Array = self.gamma / (self.beta * dt)
    self.beta_prime: Array = (1.0 - self.alpha_m) / (
        self.beta * dt * dt * (1.0 - self.alpha_f)
    )

compute_a_n

compute_a_n(
    v_dot_nm1: Array, v_dot_n: Array, a_nm1: Array
) -> Array

Calculate the pseudo-acceleration at the next time step.

Parameters:

Name	Type	Description	Default
v_dot_nm1	Array	Previous acceleration, (n_nodes, 6).	required
v_dot_n	Array	Next acceleration, (n_nodes, 6).	required
a_nm1	Array	Previous pseudo-acceleration, (n_nodes, 6).	required

Returns:

Type	Description
Array	pseudo-acceleration at next time step, (n_nodes, 6).

Source code in src/flapjax/structure/time_integration.py

def compute_a_n(self, v_dot_nm1: Array, v_dot_n: Array, a_nm1: Array) -> Array:
    r"""
    Calculate the pseudo-acceleration at the next time step.
    :param v_dot_nm1: Previous acceleration, ``(n_nodes, 6)``.
    :param v_dot_n: Next acceleration, ``(n_nodes, 6)``.
    :param a_nm1: Previous pseudo-acceleration, ``(n_nodes, 6)``.
    :return: pseudo-acceleration at next time step, ``(n_nodes, 6)``.
    """
    return (
        1.0
        / (1.0 - self.alpha_m)
        * (
            (1.0 - self.alpha_f) * v_dot_n
            + self.alpha_f * v_dot_nm1
            - self.alpha_m * a_nm1
        )
    )

predict_q

predict_q(
    q_nm1: StructureMinimalStates,
) -> tuple[Array, StructureMinimalStates]

Predict the current state based upon the previous state.

Parameters:

Name	Type	Description	Default
q_nm1	StructureMinimalStates	State at timestep n	required

Returns:

Type	Description
tuple[Array, StructureMinimalStates]	Predicted state at timestep n+1

Source code in src/flapjax/structure/time_integration.py

def predict_q(
    self, q_nm1: StructureMinimalStates
) -> tuple[Array, StructureMinimalStates]:
    r"""
    Predict the current state based upon the previous state.
    :param q_nm1: State at timestep n
    :return: Predicted state at timestep n+1
    """
    a_n = (self.alpha_f * q_nm1.v_dot - self.alpha_m * q_nm1.a) / (
        1.0 - self.alpha_m
    )

    phi_n = self.dt * q_nm1.v + self.dt * self.dt * (
        (0.5 - self.beta) * q_nm1.a + self.beta * a_n
    )

    varphi_n = vmap(
        lambda varphi_, phi_: log_se3(exp_se3(varphi_) @ exp_se3(phi_)), 0, 0
    )(q_nm1.varphi, phi_n)

    v_n = (
        q_nm1.v
        + (1.0 - self.gamma) * self.dt * q_nm1.a
        + self.gamma * self.dt * a_n
    )

    v_dot_n = (
        (1.0 - self.alpha_m) * a_n
        + self.alpha_m * q_nm1.a
        - self.alpha_f * q_nm1.v_dot
    ) / (1.0 - self.alpha_f)

    return phi_n, StructureMinimalStates(
        varphi=varphi_n, v=v_n, v_dot=v_dot_n, a=a_n
    )

compute_phi_from_phi_alpha

compute_phi_from_phi_alpha(phi_alpha: Array) -> Array

Obtain the full timestep increment from the alpha increment.

Parameters:

Name	Type	Description	Default
phi_alpha	Array	Increment from timestep n-1 to alpha, (n_nodes, 6).	required

Returns:

Type	Description
Array	Increment for timestep n, (n_nodes, 6).

Source code in src/flapjax/structure/time_integration.py

def compute_phi_from_phi_alpha(self, phi_alpha: Array) -> Array:
    r"""
    Obtain the full timestep increment from the alpha increment.
    :param phi_alpha: Increment from timestep n-1 to alpha, ``(n_nodes, 6)``.
    :return: Increment for timestep n, ``(n_nodes, 6)``.
    """
    return phi_alpha / (1.0 - self.alpha_f)

compute_v_from_v_alpha

compute_v_from_v_alpha(
    v_alpha: Array, v_nm1: Array
) -> Array

Obtain the full timestep velocity from the alpha increment and the previous velocity.

Parameters:

Name	Type	Description	Default
v_alpha	Array	Velocity at alpha step, (n_nodes, 6).	required
v_nm1	Array	Velocity at timestep n-1, (n_nodes, 6).	required

Returns:

Type	Description
Array	Velocity at timestep n, (n_nodes, 6).

Source code in src/flapjax/structure/time_integration.py

def compute_v_from_v_alpha(self, v_alpha: Array, v_nm1: Array) -> Array:
    r"""
    Obtain the full timestep velocity from the alpha increment and the previous velocity.
    :param v_alpha: Velocity at alpha step, ``(n_nodes, 6)``.
    :param v_nm1: Velocity at timestep n-1, ``(n_nodes, 6)``.
    :return: Velocity at timestep n, ``(n_nodes, 6)``.
    """

    return (v_alpha - self.alpha_f * v_nm1) / (1.0 - self.alpha_f)

utils

get_solve_dofs

get_solve_dofs(
    n_dof: int, prescribed_dofs: tuple[int, ...]
) -> tuple[int, ...]

Obtain the index of degrees of freedom to solve for, given the index of prescribed degrees of freedom.

Parameters:

Name	Type	Description	Default
n_dof	int	Total number of degrees of freedom	required
prescribed_dofs	tuple[int, ...]	Index of prescribed degrees of freedom	required

Returns:

Type	Description
tuple[int, ...]	Index of degrees of freedom to solve for

Source code in src/flapjax/structure/utils.py

def get_solve_dofs(n_dof: int, prescribed_dofs: tuple[int, ...]) -> tuple[int, ...]:
    r"""
    Obtain the index of degrees of freedom to solve for, given the index of prescribed degrees of freedom.
    :param n_dof: Total number of degrees of freedom
    :param prescribed_dofs: Index of prescribed degrees of freedom
    :return: Index of degrees of freedom to solve for
    """
    return tuple(sorted(set(range(n_dof)) - set(prescribed_dofs)))

transform_nodal_vect

transform_nodal_vect(vect: Array, rmat: Array) -> Array

Rotate a nodal vector quantity.

Parameters:

Name	Type	Description	Default
vect	Array	Nodal vectors, (..., n_nodes, 6)	required
rmat	Array	Rotation matrix, (..., n_nodes, 3, 3)	required

Returns:

Type	Description
Array	Rotated vectors, (..., n_nodes, 6)

Source code in src/flapjax/structure/utils.py

def transform_nodal_vect(vect: Array, rmat: Array) -> Array:
    r"""
    Rotate a nodal vector quantity.
    :param vect: Nodal vectors, ``(..., n_nodes, 6)``
    :param rmat: Rotation matrix, ``(..., n_nodes, 3, 3)``
    :return: Rotated vectors, ``(..., n_nodes, 6)``
    """

    idx = "...ijk,...ik->...ij"
    vect_lin = jnp.einsum(idx, rmat, vect[..., :3])
    vect_rot = jnp.einsum(idx, rmat, vect[..., 3:])
    return jnp.concatenate((vect_lin, vect_rot), axis=-1)

apply_frame_transform

apply_frame_transform(
    obj, rmat: Array, extra_fields: tuple[str, ...] = ()
) -> None

Rotate all force / velocity fields of a structure state object in place.

Source code in src/flapjax/structure/utils.py

def apply_frame_transform(obj, rmat: Array, extra_fields: tuple[str, ...] = ()) -> None:
    r"""
    Rotate all force / velocity fields of a structure state object in place.
    """
    for name in _OPTIONAL_FORCE_FIELDS:
        if getattr(obj, name) is not None:
            setattr(obj, name, transform_nodal_vect(getattr(obj, name), rmat))
    for name in ("f_int", "f_res", *extra_fields):
        setattr(obj, name, transform_nodal_vect(getattr(obj, name), rmat))