Skip to content

Coupled

flapjax.coupled

AeroelasticCase 

AeroelasticCase(structure: StructureCase, aero: AeroCase)

Coupled aeroelastic case, which is a wrapper around a StructureCase and a AeroCase pair.

A single instance may represent any of three flavours (derived from the wrapped structure):

  • Static: static snapshot structure and snapshot aero.
  • Dynamic snapshot: dynamic snapshotstructure and snapshot aero.
  • Dynamic trajectory: batched structure and batched aero.

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

Source code in src/flapjax/coupled/data_structures.py 
46
47
48
def __init__(self, structure: StructureCase, aero: AeroCase):
    self.structure: StructureCase = structure
    self.aero: AeroCase = aero

to_dynamic 

to_dynamic() -> AeroelasticCase

to_dynamic(t: None) -> AeroelasticCase

to_dynamic(t: Array) -> AeroelasticCase

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

Convert a static snapshot to a dynamic snapshot (t=None) or a batched trajectory (t provided). Calling on an already-dynamic case returns self.

Source code in src/flapjax/coupled/data_structures.py 
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
def to_dynamic(self, t: Array | None = None) -> AeroelasticCase:
    """Convert a static snapshot to a dynamic snapshot (``t=None``) or a batched
    trajectory (``t`` provided). Calling on an already-dynamic case returns
    ``self``.
    """
    if self.is_dynamic:
        return self
    if t is None:
        return AeroelasticCase(
            structure=self.structure.to_dynamic(), aero=self.aero
        )
    return AeroelasticCase(
        structure=self.structure.to_dynamic(t),
        aero=self.aero.to_dynamic(i_ts=0, n_tstep=len(t)),
    )

to_static 

to_static() -> AeroelasticCase

Return a static AeroelasticCase, dropping the structure's velocity/acceleration fields.

Source code in src/flapjax/coupled/data_structures.py 
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
def to_static(self) -> AeroelasticCase:
    """Return a static AeroelasticCase, dropping the structure's
    velocity/acceleration fields.
    """
    if not self.is_dynamic:
        return self
    if self.is_batched:
        raise ValueError(
            "to_static() on a batched AeroelasticCase is ambiguous; index a "
            "single time step first (e.g. `case[i_ts].to_static()`)."
        )
    return AeroelasticCase(
        structure=self.structure.to_static(),
        aero=self.aero,
    )

initialise classmethod 

initialise(
    initial_snapshot: AeroelasticCase,
    t: Array,
    use_f_ext_follower: bool,
    use_f_ext_dead: bool,
    structure: BeamStructure,
    x0_aero: ArrayList,
) -> AeroelasticCase

Build a batched AeroelasticCase from any single-timestep case.

Source code in src/flapjax/coupled/data_structures.py 
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
@classmethod
def initialise(
    cls,
    initial_snapshot: AeroelasticCase,
    t: Array,
    use_f_ext_follower: bool,
    use_f_ext_dead: bool,
    structure: BeamStructure,
    x0_aero: ArrayList,
) -> AeroelasticCase:
    """Build a batched AeroelasticCase from any single-timestep case."""
    if initial_snapshot.is_batched:
        if initial_snapshot.structure.n_tstep != 1:
            raise ValueError("initial_snapshot.structure.n_tstep != 1")
        if initial_snapshot.aero.n_tstep != 1:
            raise ValueError("initial_snapshot.aero.n_tstep != 1")
        init_struct: StructureCase = initial_snapshot.structure[0]
        init_aero: AeroCase = initial_snapshot.aero[0]
    elif not initial_snapshot.is_dynamic:
        init_struct = initial_snapshot.structure.to_dynamic(t=None)
        init_aero = initial_snapshot.aero
    else:
        init_struct = initial_snapshot.structure
        init_aero = initial_snapshot.aero

    struct_case = StructureCase.initialise(
        initial_snapshot=init_struct,
        t=t,
        use_f_ext_aero=True,
        use_f_ext_follower=use_f_ext_follower,
        use_f_ext_dead=use_f_ext_dead,
    )
    aero_case = AeroCase.initialise(initial_snapshot=init_aero, n_tstep=len(t))

    # compute aerodynamic forcing at timestep 0
    f_aero_init = aero_case.project_forcing_to_beam(
        i_ts=0,
        rmat=struct_case.hg[0, :, :3, :3],
        x0_aero=x0_aero,
        include_unsteady=False,
    )
    f_aero_local = structure.make_f_dead_ext(
        f_ext=f_aero_init, rmat=struct_case.hg[0, :, :3, :3]
    )

    if struct_case.f_ext_aero is None:
        raise ValueError("f_ext_aero cannot be None")

    struct_case.f_ext_aero = struct_case.f_ext_aero.at[0, ...].set(f_aero_local)

    return AeroelasticCase(structure=struct_case, aero=aero_case)

get_full_states 

get_full_states(
    i_ts: int | Array | None = None,
) -> AeroelasticFullStates

Get the full aeroelastic states. For a batched case, i_ts selects the timestep; for a snapshot, i_ts is ignored.

Source code in src/flapjax/coupled/data_structures.py 
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
def get_full_states(self, i_ts: int | Array | None = None) -> AeroelasticFullStates:
    r"""
    Get the full aeroelastic states. For a batched case, ``i_ts`` selects the
    timestep; for a snapshot, ``i_ts`` is ignored.
    """
    if self.is_batched:
        if i_ts is None:
            raise ValueError("i_ts must be provided for batched AeroelasticCase")
        return AeroelasticFullStates(
            structure=self.structure.get_full_states(i_ts=i_ts),
            aero=self.aero.get_states(i_ts=i_ts),
        )
    if self.structure.f_ext_aero is None:
        raise ValueError("f_ext_aero is None")
    return AeroelasticFullStates(
        structure=self.structure.get_full_states(),
        aero=self.aero.get_states(i_ts=0 if self.is_dynamic else None),
    )

get_minimal_states 

get_minimal_states(
    i_ts: int | Array,
) -> AeroelasticMinimalStates

Get minimal aeroelastic states at the given timestep (batched only).

Source code in src/flapjax/coupled/data_structures.py 
176
177
178
179
180
181
182
183
184
185
def get_minimal_states(self, i_ts: int | Array) -> AeroelasticMinimalStates:
    r"""Get minimal aeroelastic states at the given timestep (batched only)."""
    if not self.is_batched:
        raise TypeError(
            "get_minimal_states only supported for batched AeroelasticCase"
        )
    return AeroelasticMinimalStates(
        structure=self.structure.get_minimal_states(i_ts=i_ts),
        aero=self.aero.get_states(i_ts=i_ts),
    )

plot 

plot(
    directory: PathLike | str,
    index: int
    | Sequence[int]
    | Array
    | slice
    | None = None,
    n_interp: int = 0,
    plot_bound: bool = True,
    plot_wake: bool = True,
) -> tuple[Path, Sequence[Path]]

Plot the aeroelastic case.

Parameters:

Name Type Description Default
directory PathLike | str

Directory to save the plots to.

 required
index int | Sequence[int] | Array | slice | None

For batched cases, timestep indices to plot; ignored for snapshots.

 None
n_interp int

Number of interpolation points for plotting the structure.

 0
plot_bound bool

Whether to plot the bound aerodynamic panels.

 True
plot_wake bool

Whether to plot the wake aerodynamic panels.

 True

Returns:

Type Description
tuple[Path, Sequence[Path]]

(structure_path, aero_paths). For batched cases the structure path is a PVD; for snapshots it is a single VTU.
Source code in src/flapjax/coupled/data_structures.py 
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
def plot(
    self,
    directory: os.PathLike | str,
    index: int | Sequence[int] | Array | slice | None = None,
    n_interp: int = 0,
    plot_bound: bool = True,
    plot_wake: bool = True,
) -> tuple[Path, Sequence[Path]]:
    r"""Plot the aeroelastic case.
    :param directory: Directory to save the plots to.
    :param index: For batched cases, timestep indices to plot; ignored for
        snapshots.
    :param n_interp: Number of interpolation points for plotting the structure.
    :param plot_bound: Whether to plot the bound aerodynamic panels.
    :param plot_wake: Whether to plot the wake aerodynamic panels.
    :return: ``(structure_path, aero_paths)``. For batched cases the
        structure path is a PVD; for snapshots it is a single VTU.
    """
    if self.is_batched:
        struct_out: Path = self.structure.plot(
            directory=directory, n_interp=n_interp, index=index
        )
    else:
        struct_out = self.structure.plot(directory=directory, n_interp=n_interp)
    aero_out: Sequence[Path] = self.aero.plot(
        directory=directory,  # type: ignore
        plot_bound=plot_bound,
        plot_wake=plot_wake,
        index=index,
    )
    return struct_out, aero_out

CoupledAeroelastic 

CoupledAeroelastic(
    structure: BeamStructure,
    aero: UVLM,
    fsi_convergence_settings: ConvergenceSettings = DEFAULT_FSI_CONVERGENCE_SETTINGS,
)

Bases: BaseCoupledAeroelastic

Source code in src/flapjax/coupled/coupled.py 
53
54
55
56
57
58
59
60
61
def __init__(
    self,
    structure: BeamStructure,
    aero: UVLM,
    fsi_convergence_settings: ConvergenceSettings = DEFAULT_FSI_CONVERGENCE_SETTINGS,
):
    self.structure: BeamStructure = structure
    self.aero: UVLM = aero
    self.fsi_convergence_settings: ConvergenceSettings = fsi_convergence_settings

get_design_variables 

get_design_variables(
    case: AeroelasticCase,
    grads_to_compute: AeroelasticGradsToCompute | None,
) -> AeroelasticDesignVariables

Obtain the design variables describing the wing.

Parameters:

Name Type Description Default
case AeroelasticCase
required
grads_to_compute AeroelasticGradsToCompute | None

Data structure which describes which design variables should be obtained. If none, all variables are obtained.

 required

Returns:

Type Description
AeroelasticDesignVariables

AeroelasticDesignVariables object.
Source code in src/flapjax/coupled/coupled.py 
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
def get_design_variables(
    self,
    case: AeroelasticCase,
    grads_to_compute: AeroelasticGradsToCompute | None,
) -> AeroelasticDesignVariables:
    r"""
    Obtain the design variables describing the wing.
    :param case:
    :param grads_to_compute: Data structure which describes which design variables should be obtained. If none, all
    variables are obtained.
    :return: AeroelasticDesignVariables object.
    """
    return AeroelasticDesignVariables(
        structure_dv=self.structure.get_design_variables(
            struct_case=case.structure,
            thrust_t=case.structure.thrust,
            grads_to_compute=grads_to_compute.structure
            if grads_to_compute is not None
            else None,
        ),
        aero_dv=self.aero.get_design_variables(
            cs_ang_t=case.aero.cs_ang,
            cs_vel_t=case.aero.cs_vel,
            grads_to_compute=grads_to_compute.aero
            if grads_to_compute is not None
            else None,
        ),
    )

reference_configuration 

reference_configuration(
    prescribed_dofs: Sequence[int]
    | Array
    | slice
    | int = (),
    horseshoe: bool = False,
    use_f_ext_follower: bool = False,
    use_f_ext_dead: bool = False,
    t_init: float | Array = 0.0,
) -> AeroelasticCase

Obtain the static aeroelastic object describing the undeformed wing.

Parameters:

Name Type Description Default
prescribed_dofs Sequence[int] | Array | slice | int

Prescribed dofs for the structure. Defaults to no prescribed dofs.

 ()
horseshoe bool

Horseshoe flag.

 False
use_f_ext_follower bool

If true, allocate an array for follower forces.

 False
use_f_ext_dead bool

If true, allocate an array for dead forces.

 False
t_init float | Array

Initial time

 0.0

Returns:

Type Description
AeroelasticCase

Static aeroelastic object for undeformed wing
Source code in src/flapjax/coupled/coupled.py 
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
def reference_configuration(
    self,
    prescribed_dofs: Sequence[int] | Array | slice | int = (),
    horseshoe: bool = False,
    use_f_ext_follower: bool = False,
    use_f_ext_dead: bool = False,
    t_init: float | Array = 0.0,
) -> AeroelasticCase:
    r"""
    Obtain the static aeroelastic object describing the undeformed wing.
    :param prescribed_dofs: Prescribed dofs for the structure. Defaults to no prescribed dofs.
    :param horseshoe: Horseshoe flag.
    :param use_f_ext_follower: If true, allocate an array for follower forces.
    :param use_f_ext_dead: If true, allocate an array for dead forces.

    :param t_init: Initial time
    :return: Static aeroelastic object for undeformed wing
    """
    prescribed_dofs = self.structure.make_prescribed_dofs_tuple(prescribed_dofs)
    return AeroelasticCase(
        structure=self.structure.reference_configuration(
            use_f_grav=self.structure.use_gravity,
            use_f_ext_dead=use_f_ext_dead,
            use_f_ext_follower=use_f_ext_follower,
            use_f_aero=True,
            prescribed_dofs=prescribed_dofs,
        ),
        aero=self.aero.static_solve(
            t=t_init, hg=self.structure.hg0, horseshoe=horseshoe
        ),
    )

initialise_dynamic 

initialise_dynamic(
    static_case: AeroelasticCase,
    prescribed_dofs: Sequence[int] | Array | slice | int,
) -> AeroelasticCase

Initialise a dynamic aeroelastic snapshot from a static aeroelastic case. This takes a static aeroelastic case obtained under clamped conditions (i.e. relative_motion is True for the free stream), and sets the structural velocity to be that of the freestream.

Parameters:

Name Type Description Default
static_case AeroelasticCase

Static aeroelastic case.

 required
prescribed_dofs Sequence[int] | Array | slice | int

Prescribed dofs. This is often useful for updating from a clamped trim to a free-flying dynamic case.

 required

Returns:

Type Description
AeroelasticCase

Dynamic aeroelastic snapshot.
Source code in src/flapjax/coupled/coupled.py 
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
def initialise_dynamic(
    self,
    static_case: AeroelasticCase,
    prescribed_dofs: Sequence[int] | Array | slice | int,
) -> AeroelasticCase:
    r"""
    Initialise a dynamic aeroelastic snapshot from a static aeroelastic case. This takes a static aeroelastic case
    obtained under clamped conditions (i.e. `relative_motion` is True for the free stream), and sets the structural
    velocity to be that of the freestream.
    :param static_case: Static aeroelastic case.
    :param prescribed_dofs: Prescribed dofs. This is often useful for updating from a clamped trim to a free-flying
    dynamic case.
    :return: Dynamic aeroelastic snapshot.
    """
    u_inf = self.aero.flowfield.u_inf  # flowfield velocity to set
    self.aero.flowfield.relative_motion = (
        False  # the output will have relative motion disabled
    )

    rmat_struct = static_case.structure.hg[:, :3, :3]  # deformed rotations
    v_local = jnp.einsum(
        "ijk,j->ik", rmat_struct, -u_inf
    )  # local frame velocity, (n_nodes, 3)

    dynamic_case = static_case.to_dynamic(t=None)
    dynamic_case.structure.v = dynamic_case.structure.v.at[:, :3].set(v_local)
    dynamic_case.structure.prescribed_dofs = (
        self.structure.make_prescribed_dofs_tuple(prescribed_dofs)
    )
    dynamic_case.structure.free_dofs = get_solve_dofs(
        n_dof=self.structure.n_dof,
        prescribed_dofs=dynamic_case.structure.prescribed_dofs,
    )

    return dynamic_case

aeroelastic_states_res_from_dv_varphi 

aeroelastic_states_res_from_dv_varphi(
    dv: AeroelasticDesignVariables,
    varphi: Array,
    thrust: dict[str, Array],
    i_ts: int,
    t: Array,
    use_horseshoe: bool,
) -> tuple[AeroelasticFullStates, Array]

Obtain useful states and forcing residual from design variables and a minimal configuration vector.

Source code in src/flapjax/coupled/gradients/coupled.py 
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
def aeroelastic_states_res_from_dv_varphi(
    self,
    dv: AeroelasticDesignVariables,
    varphi: Array,
    thrust: dict[str, Array],
    i_ts: int,
    t: Array,
    use_horseshoe: bool,
) -> tuple[AeroelasticFullStates, Array]:
    r"""
    Obtain useful states and forcing residual from design variables and a minimal configuration vector.
    """

    # make a copy of the structure object to prevent modifying the original states
    inner_case = pytree_clone(self)

    struct_dv = dv.structure
    aero_dv = dv.aero
    struct = self.structure
    aero = self.aero
    flowfield = (
        aero.flowfield.from_design_variables(design_variables=aero_dv.flowfield)
        if aero_dv.flowfield is not None
        else aero.flowfield
    )
    inner_case.set_design_variables(
        coords=dv_or(struct_dv.x0, struct.x0),
        k_cs=dv_or(struct_dv.k_cs, struct.k_cs),
        m_cs=dv_or(struct_dv.m_cs, struct.m_cs),
        m_lumped=dv_or(struct_dv.m_lumped, struct.m_lumped)
        if struct.use_lumped_mass
        else None,
        thrust_reference=dv_or(struct_dv.thrust_t, struct.thrust_reference),
        flowfield=flowfield,
        delta_w=aero.delta_w,
        dt=aero.dt,
        x0_aero=dv_or(aero_dv.zeta_b0, aero.zeta_b0),
        orientation_euler=dv_or(
            struct_dv.orientation_euler, struct.orientation_euler
        ),
        cs_angles_reference=dv_or(aero_dv.cs_ang_t, aero.cs_ang0),
        remove_checks=True,
    )

    exp_varphi = vmap(exp_se3)(varphi.reshape(-1, 6))  # (n_nodes, 4, 4)
    hg = jnp.einsum(
        "ijk,ikl->ijl", inner_case.structure.hg0, exp_varphi
    )  # (n_nodes, 4, 4)

    # evaluate aero forcing and project to beam nodes
    aero_sol = inner_case.aero.static_solve(hg=hg, t=t, horseshoe=use_horseshoe)
    f_ext_aero_global = aero_sol.project_forcing_to_beam(
        i_ts=0,
        rmat=hg[:, :3, :3],
        x0_aero=self.aero.zeta_b0,
        include_unsteady=False,
    )

    d = inner_case.structure.make_d(hg)
    p_d = inner_case.structure.make_p_d(d)
    eps = inner_case.structure.make_eps(d)
    f_elem = inner_case.structure.make_f_elem(eps=eps)

    if inner_case.structure.use_gravity:
        m_t = inner_case.structure.make_m_t(d)
    else:
        m_t = None

    if dv.structure.f_ext_dead is not None:
        f_ext_dead = inner_case.structure.make_f_dead_ext(
            dv.structure.f_ext_dead, hg[:, :3, :3]
        )
    else:
        f_ext_dead = None

    f_dead_total = inner_case.structure.make_f_ext_dead_tot(
        f_ext_dead, f_ext_aero_global, i_load_step=None
    )

    f_res = inner_case.structure.make_f_res(
        solve_dofs=None,
        p_d=p_d,
        eps=eps,
        hg=hg,
        f_ext_follower_n=dv.structure.f_ext_follower,
        f_ext_dead_n=f_dead_total,
        thrust_n=dv.structure.thrust_t
        if dv.structure.thrust_t is not None
        else thrust,
        dynamic=False,
        m_t=m_t,
        c_l=None,
        c_l_lumped=None,
        v=None,
        v_dot=None,
    )[0]

    struct_states = StructureFullStates(
        hg=hg,
        varphi=varphi,
        eps=eps,
        f_elem=f_elem,
        f_res=f_res.reshape(-1, 6),
        v=None,
        v_dot=None,
    )

    aero_states = aero_sol.get_states(i_ts=i_ts)

    return AeroelasticFullStates(structure=struct_states, aero=aero_states), f_res

static_adjoint 

static_adjoint(
    case: AeroelasticCase,
    objective: AeroelasticObjectiveFunction,
    grads_to_compute: AeroelasticGradsToCompute = DEFAULT_GRADS_TO_COMPUTE,
    optional_jacobians: OptionalJacobians
    | None = DEFAULT_OPTIONAL_JACOBIANS,
    ad_mode: ADMode = "forward",
    batch_size: int | None = 32,
) -> tuple[AeroelasticDesignVariables, 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
case AeroelasticCase

AeroelasticCase containing the current state of the aeroelastic system.

 required
objective AeroelasticObjectiveFunction

Objective function that takes the structure and design variables and returns an array

 required
grads_to_compute AeroelasticGradsToCompute

Data structure which specifies which gradients to compute. This is used to speed up the adjoint solve by only computing the necessary Jacobian blocks.

 DEFAULT_GRADS_TO_COMPUTE
optional_jacobians OptionalJacobians | None

OptionalJacobians object specifying which Jacobians to compute.

 DEFAULT_OPTIONAL_JACOBIANS
ad_mode ADMode

Optional use of either forward or reverse adjoint. For passing the initial state sensitivities to a dynamic solve, only forward mode can be used to give the required adjoint.

 'forward'
batch_size int | None

Batch size for computing p_res_p_varphi to reduce memory usage. Ignored when matrix_free is True.

 32

Returns:

Type Description
tuple[AeroelasticDesignVariables, Array]

Gradient of objective function output with respect to design variables.
Source code in src/flapjax/coupled/gradients/coupled.py 
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
@jax.jit(static_argnums=(0, 1, 2, 3, 4, 5, 6))
def static_adjoint(
    self,
    case: AeroelasticCase,
    objective: AeroelasticObjectiveFunction,
    grads_to_compute: AeroelasticGradsToCompute = DEFAULT_GRADS_TO_COMPUTE,
    optional_jacobians: OptionalJacobians | None = DEFAULT_OPTIONAL_JACOBIANS,
    ad_mode: ADMode = "forward",
    batch_size: int | None = 32,
) -> tuple[AeroelasticDesignVariables, 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 case: AeroelasticCase containing the current state of the aeroelastic system.
    :param objective: Objective function that takes the structure and design variables and returns an array
    :param grads_to_compute: Data structure which specifies which gradients to compute. This is used to speed up the
    adjoint solve by only computing the necessary Jacobian blocks.
    :param optional_jacobians: OptionalJacobians object specifying which Jacobians to compute.
    :param ad_mode: Optional use of either forward or reverse adjoint. For passing the initial state sensitivities
    to a dynamic solve, only forward mode can be used to give the required adjoint.
    :param batch_size: Batch size for computing p_res_p_varphi to reduce memory usage. Ignored when
    ``matrix_free`` is True.
    :return: Gradient of objective function output with respect to design variables.
    """

    if ad_mode not in ("forward", "reverse"):
        raise ValueError("ad_mode must be either 'forward' or 'reverse'")

    jax_print("Computing static adjoint", verbose_level="normal")

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

    if optional_jacobians is not None:
        self.structure.optional_jacobians = optional_jacobians

    dv = self.get_design_variables(case=case, grads_to_compute=grads_to_compute)
    states = case.get_full_states()

    # find shape of objective function output without evaluating function
    f_properties = jax.eval_shape(lambda: objective(states, dv, None))
    f_shape = f_properties.shape
    j0_shape = f_shape if len(f_shape) > 0 else (1,)
    n_f = f_properties.size
    n_x = dv.structure.n_x + dv.aero.n_x
    n_u_full = self.structure.n_dof

    varphi = case.structure.varphi

    if case.aero.static_horseshoe is None:
        raise ValueError("static_horseshoe not defined")
    static_horseshoe: bool = case.aero.static_horseshoe

    # function for computing sensitivity of objective to design variables and degrees of freedom
    # to obtain the actual Jacobian we must pull back the identity through it
    vjp_fn = self.compute_p_j0_p_x(
        case=case,
        objective=objective,
        grads_to_compute=grads_to_compute,
        horseshoe=static_horseshoe,
    )

    cot_j0 = jnp.eye(n_f).reshape(n_f, *j0_shape)  # seed for backpropogation
    p_j_p_varphi_raw, p_j_p_x_raw = jax.vmap(vjp_fn)(cot_j0)  # sensitivities

    p_j_p_varphi_flat = p_j_p_varphi_raw.reshape(n_f, -1)  # (n_f, n_dof)
    p_j_p_x_flat = AeroelasticDesignVariables(
        structure_dv=StructureDesignVariables(
            **{
                k: getattr(p_j_p_x_raw.structure, k) for k in dv.structure.to_dict()
            },
            f_shape=(n_f,),
        ),
        aero_dv=AeroDesignVariables(
            **{k: getattr(p_j_p_x_raw.aero, k) for k in dv.aero.to_dict()},
            f_shape=(n_f,),
        ),
    ).ravel_jacobian(f_size=n_f, x_size=n_x)

    def _residual(varphi_vec: Array, dv_: AeroelasticDesignVariables) -> Array:
        r"""
        Helper function to give the static aeroelastic residual for a given deformation and design variables.
        """
        return self.aeroelastic_states_res_from_dv_varphi(
            dv=dv_,
            varphi=varphi_vec.reshape(self.structure.n_nodes, 6),
            thrust=case.structure.thrust,
            t=case.aero.t,
            i_ts=0,
            use_horseshoe=static_horseshoe,
        )[1]

    if ad_mode == "forward":
        # single joint VJP shared between p_res_p_varphi and p_res_p_x
        _, vjp_res_both = jax.vjp(_residual, varphi.ravel(), dv)
        p_res_p_varphi, p_res_p_x = jax.lax.map(
            vjp_res_both,
            jnp.eye(n_u_full),
            batch_size=batch_size,
        )
        # solve for adjoint
        adj = jnp.linalg.solve(
            p_res_p_varphi[jnp.ix_(solve_dofs, solve_dofs)],
            p_res_p_x.ravel_jacobian(f_size=n_u_full, x_size=n_x)[solve_dofs, :],
        )

        d_f_d_x_dict = dv.from_adjoint(
            f_shape,
            p_j_p_x_flat - p_j_p_varphi_flat[:, solve_dofs] @ adj,
        )
    else:
        # construct residual Jacobian
        _, vjp_res_varphi = jax.vjp(lambda v: _residual(v, dv), varphi.ravel())
        p_res_p_varphi = jax.lax.map(
            lambda cot: vjp_res_varphi(cot)[0],
            jnp.eye(n_u_full),
            batch_size=batch_size,
        )
        adj = jnp.linalg.solve(
            p_res_p_varphi[jnp.ix_(solve_dofs, solve_dofs)].T,
            p_j_p_varphi_flat[:, solve_dofs].T,
        ).T  # (n_f, n_solve_dofs)

        adj_full = (
            jnp.zeros((n_f, n_u_full), dtype=adj.dtype).at[:, solve_dofs].set(adj)
        )

        # sensitivity of residual w.r.t. design variables
        _, vjp_res_dv = jax.vjp(lambda dv_: _residual(varphi.ravel(), dv_), dv)
        (adj_p_res_p_x_raw,) = jax.vmap(vjp_res_dv)(adj_full)

        adj_p_res_p_x_flat = AeroelasticDesignVariables(
            structure_dv=StructureDesignVariables(
                **{
                    k: getattr(adj_p_res_p_x_raw.structure, k)
                    for k in dv.structure.to_dict()
                },
                f_shape=(n_f,),
            ),
            aero_dv=AeroDesignVariables(
                **{
                    k: getattr(adj_p_res_p_x_raw.aero, k) for k in dv.aero.to_dict()
                },
                f_shape=(n_f,),
            ),
        ).ravel_jacobian(f_size=n_f, x_size=n_x)

        d_f_d_x_dict = dv.from_adjoint(
            f_shape,
            p_j_p_x_flat - adj_p_res_p_x_flat,
        )

    return dv.split_adjoint(d_f_d_x=d_f_d_x_dict, f_shape=f_shape), adj

compute_p_j0_p_x 

compute_p_j0_p_x(
    case: AeroelasticCase,
    objective: AeroelasticObjectiveFunction,
    grads_to_compute: AeroelasticGradsToCompute | None,
    horseshoe: bool = False,
    include_q0: bool = False,
) -> Callable[
    ..., tuple[Array, AeroelasticDesignVariables]
]

Build the VJP of the initial-timestep objective for pertubations in the design variables, and optionally the initial states.

Parameters:

Name Type Description Default
case AeroelasticCase

AeroelasticCase solution for the initial timestep.

 required
objective AeroelasticObjectiveFunction

Objective function which takes the system full states, design variables and timestep index.

 required
grads_to_compute AeroelasticGradsToCompute | None

Grads to compute when computing design gradient.

 required
horseshoe bool

Flag for using horseshoe wake.

 False
include_q0 bool

If True, the returned VJP also propagates a cotangent through q0.

 False

Returns:

Type Description
Callable[..., tuple[Array, AeroelasticDesignVariables]]

VJP for cotangents of the design variables, and optionally the initial states.
Source code in src/flapjax/coupled/gradients/coupled.py 
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
def compute_p_j0_p_x(
    self,
    case: AeroelasticCase,
    objective: AeroelasticObjectiveFunction,
    grads_to_compute: AeroelasticGradsToCompute | None,
    horseshoe: bool = False,
    include_q0: bool = False,
) -> Callable[
    ...,
    tuple[Array, AeroelasticDesignVariables],
]:
    r"""
    Build the VJP of the initial-timestep objective for pertubations in the design variables, and optionally the
    initial states.
    :param case: AeroelasticCase solution for the initial timestep.
    :param objective: Objective function which takes the system full states, design variables and timestep index.
    :param grads_to_compute: Grads to compute when computing design gradient.
    :param horseshoe: Flag for using horseshoe wake.
    :param include_q0: If True, the returned VJP also propagates a cotangent through ``q0``.
    :return: VJP for cotangents of the design variables, and optionally the initial states.
    """

    # design variables with variables that we don't require gradients omitted to speed up computations.
    dv = self.get_design_variables(case=case, grads_to_compute=grads_to_compute)

    # design variables with no omissions
    dv_full = self.get_design_variables(case=case, grads_to_compute=None)

    varphi = case.structure.varphi
    n_dof = self.structure.n_dof

    @jax.checkpoint
    def objective_from_varphi(
        varphi_: Array,
        dv_: AeroelasticDesignVariables,
    ) -> Array | tuple[Array, Array]:
        inner_case = self.case_from_dv(dv=dv_)

        assert (
            dv_full.aero.cs_ang_t is not None and dv_full.aero.cs_vel_t is not None
        )

        # solve aero problem
        hg = inner_case.structure.compute_hg_from_varphi(varphi=varphi_)
        _, _, gamma_b, gamma_w, _, _, zeta_w, _, f_steady, _, _, _, _, _ = (
            inner_case.aero.base_solve(
                q_nm1=None,
                t_n=case.aero.t,
                hg_n=hg,
                hg_nm1=None,
                hg_dot_n=None,
                static=True,
                horseshoe=horseshoe,
                cs_ang_n={
                    k: jnp.atleast_1d(v)[0]
                    for k, v in (
                        dv_.aero.cs_ang_t
                        if dv_.aero.cs_ang_t is not None
                        else dv_full.aero.cs_ang_t
                    ).items()
                },
                cs_ang_nm1=None,
                cs_vel_n={
                    k: jnp.atleast_1d(v)[0]
                    for k, v in (
                        dv_.aero.cs_vel_t
                        if dv_.aero.cs_vel_t is not None
                        else dv_full.aero.cs_vel_t
                    ).items()
                },
            )
        )

        f_aero_beam_global = project_forcing_to_beam(
            f_total=f_steady,
            rmat=hg[:, :3, :3],
            dof_mapping=inner_case.aero.dof_mapping,
            x0_aero=inner_case.aero.zeta_b0,
            mirror_edge_low=inner_case.aero.mirror_edge_low,
            mirror_edge_high=inner_case.aero.mirror_edge_high,
        )

        f_aero_beam_local = transform_nodal_vect(
            vect=f_aero_beam_global, rmat=jnp.transpose(hg[:, :3, :3], (0, 2, 1))
        )

        q_aero = AeroFullStates(
            gamma_b=gamma_b,
            gamma_w=gamma_w,
            zeta_w=zeta_w,
            gamma_b_dot=ArrayList.zeros_like(gamma_b),
        )

        # assume initial velocities and accelerations are zero
        q_structure = StructureMinimalStates(
            varphi=varphi_,
            v=jnp.zeros_like(varphi_),
            v_dot=jnp.zeros_like(varphi_),
            a=jnp.zeros_like(varphi_),
            f_ext_aero=f_aero_beam_local,
        )

        q0 = AeroelasticMinimalStates(structure=q_structure, aero=q_aero).ravel()
        q0_full = self.minimal_states_to_full_states(
            i_ts=0,
            q=AeroelasticMinimalStates.from_vector(
                vect=q0,
                n_dof=n_dof,
                aero_shapes=q_aero.shapes(),
            ),
            dv=dv_,
            dv_full=dv_full,
        )
        j0 = jnp.atleast_1d(objective(q0_full, dv_, 0))
        if include_q0:
            return j0, q0
        return j0

    _, vjp_fn = jax.vjp(objective_from_varphi, varphi, dv)
    return vjp_fn

timestep_residual 

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

Compute the coupled aeroelastic residual vector. This is used for the matrix-free time domain case by applying VJP to this function.

Parameters:

Name Type Description Default
i_ts int | Array

Time step index.

 required
t Array

Time at step n.

 required
q_nm1 AeroelasticMinimalStates

Minimal states at step n-1.

 required
q_n AeroelasticMinimalStates

Minimal states at step n.

 required
dv_ AeroelasticDesignVariables

Aeroelastic design variables (may have some fields omitted).

 required
dv_full AeroelasticDesignVariables

Aeroelastic design variables without omissions.

 required
thrust_t dict[str, Array]

Thrust time history, {key: [n_tstep]}.

 required
solve_dofs tuple[int, ...]

Structural degrees of freedom which are solved for.

 required
approx_grads bool

If True, remove some gradient terms which are generally small.

 required

Returns:

Type Description
Array

Coupled residual (n_adj_dof, ).
Source code in src/flapjax/coupled/gradients/coupled.py 
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
def timestep_residual(
    self,
    i_ts: int | Array,
    t: Array,
    q_nm1: AeroelasticMinimalStates,
    q_n: AeroelasticMinimalStates,
    dv_: AeroelasticDesignVariables,
    dv_full: AeroelasticDesignVariables,
    thrust_t: dict[str, Array],
    solve_dofs: tuple[int, ...],
    approx_grads: bool,
) -> Array:
    r"""
    Compute the coupled aeroelastic residual vector. This is used for the matrix-free time domain case by applying
    VJP to this function.
    :param i_ts: Time step index.
    :param t: Time at step n.
    :param q_nm1: Minimal states at step n-1.
    :param q_n: Minimal states at step n.
    :param dv_: Aeroelastic design variables (may have some fields omitted).
    :param dv_full: Aeroelastic design variables without omissions.
    :param thrust_t: Thrust time history, {key: [n_tstep]}.
    :param solve_dofs: Structural degrees of freedom which are solved for.
    :param approx_grads: If True, remove some gradient terms which are generally small.
    :return: Coupled residual ``(n_adj_dof, )``.
    """
    assert (
        q_n.structure.f_ext_aero is not None
        and q_nm1.structure.f_ext_aero is not None
    )

    struct_res = self.structure.timestep_residual(
        i_ts=i_ts,
        q_nm1=q_nm1.structure,
        q_n=q_n.structure,
        dv_=dv_.structure,
        thrust_t=thrust_t,
        solve_dofs=solve_dofs,
        approx_grads=approx_grads,
    )

    # Rematerialise the aero pass to reduce memory usage
    @jax.checkpoint
    def _aero_forward(
        varphi_nm1_: Array,
        varphi_n_: Array,
        v_n_: Array,
        t_n_: Array,
        q_nm1_aero_: AeroFullStates,
        q_n_aero_: AeroFullStates,
        dv__: AeroelasticDesignVariables,
        dv_full_: AeroelasticDesignVariables,
        f_aero_beam_n_: Array,
    ) -> Array:
        return self.aero.timestep_residual(
            i_ts=i_ts,
            varphi_nm1=varphi_nm1_,
            varphi_n=varphi_n_,
            v_n=v_n_,
            t_n=t_n_,
            q_n=q_n_aero_,
            q_nm1=q_nm1_aero_,
            dv=dv__,
            dv_full=dv_full_,
            f_aero_beam_n=f_aero_beam_n_,
            struct_obj=self.structure,
            approx_grads=approx_grads,
        )

    # evaluate checkpointed function
    aero_res = _aero_forward(
        varphi_nm1_=q_nm1.structure.varphi,
        varphi_n_=q_n.structure.varphi,
        v_n_=q_n.structure.v,
        t_n_=t,
        q_nm1_aero_=q_nm1.aero,
        q_n_aero_=q_n.aero,
        dv__=dv_,
        dv_full_=dv_full,
        f_aero_beam_n_=q_n.structure.f_ext_aero,
    )

    # remove forces from degrees of freedom which are not solved for
    n_aero_states: int = q_n.aero.n_states
    solve_dofs_arr = jnp.array(solve_dofs)
    aero_res_solve = jnp.concatenate(
        (aero_res[:n_aero_states], aero_res[n_aero_states:][solve_dofs_arr])
    )

    return jnp.concatenate((struct_res, aero_res_solve))

timestep_residual_jacobians 

timestep_residual_jacobians(
    i_ts: int | Array,
    t: Array,
    q_nm1: AeroelasticMinimalStates,
    q_n: AeroelasticMinimalStates,
    dv_: AeroelasticDesignVariables,
    dv_full: AeroelasticDesignVariables,
    thrust_t: dict[str, Array],
    solve_dofs: tuple[int, ...],
    approx_grads: bool,
    n_profile_loops: int | None,
    jac_options: dict,
    mode: ADMode = "reverse",
    map_batch_size: int | None = None,
) -> tuple[
    Array,
    Array,
    StructureDesignVariables,
    AeroelasticDesignVariables,
    dict[str, dict[str, float]] | None,
    dict[str, dict[str, float]] | None,
]

Compute the required time-domain residual Jacobians for the adjoint solution.

Parameters:

Name Type Description Default
i_ts int | Array

Time step index.

 required
t Array

Time.

 required
q_nm1 AeroelasticMinimalStates

Minimal degrees of freedom at timestep n-1.

 required
q_n AeroelasticMinimalStates

Minimal degrees of freedom at timestep n.

 required
dv_ AeroelasticDesignVariables

Design variables for which to obtain gradients.

 required
dv_full AeroelasticDesignVariables

All design variables.

 required
thrust_t dict[str, Array]

Thrust at each time step.

 required
solve_dofs tuple[int, ...]

Degrees of freedom to solve for.

 required
approx_grads bool

Whether to use approximate gradients for the structural dynamic subproblem.

 required
n_profile_loops int | None

Number of profile loops. Used for profiling routines only.

 required
jac_options dict

Options for Jacobian computation, allowing for approximations to be introduced.

 required
mode ADMode

Mode for automatic differentiation, either forward or reverse.

 'reverse'
map_batch_size int | None

Batch size used for vectorising Jacobian construction.

 None

Returns:

Type Description
tuple[Array, Array, StructureDesignVariables, AeroelasticDesignVariables, dict[str, dict[str, float]] | None, dict[str, dict[str, float]] | None]

Jacobian of residual with respect to previous degrees of freedom, Jacobian of residual with respect to current degrees of freedom. Jacobian of v_dot residual with respect to structural design variables, Jacobian of aero residual with respect to design variables. Can also include compile time and run time when profiling is used.
Source code in src/flapjax/coupled/gradients/coupled.py 
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
def timestep_residual_jacobians(
    self,
    i_ts: int | Array,
    t: Array,
    q_nm1: AeroelasticMinimalStates,
    q_n: AeroelasticMinimalStates,
    dv_: AeroelasticDesignVariables,
    dv_full: AeroelasticDesignVariables,
    thrust_t: dict[str, Array],
    solve_dofs: tuple[int, ...],
    approx_grads: bool,
    n_profile_loops: int | None,
    jac_options: dict,
    mode: ADMode = "reverse",
    map_batch_size: int | None = None,
) -> tuple[
    Array,
    Array,
    StructureDesignVariables,
    AeroelasticDesignVariables,
    dict[str, dict[str, float]] | None,
    dict[str, dict[str, float]] | None,
]:
    r"""
    Compute the required time-domain residual Jacobians for the adjoint solution.
    :param i_ts: Time step index.
    :param t: Time.
    :param q_nm1: Minimal degrees of freedom at timestep n-1.
    :param q_n: Minimal degrees of freedom at timestep n.
    :param dv_: Design variables for which to obtain gradients.
    :param dv_full: All design variables.
    :param thrust_t: Thrust at each time step.
    :param solve_dofs: Degrees of freedom to solve for.
    :param approx_grads: Whether to use approximate gradients for the structural dynamic subproblem.
    :param n_profile_loops: Number of profile loops. Used for profiling routines only.
    :param jac_options: Options for Jacobian computation, allowing for approximations to be introduced.
    :param mode: Mode for automatic differentiation, either ``forward`` or ``reverse``.
    :param map_batch_size: Batch size used for vectorising Jacobian construction.
    :return: Jacobian of residual with respect to previous degrees of freedom, Jacobian of residual with respect to
    current degrees of freedom. Jacobian of v_dot residual with respect to structural design variables, Jacobian of
    aero residual with respect to design variables. Can also include compile time and run time when profiling is
    used.
    """

    assert (
        q_n.structure.f_ext_aero is not None
        and q_nm1.structure.f_ext_aero is not None
    )

    (
        p_aero_res_p_q_aero_nm1,
        p_aero_res_p_q_aero_n,
        p_aero_res_d_dv,
        p_aero_res_p_q_struct_nm1,
        p_aero_res_p_q_struct_n,
        aero_compile_time,
        aero_run_time,
    ) = self.aero.timestep_residual_jacobians(
        i_ts=i_ts,
        varphi_nm1=q_nm1.structure.varphi,
        varphi_n=q_n.structure.varphi,
        v_n=q_n.structure.v,
        t_n=t,
        q_n=q_n.aero,
        q_nm1=q_nm1.aero,
        dv=dv_,
        dv_full=dv_full,
        f_aero_beam_n=q_n.structure.f_ext_aero,
        struct_obj=self.structure,
        approx_grads=approx_grads,
        solve_dofs=solve_dofs,
        n_profile_loops=n_profile_loops,
        jac_options=jac_options,
        mode=mode,
        map_batch_size=map_batch_size,
    )

    n_aero_dof, _ = p_aero_res_p_q_struct_nm1.shape

    (
        p_struct_res_p_q_struct_nm1,
        p_struct_res_p_q_struct_n,
        p_v_dot_res_p_struct_dv,
        p_v_dot_res_p_f_ext_nm1,
        p_v_dot_res_p_f_ext_n,
        struct_compile_time,
        struct_run_time,
    ) = self.structure.timestep_residual_jacobians(
        i_ts=i_ts,
        q_nm1=q_nm1.structure,
        q_n=q_n.structure,
        f_ext_aero_n=q_n.structure.f_ext_aero,
        f_ext_aero_nm1=q_nm1.structure.f_ext_aero,
        thrust_t=thrust_t,
        dv=dv_.structure,
        solve_dofs=solve_dofs,
        approx_grads=approx_grads,
        n_profile_loops=n_profile_loops,
        jac_options=jac_options,
        mode=mode,
    )

    assert p_v_dot_res_p_f_ext_nm1 is not None and p_v_dot_res_p_f_ext_n is not None

    # reduce aero-to-struct Jacobian columns to solve_dofs
    n_solve = len(solve_dofs)
    n_struct_res = p_struct_res_p_q_struct_nm1.shape[0]
    solve_dofs_arr = jnp.array(solve_dofs)
    struct_col_ix = jnp.concatenate(
        [solve_dofs_arr + i * self.structure.n_dof for i in range(4)]
    )
    p_aero_res_p_q_struct_nm1 = p_aero_res_p_q_struct_nm1[:, struct_col_ix]
    p_aero_res_p_q_struct_n = p_aero_res_p_q_struct_n[:, struct_col_ix]

    # create struct-to-aero cross-coupling (only v_dot residual depends on f_ext_aero)
    # f_aero block is n_solve-wide (last n_solve cols of aero state), so slice Jacobian to solve_dofs cols
    p_struct_res_p_q_aero_nm1 = jnp.zeros((n_struct_res, n_aero_dof))
    p_struct_res_p_q_aero_nm1 = p_struct_res_p_q_aero_nm1.at[
        jnp.arange(n_solve) + 2 * n_solve, -n_solve:
    ].set(p_v_dot_res_p_f_ext_nm1[:, solve_dofs_arr])

    p_struct_res_p_q_aero_n = jnp.zeros((n_struct_res, n_aero_dof))
    p_struct_res_p_q_aero_n = p_struct_res_p_q_aero_n.at[
        jnp.arange(n_solve) + 2 * n_solve, -n_solve:
    ].set(p_v_dot_res_p_f_ext_n[:, solve_dofs_arr])

    p_res_p_q_nm1 = jnp.block(
        [
            [p_struct_res_p_q_struct_nm1, p_struct_res_p_q_aero_nm1],
            [p_aero_res_p_q_struct_nm1, p_aero_res_p_q_aero_nm1],
        ]
    )

    p_res_p_q_n = jnp.block(
        [
            [p_struct_res_p_q_struct_n, p_struct_res_p_q_aero_n],
            [p_aero_res_p_q_struct_n, p_aero_res_p_q_aero_n],
        ]
    )

    if n_profile_loops is not None:
        assert (
            aero_compile_time is not None
            and aero_run_time is not None
            and struct_compile_time is not None
            and struct_run_time is not None
        )
        compile_time = aero_compile_time | struct_compile_time
        run_time = aero_run_time | struct_run_time
    else:
        compile_time = None
        run_time = None
    return (
        p_res_p_q_nm1,
        p_res_p_q_n,
        p_v_dot_res_p_struct_dv,
        p_aero_res_d_dv,
        compile_time,
        run_time,
    )

construct_approximate_jacobians 

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

Compute approximations for Jacobians which are specified in the jacobian_approximations data structure. The aerodynamic residual approximations are delegated to UVLM.construct_approximate_jacobians, and the structural residual approximations to BeamStructure.construct_approximate_jacobians. The two dictionaries are merged into a single result.

Parameters:

Name Type Description Default
sol AeroelasticCase

Solution for which approximations will be created for the initial time step.

 required
jacobian_approximations AeroelasticJacobianApproximations

Data structure which defines which approximations to create.

 required

Returns:

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

Dictionary of approximations keyed by residual name.
Source code in src/flapjax/coupled/gradients/coupled.py 
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
def construct_approximate_jacobians(
    self,
    sol: AeroelasticCase,
    jacobian_approximations: AeroelasticJacobianApproximations,
) -> dict[str, dict[str, Callable[..., Any] | None]]:
    r"""
    Compute approximations for Jacobians which are specified in the jacobian_approximations data structure. The
    aerodynamic residual approximations are delegated to ``UVLM.construct_approximate_jacobians``, and the
    structural residual approximations to ``BeamStructure.construct_approximate_jacobians``. The two
    dictionaries are merged into a single result.
    :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 keyed by residual name.
    """
    dv = self.get_design_variables(case=sol, grads_to_compute=None)
    solve_dofs = get_solve_dofs(
        n_dof=self.structure.n_dof,
        prescribed_dofs=sol.structure.prescribed_dofs,
    )

    aero_options = self.aero.construct_approximate_jacobians(
        aero_sol=sol.aero,
        structure_sol=sol.structure,
        struct_obj=self.structure,
        dv=dv,
        dv_full=dv,
        solve_dofs=solve_dofs,
        jacobian_approximations=jacobian_approximations.aero,
    )

    struct_options = self.structure.construct_approximate_jacobians(
        sol=sol.structure,
        jacobian_approximations=jacobian_approximations.structure,
    )

    return aero_options | struct_options

evaluate_dynamic_objective 

evaluate_dynamic_objective(
    case: AeroelasticCase,
    objective: AeroelasticObjectiveFunction,
) -> Array

Evaluate the dynamic objective for a given case.

Parameters:

Name Type Description Default
case AeroelasticCase

Dynamic aeroelastic case object.

 required
objective AeroelasticObjectiveFunction

Objective function to be evaluated.

 required

Returns:

Type Description
Array

Value of subobjective at every time step, [n_tstep].
Source code in src/flapjax/coupled/gradients/coupled.py 
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
def evaluate_dynamic_objective(
    self, case: AeroelasticCase, objective: AeroelasticObjectiveFunction
) -> Array:
    r"""
    Evaluate the dynamic objective for a given case.
    :param case: Dynamic aeroelastic case object.
    :param objective: Objective function to be evaluated.
    :return: Value of subobjective at every time step, [n_tstep].
    """
    n_tstep = case.structure.n_tstep

    dv = self.get_design_variables(case=case, grads_to_compute=None)

    return jax.vmap(
        lambda i_ts: jnp.atleast_1d(
            objective(case.get_full_states(i_ts=i_ts), dv, i_ts)
        )
    )(jnp.arange(n_tstep)).reshape(n_tstep, -1)

make_frozen_wake_preconditioner 

make_frozen_wake_preconditioner(
    case: AeroelasticCase,
    dv_full: AeroelasticDesignVariables,
    solve_dofs: tuple[int, ...],
    approx_grads: bool = False,
    precond_i_ts: int = 0,
    batch_size: int | None = 32,
) -> Callable[[Array], Array]

Build a preconditioner for the coupled aeroelastic system which skips the wake grid and circulation. When applied to the matrix-free system GMRES, this found a good reduction in the number of iterations required whilst avoiding the large memory overhead involved in computing the full Jacobian due to the large number of wake states.

Parameters:

Name Type Description Default
case AeroelasticCase

Dynamic aeroelastic case object.

 required
dv_full AeroelasticDesignVariables

Full dynamic aeroelastic design variables.

 required
solve_dofs tuple[int, ...]

Solve degree of freedom index.

 required
approx_grads bool

Approximate gradient of the coupled aeroelastic system, removing some negligible terms.

 False
precond_i_ts int

Time step index for which to create the preconditioner. Defaults to 0.

 0
batch_size int | None

Batch size for mapping the Jacobian construction on the aerodynamic system.

 32

Returns:

Type Description
Callable[[Array], Array]

Preconditioner function.
Source code in src/flapjax/coupled/gradients/coupled.py 
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
def make_frozen_wake_preconditioner(
    self,
    case: AeroelasticCase,
    dv_full: AeroelasticDesignVariables,
    solve_dofs: tuple[int, ...],
    approx_grads: bool = False,
    precond_i_ts: int = 0,
    batch_size: int | None = 32,
) -> Callable[[Array], Array]:
    r"""
    Build a preconditioner for the coupled aeroelastic system which skips the wake grid and circulation. When
    applied to the matrix-free system GMRES, this found a good reduction in the number of iterations required whilst
    avoiding the large memory overhead involved in computing the full Jacobian due to the large number of wake
    states.
    :param case: Dynamic aeroelastic case object.
    :param dv_full: Full dynamic aeroelastic design variables.
    :param solve_dofs: Solve degree of freedom index.
    :param approx_grads: Approximate gradient of the coupled aeroelastic system, removing some negligible terms.
    :param precond_i_ts: Time step index for which to create the preconditioner. Defaults to 0.
    :param batch_size: Batch size for mapping the Jacobian construction on the aerodynamic system.
    :return: Preconditioner function.
    """
    precond_q_nm1 = case.get_minimal_states(i_ts=max(precond_i_ts - 1, 0))
    precond_q_n = case.get_minimal_states(i_ts=precond_i_ts)
    precond_t_n = case.structure.t[precond_i_ts]

    dv_precond = self.get_design_variables(case=case, grads_to_compute=None)

    assert precond_q_n.structure.f_ext_aero is not None
    assert precond_q_nm1.structure.f_ext_aero is not None

    # populate jac_options with the expected residual/argname keys, all mapped to None so AD is used everywhere
    jac_options = self.construct_approximate_jacobians(
        sol=case,
        jacobian_approximations=AeroelasticJacobianApproximations(),
    )

    # compute structural Jacobians
    (
        _,
        p_struct_res_p_q_struct_n,
        _,
        _,
        p_v_dot_res_p_f_ext_n,
        *_,
    ) = self.structure.timestep_residual_jacobians(
        i_ts=precond_i_ts,
        q_nm1=precond_q_nm1.structure,
        q_n=precond_q_n.structure,
        f_ext_aero_n=precond_q_n.structure.f_ext_aero,
        f_ext_aero_nm1=precond_q_nm1.structure.f_ext_aero,
        thrust_t=case.structure.thrust,
        dv=dv_precond.structure,
        solve_dofs=solve_dofs,
        approx_grads=approx_grads,
        n_profile_loops=None,
        jac_options=jac_options,
    )

    # compute aero Jacobians
    (
        _,
        p_aero_res_p_q_aero_n,
        _,
        _,
        p_aero_res_p_q_struct_n,
        *_,
    ) = self.aero.timestep_residual_jacobians(
        i_ts=precond_i_ts,
        varphi_nm1=precond_q_nm1.structure.varphi,
        varphi_n=precond_q_n.structure.varphi,
        v_n=precond_q_n.structure.v,
        t_n=precond_t_n,
        q_n=precond_q_n.aero,
        q_nm1=precond_q_nm1.aero,
        dv=dv_precond,
        dv_full=dv_full,
        f_aero_beam_n=precond_q_n.structure.f_ext_aero,
        struct_obj=self.structure,
        approx_grads=approx_grads,
        solve_dofs=solve_dofs,
        n_profile_loops=None,
        jac_options=jac_options,
        compute_wake_gradients=False,
        map_batch_size=batch_size,
    )

    assert p_v_dot_res_p_f_ext_n is not None

    solve_dofs_arr = jnp.array(solve_dofs)
    struct_col_ix = jnp.concatenate(
        [solve_dofs_arr + i * self.structure.n_dof for i in range(4)]
    )
    p_aero_res_p_q_struct_n = p_aero_res_p_q_struct_n[:, struct_col_ix]

    # assemble Jacobians
    n_solve = len(solve_dofs)
    n_struct_res = p_struct_res_p_q_struct_n.shape[0]
    n_aero_reduced = p_aero_res_p_q_aero_n.shape[0]
    p_struct_res_p_q_aero_n = jnp.zeros((n_struct_res, n_aero_reduced))
    p_struct_res_p_q_aero_n = p_struct_res_p_q_aero_n.at[
        jnp.arange(n_solve) + 2 * n_solve, -n_solve:
    ].set(p_v_dot_res_p_f_ext_n[:, solve_dofs_arr])

    p_res_p_q_n_reduced = jnp.block(
        [
            [p_struct_res_p_q_struct_n, p_struct_res_p_q_aero_n],
            [p_aero_res_p_q_struct_n, p_aero_res_p_q_aero_n],
        ]
    )

    # compute the LU decomposition for fast reuse when solving
    precond_lu = jax.scipy.linalg.lu_factor(p_res_p_q_n_reduced.T)

    # adjoint state counts and placement
    n_struct = 4 * n_solve
    n_gamma_b = int(precond_q_n.aero.gamma_b.ravel().size)
    n_gamma_w = int(precond_q_n.aero.gamma_w.ravel().size)
    n_gamma_b_dot = int(precond_q_n.aero.gamma_b_dot.ravel().size)
    n_zeta_w = int(precond_q_n.aero.zeta_w.ravel().size)

    gamma_w_start = n_struct + n_gamma_b
    gamma_w_end = gamma_w_start + n_gamma_w
    zeta_w_start = gamma_w_end + n_gamma_b_dot
    zeta_w_end = zeta_w_start + n_zeta_w

    def apply_precond(vec: Array) -> Array:
        # split the system to remove gamma_w and zeta_w. The removed blocks use the negative identity as preconditioner.
        vec_pre_gw = vec[:gamma_w_start]  # struct + gamma_b
        vec_gw = vec[gamma_w_start:gamma_w_end]  # gamma_w (identity)
        vec_mid = vec[gamma_w_end:zeta_w_start]  # gamma_b_dot
        vec_zw = vec[zeta_w_start:zeta_w_end]  # zeta_w (identity)
        vec_post_zw = vec[zeta_w_end:]  # f_ext_aero
        vec_reduced = jnp.concatenate([vec_pre_gw, vec_mid, vec_post_zw])
        x_reduced = jax.scipy.linalg.lu_solve(precond_lu, vec_reduced)

        x_pre_gw = x_reduced[:gamma_w_start]
        x_mid = x_reduced[gamma_w_start : gamma_w_start + n_gamma_b_dot]
        x_post_zw = x_reduced[gamma_w_start + n_gamma_b_dot :]

        return jnp.concatenate([x_pre_gw, -vec_gw, x_mid, -vec_zw, x_post_zw])

    return apply_precond

dynamic_adjoint 

dynamic_adjoint(
    case: AeroelasticCase,
    objective: AeroelasticObjectiveFunction,
    matrix_free: bool = True,
    jacobian_approximations: AeroelasticJacobianApproximations
    | None = None,
    grads_to_compute: AeroelasticGradsToCompute
    | None = DEFAULT_GRADS_TO_COMPUTE,
    p_varphi_p_x: Array | None = None,
    save_adjoint: bool = False,
    approx_grads: bool = True,
    i_ts_adjoint_range: tuple[int | None, int | None] = (
        None,
        None,
    ),
    include_initial_state_grad: bool = True,
    gmres_mode: Literal[
        "batched", "incremental"
    ] = "incremental",
    gmres_warm_start: bool = True,
    gmres_precond: bool = True,
    gmres_restart: int = 50,
    i_ts_preconditioner: int = 0,
    preconditioner_batch_size: int | None = 16,
    preconditioner: Callable[[Array], Array] | None = None,
) -> tuple[AeroelasticDesignVariables, Array, Array | None]

Compute the adjoint of a coupled dynamic aeroelastic system.

Parameters:

Name Type Description Default
case AeroelasticCase

Dynamic aeroelastic case

 required
objective AeroelasticObjectiveFunction

Objective function that takes the system full states, design variables and timestep index, and returns an array.

 required
matrix_free bool

If true, do not explicitly compute the residual Jacobians and instead use the VJP and GMRES to solve.

 True
jacobian_approximations AeroelasticJacobianApproximations | None

Data structure which specifies Jacobian approximations to use for each part of the problem.

 None
grads_to_compute AeroelasticGradsToCompute | None

Specify which design variables for which to compute gradients for. If None, all available gradients are computed.

 DEFAULT_GRADS_TO_COMPUTE
p_varphi_p_x Array | None

Gradient of initial twists with respect to design variables. In practice, this is found from the static solve.

 None
save_adjoint bool

Whether to save the adjoint of the dynamic aeroelastic system.

 False
approx_grads bool

Whether to use gradient approximation or not. This removes some negligible contributions in the structural dynamic system.

 True
i_ts_adjoint_range tuple[int | None, int | None]

Optional (start, end) window of time steps for which to compute the adjoint. Either entry may be None to leave that side untruncated. When start > 1 the initial-state gradient contribution is automatically skipped.

 (None, None)
include_initial_state_grad bool

If False, skip the _initial_timestep_grad_contribution call that solves the static adjoint at t = 0 and propagates p_varphi_p_x. Intended for profiling only.

 True
gmres_mode Literal['batched', 'incremental']

If using matrix free, sets the mode for GMRES. Batched is preferred for GPU, whereas incremental may be preferred on CPU.

 'incremental'
gmres_warm_start bool

If True, use the previous timestep adjoint vector as the first guess for the current value. Otherwise, initialise with the zero vector.

 True
gmres_precond bool

If True and preconditioner is None, build the frozen-wake preconditioner internally. Ignored when preconditioner is supplied.

 True
gmres_restart int

Number of times to restart the GMRES algorithm.

 50
i_ts_preconditioner int

Timestep at which to build the preconditioner if requested.

 0
preconditioner_batch_size int | None

Batch size for creating the Jacobians for the preconditioner. Ignored if preconditioner is supplied.

 16
preconditioner Callable[[Array], Array] | None

Optional pass a prebuild preconditioner. Useful for profiling.

 None

Returns:

Type Description
tuple[AeroelasticDesignVariables, Array, Array | None]

Gradient of sum of objective across timesteps with respect to design variables, objective at each time step, and optional adjoint states.
Source code in src/flapjax/coupled/gradients/coupled.py 
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
@jax.jit(static_argnums=(0, 1, 2, 3, 4, 5, 7, 8, 9, 10, 11, 12, 13, 14, 15, 17))
def dynamic_adjoint(
    self,
    case: AeroelasticCase,
    objective: AeroelasticObjectiveFunction,
    matrix_free: bool = True,
    jacobian_approximations: AeroelasticJacobianApproximations | None = None,
    grads_to_compute: AeroelasticGradsToCompute | None = DEFAULT_GRADS_TO_COMPUTE,
    p_varphi_p_x: Array | None = None,
    save_adjoint: bool = False,
    approx_grads: bool = True,
    i_ts_adjoint_range: tuple[int | None, int | None] = (None, None),
    include_initial_state_grad: bool = True,
    gmres_mode: Literal["batched", "incremental"] = "incremental",
    gmres_warm_start: bool = True,
    gmres_precond: bool = True,
    gmres_restart: int = 50,
    i_ts_preconditioner: int = 0,
    preconditioner_batch_size: int | None = 16,
    preconditioner: Callable[[Array], Array] | None = None,
) -> tuple[AeroelasticDesignVariables, Array, Array | None]:
    r"""
    Compute the adjoint of a coupled dynamic aeroelastic system.
    :param case: Dynamic aeroelastic case
    :param objective: Objective function that takes the system full states, design variables and timestep index,
    and returns an array.
    :param matrix_free: If true, do not explicitly compute the residual Jacobians and instead use the VJP and GMRES
    to solve.
    :param jacobian_approximations: Data structure which specifies Jacobian approximations to use for each part of
    the problem.
    :param grads_to_compute: Specify which design variables for which to compute gradients for. If None, all
    available gradients are computed.
    :param p_varphi_p_x: Gradient of initial twists with respect to design variables. In practice, this is found
    from the static solve.
    :param save_adjoint: Whether to save the adjoint of the dynamic aeroelastic system.
    :param approx_grads: Whether to use gradient approximation or not. This removes some negligible contributions
    in the structural dynamic system.
    :param i_ts_adjoint_range: Optional ``(start, end)`` window of time steps for which to compute the adjoint.
     Either entry may be ``None`` to leave that side untruncated. When ``start > 1`` the initial-state gradient
     contribution is automatically skipped.
    :param include_initial_state_grad: If False, skip the ``_initial_timestep_grad_contribution`` call that solves
    the static adjoint at ``t = 0`` and propagates ``p_varphi_p_x``. Intended for profiling only.
    :param gmres_mode: If using matrix free, sets the mode for GMRES. Batched is preferred for GPU, whereas
    incremental may be preferred on CPU.
    :param gmres_warm_start: If True, use the previous timestep adjoint vector as the first guess for the current
    value. Otherwise, initialise with the zero vector.
    :param gmres_precond: If True and ``preconditioner`` is None, build the frozen-wake preconditioner internally.
    Ignored when ``preconditioner`` is supplied.
    :param gmres_restart: Number of times to restart the GMRES algorithm.
    :param i_ts_preconditioner: Timestep at which to build the preconditioner if requested.
    :param preconditioner_batch_size: Batch size for creating the Jacobians for the preconditioner. Ignored if
    ``preconditioner`` is supplied.
    :param preconditioner: Optional pass a prebuild preconditioner. Useful for profiling.
    :return: Gradient of sum of objective across timesteps with respect to design variables, objective at each time
    step, and optional adjoint states.
    """

    # make copies to prevent contaminating input object with tracer
    case = deepcopy(case)
    p_varphi_p_x = deepcopy(p_varphi_p_x)

    solve_dofs: tuple[int, ...] = get_solve_dofs(
        n_dof=self.structure.n_dof,
        prescribed_dofs=case.structure.prescribed_dofs,
    )

    solve_dofs_arr: Array = jnp.array(solve_dofs)

    n_tstep = case.structure.n_tstep

    dv = self.get_design_variables(case=case, grads_to_compute=grads_to_compute)

    dv_full = self.get_design_variables(case=case, grads_to_compute=None)

    assert case.aero.static_horseshoe is not None
    static_horseshoe: bool = case.aero.static_horseshoe

    full_states_init = case.get_full_states(i_ts=0)
    minimal_states_init = case.get_minimal_states(i_ts=0)

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

    n_solve = len(solve_dofs)

    j_eval = jax.vmap(
        lambda i_ts: jnp.atleast_1d(
            objective(case.get_full_states(i_ts=i_ts), dv, i_ts)
        )
    )(jnp.arange(n_tstep)).reshape(n_tstep, n_j)

    jac_options = self.construct_approximate_jacobians(
        sol=case,
        jacobian_approximations=jacobian_approximations
        if jacobian_approximations is not None
        else AeroelasticJacobianApproximations(),
    )

    # define 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 = 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_ > n_tstep - 1:
        raise ValueError(
            f"i_ts_adjoint_range end must be <= n_tstep - 1 = {n_tstep - 1}, got "
            f"{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

    @jax.jit
    def objective_jacobians(
        i_ts: int, q_n: AeroelasticMinimalStates
    ) -> tuple[Array, AeroelasticDesignVariables]:
        # function to obtain the Jacobians of the objective w.r.t. the minimal states and the design variables
        p_j_n_p_q_n, p_j_n_p_x = jax.jacrev(
            lambda q_free, dv__: jnp.atleast_1d(
                objective(
                    self.minimal_states_to_full_states(
                        i_ts=i_ts,
                        q=AeroelasticMinimalStates.from_vector(
                            vect=q_n.ravel().at[free_state_ix].set(q_free),
                            n_dof=self.structure.n_dof,
                            aero_shapes=minimal_states_init.aero.shapes(),
                        ),
                        dv=dv__,
                        dv_full=dv_full,
                    ),
                    dv__,
                    i_ts,
                )
            ),
            argnums=(0, 1),
            allow_int=True,
        )(q_n.ravel()[free_state_ix], dv)
        return p_j_n_p_q_n, p_j_n_p_x

    assert dv_full.aero.cs_ang_t is not None and dv_full.aero.cs_vel_t is not None

    # create the initial d_j_d_x sensitivites which are accumulated though the solve process
    dv_grad_init = AeroelasticDesignVariables.zeros(
        system=self, case=case, grads_to_compute=grads_to_compute, j_shape=j_shape
    )

    n_dof: int = self.structure.n_dof
    n_aero_states: int = minimal_states_init.aero.n_states
    free_state_ix: Array = jnp.concatenate(
        [solve_dofs_arr + i * n_dof for i in range(4)]
        + [jnp.arange(5 * n_dof, 5 * n_dof + n_aero_states)]
        + [solve_dofs_arr + 4 * n_dof]
    )
    n_adj_dof = 4 * n_solve + n_aero_states + n_solve

    adj_full_init: Array | None = (
        jnp.zeros((case.structure.n_tstep + 1, n_j, n_adj_dof))
        if save_adjoint
        else None
    )

    d_j_d_x: AeroelasticDesignVariables
    if matrix_free:
        if preconditioner is not None:
            apply_precond = preconditioner
        elif gmres_precond:
            apply_precond = self.make_frozen_wake_preconditioner(
                case=case,
                dv_full=dv_full,
                solve_dofs=solve_dofs,
                approx_grads=approx_grads,
                precond_i_ts=i_ts_preconditioner,
                batch_size=preconditioner_batch_size,
            )
            jax_print(
                "Built frozen-wake preconditioner",
                verbose_level="normal",
            )
        else:
            apply_precond = None

        def matrix_free_body(
            rev_i_ts_: int,
            carry: tuple[AeroelasticDesignVariables, Array, Array, Array],
        ) -> tuple[AeroelasticDesignVariables, Array, Array, Array]:
            d_j_d_x_, adj_np1, adj_t_p_r_np1_p_q_n, adj_full_ = carry

            i_ts = i_ts_end_adj_ - rev_i_ts_
            i_ts_nm1 = jnp.maximum(i_ts - 1, 0)
            q_nm1 = case.get_minimal_states(i_ts=i_ts_nm1)
            q_n = case.get_minimal_states(i_ts=i_ts)
            t_n = case.structure.t[i_ts]

            p_j_n_p_q_n, p_j_n_p_x = objective_jacobians(i_ts=i_ts, q_n=q_n)

            def _residual_all(
                q_n_: AeroelasticMinimalStates,
                q_nm1_: AeroelasticMinimalStates,
                dv_: AeroelasticDesignVariables,
            ) -> Array:
                return self.timestep_residual(
                    i_ts=i_ts,
                    t=t_n,
                    q_nm1=q_nm1_,
                    q_n=q_n_,
                    dv_=dv_,
                    dv_full=dv_full,
                    thrust_t=case.structure.thrust,
                    solve_dofs=solve_dofs,
                    approx_grads=approx_grads,
                )

            # single VJP shared between the GMRES matvec, the coupling term and the design-variable pull
            _, pull_all = jax.vjp(_residual_all, q_n, q_nm1, dv)

            def matvec_qn_t(v: Array) -> Array:
                if map_verbosity_level(get_verbosity()) >= map_verbosity_level(
                    "normal"
                ):
                    # print a dot for every GMRES iteration. Due to the jax GMRES function not returning the number
                    # of iterations, this at least allows us to count the dots!
                    def _print_gmres_dot() -> None:
                        sys.stdout.write(".")
                        sys.stdout.flush()

                    jax.debug.callback(_print_gmres_dot, ordered=True)

                d_q: AeroelasticMinimalStates = pull_all(v)[0]
                return d_q.to_free_dofs(solve_dofs_arr=solve_dofs_arr)

            b_rhs = -(p_j_n_p_q_n.reshape(n_j, -1) + adj_t_p_r_np1_p_q_n)

            def _solve_row(b_row: Array, x0_row: Array) -> tuple[Array, Array]:
                # noinspection PyTypeChecker
                x, info = jax.scipy.sparse.linalg.gmres(
                    matvec_qn_t,
                    b_row,
                    x0=x0_row if gmres_warm_start else None,
                    tol=1e-6,
                    atol=1e-6,
                    restart=gmres_restart,
                    maxiter=50,
                    M=apply_precond,
                    solve_method=gmres_mode,
                )
                return x, info

            adj_n, gmres_info = jax.vmap(_solve_row)(b_rhs, adj_np1)

            def _pull_row(
                a: Array,
            ) -> tuple[Array, AeroelasticDesignVariables]:
                q_nm1_cot: AeroelasticMinimalStates
                _, q_nm1_cot, dv_cot = pull_all(a)

                return q_nm1_cot.to_free_dofs(solve_dofs_arr=solve_dofs_arr), dv_cot

            adj_t_p_r_n_p_q_nm1, dv_grads = jax.vmap(_pull_row)(adj_n)

            d_j_d_x_ += dv_grads
            d_j_d_x_ += p_j_n_p_x

            jax_print(
                "\nSolved adjoint for timestep {i_ts} (GMRES converged={converged}, max|adj|={ma:.2e})",
                i_ts=i_ts,
                converged=jnp.max(gmres_info) == 0,
                ma=jnp.max(jnp.abs(adj_n)),
                verbose_level="normal",
            )

            if save_adjoint:
                adj_full_ = adj_full_.at[i_ts].set(adj_n)

            return d_j_d_x_, adj_n, adj_t_p_r_n_p_q_nm1, adj_full_

        d_j_d_x, _, future_row, adj_full = jax.lax.fori_loop(
            lower=0,
            upper=n_adj_iters,
            body_fun=matrix_free_body,
            init_val=(
                dv_grad_init,
                jnp.zeros((n_j, n_adj_dof)),
                jnp.zeros((n_j, n_adj_dof)),
                adj_full_init,
            ),
        )
    else:

        def step_body(
            rev_i_ts_: int,
            carry: tuple[AeroelasticDesignVariables, Array, Array, Array],
        ) -> tuple[AeroelasticDesignVariables, Array, Array, Array]:
            d_j_d_x_, adj_np1, p_r_np1_p_q_n, adj_full_ = carry

            i_ts = i_ts_end_adj_ - rev_i_ts_
            i_ts_nm1 = jnp.maximum(i_ts - 1, 0)
            q_nm1 = case.get_minimal_states(i_ts=i_ts_nm1)
            q_n = case.get_minimal_states(i_ts=i_ts)
            (
                p_res_p_q_nm1,
                p_res_p_q_n,
                p_v_dot_res_p_struct_dv_,
                p_aero_res_d_dv_,
                *_,
            ) = self.timestep_residual_jacobians(
                i_ts=i_ts,
                t=case.structure.t[i_ts],
                q_nm1=q_nm1,
                q_n=q_n,
                dv_=dv,
                dv_full=dv_full,
                thrust_t=case.structure.thrust,
                solve_dofs=solve_dofs,
                approx_grads=approx_grads,
                n_profile_loops=None,
                jac_options=jac_options,
            )
            p_j_n_p_q_n_, p_j_n_p_x_ = objective_jacobians(i_ts=i_ts, q_n=q_n)

            # solve adjoint step
            b = -(p_j_n_p_q_n_.reshape(n_j, -1) + adj_np1 @ p_r_np1_p_q_n).T
            adj_n = jnp.linalg.solve(p_res_p_q_n.T, b).T

            jax_print(
                "Solved adjoint for timestep {i_ts}",
                i_ts=i_ts,
                verbose_level="normal",
            )

            # add sentitivity of aerodynamic problem through full aero residual
            d_j_d_x_ += p_aero_res_d_dv_.premultiply_adj(adj_n[:, 4 * n_solve :])

            # add sensitivity of structural problem through v_dot residual
            d_j_d_x_.structure += p_v_dot_res_p_struct_dv_.premultiply_adj(
                adj_n[:, 2 * n_solve : 3 * n_solve]
            )
            d_j_d_x_ += p_j_n_p_x_

            if save_adjoint:
                adj_full_ = adj_full_.at[i_ts].set(adj_n)

            return d_j_d_x_, adj_n, p_res_p_q_nm1, adj_full_

        d_j_d_x, adj_last, p_r1_p_q0, adj_full = jax.lax.fori_loop(
            lower=0,
            upper=n_adj_iters,
            body_fun=step_body,
            init_val=(
                dv_grad_init,
                jnp.zeros((n_j, n_adj_dof)),
                jnp.zeros((n_adj_dof, n_adj_dof)),
                adj_full_init,
            ),
        )
        future_row = adj_last @ p_r1_p_q0

    # solve initial timestep adjoint, as there is no r0. Skipped when the adjoint window truncates early time steps
    if include_initial_state_grad and i_ts_start_adj_ <= 1:
        future_cot_q0_full = jnp.zeros((n_j, minimal_states_init.n_states))
        if case.structure.n_tstep > 1:
            future_cot_q0_full = future_cot_q0_full.at[:, free_state_ix].set(
                future_row
            )

        d_j_d_x += self._initial_timestep_grad_contribution(
            case=case[0].to_static(),
            objective=objective,
            grads_to_compute=grads_to_compute,
            p_varphi_p_x=p_varphi_p_x,
            solve_dofs=solve_dofs_arr,
            adj_t_p_r1_p_q0=future_cot_q0_full,
            horseshoe=static_horseshoe,
        )

    # restore original shape of j, and cut off zeros for past-end timestep and initial timestep which are always 0
    adj = (
        adj_full.reshape(adj_full.shape[0], *j_shape, *adj_full.shape[2:])[1:-1]
        if save_adjoint
        else None
    )

    d_j_d_x.mapping = dv.mapping

    return d_j_d_x, j_eval, adj

dynamic_adjoint_profile 

dynamic_adjoint_profile(
    case: AeroelasticCase,
    approx_grads: bool,
    jacobian_approximations: AeroelasticJacobianApproximations
    | None = None,
    grads_to_compute: AeroelasticGradsToCompute
    | None = None,
    i_ts: int = 1,
    n_profile_loops: int = 10,
) -> tuple[
    dict[str, dict[str, float]], dict[str, dict[str, float]]
]

Function to time evaluation of the Jacobians used for the coupled aeroelastic adjoint solution for the case where the full Jacobian is computed.

Parameters:

Name Type Description Default
case AeroelasticCase

Dynamic aeroelastic case from which to extract states.

 required
approx_grads bool

If True, neglect small gradient terms.

 required
jacobian_approximations AeroelasticJacobianApproximations | None

Define which blocks of the adjoint Jacobians will be substituted for approximations.

 None
grads_to_compute AeroelasticGradsToCompute | None

AeroelasticGradsToCompute object describing which design gradients to compute. If None, all gradients will be computed.

 None
i_ts int

Time step index where to evaluate residual Jacobians.

 1
n_profile_loops int

Number of times to loop the Jacobian evaluation time for averaging the runtime.

 10

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/coupled/gradients/coupled.py 
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
def dynamic_adjoint_profile(
    self,
    case: AeroelasticCase,
    approx_grads: bool,
    jacobian_approximations: AeroelasticJacobianApproximations | None = None,
    grads_to_compute: AeroelasticGradsToCompute | None = None,
    i_ts: int = 1,
    n_profile_loops: int = 10,
) -> tuple[dict[str, dict[str, float]], dict[str, dict[str, float]]]:
    r"""
    Function to time evaluation of the Jacobians used for the coupled aeroelastic adjoint solution for the case
    where the full Jacobian is computed.
    :param case: Dynamic aeroelastic case from which to extract states.
    :param approx_grads: If True, neglect small gradient terms.
    :param jacobian_approximations: Define which blocks of the adjoint Jacobians will be substituted for
    approximations.
    :param grads_to_compute: AeroelasticGradsToCompute object describing which design gradients to compute. If
    None, all gradients will be computed.
    :param i_ts: Time step index where to evaluate residual Jacobians.
    :param n_profile_loops: Number of times to loop the Jacobian evaluation time for averaging the runtime.
    :return: Dictionary of {residual_name: {gradient_argument: val}} for compile time and run time respectively.
    """

    print_table_title(inner_width=95, title="Aeroelastic Adjoint Profile")

    jac_options = self.construct_approximate_jacobians(
        sol=case,
        jacobian_approximations=jacobian_approximations
        if jacobian_approximations is not None
        else AeroelasticJacobianApproximations(),
    )

    *_, compile_time, run_time = self.timestep_residual_jacobians(
        i_ts=i_ts,
        t=case.aero.t[i_ts],
        q_nm1=case.get_minimal_states(i_ts=i_ts - 1),
        q_n=case.get_minimal_states(i_ts=i_ts),
        dv_=self.get_design_variables(case=case, grads_to_compute=grads_to_compute),
        dv_full=self.get_design_variables(case=case, grads_to_compute=None),
        thrust_t=case.structure.thrust,
        solve_dofs=get_solve_dofs(
            n_dof=self.structure.n_dof,
            prescribed_dofs=case.structure.prescribed_dofs,
        ),
        approx_grads=approx_grads,
        n_profile_loops=n_profile_loops,
        jac_options=jac_options,
    )

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

    print_table_line(inner_width=95)

    return compile_time, run_time

trim 

trim(
    prescribed_dofs: Sequence[int] | Array | slice | int,
    zero_force_dofs: Sequence[int] | Array | slice | int,
    trim_cs: Sequence[str | Sequence[str]] | str | None,
    thrust_nodes: Sequence[str | Sequence[str]]
    | str
    | None,
    trim_orientation: str | Sequence[str] | None = "x",
    trim_hinges: Sequence[str | Sequence[str]]
    | str
    | None = None,
    trim_f_abs_tolerance: float = 0.01,
    f_ext_follower: Array | None = None,
    f_ext_dead: Array | None = None,
    t: float | Array = 0.0,
    load_steps: int = 1,
    trim_relaxation: float = 0.9,
    horseshoe: bool = False,
    method: Literal[
        "adjoint", "finite_difference"
    ] = "finite_difference",
    broyden_fd_step: float = 0.001,
    max_iter: int = 100,
) -> tuple[AeroelasticCase, TrimVariables]

Trim an aircraft such that the resulting sum of forces on the aircraft is zero without any supports.

Parameters:

Name Type Description Default
prescribed_dofs Sequence[int] | Array | slice | int

Degrees of freedom which are clamped for the trim process.

 required
zero_force_dofs Sequence[int] | Array | slice | int

Degrees freedom where we wish to drive the clamping force to zero. This is not necessarily the same as prescribed_dofs, as in some cases there are degrees of freedom we will allow to have a non-zero clamping force. For example in the case of a clamped cantilever wing where we wish to find the angle of attack that gives lift equal to the weight, there would be a nonzero pitching moment.

 required
trim_cs Sequence[str | Sequence[str]] | str | None

Keys of control surfaces which are to be used to trim the aircraft. Each element may either be a single control-surface key (that surface gets its own independent deflection) or a sequence of keys (those surfaces are tied together and share a single deflection).

 required
thrust_nodes Sequence[str | Sequence[str]] | str | None

Keys of thrust nodes which are to be used to trim the aircraft. Each element may either be a single node key (that node gets its own independent thrust value) or a sequence of node keys (those nodes are tied together and share a single thrust value).

 required
trim_orientation str | Sequence[str] | None

Inertial axis (or axes if a sequence is provided) around which the aircraft is rotated about at the clamp to achieve trim.

 'x'
trim_hinges Sequence[str | Sequence[str]] | str | None

Names of MultibodyHinge constraints (keys into the structure's named constraints) whose rotation should be solved as a trim variable.

 None
trim_f_abs_tolerance float

Absolute maximum force residual at the clamped nodes for convergence to be achieved.

 0.01
f_ext_follower Array | None

External follower forces, [n_nodes, 6].

 None
f_ext_dead Array | None

external dead forces, [n_nodes, 6].

 None
t float | Array

Time at which to trim the aircraft, default zero.

 0.0
load_steps int

Number of load steps used for the static solution.

 1
trim_relaxation float

Relaxation factor for updates to degrees of freedom used to achieve trim.

 0.9
horseshoe bool

If true, use a horseshoe wake formulation.

 False
method Literal['adjoint', 'finite_difference']

"adjoint" rebuilds the trim Jacobian each iteration via the adjoint method. "finite_difference" approximates the Jacobian with a forward finite-difference sweep. Usually the latter is faster assuming a small number of trim variables.

 'finite_difference'
broyden_fd_step float

Step size used for the finite-difference bootstrap of the Broyden Jacobian.

 0.001
max_iter int

Maximum number of trim iterations. A warning is emitted if the routine fails to converge within this limit.

 100

Returns:

Type Description
tuple[AeroelasticCase, TrimVariables]

Aeroelastic solution object for the trimmed aircraft.
Source code in src/flapjax/coupled/gradients/coupled.py 
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
def trim(
    self,
    prescribed_dofs: Sequence[int] | Array | slice | int,
    zero_force_dofs: Sequence[int] | Array | slice | int,
    trim_cs: Sequence[str | Sequence[str]] | str | None,
    thrust_nodes: Sequence[str | Sequence[str]] | str | None,
    trim_orientation: str | Sequence[str] | None = "x",
    trim_hinges: Sequence[str | Sequence[str]] | str | None = None,
    trim_f_abs_tolerance: float = 1e-2,
    f_ext_follower: Array | None = None,
    f_ext_dead: Array | None = None,
    t: float | Array = 0.0,
    load_steps: int = 1,
    trim_relaxation: float = 0.9,
    horseshoe: bool = False,
    method: Literal["adjoint", "finite_difference"] = "finite_difference",
    broyden_fd_step: float = 1e-3,
    max_iter: int = 100,
) -> tuple[AeroelasticCase, TrimVariables]:
    r"""
    Trim an aircraft such that the resulting sum of forces on the aircraft is zero without any supports.
    :param prescribed_dofs: Degrees of freedom which are clamped for the trim process.
    :param zero_force_dofs: Degrees freedom where we wish to drive the clamping force to zero. This is not
    necessarily the same as prescribed_dofs, as in some cases there are degrees of freedom we will allow to have a
    non-zero clamping force. For example in the case of a clamped cantilever wing where we wish to find the angle
    of attack that gives lift equal to the weight, there would be a nonzero pitching moment.
    :param trim_cs: Keys of control surfaces which are to be used to trim the aircraft. Each element may either be
    a single control-surface key (that surface gets its own independent deflection) or a sequence of keys (those
    surfaces are tied together and share a single deflection).
    :param thrust_nodes: Keys of thrust nodes which are to be used to trim the aircraft. Each element may either be
    a single node key (that node gets its own independent thrust value) or a sequence of node keys (those nodes are
    tied together and share a single thrust value).
    :param trim_orientation: Inertial axis (or axes if a sequence is provided) around which the aircraft is
    rotated about at the clamp to achieve trim.
    :param trim_hinges: Names of ``MultibodyHinge`` constraints (keys into the structure's named constraints)
    whose rotation should be solved as a trim variable.
    :param trim_f_abs_tolerance: Absolute maximum force residual at the clamped nodes for convergence to be achieved.
    :param f_ext_follower: External follower forces, [n_nodes, 6].
    :param f_ext_dead: external dead forces, [n_nodes, 6].
    :param t: Time at which to trim the aircraft, default zero.
    :param load_steps: Number of load steps used for the static solution.
    :param trim_relaxation: Relaxation factor for updates to degrees of freedom used to achieve trim.
    :param horseshoe: If true, use a horseshoe wake formulation.
    :param method: "adjoint" rebuilds the trim Jacobian each iteration via the adjoint method. "finite_difference"
    approximates the Jacobian with a forward finite-difference sweep. Usually the latter is faster assuming a small
    number of trim variables.
    :param broyden_fd_step: Step size used for the finite-difference bootstrap of the Broyden Jacobian.
    :param max_iter: Maximum number of trim iterations. A warning is emitted if the routine fails to converge within
    this limit.
    :return: Aeroelastic solution object for the trimmed aircraft.
    """

    # parse groups for paired thrust nodes/control surfaces/hinges
    cs_groups: list[tuple[str, ...]] = parse_groups(trim_cs, "trim_cs")
    thrust_groups: list[tuple[str, ...]] = parse_groups(
        thrust_nodes, "thrust_nodes"
    )
    hinge_groups: list[tuple[str, ...]] = parse_groups(trim_hinges, "trim_hinges")

    check_unique_members(cs_groups, "trim_cs")
    check_unique_members(thrust_groups, "thrust_nodes")
    check_unique_members(hinge_groups, "trim_hinges")

    if hinge_groups and method == "adjoint":
        raise ValueError(
            "trim_hinges is only supported with method='finite_difference'."
        )

    trim_orientation_: Sequence[str] = (
        [trim_orientation]
        if isinstance(trim_orientation, str)
        else trim_orientation
        if trim_orientation is not None
        else []
    )

    zero_force_dofs_: tuple[int, ...] = self.structure.make_prescribed_dofs_tuple(
        zero_force_dofs
    )

    prescribed_dofs_: tuple[int, ...] = self.structure.make_prescribed_dofs_tuple(
        prescribed_dofs
    )

    if not self.structure.use_gravity:
        warn("Gravity is not enabled. Trim may result in unexpected behaviour.")

    # initial set of variables. Tied members share a single value, initialised from the first member.
    trim_variables_init: TrimVariables = TrimVariables(
        cs_ang={group_key(g): self.aero.cs_ang0[g[0]] for g in cs_groups},
        thrust={
            group_key(g): self.structure.thrust_reference[g[0]]
            for g in thrust_groups
        },
        trim_angles={
            k: self.structure.orientation_euler[ORIENTATION_DICT[k]]
            for k in trim_orientation_
        },
        hinge_angle={group_key(g): jnp.array(0.0) for g in hinge_groups},
    )

    ae_sol_init = self.reference_configuration(
        horseshoe=horseshoe,
        prescribed_dofs=prescribed_dofs_,
        use_f_ext_dead=f_ext_dead is not None,
        use_f_ext_follower=f_ext_follower is not None,
    )

    inner_case = deepcopy(self)

    if method == "adjoint":

        def trim_body(
            i_iter: int,
            trim_variables_: TrimVariables,
            sol_: AeroelasticCase,
            f_clamp_: Array,
        ) -> tuple[int, TrimVariables, AeroelasticCase, Array]:
            i_iter, _, tv, sol, fc = self.trim_iter(
                i_iter,
                inner_case,
                trim_variables_,
                sol_,
                f_clamp_,
                prescribed_dofs=prescribed_dofs_,
                zero_force_dofs=zero_force_dofs_,
                f_ext_dead=f_ext_dead,
                f_ext_follower=f_ext_follower,
                t=jnp.array(t),
                load_steps=load_steps,
                horseshoe=horseshoe,
                cs_groups=cs_groups,
                thrust_groups=thrust_groups,
                trim_orientation=trim_orientation_,
                trim_relaxation=trim_relaxation,
            )
            return i_iter, tv, sol, fc

        f_clamp_init = jnp.full((len(zero_force_dofs_)), 1e10)
        print_table_title(title="Trim (Adjoint)", inner_width=104)
        trim_variables_init.print_header(f_clamp=f_clamp_init)
        _, trim_variables, ae_sol, f_clamp_final = jax.lax.while_loop(
            lambda args_: jnp.logical_and(
                jnp.any(jnp.abs(args_[3]) >= trim_f_abs_tolerance),
                args_[0] < max_iter,
            ),
            body_fun=lambda args_: trim_body(*args_),
            init_val=(
                0,
                trim_variables_init,
                ae_sol_init,
                f_clamp_init,
            ),
        )
        print_table_line(inner_width=104)
    elif method == "finite_difference":
        print_table_title(title="Trim (Finite Difference)", inner_width=104)

        b_approx_init, f_clamp_init, ae_sol_bootstrap = self._trim_fd_jacobian(
            inner_case=inner_case,
            trim_variables=trim_variables_init,
            fd_step=broyden_fd_step,
            prescribed_dofs=prescribed_dofs_,
            zero_force_dofs=zero_force_dofs_,
            f_ext_follower=f_ext_follower,
            f_ext_dead=f_ext_dead,
            t=jnp.array(t),
            load_steps=load_steps,
            horseshoe=horseshoe,
            cs_groups=cs_groups,
            thrust_groups=thrust_groups,
            hinge_groups=hinge_groups,
            trim_orientation=trim_orientation_,
        )

        def trim_body_broyden(
            i_iter: int,
            trim_variables_: TrimVariables,
            sol_: AeroelasticCase,
            f_clamp_: Array,
            b_approx_: Array,
        ) -> tuple[int, TrimVariables, AeroelasticCase, Array, Array]:
            i_iter, _, tv, sol, fc, b = self._trim_iter_fd(
                i_iter,
                inner_case,
                trim_variables_,
                sol_,
                f_clamp_,
                b_approx_,
                prescribed_dofs=prescribed_dofs_,
                zero_force_dofs=zero_force_dofs_,
                f_ext_dead=f_ext_dead,
                f_ext_follower=f_ext_follower,
                t=jnp.array(t),
                load_steps=load_steps,
                horseshoe=horseshoe,
                cs_groups=cs_groups,
                thrust_groups=thrust_groups,
                hinge_groups=hinge_groups,
                trim_orientation=trim_orientation_,
                trim_relaxation=trim_relaxation,
            )
            return i_iter, tv, sol, fc, b

        trim_variables_init.print_header(f_clamp=f_clamp_init)
        _, trim_variables, ae_sol, f_clamp_final, _ = jax.lax.while_loop(
            lambda args_: jnp.logical_and(
                jnp.any(jnp.abs(args_[3]) >= trim_f_abs_tolerance),
                args_[0] < max_iter,
            ),
            body_fun=lambda args_: trim_body_broyden(*args_),
            init_val=(
                0,
                trim_variables_init,
                ae_sol_bootstrap,
                f_clamp_init,
                b_approx_init,
            ),
        )
        print_table_line(inner_width=104)
    else:
        raise ValueError(f"Unknown trim method: {method!r}.")

    if bool(jnp.any(jnp.isnan(f_clamp_final))):
        warn("Trim residual is NaN - solution diverged")
    elif bool(jnp.any(jnp.abs(f_clamp_final) >= trim_f_abs_tolerance)):
        warn(f"Trim did not converge within max_iter={max_iter} iterations ")

    new_orientation: Array = self.structure.orientation_euler
    for k, v in trim_variables.trim_angles.items():
        new_orientation = new_orientation.at[ORIENTATION_DICT[k]].set(v)

    # set solutions into case object
    self.set_design_variables(
        coords=self.structure.x0_reference,
        k_cs=self.structure.k_cs,
        m_cs=self.structure.m_cs,
        m_lumped=self.structure.m_lumped
        if self.structure.use_lumped_mass
        else None,
        dt=self.aero.dt,
        flowfield=self.aero.flowfield,
        delta_w=self.aero.delta_w,
        x0_aero=self.aero.zeta_b0,
        thrust_reference=self.structure.thrust_reference
        | expand_groups(trim_variables.thrust, thrust_groups),
        orientation_euler=new_orientation,
        cs_angles_reference=self.aero.cs_ang0
        | expand_groups(trim_variables.cs_ang, cs_groups),
        remove_checks=True,
    )

    if hinge_groups:
        # release the angle-pinning constraint used to condition the trim solve
        self._revert_hinge_trim(hinge_groups)

    return ae_sol, trim_variables

trim_angles_to_euler 

trim_angles_to_euler(
    trim_angles: dict[str, Array],
) -> Array

Find the 3 Euler angles describing the aircraft orientation. This allows for any combination to be set by the trim routine, with values not passed using the fixed values provided in the reference orientation.

Parameters:

Name Type Description Default
trim_angles dict[str, Array]

Dictionary of axis-angle pairs.

 required

Returns:

Type Description
Array

Euler angles describing the aircraft orientation, (3, ).
Source code in src/flapjax/coupled/gradients/coupled.py 
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
def trim_angles_to_euler(self, trim_angles: dict[str, Array]) -> Array:
    r"""
    Find the 3 Euler angles describing the aircraft orientation. This allows for any combination to be set by
    the trim routine, with values not passed using the fixed values provided in the reference orientation.
    :param trim_angles: Dictionary of axis-angle pairs.
    :return: Euler angles describing the aircraft orientation, ``(3, )``.
    """

    orientation_euler = self.structure.orientation_euler
    for k, v in trim_angles.items():
        orientation_euler = orientation_euler.at[ORIENTATION_DICT[k]].set(v)
    return orientation_euler

TrimVariables 

TrimVariables(
    cs_ang: dict[str, Array],
    thrust: dict[str, Array],
    trim_angles: dict[str, Array],
    hinge_angle: dict[str, Array] | None = None,
)
Source code in src/flapjax/coupled/gradients/data_structures.py 
41
42
43
44
45
46
47
48
49
50
51
def __init__(
    self,
    cs_ang: dict[str, Array],
    thrust: dict[str, Array],
    trim_angles: dict[str, Array],
    hinge_angle: dict[str, Array] | None = None,
):
    self.cs_ang: dict[str, Array] = cs_ang
    self.thrust: dict[str, Array] = thrust
    self.trim_angles: dict[str, Array] = trim_angles
    self.hinge_angle: dict[str, Array] = {} if hinge_angle is None else hinge_angle

print_header 

print_header(f_clamp: Array | None) -> None

Print the column-header row. Call once before the iteration loop.

Source code in src/flapjax/coupled/gradients/data_structures.py 
71
72
73
74
75
76
77
78
79
80
81
82
83
84
def print_header(self, f_clamp: Array | None) -> None:
    """Print the column-header row. Call once before the iteration loop."""
    specs = self._column_specs(f_clamp)
    col_widths = [max(len(self._header_label(k, u)), vw) for k, u, vw, _ in specs]

    cells = ["iter".rjust(self._ITER_W)]
    cells += [
        self._header_label(k, u).rjust(cw)
        for (k, u, _, _), cw in zip(specs, col_widths)
    ]
    inner = " | ".join(cells)
    padding = self._INNER_WIDTH - len(inner) - 2
    jax_print("| " + inner + " " * padding + " |", verbose_level="normal")
    print_table_line(inner_width=self._INNER_WIDTH)

print_values 

print_values(i_iter: int, f_clamp: Array | None) -> None

Print one row of numeric values, aligned with the header.

Source code in src/flapjax/coupled/gradients/data_structures.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
def print_values(self, i_iter: int, f_clamp: Array | None) -> None:
    """Print one row of numeric values, aligned with the header."""
    specs = self._column_specs(f_clamp)
    col_widths = [max(len(self._header_label(k, u)), vw) for k, u, vw, _ in specs]

    def _scalar(x: Array) -> Array:
        return jnp.ravel(x)[0]

    values: list[Array] = []
    values.extend(jnp.rad2deg(_scalar(v)) for v in self.cs_ang.values())
    values.extend(_scalar(v) for v in self.thrust.values())
    values.extend(jnp.rad2deg(_scalar(v)) for v in self.trim_angles.values())
    values.extend(jnp.rad2deg(_scalar(v)) for v in self.hinge_angle.values())
    if f_clamp is not None:
        values.extend(_scalar(f_clamp[i]) for i in range(f_clamp.shape[0]))

    placeholders = [f"c{i}" for i in range(len(specs))]
    cells = [f"{{i_iter:>{self._ITER_W}}}"]
    cells += [
        f"{{{ph}:>{cw}{fmt}}}"
        for ph, (_, _, _, fmt), cw in zip(placeholders, specs, col_widths)
    ]
    inner = " | ".join(cells)
    rendered_len = self._ITER_W + sum(col_widths) + 3 * len(specs)
    padding = self._INNER_WIDTH - rendered_len - 2

    kwargs = {"i_iter": i_iter, **dict(zip(placeholders, values))}
    jax_print("| " + inner + " " * padding + " |", **kwargs, verbose_level="normal")

LinearCoupled 

LinearCoupled(
    case: BaseCoupledAeroelastic,
    reference: AeroelasticCase,
    batch_size: int | None,
    n_struct_modes: int | None,
    wake_type: LinearWakeType = "frozen",
    bound_upwash: bool = True,
    wake_upwash: bool = False,
    unsteady_force: bool = True,
    int_order: Literal[3, 4, 5] = BASE_LOBATTO_ORDER,
    *,
    skip_checks: bool = False,
    prescribed_dofs: Sequence[int]
    | Array
    | slice
    | int
    | None = None,
)

Bases: LinearModel[AeroelasticCase, AeroelasticInputUnflattened, AeroelasticStateUnflattened, AeroelasticOutputUnflattened, AeroelasticLinearResult]

Source code in src/flapjax/coupled/linear/linear_coupled.py 
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
def __init__(
    self,
    case: BaseCoupledAeroelastic,
    reference: AeroelasticCase,
    batch_size: int | None,
    n_struct_modes: int | None,
    wake_type: LinearWakeType = "frozen",
    bound_upwash: bool = True,
    wake_upwash: bool = False,
    unsteady_force: bool = True,
    int_order: Literal[3, 4, 5] = BASE_LOBATTO_ORDER,
    *,
    skip_checks: bool = False,
    prescribed_dofs: Sequence[int] | Array | slice | int | None = None,
):
    if prescribed_dofs is not None:
        prescribed_dofs = case.structure.make_prescribed_dofs_tuple(prescribed_dofs)

    self.aero = LinearUVLM(
        case=case.aero,
        reference=reference.aero,
        wake_type=wake_type,
        bound_upwash=bound_upwash,
        wake_upwash=wake_upwash,
        unsteady_force=unsteady_force,
        skip_linearisation=True,
        skip_checks=skip_checks,
    )
    self.structure = LinearBeam(
        beam=case.structure,
        reference=reference.structure,
        dt=case.aero.dt,
        n_modes=n_struct_modes,
        int_order=int_order,
        prescribed_dofs=prescribed_dofs,
    )

    effective_prescribed = (
        prescribed_dofs
        if prescribed_dofs is not None
        else reference.structure.prescribed_dofs
    )
    self.n_beam_nodal_dof: int = case.structure.n_dof - len(effective_prescribed)
    self.n_beam_input_dof: int = (
        self.structure.n_modes
        if self.structure.modal_inputs
        else self.n_beam_nodal_dof
    )
    self.n_beam_state_dof: int = (
        self.structure.n_modes
        if self.structure.modal_states
        else self.n_beam_nodal_dof
    )
    self.n_beam_output_dof: int = (
        self.structure.n_modes
        if self.structure.modal_outputs
        else self.n_beam_nodal_dof
    )

    self.n_nodes: int = case.structure.n_nodes
    self.free_dofs: Array = jnp.array(
        get_solve_dofs(
            n_dof=case.structure.n_dof,
            prescribed_dofs=effective_prescribed,
        )
    )

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

    self._case: BaseCoupledAeroelastic = case
    self.unsteady_force: bool = unsteady_force
    if batch_size is not False:
        self.sys = self.linearise(batch_size=batch_size)

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 
286
287
288
289
290
291
292
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 
294
295
296
297
298
299
300
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 
362
363
364
365
366
367
368
369
370
371
372
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 
374
375
376
377
378
379
380
381
382
383
384
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 
386
387
388
389
390
391
392
393
394
395
396
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 
522
523
524
525
526
527
528
529
530
531
532
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 
534
535
536
537
538
539
540
541
542
543
544
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 
546
547
548
549
550
551
552
553
554
555
556
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 
558
559
560
561
562
563
564
565
566
567
568
569
570
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 
572
573
574
575
576
577
578
579
580
581
582
583
584
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 
586
587
588
589
590
591
592
593
594
595
596
597
598
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 
627
628
629
630
631
632
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 
634
635
636
637
638
639
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 
641
642
643
644
645
646
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)))

step 

step(
    gamma_b_vec: Array | None = None,
    gamma_w_vec: Array | None = None,
    gamma_b_nm1_vec: Array | None = None,
    zeta_w_vec: Array | None = None,
    nu_b_vec: Array | None = None,
    nu_w_vec: Array | None = None,
    f_ext: Array | None = None,
    q_nodal: Array | None = None,
    q_dot_nodal: Array | None = None,
) -> tuple[
    AeroelasticStateUnflattened,
    AeroelasticOutputUnflattened,
]

Step solution from states at timestep n and inputs at timestep n+1 to give states at timestep n+1 and outputs at timestep n.

Source code in src/flapjax/coupled/linear/linear_coupled.py 
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
def step(
    self,
    gamma_b_vec: Array | None = None,
    gamma_w_vec: Array | None = None,
    gamma_b_nm1_vec: Array | None = None,
    zeta_w_vec: Array | None = None,
    nu_b_vec: Array | None = None,
    nu_w_vec: Array | None = None,
    f_ext: Array | None = None,
    q_nodal: Array | None = None,
    q_dot_nodal: Array | None = None,
) -> tuple[AeroelasticStateUnflattened, AeroelasticOutputUnflattened]:
    r"""
    Step solution from states at timestep n and inputs at timestep n+1 to give states at timestep n+1 and outputs
    at timestep n.
    """
    ref = self.reference

    # unravel vector inputs, falling back to reference values when None
    gamma_b = (
        ArrayList.from_vector(vect=gamma_b_vec, shapes=ref.aero.gamma_b.shape)
        if gamma_b_vec is not None
        else ref.aero.gamma_b
    )
    gamma_w = (
        ArrayList.from_vector(vect=gamma_w_vec, shapes=ref.aero.gamma_w.shape)
        if gamma_w_vec is not None
        else ref.aero.gamma_w
    )
    gamma_b_nm1 = (
        (
            ArrayList.from_vector(
                vect=gamma_b_nm1_vec, shapes=ref.aero.gamma_b.shape
            )
            if gamma_b_nm1_vec is not None
            else ref.aero.gamma_b
        )
        if self.aero.unsteady_force
        else None
    )
    zeta_w = (
        ArrayList.from_vector(vect=zeta_w_vec, shapes=ref.aero.zeta_w.shape)
        if zeta_w_vec is not None
        else ref.aero.zeta_w
    )
    nu_b = (
        (
            ArrayList.from_vector(vect=nu_b_vec, shapes=ref.aero.zeta_b.shape)
            if nu_b_vec is not None
            else ArrayList.zeros_like(ref.aero.zeta_b)
        )
        if self.aero.bound_upwash
        else None
    )
    nu_w = (
        (
            ArrayList.from_vector(vect=nu_w_vec, shapes=ref.aero.zeta_w.shape)
            if nu_w_vec is not None
            else ArrayList.zeros_like(ref.aero.zeta_w)
        )
        if self.aero.wake_upwash
        else None
    )
    q_nodal = (
        q_nodal if q_nodal is not None else jnp.zeros(len(self.structure.free_dofs))
    )
    q_dot_nodal = (
        q_dot_nodal
        if q_dot_nodal is not None
        else jnp.zeros(len(self.structure.free_dofs))
    )

    # fill in prescribed dofs with zeros
    q_full = (
        jnp.zeros(self.n_nodes * 6)
        .at[self.free_dofs]
        .set(q_nodal)
        .reshape(self.n_nodes, 6)
    )
    q_dot_full = (
        jnp.zeros(self.n_nodes * 6)
        .at[self.free_dofs]
        .set(q_dot_nodal)
        .reshape(self.n_nodes, 6)
    )

    # total perturbed coordinates and time derivative
    hg = jnp.einsum("ijk,ikl->ijl", ref.structure.hg, vmap(exp_se3)(q_full))
    hg_dot = jnp.einsum(
        "ijk,ikl->ijl",
        ref.structure.hg,
        vmap(ha_to_ha_tilde)(q_dot_full),
    )

    # aerodynamic grid
    zeta_b = self.aero.case.hg_to_zeta_b(hg_n=hg, cs_ang_n=self.aero.case.cs_ang0)
    zeta_b_dot = self.aero.case.hg_dot_to_zeta_b_dot(
        hg_n=hg,
        hg_dot_n=hg_dot,
        cs_ang_n=self.aero.case.cs_ang0,
        cs_vel_n=self.aero.case.cs_vel0,
    )

    # pass through aerodynamic system
    u_n_aero = AeroInputUnflattened(
        zeta_b=zeta_b, zeta_b_dot=zeta_b_dot, nu_b=nu_b, nu_w=nu_w
    )

    x_n_aero = AeroStateUnflattened(
        gamma_b=gamma_b,
        gamma_w=gamma_w,
        gamma_b_nm1=gamma_b_nm1,
        zeta_w=zeta_w if self.aero.prescribed_wake else None,
        zeta_b=zeta_b if self.aero.prescribed_wake else None,
    )

    u_n_aero_vec = self.aero.pack_input_vector(u_n_aero)
    x_n_aero_vec = self.aero.pack_state_vector(x_n_aero)

    x_np1_aero_vec, y_np1_aero_vec = self.aero.step_vec(
        x_vec=x_n_aero_vec, u_vec=u_n_aero_vec
    )

    x_np1_aero = self.aero.unpack_state_vector(x=x_np1_aero_vec)
    y_np1_aero = self.aero.unpack_output_vector(y=y_np1_aero_vec)
    assert isinstance(x_np1_aero, AeroStateUnflattened) and isinstance(
        y_np1_aero, AeroOutputUnflattened
    ), (
        "Unpacked aero state and output must be of type AeroStateUnflattened and AeroOutputUnflattened."
    )

    # total aero forces on the grid, from the aero step (returns totals)
    f_aero_np1 = y_np1_aero.f_steady
    if self.unsteady_force:
        assert y_np1_aero.f_unsteady is not None
        # don't want to mutate in place
        # noinspection augment-assignment
        f_aero_np1 = f_aero_np1 + y_np1_aero.f_unsteady

    # project total aero forces onto the beam under the current (perturbed) rotation
    rmat = hg[:, :3, :3]
    f_aero_beam_total = project_forcing_to_beam(
        f_total=f_aero_np1,
        rmat=rmat,
        dof_mapping=self.aero.case.dof_mapping,
        x0_aero=self.aero.case.zeta_b0,
        mirror_edge_low=self.aero.case.mirror_edge_low,
        mirror_edge_high=self.aero.case.mirror_edge_high,
    )

    # subtract the reference contribution so the aero forcing fed into the beam operator
    # is a pure perturbation (the beam sys.a / sys.b operate on perturbations)
    f_aero_ref_total = ref.aero.f_steady
    if self.aero.unsteady_force:
        # noinspection augment-assignment
        f_aero_ref_total = f_aero_ref_total + ref.aero.f_unsteady
    f_aero_beam_ref = project_forcing_to_beam(
        f_total=f_aero_ref_total,
        rmat=ref.structure.hg[:, :3, :3],
        dof_mapping=self.aero.case.dof_mapping,
        x0_aero=self.aero.case.zeta_b0,
        mirror_edge_low=self.aero.case.mirror_edge_low,
        mirror_edge_high=self.aero.case.mirror_edge_high,
    )
    delta_f_aero_beam = f_aero_beam_total - f_aero_beam_ref

    f_ext_: Array = f_ext if f_ext is not None else jnp.zeros(self.free_dofs.size)

    # scatter free-dof forcing to full (n_nodes * 6) vectors expected by the beam B operator, and add aero
    # forcing (global frame, perturbation) as an external force
    f_ext_full = (
        jnp.zeros(self.n_nodes * 6).at[self.free_dofs].set(f_ext_)
        + delta_f_aero_beam.ravel()
    )  # [n_nodes * 6]

    if self.structure.modal_inputs:
        f_ext_full = self.structure.nodal_to_modal(f_ext_full[self.free_dofs])

    # beam step in perturbation form (discrete-time Tustin)
    x_beam_n = jnp.concatenate(
        [
            self.structure.nodal_to_modal(q_nodal)
            if self.structure.modal_states
            else q_nodal,
            self.structure.nodal_to_modal(q_dot_nodal)
            if self.structure.modal_states
            else q_dot_nodal,
        ]
    )

    x_beam_np1 = (
        self.structure.sys.a @ x_beam_n
        + self.structure.sys.b[:, self.structure.input_slices["f_ext"].slices]
        @ f_ext_full
    )

    q_np1 = x_beam_np1[: self.n_beam_state_dof]
    q_dot_np1 = x_beam_np1[self.n_beam_state_dof :]

    state_np1 = AeroelasticStateUnflattened(
        gamma_b=x_np1_aero.gamma_b,
        gamma_w=x_np1_aero.gamma_w,
        gamma_b_nm1=x_np1_aero.gamma_b_nm1,
        zeta_w=x_np1_aero.zeta_w,
        q=q_np1,
        q_dot=q_dot_np1,
    )
    output_n = AeroelasticOutputUnflattened(q=q_np1, q_dot=q_dot_np1)

    return state_np1, output_n

gamma_b_step 

gamma_b_step(
    gamma_b_n_vec: Array,
    gamma_w_n_vec: Array,
    q_n: Array,
    q_dot_n: Array,
    zeta_w_n_vec: Array | None = None,
    nu_b_n_vec: Array | None = None,
) -> Array

Bound circulation at timestep n+1 as a function of states at n and inputs at n+1.

Source code in src/flapjax/coupled/linear/linear_coupled.py 
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
def gamma_b_step(
    self,
    gamma_b_n_vec: Array,
    gamma_w_n_vec: Array,
    q_n: Array,
    q_dot_n: Array,
    zeta_w_n_vec: Array | None = None,
    nu_b_n_vec: Array | None = None,
) -> Array:
    r"""
    Bound circulation at timestep n+1 as a function of states at n and inputs at n+1.
    """
    x_np1, _ = self.step(
        nu_b_vec=nu_b_n_vec,
        gamma_b_vec=gamma_b_n_vec,
        gamma_w_vec=gamma_w_n_vec,
        zeta_w_vec=zeta_w_n_vec,
        q_nodal=self.structure.modal_to_nodal(q_n)
        if self.structure.modal_states
        else q_n,
        q_dot_nodal=self.structure.modal_to_nodal(q_dot_n)
        if self.structure.modal_states
        else q_dot_n,
    )
    return x_np1.gamma_b.ravel()

wake_prop_step 

wake_prop_step(
    gamma_b_n_vec: Array,
    gamma_w_n_vec: Array,
    q_n: Array,
    zeta_w_n_vec: Array | None = None,
    nu_w_n_vec: Array | None = None,
) -> tuple[Array | None, Array]

Wake propagation as a function of states at n and inputs at n+1.

Source code in src/flapjax/coupled/linear/linear_coupled.py 
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
def wake_prop_step(
    self,
    gamma_b_n_vec: Array,
    gamma_w_n_vec: Array,
    q_n: Array,
    zeta_w_n_vec: Array | None = None,
    nu_w_n_vec: Array | None = None,
) -> tuple[Array | None, Array]:
    r"""
    Wake propagation as a function of states at n and inputs at n+1.
    """
    x_new, _ = self.step(
        nu_w_vec=nu_w_n_vec,
        gamma_b_vec=gamma_b_n_vec,
        gamma_w_vec=gamma_w_n_vec,
        zeta_w_vec=zeta_w_n_vec,
        q_nodal=self.structure.modal_to_nodal(q_n)
        if self.structure.modal_states
        else q_n,
    )
    assert not ((x_new.zeta_w is not None) ^ self.aero.prescribed_wake), (
        "zeta_w should be None only if prescribed_wake is False."
    )
    return (
        x_new.zeta_w.ravel() if x_new.zeta_w is not None else None,
        x_new.gamma_w.ravel(),
    )

q_step 

q_step(
    q_n: Array,
    q_dot_n: Array,
    gamma_b_n_vec: Array,
    gamma_w_n_vec: Array,
    f_ext: Array | None,
    zeta_w_n_vec: Array | None = None,
    gamma_b_nm1_vec: Array | None = None,
    nu_b_n_vec: Array | None = None,
    nu_w_n_vec: Array | None = None,
) -> Array

Beam displacement at timestep n+1 as a function of states at n and external forcing.

Source code in src/flapjax/coupled/linear/linear_coupled.py 
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
def q_step(
    self,
    q_n: Array,
    q_dot_n: Array,
    gamma_b_n_vec: Array,
    gamma_w_n_vec: Array,
    f_ext: Array | None,
    zeta_w_n_vec: Array | None = None,
    gamma_b_nm1_vec: Array | None = None,
    nu_b_n_vec: Array | None = None,
    nu_w_n_vec: Array | None = None,
) -> Array:
    r"""
    Beam displacement at timestep n+1 as a function of states at n and external forcing.
    """
    x_new, _ = self.step(
        f_ext=f_ext,
        gamma_b_vec=gamma_b_n_vec,
        gamma_w_vec=gamma_w_n_vec,
        gamma_b_nm1_vec=gamma_b_nm1_vec,
        zeta_w_vec=zeta_w_n_vec,
        nu_b_vec=nu_b_n_vec,
        nu_w_vec=nu_w_n_vec,
        q_nodal=self.structure.modal_to_nodal(q_n)
        if self.structure.modal_states
        else q_n,
        q_dot_nodal=self.structure.modal_to_nodal(q_dot_n)
        if self.structure.modal_states
        else q_dot_n,
    )
    return x_new.q.ravel()

q_dot_step 

q_dot_step(
    q_n: Array,
    q_dot_n: Array,
    gamma_b_n_vec: Array,
    gamma_w_n_vec: Array,
    f_ext: Array | None,
    zeta_w_n_vec: Array | None,
    gamma_b_nm1_vec: Array | None = None,
    nu_b_n_vec: Array | None = None,
    nu_w_n_vec: Array | None = None,
) -> Array

Beam velocity at timestep n+1 as a function of states at n and external forcing.

Source code in src/flapjax/coupled/linear/linear_coupled.py 
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
def q_dot_step(
    self,
    q_n: Array,
    q_dot_n: Array,
    gamma_b_n_vec: Array,
    gamma_w_n_vec: Array,
    f_ext: Array | None,
    zeta_w_n_vec: Array | None,
    gamma_b_nm1_vec: Array | None = None,
    nu_b_n_vec: Array | None = None,
    nu_w_n_vec: Array | None = None,
) -> Array:
    r"""
    Beam velocity at timestep n+1 as a function of states at n and external forcing.
    """
    x_new, _ = self.step(
        f_ext=f_ext,
        gamma_b_vec=gamma_b_n_vec,
        gamma_w_vec=gamma_w_n_vec,
        gamma_b_nm1_vec=gamma_b_nm1_vec,
        zeta_w_vec=zeta_w_n_vec,
        nu_b_vec=nu_b_n_vec,
        nu_w_vec=nu_w_n_vec,
        q_nodal=self.structure.modal_to_nodal(q_n)
        if self.structure.modal_states
        else q_n,
        q_dot_nodal=self.structure.modal_to_nodal(q_dot_n)
        if self.structure.modal_states
        else q_dot_n,
    )
    return x_new.q_dot.ravel()

compute_jacobians 

compute_jacobians() -> tuple[
    dict[
        str,
        tuple[
            Callable[..., Any],
            dict[str, Any],
            Sequence[str],
        ],
    ],
    dict[str, dict[str, Callable[..., Array]]],
]

Returns:

Type Description
tuple[dict[str, tuple[Callable[..., Any], dict[str, Any], Sequence[str]]], dict[str, dict[str, Callable[..., Array]]]]

Tuple. First entry is dictionaries with keys being the function name (e.g., gamma_b, gamma_w), with each entry containing the relevant stepping function, the arguments for the function, and the name of the arguments for which to obtain derivatives. Second entry is a dictionary of explicit Jacobians functions that take the same arguments as the first output.
Source code in src/flapjax/coupled/linear/linear_coupled.py 
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
def compute_jacobians(
    self,
) -> tuple[
    dict[str, tuple[Callable[..., Any], dict[str, Any], Sequence[str]]],
    dict[str, dict[str, Callable[..., Array]]],
]:
    r"""
    :return: Tuple. First entry is dictionaries with keys being the function name (e.g., gamma_b, gamma_w), with
    each entry containing the relevant stepping function, the arguments for the function, and the name of the
    arguments for which to obtain derivatives. Second entry is a dictionary of explicit Jacobians functions that
    take the same arguments as the first output.
    """
    ref = self.reference

    # bound circulation
    gamma_b_args: dict[str, Any] = {
        "gamma_b_n_vec": ref.aero.gamma_b.ravel(),
        "gamma_w_n_vec": ref.aero.gamma_w.ravel(),
        "q_n": jnp.zeros((self.n_beam_state_dof,)),
        "q_dot_n": jnp.zeros((self.n_beam_state_dof,)),
        "zeta_w_n_vec": ref.aero.zeta_w.ravel(),
        "nu_b_n_vec": jnp.zeros(ref.aero.zeta_b.size)
        if self.aero.bound_upwash
        else None,
    }
    gamma_b_diff = ["gamma_w_n_vec", "q_n", "q_dot_n"]
    if self.aero.prescribed_wake:
        gamma_b_diff.append("zeta_w_n_vec")
    if self.aero.bound_upwash:
        gamma_b_diff.append("nu_b_n_vec")

    # wake
    wake_args: dict[str, Any] = {
        "gamma_b_n_vec": ref.aero.gamma_b.ravel(),
        "gamma_w_n_vec": ref.aero.gamma_w.ravel(),
        "q_n": jnp.zeros((self.n_beam_state_dof,)),
        "zeta_w_n_vec": ref.aero.zeta_w.ravel(),
        "nu_w_n_vec": jnp.zeros(ref.aero.zeta_w.size)
        if self.aero.wake_upwash
        else None,
    }
    gamma_w_diff = ["gamma_b_n_vec", "gamma_w_n_vec"]
    zeta_w_diff = ["q_n"]
    if self.aero.prescribed_wake:
        zeta_w_diff.append("zeta_w_n_vec")
    if self.aero.wake_upwash:
        zeta_w_diff.append("nu_w_n_vec")
    if self.aero.free_wake:
        zeta_w_diff.extend(["gamma_b_n_vec", "gamma_w_n_vec"])

    q_args: dict[str, Any] = {
        "q_n": jnp.zeros((self.n_beam_state_dof,)),
        "q_dot_n": jnp.zeros((self.n_beam_state_dof,)),
        "gamma_b_n_vec": ref.aero.gamma_b.ravel(),
        "gamma_w_n_vec": ref.aero.gamma_w.ravel(),
        "f_ext": jnp.zeros((self.n_beam_input_dof,)),
        "zeta_w_n_vec": ref.aero.zeta_w.ravel(),
        "gamma_b_nm1_vec": ref.aero.gamma_b.ravel()
        if self.aero.unsteady_force
        else None,
        "nu_b_n_vec": jnp.zeros(ref.aero.zeta_b.size)
        if self.aero.bound_upwash
        else None,
        "nu_w_n_vec": jnp.zeros(ref.aero.zeta_w.size)
        if self.aero.wake_upwash
        else None,
    }
    q_diff = [
        "q_n",
        "q_dot_n",
        "gamma_b_n_vec",
        "gamma_w_n_vec",
        "f_ext",
    ]

    if self.aero.prescribed_wake:
        q_diff.append("zeta_w_n_vec")
    if self.aero.unsteady_force:
        q_diff.append("gamma_b_nm1_vec")
    if self.aero.bound_upwash:
        q_diff.append("nu_b_n_vec")
    if self.aero.wake_upwash:
        q_diff.append("nu_w_n_vec")

    linear_args: dict[
        str, tuple[Callable[..., Any], dict[str, Any], Sequence[str]]
    ] = {
        "gamma_b": (self.gamma_b_step, gamma_b_args, gamma_b_diff),
        "gamma_w": (
            lambda *args, **kwargs: self.wake_prop_step(**kwargs)[1],
            wake_args,
            gamma_w_diff,
        ),
        "gamma_b_nm1": (
            lambda *args, **kwargs: None,
            {
                "gamma_b_n_vec": ref.aero.gamma_b.ravel(),
            },
            ["gamma_b_n_vec"],
        ),
        "q": (self.q_step, q_args, q_diff),
        "q_dot": (self.q_dot_step, q_args, q_diff),
    }

    # add zeta_w for linearisation
    if self.aero.prescribed_wake:
        linear_args["zeta_w"] = (
            lambda *args, **kwargs: self.wake_prop_step(**kwargs)[0],
            wake_args,
            zeta_w_diff,
        )

    # define Jacobians we know nicely as jac_options
    jac_options = {
        "q": {
            "q_n": lambda *args, **kwargs: self.structure.sys.a[
                : self.n_beam_state_dof, : self.n_beam_state_dof
            ],
            "q_dot_n": lambda *args, **kwargs: self.structure.sys.a[
                : self.n_beam_state_dof, self.n_beam_state_dof :
            ],
        },
        "q_dot": {
            "q_n": lambda *args, **kwargs: self.structure.sys.a[
                self.n_beam_state_dof :, : self.n_beam_state_dof
            ],
            "q_dot_n": lambda *args, **kwargs: self.structure.sys.a[
                self.n_beam_state_dof :, self.n_beam_state_dof :
            ],
        },
    }
    if self.aero.unsteady_force:
        jac_options["gamma_b_nm1"] = {
            "gamma_b_n_vec": lambda *args, **kwargs: jnp.eye(ref.aero.gamma_b.size)
        }

    return linear_args, jac_options

create_jacobians 

create_jacobians(
    mode: ADMode | dict[str, ADMode] = "reverse",
    batch_size: int | None = None,
    n_profile_loops: int | None = None,
    jac_options: dict[
        str, dict[str, Callable[..., Any] | None]
    ]
    | None = None,
) -> tuple[
    dict[str, dict[str, Array]],
    dict[str, dict[str, float]] | None,
    dict[str, dict[str, float]] | None,
]

Assemble the per-residual Jacobians for the coupled linearisation. When n_profile_loops is set, all Jacobian constructions are looped locally so that they can be timed.

Source code in src/flapjax/coupled/linear/linear_coupled.py 
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
def create_jacobians(
    self,
    mode: ADMode | dict[str, ADMode] = "reverse",
    batch_size: int | None = None,
    n_profile_loops: int | None = None,
    jac_options: dict[str, dict[str, Callable[..., Any] | None]] | None = None,
) -> tuple[
    dict[str, dict[str, Array]],
    dict[str, dict[str, float]] | None,
    dict[str, dict[str, float]] | None,
]:
    """
    Assemble the per-residual Jacobians for the coupled linearisation. When
    ``n_profile_loops`` is set, all Jacobian constructions are looped locally so that they can be timed.
    """
    res_args, jac_options_exp = self.compute_jacobians()
    jac_options_total: dict[str, dict[str, Callable[..., Any] | None]] = (
        jac_options if jac_options is not None else {}
    ) | jac_options_exp

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

    aero_residual_names = {"gamma_b", "gamma_w", "gamma_b_nm1", "zeta_w"}
    delegate_aero = n_profile_loops is None

    for res_name, (res_func, args, diff_arg_names) in res_args.items():
        if delegate_aero and res_name in aero_residual_names:
            continue

        res_jac_options: dict[str, Callable[..., Any] | None] = {
            arg: None for arg in diff_arg_names
        }
        if res_name in jac_options_total:
            for arg, entry in jac_options_total[res_name].items():
                if arg in res_jac_options:
                    res_jac_options[arg] = entry

        if isinstance(mode, str):
            res_mode: ADMode = mode
        elif isinstance(mode, dict):
            try:
                res_mode = mode[res_name]
            except KeyError:
                res_mode = "reverse"
        else:
            raise NotImplementedError

        jacs, res_compile_time, res_run_time = jacrev_custom(
            func=res_func,
            jac_options=res_jac_options,
            n_profile_loops=n_profile_loops,
            func_name=res_name,
            map_batch_size=batch_size,
            mode=res_mode,
        )(**args)

        jacobians[res_name] = jacs
        if n_profile_loops is not None:
            assert res_compile_time is not None and res_run_time is not None
            compile_time[res_name] = res_compile_time
            run_time[res_name] = res_run_time

    if delegate_aero:
        # delegate aero linearisation to the aero system, with the beam kinematics wrapped as an input projection
        beam_proj = self._build_beam_projection()
        aero_mode: ADMode | dict[str, ADMode]
        if isinstance(mode, dict):
            aero_mode = {k: mode[k] for k in aero_residual_names if k in mode}
        else:
            aero_mode = mode
        needed_aero_residuals = {"gamma_b", "gamma_w"}
        if self.aero.unsteady_force:
            needed_aero_residuals.add("gamma_b_nm1")
        if self.aero.prescribed_wake:
            needed_aero_residuals.add("zeta_w")
        aero_jacs = self.aero.create_jacobians(
            mode=aero_mode,
            batch_size=batch_size,
            input_projection=beam_proj,
            residual_names=tuple(needed_aero_residuals),
        )
        jacobians.update(aero_jacs)

    return (
        jacobians,
        compile_time if n_profile_loops is not None else None,
        run_time if n_profile_loops is not None else None,
    )

linearise_profile 

linearise_profile(
    n_profile_loops: int = 3,
) -> tuple[
    dict[str, dict[str, float]], dict[str, dict[str, float]]
]

Profile forming the Jacobians required for the linearised model.

Parameters:

Name Type Description Default
n_profile_loops int

Number of times to loop Jacobian creation for averaging.

 3

Returns:

Type Description
tuple[dict[str, dict[str, float]], dict[str, dict[str, float]]]

Dictionaries of compile and run times for each sub function.
Source code in src/flapjax/coupled/linear/linear_coupled.py 
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
def linearise_profile(
    self,
    n_profile_loops: int = 3,
) -> tuple[dict[str, dict[str, float]], dict[str, dict[str, float]]]:
    r"""
    Profile forming the Jacobians required for the linearised model.
    :param n_profile_loops: Number of times to loop Jacobian creation for averaging.
    :return: Dictionaries of compile and run times for each sub function.
    """

    print_table_title(inner_width=95, title="Aeroelastic Adjoint Profile")

    _, compile_time, run_time = self.create_jacobians(
        n_profile_loops=n_profile_loops,
        jac_options=None,
        mode={"gamma_b": "forward"} if self.structure.modal_states else {},  # type: ignore
    )

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

    print_table_line(inner_width=95)

    return compile_time, run_time

modal 

modal(
    n_modes: int | None = None,
    freq_range: tuple[float | Array, float | Array] = (
        0.0,
        jnp.inf,
    ),
    damp_range: tuple[float | Array, float | Array] = (
        -jnp.inf,
        jnp.inf,
    ),
    min_struct_content: float | Array = 0.0,
    remove_complex_conjugate: bool = True,
    plot_eigvals: bool = False,
    sort: Literal["frequency", "damping"] = "frequency",
    plot_xlim: tuple[float, float] = (-500.0, 50.0),
    plot_ylim: tuple[float, float] = (-400.0, 400.0),
    n_plot_vtk: int = 0,
    vtu_directory: PathLike | str = "./modal",
    n_phase: int = 8,
    n_interp: int = 0,
    max_disp: float = 0.2,
    max_ang: float = 0.2,
    max_gamma: float = 100.0,
) -> Array

Compute stability eigenvalues of the linear system A matrix.

Parameters:

Name Type Description Default
n_modes int | None

Number of modes to be kept. If None, all eigenvalues are returned.

 None
freq_range tuple[float | Array, float | Array]

(min, max) natural frequency window in Hz. Modes outside are pushed past the truncation and dropped when n_modes is set.

 (0.0, inf)
damp_range tuple[float | Array, float | Array]

(min, max) damping-ratio window. Modes outside are pushed past the truncation and dropped when n_modes is set.

 (-inf, inf)
min_struct_content float | Array

Minimum fraction of eigenvector energy that must live in the beam (q, q_dot) states in range [0, 1]. Modes below this threshold (typically wake convection modes) are pushed past the truncation.

 0.0
remove_complex_conjugate bool

If true, one mode from each complex-conjugate pair is dropped.

 True
plot_eigvals bool

If true, plot the eigenvalues with Matplotlib.

 False
sort Literal['frequency', 'damping']

Method for sorting eigenvalues before truncation, can be either "frequency" or "damping".

 'frequency'
plot_xlim tuple[float, float]

Range of real component to be used for plotting.

 (-500.0, 50.0)
plot_ylim tuple[float, float]

Range of imaginary component to be used for plotting.

 (-400.0, 400.0)
n_plot_vtk int

Number of modes (starting from the most damped) to write to VTK for visualisation. Set to 0 to skip plotting.

 0
vtu_directory PathLike | str

Directory for plotting vtu files, defaults to "./modal".

 './modal'
n_phase int

Number of phase samples of the complex eigenvector to plot per mode.

 8
n_interp int

Number of interpolation points to add along each beam element in the beam VTU output.

 0
max_disp float

Maximum linear displacement used to normalise the plotted mode shape (in reference units).

 0.2
max_ang float

Maximum angular displacement used to normalise the plotted mode shape.

 0.2
max_gamma float

Maximum circulation used to normalise the plotted mode shape.

 100.0

Returns:

Type Description
Array

Continuous-time eigenvalues of the system A matrix, (n_states, ) or (n_states, 2) if to_components=True.
Source code in src/flapjax/coupled/linear/linear_coupled.py 
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
def modal(
    self,
    n_modes: int | None = None,
    freq_range: tuple[float | Array, float | Array] = (0.0, jnp.inf),
    damp_range: tuple[float | Array, float | Array] = (-jnp.inf, jnp.inf),
    min_struct_content: float | Array = 0.0,
    remove_complex_conjugate: bool = True,
    plot_eigvals: bool = False,
    sort: Literal["frequency", "damping"] = "frequency",
    plot_xlim: tuple[float, float] = (-500.0, 50.0),
    plot_ylim: tuple[float, float] = (-400.0, 400.0),
    n_plot_vtk: int = 0,
    vtu_directory: os.PathLike | str = "./modal",
    n_phase: int = 8,
    n_interp: int = 0,
    max_disp: float = 0.2,
    max_ang: float = 0.2,
    max_gamma: float = 100.0,
) -> Array:
    r"""
    Compute stability eigenvalues of the linear system A matrix.
    :param n_modes: Number of modes to be kept. If None, all eigenvalues are returned.
    :param freq_range: (min, max) natural frequency window in Hz. Modes outside are pushed past the
    truncation and dropped when n_modes is set.
    :param damp_range: (min, max) damping-ratio window. Modes outside are pushed past the truncation and
    dropped when n_modes is set.
    :param min_struct_content: Minimum fraction of eigenvector energy that must live in the beam
    ``(q, q_dot)`` states in range `[0, 1]`. Modes below this threshold (typically wake convection modes) are pushed
    past the truncation.
    :param remove_complex_conjugate: If true, one mode from each complex-conjugate pair is dropped.
    :param plot_eigvals: If true, plot the eigenvalues with Matplotlib.
    :param sort: Method for sorting eigenvalues before truncation, can be either "frequency" or "damping".
    :param plot_xlim: Range of real component to be used for plotting.
    :param plot_ylim: Range of imaginary component to be used for plotting.
    :param n_plot_vtk: Number of modes (starting from the most damped) to write to VTK for visualisation. Set to 0
    to skip plotting.
    :param vtu_directory: Directory for plotting vtu files, defaults to "./modal".
    :param n_phase: Number of phase samples of the complex eigenvector to plot per mode.
    :param n_interp: Number of interpolation points to add along each beam element in the beam VTU output.
    :param max_disp: Maximum linear displacement used to normalise the plotted mode shape (in reference units).
    :param max_ang: Maximum angular displacement used to normalise the plotted mode shape.
    :param max_gamma: Maximum circulation used to normalise the plotted mode shape.
    :return: Continuous-time eigenvalues of the system A matrix, ``(n_states, )`` or ``(n_states, 2)`` if ``to_components=True``.
    """

    evals_d, evecs = jnp.linalg.eig(self.sys.a)
    evals = jnp.log(evals_d) / self.dt  # convert to continuous time

    # order from most to least damped and truncate
    omega_damped = jnp.abs(evals.imag)
    damping = -evals.real / jnp.abs(evals)
    omega_natural = omega_damped / jnp.sqrt(1.0 - damping**2)

    freq_natural_hz = omega_natural / (2.0 * jnp.pi)

    match sort:
        case "frequency":
            idx = omega_natural.argsort()
        case "damping":
            idx = damping.argsort()

    # push conjugate partners past the truncation point (indexed by original position, then re-sorted so it
    # aligns with `idx`). Stable-argsort preserves the primary sort within each group.
    if remove_complex_conjugate:
        partner = conjugate_partner_mask(
            freq_hz=freq_natural_hz, damping=damping, tiebreaker=evals.real
        )
        idx = idx[jnp.argsort(partner[idx], stable=True)]

    # fraction of eigenvector energy in the structural states — used to reject
    # aero-only modes (e.g. wake convection)
    q_slice = self.state_slices["q"].slices
    q_dot_slice = self.state_slices["q_dot"].slices
    evec_sq = jnp.abs(evecs) ** 2
    struct_content = (
        evec_sq[q_slice].sum(axis=0) + evec_sq[q_dot_slice].sum(axis=0)
    ) / evec_sq.sum(axis=0)

    # push modes outside the requested natural-frequency / damping window to the back so truncation to
    # n_modes keeps only the in-range ones.
    in_range = (
        (freq_natural_hz[idx] >= freq_range[0])
        & (freq_natural_hz[idx] <= freq_range[1])
        & (damping[idx] >= damp_range[0])
        & (damping[idx] <= damp_range[1])
        & (struct_content[idx] >= min_struct_content)
    )
    idx = idx[jnp.argsort(~in_range, stable=True)]

    if n_modes is not None:
        idx = idx[:n_modes]

    freq_damped_ordered = omega_damped[idx] / (2.0 * jnp.pi)
    freq_natural_ordered = omega_natural[idx] / (2.0 * jnp.pi)
    damping_ordered = damping[idx]

    # write to console
    if n_modes is not None:
        print_table_line(inner_width=71)
        jax_print(
            "| Mode | Damped Frequency [Hz] | Natural Frequency [Hz] | Damping Ratio |",
            verbose_level="normal",
        )
        print_table_line(inner_width=71)
        for i_mode in range(n_modes):
            jax_print(
                "| {mode:>4d} | {freq_damped:>21.3f} | {freq_natural:>22.3f} | {damp:>13.6f} |",
                mode=i_mode + 1,
                freq_damped=freq_damped_ordered[i_mode],
                freq_natural=freq_natural_ordered[i_mode],
                damp=damping_ordered[i_mode],
                verbose_level="normal",
            )
        print_table_line(inner_width=71)

    if plot_eigvals:
        _, ax = plt.subplots()
        ax.scatter(
            evals.real,
            evals.imag,
        )
        ax.set_xlim(*plot_xlim)
        ax.set_ylim(*plot_ylim)
        ax.set_xlabel("Re(eig) [1/s]")
        ax.set_ylabel("Im(eig) [1/s]")
        ax.set_title("Eigenvalues")
        plt.show()

    if n_plot_vtk > 0:
        evecs_ordered = evecs[:, idx[:n_plot_vtk]].T  # (m, n_states)

        q_mode = evecs_ordered[
            :, self.state_slices["q"].slices
        ]  # (m, n_free_dof | n_modes)
        if self.structure.modal_states:
            q_mode = self.structure.modal_to_nodal(q_mode)  # (m, n_free_dof)
        q_full = (
            jnp.zeros((n_plot_vtk, self.n_nodes * 6), dtype=complex)
            .at[:, self.free_dofs]
            .set(q_mode)
        )

        def _extract_array_list(name: str) -> ArrayList | None:
            component = self.state_slices[name]
            if not component.enabled:
                return None
            return ArrayList(
                [
                    evecs_ordered[:, s].reshape((n_plot_vtk, *shape))
                    for s, shape in zip(component.slices, component.shapes)
                ]
            )

        gamma_b_full = _extract_array_list("gamma_b")
        gamma_w_full = _extract_array_list("gamma_w")
        zeta_w_full = _extract_array_list("zeta_w")

        plot_modes_vtu(
            reference=self.reference,
            directory=vtu_directory,
            q_full=q_full.reshape(n_plot_vtk, self.n_nodes, 6),
            freqs=freq_damped_ordered,
            dampings=damping_ordered,
            gamma_b_full=gamma_b_full,
            gamma_w_full=gamma_w_full,
            zeta_w_full=zeta_w_full,
            uvlm=self.aero.case,
            n_phase=n_phase,
            n_interp=n_interp,
            max_disp=max_disp,
            max_ang=max_ang,
            max_gamma=max_gamma,
        )

    return evals[idx]

modal_rescaled 

modal_rescaled(
    velocity: float | Array,
    density: float | Array,
    chord: float | Array,
    n_modes: int | None = None,
    freq_range: tuple[float | Array, float | Array] = (
        0.0,
        jnp.inf,
    ),
    damp_range: tuple[float | Array, float | Array] = (
        -jnp.inf,
        jnp.inf,
    ),
    min_struct_content: float | Array = 0.0,
    remove_complex_conjugate: bool = True,
    sort: Literal["frequency", "damping"] = "frequency",
) -> Array

Compute eigenvalues of the rescaled linear system at a new velocity, density, and chord length without re-linearising.

Parameters:

Name Type Description Default
velocity float | Array

Freestream velocity magnitude(s) at the new condition(s), scalar or (*n_points,).

 required
density float | Array

Flow density(s) at the new condition(s), scalar or broadcastable with velocity.

 required
chord float | Array

Reference chord length(s) at the new condition(s), scalar or broadcastable with velocity.

 required
n_modes int | None

Number of modes to keep. If None, all eigenvalues are returned.

 None
freq_range tuple[float | Array, float | Array]

(min, max) natural frequency window in Hz.

 (0.0, inf)
damp_range tuple[float | Array, float | Array]

(min, max) damping-ratio window.

 (-inf, inf)
min_struct_content float | Array

Minimum structural eigenvector energy fraction in [0, 1].

 0.0
remove_complex_conjugate bool

Drop one partner from each conjugate pair.

 True
sort Literal['frequency', 'damping']

Sort eigenvalues by "frequency" or "damping".

 'frequency'

Returns:

Type Description
Array

Continuous-time eigenvalues of the rescaled system(s), (n_out,) for scalar inputs or (*n_points, n_out) for batched inputs.
Source code in src/flapjax/coupled/linear/linear_coupled.py 
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
def modal_rescaled(
    self,
    velocity: float | Array,
    density: float | Array,
    chord: float | Array,
    n_modes: int | None = None,
    freq_range: tuple[float | Array, float | Array] = (0.0, jnp.inf),
    damp_range: tuple[float | Array, float | Array] = (-jnp.inf, jnp.inf),
    min_struct_content: float | Array = 0.0,
    remove_complex_conjugate: bool = True,
    sort: Literal["frequency", "damping"] = "frequency",
) -> Array:
    r"""
    Compute eigenvalues of the rescaled linear system at a new velocity, density, and chord length without
    re-linearising.
    :param velocity: Freestream velocity magnitude(s) at the new condition(s), scalar or ``(*n_points,)``.
    :param density: Flow density(s) at the new condition(s), scalar or broadcastable with ``velocity``.
    :param chord: Reference chord length(s) at the new condition(s), scalar or broadcastable with ``velocity``.
    :param n_modes: Number of modes to keep. If ``None``, all eigenvalues are returned.
    :param freq_range: (min, max) natural frequency window in Hz.
    :param damp_range: (min, max) damping-ratio window.
    :param min_struct_content: Minimum structural eigenvector energy fraction in ``[0, 1]``.
    :param remove_complex_conjugate: Drop one partner from each conjugate pair.
    :param sort: Sort eigenvalues by ``"frequency"`` or ``"damping"``.
    :return: Continuous-time eigenvalues of the rescaled system(s), ``(n_out,)`` for scalar
        inputs or ``(*n_points, n_out)`` for batched inputs.
    """
    velocity, density, chord = jnp.broadcast_arrays(
        jnp.asarray(velocity, dtype=float),
        jnp.asarray(density, dtype=float),
        jnp.asarray(chord, dtype=float),
    )
    batch_shape = velocity.shape

    def _single(velocity_: Array, density_: Array, chord_: Array) -> Array:
        r"""
        Rescale and diagonalise the linear system at a single (velocity, density, chord) point.
        """
        # reference conditions
        u_ref = self._case.aero.flowfield.u_inf_mag
        rho_ref = self._case.aero.flowfield.rho
        m_chord = self.reference.aero.gamma_b[0].shape[0]

        dt_new = chord_ / (m_chord * velocity_)

        # discretise structural system at new dt
        struct_cont = self.structure.linearise_continuous()
        n_struct = struct_cont.a.shape[0]
        eye_s = jnp.eye(n_struct)

        mat_inv_new = jnp.linalg.inv(eye_s - struct_cont.a * 0.5 * dt_new)
        a_struct_new = mat_inv_new @ (eye_s + struct_cont.a * 0.5 * dt_new)
        b_struct_new = mat_inv_new @ (struct_cont.b * dt_new)

        f_ext_slice = self.structure.input_slices["f_ext"].slices
        b_d_f_ref = self.structure.sys.b[:, f_ext_slice]
        b_d_f_new = b_struct_new[:, f_ext_slice]

        # rescaled A matrix
        a_ref = self.sys.a

        # structural and aero index ranges in the coupled state vector
        q_start = self.state_slices["q"].slices.start
        q_dot_end = self.state_slices["q_dot"].slices.stop
        struct_slice = slice(q_start, q_dot_end)

        # reference structural discrete-time A_d
        a_struct_ref = self.structure.sys.a

        # aero-induced contribution in the structural rows
        struct_rows = a_ref[struct_slice, :]
        delta = struct_rows.at[:, struct_slice].add(-a_struct_ref)

        # extract force Jacobian
        force_jac = jnp.linalg.pinv(b_d_f_ref) @ delta

        # per-column force scaling
        n_total = a_ref.shape[0]
        q_state_slice = self.state_slices["q"].slices
        base_force_scale = (density_ * velocity_) / (rho_ref * u_ref)
        force_col_scale = jnp.ones(n_total) * base_force_scale
        force_col_scale = force_col_scale.at[q_state_slice].set(
            (density_ * velocity_**2) / (rho_ref * u_ref**2)
        )
        delta_new = b_d_f_new @ (force_jac * force_col_scale[None, :])

        a_new = a_ref.at[struct_slice, :].set(delta_new)
        a_new = a_new.at[struct_slice, struct_slice].add(a_struct_new)

        vel_ratio = velocity_ / u_ref
        a_new = a_new.at[:q_start, q_state_slice].multiply(vel_ratio)

        # eigenvalue computation
        evals_d, evecs = jnp.linalg.eig(a_new)
        evals = jnp.log(evals_d) / dt_new

        omega_damped = jnp.abs(evals.imag)
        damping = -evals.real / jnp.abs(evals)
        omega_natural = omega_damped / jnp.sqrt(1.0 - damping**2)
        freq_natural_hz = omega_natural / (2.0 * jnp.pi)

        match sort:
            case "frequency":
                idx = omega_natural.argsort()
            case "damping":
                idx = damping.argsort()

        if remove_complex_conjugate:
            partner = conjugate_partner_mask(
                freq_hz=freq_natural_hz, damping=damping, tiebreaker=evals.real
            )
            idx = idx[jnp.argsort(partner[idx], stable=True)]

        q_slice = self.state_slices["q"].slices
        q_dot_slice = self.state_slices["q_dot"].slices
        evec_sq = jnp.abs(evecs) ** 2
        struct_content = (
            evec_sq[q_slice].sum(axis=0) + evec_sq[q_dot_slice].sum(axis=0)
        ) / evec_sq.sum(axis=0)

        in_range = (
            (freq_natural_hz[idx] >= freq_range[0])
            & (freq_natural_hz[idx] <= freq_range[1])
            & (damping[idx] >= damp_range[0])
            & (damping[idx] <= damp_range[1])
            & (struct_content[idx] >= min_struct_content)
        )
        idx = idx[jnp.argsort(~in_range, stable=True)]

        if n_modes is not None:
            idx = idx[:n_modes]

        return evals[idx]

    if batch_shape == ():
        return _single(velocity, density, chord)

    n_points = velocity.size

    # map scaling across multiple points
    evals_flat = vmap(_single)(
        velocity.reshape(n_points),
        density.reshape(n_points),
        chord.reshape(n_points),
    )  # (n_points, n_out)
    return evals_flat.reshape(*batch_shape, evals_flat.shape[-1])

frf 

frf(
    omega: Array, flowfield: FrequencyFlowField | None
) -> AeroelasticOutputUnflattened

Compute the frequency response function for a gust input.

Parameters:

Name Type Description Default
omega Array

Frequencies in rad/s, (n_freq,).

 required
flowfield FrequencyFlowField | None

Frequency-domain turbulence spectrum. If None, the raw transfer function is returned.

 required

Returns:

Type Description
AeroelasticOutputUnflattened

Gust FRF with q and q_dot fields, each (n_freq, n_dof).
Source code in src/flapjax/coupled/linear/linear_coupled.py 
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
def frf(
    self,
    omega: Array,
    flowfield: FrequencyFlowField | None,
) -> AeroelasticOutputUnflattened:
    r"""
    Compute the frequency response function for a gust input.
    :param omega: Frequencies in rad/s, ``(n_freq,)``.
    :param flowfield: Frequency-domain turbulence spectrum. If ``None``,
        the raw transfer function is returned.
    :return: Gust FRF with ``q`` and ``q_dot`` fields, each ``(n_freq, n_dof)``.
    """
    h = self.frf_base(omega, flowfield)
    n_out = self.n_beam_output_dof
    return AeroelasticOutputUnflattened(q=h[:, :n_out], q_dot=h[:, n_out:])

frf_base 

frf_base(
    omega: Array, flowfield: FrequencyFlowField | None
) -> Array

Compute the gust FRF as a flat vector.

Parameters:

Name Type Description Default
omega Array

Frequencies in rad/s, (n_freq,).

 required
flowfield FrequencyFlowField | None

Frequency-domain turbulence spectrum. If None, the raw transfer function is returned.

 required

Returns:

Type Description
Array

Complex FRF, (n_freq, n_outputs).
Source code in src/flapjax/coupled/linear/linear_coupled.py 
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
def frf_base(
    self,
    omega: Array,
    flowfield: FrequencyFlowField | None,
) -> Array:
    r"""
    Compute the gust FRF as a flat vector.
    :param omega: Frequencies in rad/s, ``(n_freq,)``.
    :param flowfield: Frequency-domain turbulence spectrum. If ``None``,
        the raw transfer function is returned.
    :return: Complex FRF, ``(n_freq, n_outputs)``.
    """
    from flapjax.coupled.linear_gradients.frf import gust_penetration_vector

    u_inf = (
        flowfield.u_inf
        if flowfield is not None
        else self._case.aero.flowfield.u_inf_mag
    )

    zeta_b0 = self._reference.aero.zeta_b
    vertex_x = jnp.concatenate([z[..., 0].ravel() for z in zeta_b0])
    g = gust_penetration_vector(omega=omega, vertex_x=vertex_x, u_inf=u_inf)

    a = self.sys.a
    c = self.sys.c
    n_states = a.shape[0]

    linear_upwash = LinearCoupled(
        case=self._case,
        reference=self._reference,
        batch_size=False,
        n_struct_modes=None,
        bound_upwash=True,
        skip_checks=True,
    )
    nu_b_zero = jnp.zeros(self._reference.aero.zeta_b.size)

    def _step_nu_b(nu_b_vec: Array) -> Array:
        state_np1, _ = linear_upwash.step(nu_b_vec=nu_b_vec)
        return linear_upwash.pack_state_vector(state_np1)

    def _b_g_single(g_k: Array) -> Array:
        _, bg_re = jax.jvp(_step_nu_b, (nu_b_zero,), (g_k.real,))
        _, bg_im = jax.jvp(_step_nu_b, (nu_b_zero,), (g_k.imag,))
        return bg_re + 1j * bg_im

    b_g = jax.vmap(_b_g_single)(g)

    z = jnp.exp(1j * omega * self.dt)
    eye = jnp.eye(n_states)

    def _solve_single(z_k: Array, bg_k: Array) -> Array:
        return z_k * c @ jnp.linalg.solve(z_k * eye - a, bg_k)

    h = jax.vmap(_solve_single)(z, b_g)

    if flowfield is not None:
        # scale with flowfield PSD if available
        h *= jnp.sqrt(flowfield.psd(omega))[:, None]

    return h

NonlinearBeamLinearAero 

NonlinearBeamLinearAero(
    structure: BeamStructure,
    aero: LinearUVLM,
    fsi_convergence_settings: ConvergenceSettings = DEFAULT_FSI_CONVERGENCE_SETTINGS,
)

Aeroelastic case coupling a nonlinear beam with a pre-built linear UVLM model. Useful for relatively cheap time-domain aeroelastic simulations where the aerodynamic nonlinearities are not significant.

Parameters:

Name Type Description Default
structure BeamStructure

Nonlinear BeamStructure with design variables set.

 required
aero LinearUVLM

Pre-built LinearUVLM linearised about some reference state.

 required
fsi_convergence_settings ConvergenceSettings

FSI iteration convergence controls.

 DEFAULT_FSI_CONVERGENCE_SETTINGS
Source code in src/flapjax/coupled/linear_aero_coupled.py 
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
def __init__(
    self,
    structure: BeamStructure,
    aero: LinearUVLM,
    fsi_convergence_settings: ConvergenceSettings = DEFAULT_FSI_CONVERGENCE_SETTINGS,
) -> None:
    r"""
    :param structure: Nonlinear ``BeamStructure`` with design variables set.
    :param aero: Pre-built ``LinearUVLM`` linearised about some reference state.
    :param fsi_convergence_settings: FSI iteration convergence controls.
    """
    self.structure: BeamStructure = structure
    self.aero: LinearUVLM = aero
    self.fsi_convergence_settings: ConvergenceSettings = fsi_convergence_settings
    self.include_unsteady_force: bool = aero.unsteady_force

get_state 

get_state(
    i_ts: int, case: AeroCase
) -> AeroStateUnflattened

Extract linear aero states at time step i_ts from the case object.

Parameters:

Name Type Description Default
i_ts int

Time step index.

 required
case AeroCase

Batched AeroCase containing the linear aero states.

 required

Returns:

Type Description
AeroStateUnflattened

Linear aero states at time step i_ts.
Source code in src/flapjax/coupled/linear_aero_coupled.py 
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
def get_state(self, i_ts: int, case: AeroCase) -> AeroStateUnflattened:
    r"""
    Extract linear aero states at time step ``i_ts`` from the case object.
    :param i_ts: Time step index.
    :param case: Batched ``AeroCase`` containing the linear aero states.
    :return: Linear aero states at time step ``i_ts``.
    """
    gamma_b = case.gamma_b.index_all(i_ts, ...)
    gamma_w = case.gamma_w.index_all(i_ts, ...)

    if self.aero.unsteady_force:
        i_prev = jnp.maximum(i_ts - 1, 0)
        gamma_b_nm1 = case.gamma_b.index_all(i_prev, ...)
    else:
        gamma_b_nm1 = None

    if self.aero.prescribed_wake:
        assert case.zeta_w is not None
        zeta_w = case.zeta_w.index_all(i_ts, ...)
        zeta_b_state = case.zeta_b.index_all(i_ts, ...)
    else:
        zeta_w = None
        zeta_b_state = None

    return AeroStateUnflattened(
        gamma_b=gamma_b,
        gamma_w=gamma_w,
        gamma_b_nm1=gamma_b_nm1,
        zeta_w=zeta_w,
        zeta_b=zeta_b_state,
    )

case_solve 

case_solve(
    case: AeroCase,
    i_ts: int,
    hg_n: Array | None,
    hg_nm1: Array | None,
    hg_dot_n: Array | None,
    static: bool,
    horseshoe: bool,
    cs_ang_n: dict[str, Array],
    cs_ang_nm1: dict[str, Array] | None,
    cs_vel_n: dict[str, Array] | None,
) -> AeroCase

Step the linear aero system one step and write the result into case. Called by BeamStructure.base_dynamic_solve inside the FSI loop. Only supports dynamic solves; hg_nm1, cs_ang_nm1 and horseshoe are part of the DynamicAeroSolver protocol but not consumed here.

Source code in src/flapjax/coupled/linear_aero_coupled.py 
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
def case_solve(
    self,
    case: AeroCase,
    i_ts: int,
    hg_n: Array | None,
    hg_nm1: Array | None,
    hg_dot_n: Array | None,
    static: bool,
    horseshoe: bool,
    cs_ang_n: dict[str, Array],
    cs_ang_nm1: dict[str, Array] | None,
    cs_vel_n: dict[str, Array] | None,
) -> AeroCase:
    r"""
    Step the linear aero system one step and write the result into ``case``. Called by
    ``BeamStructure.base_dynamic_solve`` inside the FSI loop. Only supports dynamic solves;
    ``hg_nm1``, ``cs_ang_nm1`` and ``horseshoe`` are part of the ``DynamicAeroSolver`` protocol
    but not consumed here.
    """
    del hg_nm1, cs_ang_nm1, horseshoe
    if static:
        raise NotImplementedError(
            "case_solve(static=True) is not supported for the linear aero adapter"
        )
    assert hg_n is not None and hg_dot_n is not None

    uvlm = self.aero.case
    sys = self.aero.sys
    cs_vel_n_ = cs_vel_n if cs_vel_n is not None else {}

    zeta_b_n = uvlm.hg_to_zeta_b(hg_n=hg_n, cs_ang_n=cs_ang_n)
    zeta_b_dot_n = uvlm.hg_dot_to_zeta_b_dot(
        hg_n=hg_n,
        hg_dot_n=hg_dot_n,
        cs_ang_n=cs_ang_n,
        cs_vel_n=cs_vel_n_,
    )

    # flowfield perturbation from the linearisation reference, applied as extra input upwash
    t_n = case.t[i_ts]
    t_ref = self.aero.reference.t

    if self.aero.bound_upwash:
        nu_b_n = ArrayList.zeros_like(zeta_b_n)
        nu_b_n += self.flowfield.surf_vmap_call(
            xs=self.aero.reference.zeta_b, t=t_n
        ) - self.flowfield.surf_vmap_call(xs=self.aero.reference.zeta_b, t=t_ref)
    else:
        nu_b_n = None

    if self.aero.wake_upwash:
        nu_w_n = ArrayList.zeros_like(self.aero.reference.zeta_w)
        nu_w_n += self.flowfield.surf_vmap_call(
            xs=self.aero.reference.zeta_w, t=t_n
        ) - self.flowfield.surf_vmap_call(xs=self.aero.reference.zeta_w, t=t_ref)
    else:
        nu_w_n = None

    u_n = AeroInputUnflattened(
        zeta_b=zeta_b_n,
        zeta_b_dot=zeta_b_dot_n,
        nu_b=nu_b_n,
        nu_w=nu_w_n,
    )
    u_n_vec = self.aero.pack_input_vector(u_n)
    u_ref_vec = self.aero.pack_input_vector(self.aero.reference_inputs)
    du_n = u_n_vec - u_ref_vec

    x_nm1_unflat = self.get_state(i_ts=i_ts - 1, case=case)
    x_nm1_vec = self.aero.pack_state_vector(x_nm1_unflat)
    x_ref_vec = self.aero.pack_state_vector(self.aero.reference_states)
    dx_prev = x_nm1_vec - x_ref_vec

    dx_n = sys.a @ dx_prev + sys.b @ du_n
    dy_n = sys.c @ dx_n + sys.d @ du_n

    x_n_vec = dx_n + x_ref_vec
    y_ref_vec = self.aero.pack_output_vector(self.aero.reference_outputs)
    y_n_vec = dy_n + y_ref_vec

    x_n_unflat = self.aero.unpack_state_vector(x_n_vec)
    y_n_unflat = self.aero.unpack_output_vector(y_n_vec)

    case.set_arraylist_at_ts("zeta_b", zeta_b_n, i_ts)
    case.set_arraylist_at_ts("zeta_b_dot", zeta_b_dot_n, i_ts)
    case.set_arraylist_at_ts("c", compute_c(zeta_b_n), i_ts)
    case.set_arraylist_at_ts("nc", compute_nc(zeta_b_n), i_ts)
    case.set_arraylist_at_ts("gamma_b", x_n_unflat.gamma_b, i_ts)
    case.set_arraylist_at_ts("gamma_w", x_n_unflat.gamma_w, i_ts)
    case.set_arraylist_at_ts("f_steady", y_n_unflat.f_steady, i_ts)

    if self.aero.unsteady_force:
        assert y_n_unflat.f_unsteady is not None
        gamma_b_dot_n = ArrayList(
            [
                (gb - gb_prev) / self.dt
                for gb, gb_prev in zip(x_n_unflat.gamma_b, x_nm1_unflat.gamma_b)
            ]
        )
        case.set_arraylist_at_ts("gamma_b_dot", gamma_b_dot_n, i_ts)
        case.set_arraylist_at_ts("f_unsteady", y_n_unflat.f_unsteady, i_ts)

    assert case.zeta_w is not None
    if self.aero.prescribed_wake:
        assert x_n_unflat.zeta_w is not None
        case.set_arraylist_at_ts("zeta_w", x_n_unflat.zeta_w, i_ts)
    else:
        case.set_arraylist_at_ts("zeta_w", self.aero.reference.zeta_w, i_ts)

    case.t = case.t.at[i_ts].set(t_n)

    return case

reference_configuration 

reference_configuration(
    prescribed_dofs: Sequence[int]
    | Array
    | slice
    | int = (),
    use_f_ext_follower: bool = False,
    use_f_ext_dead: bool = False,
) -> AeroelasticCase

Aeroelastic snapshot built from the beam's reference configuration and the aero linearisation reference state.

Parameters:

Name Type Description Default
prescribed_dofs Sequence[int] | Array | slice | int

Prescribed DOFs for the beam structure. Defaults to no prescribed DOFs.

 ()
use_f_ext_follower bool

Whether to include follower forces in the reference configuration.

 False
use_f_ext_dead bool

Whether to include dead forces in the reference configuration.

 False

Returns:

Type Description
AeroelasticCase

Aeroelastic snapshot at the reference configuration.
Source code in src/flapjax/coupled/linear_aero_coupled.py 
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
def reference_configuration(
    self,
    prescribed_dofs: Sequence[int] | Array | slice | int = (),
    use_f_ext_follower: bool = False,
    use_f_ext_dead: bool = False,
) -> AeroelasticCase:
    r"""
    Aeroelastic snapshot built from the beam's reference configuration and
    the aero linearisation reference state.
    :param prescribed_dofs: Prescribed DOFs for the beam structure. 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.
    :return: Aeroelastic snapshot at the reference configuration.
    """
    prescribed_dofs_tuple = self.structure.make_prescribed_dofs_tuple(
        prescribed_dofs
    )
    return AeroelasticCase(
        structure=self.structure.reference_configuration(
            use_f_grav=self.structure.use_gravity,
            use_f_ext_dead=use_f_ext_dead,
            use_f_ext_follower=use_f_ext_follower,
            use_f_aero=True,
            prescribed_dofs=prescribed_dofs_tuple,
        ),
        aero=self.aero.reference_snapshot(),
    )

dynamic_solve 

dynamic_solve(
    init_case: AeroelasticCase | None,
    prescribed_dofs: Sequence[int] | Array | slice | int,
    n_tstep: int,
    f_ext_follower: Array | None = None,
    f_ext_dead: Array | None = None,
    t_init: float = 0.0,
    load_steps: int = 1,
    thrust_t: dict[str, Array] | None = None,
    cs_ang_t: dict[str, Array] | None = None,
    cs_vel_t: dict[str, Array] | None = None,
) -> AeroelasticCase

Dynamic aeroelastic solve with nonlinear beam and linear UVLM.

Source code in src/flapjax/coupled/linear_aero_coupled.py 
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
def dynamic_solve(
    self,
    init_case: AeroelasticCase | None,
    prescribed_dofs: Sequence[int] | Array | slice | int,
    n_tstep: int,
    f_ext_follower: Array | None = None,
    f_ext_dead: Array | None = None,
    t_init: float = 0.0,
    load_steps: int = 1,
    thrust_t: dict[str, Array] | None = None,
    cs_ang_t: dict[str, Array] | None = None,
    cs_vel_t: dict[str, Array] | None = None,
) -> AeroelasticCase:
    r"""
    Dynamic aeroelastic solve with nonlinear beam and linear UVLM.
    """
    ref_cs_ang = self.aero.reference.cs_ang
    if cs_ang_t is not None:
        for key, series in cs_ang_t.items():
            if key not in ref_cs_ang:
                raise ValueError(
                    f"cs_ang_t key '{key}' not present in the linearisation reference"
                )
            if series.shape[0] != n_tstep:
                raise ValueError(
                    f"Inconsistent number of time steps for control surface input cs_ang_t['{key}']"
                )

            # TODO: implement linear control surfaces
            if not jnp.allclose(series, ref_cs_ang[key]):
                warn(f"cs_ang_t['{key}'] deviates from the linearisation reference")
    cs_ang_t_ = (
        cs_ang_t
        if cs_ang_t is not None
        else {k: jnp.full(n_tstep, v) for k, v in ref_cs_ang.items()}
    )

    if cs_vel_t is None:
        cs_vel_t_ = cs_ang_to_cs_vel(cs_ang_t=cs_ang_t_, dt=self.aero.dt)
    else:
        cs_vel_t_ = cs_vel_t

    if thrust_t is not None:
        if thrust_t.keys() != dict(self.structure.thrust_direction).keys():
            raise ValueError("Mismatch in keys for thrust")
        for k, v in thrust_t.items():
            check_arr_shape(v, (n_tstep,), name=f"thrust_t['{k}']")
        thrust_t_: dict[str, Array] = thrust_t
    else:
        thrust_t_ = {
            k: jnp.full(n_tstep, v)
            for k, v in self.structure.thrust_reference.items()
        }

    prescribed_dofs_tuple = self.structure.make_prescribed_dofs_tuple(
        prescribed_dofs
    )
    solve_dofs = get_solve_dofs(
        n_dof=self.structure.n_dof, prescribed_dofs=prescribed_dofs_tuple
    )

    t = jnp.arange(n_tstep) * self.aero.dt + t_init

    self.structure.time_integrator = TimeIntegrator(
        spectral_radius=self.structure.spectral_radius, dt=self.aero.dt
    )

    if init_case is None:
        initial_snapshot = self.reference_configuration(
            prescribed_dofs=prescribed_dofs_tuple,
            use_f_ext_follower=f_ext_follower is not None,
            use_f_ext_dead=f_ext_dead is not None,
        ).to_dynamic(t=None)
    else:
        initial_snapshot = init_case

    case = AeroelasticCase.initialise(
        initial_snapshot=initial_snapshot,
        t=t,
        use_f_ext_follower=f_ext_follower is not None,
        use_f_ext_dead=f_ext_dead is not None,
        structure=self.structure,
        x0_aero=self.aero.case.zeta_b0,
    )

    if f_ext_follower is not None and case.structure.f_ext_follower is not None:
        case.structure.f_ext_follower = case.structure.f_ext_follower.at[
            0, ...
        ].set(f_ext_follower[0, ...])
    if f_ext_dead is not None and case.structure.f_ext_dead is not None:
        case.structure.f_ext_dead = case.structure.f_ext_dead.at[0, ...].set(
            self.structure.make_f_dead_ext(
                f_ext=f_ext_dead[0, ...], rmat=case.structure.hg[0, :, :3, :3]
            )
        )

    case.structure.prescribed_dofs = prescribed_dofs_tuple

    fsi_converge_status = ConvergenceStatus(self.fsi_convergence_settings)
    fsi_converge_status.print_header(dynamic=True)

    out = self.structure.base_dynamic_solve(
        struct_case=case.structure,
        struct_convergence_status=ConvergenceStatus(
            self.structure.struct_convergence_settings
        ),
        t=t,
        solve_dofs=solve_dofs,
        load_steps=load_steps,
        f_ext_follower=f_ext_follower,
        f_ext_dead=f_ext_dead,
        aero_obj=self,
        aero_case=case.aero,
        fsi_convergence_status=fsi_converge_status,
        thrust_t=thrust_t_,
        cs_ang_t=cs_ang_t_,
        cs_vel_t=cs_vel_t_,
    )

    fsi_converge_status.print_line(dynamic=True)
    return out

compute_gust_frf 

compute_gust_frf(
    system: CoupledAeroelastic,
    dv: AeroelasticDesignVariables,
    varphi: Array,
    case: AeroelasticCase,
    omega: Array,
    frf_flowfield: FrequencyFlowField | None = None,
    batch_size: int | None = 4,
) -> AeroelasticOutputUnflattened

Compute the gust frequency response function using JVP through the coupled step.

When frf_flowfield is provided the result is PSD-weighted (each frequency row scaled by sqrt(psd(omega))). When None, the raw transfer function is returned.

Parameters:

Name Type Description Default
system CoupledAeroelastic

The coupled aeroelastic system.

 required
dv AeroelasticDesignVariables

Aeroelastic design variables.

 required
varphi Array

Structural configuration, (n_nodes, 6).

 required
case AeroelasticCase

Converged aeroelastic case around which to linearise.

 required
omega Array

Angular frequencies in rad/s, (n_freq,).

 required
frf_flowfield FrequencyFlowField | None

Frequency-domain turbulence spectrum. If None, the raw transfer function is returned.

 None
batch_size int | None

Batch size for Jacobian materialisation of the A matrix.

 4

Returns:

Type Description
AeroelasticOutputUnflattened

Gust FRF with q and q_dot fields.
Source code in src/flapjax/coupled/linear_gradients/frf.py 
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
def compute_gust_frf(
    system: CoupledAeroelastic,
    dv: AeroelasticDesignVariables,
    varphi: Array,
    case: AeroelasticCase,
    omega: Array,
    frf_flowfield: FrequencyFlowField | None = None,
    batch_size: int | None = 4,
) -> AeroelasticOutputUnflattened:
    r"""
    Compute the gust frequency response function using JVP through the coupled step.

    When ``frf_flowfield`` is provided the result is PSD-weighted (each
    frequency row scaled by ``sqrt(psd(omega))``). When ``None``, the raw
    transfer function is returned.

    :param system: The coupled aeroelastic system.
    :param dv: Aeroelastic design variables.
    :param varphi: Structural configuration, ``(n_nodes, 6)``.
    :param case: Converged aeroelastic case around which to linearise.
    :param omega: Angular frequencies in rad/s, ``(n_freq,)``.
    :param frf_flowfield: Frequency-domain turbulence spectrum. If ``None``,
        the raw transfer function is returned.
    :param batch_size: Batch size for Jacobian materialisation of the A matrix.
    :return: Gust FRF with ``q`` and ``q_dot`` fields.
    """
    ref, inner = build_reference_case(system, dv, varphi, case)
    linear = inner.linearise(
        reference=ref, skip_checks=True, batch_size=batch_size, n_struct_modes=None
    )
    return linear.frf(omega=omega, flowfield=frf_flowfield)

gust_frf_adjoint 

gust_frf_adjoint(
    system: CoupledAeroelastic,
    case: AeroelasticCase,
    omega: Array,
    objective: FRFObjective,
    frf_flowfield: FrequencyFlowField | None = None,
    grads_to_compute: AeroelasticGradsToCompute
    | None = None,
    batch_size: int = 32,
) -> tuple[Array, AeroelasticDesignVariables]

Compute sensitivities of an objective that depends on the gust FRF with respect to design variables.

Parameters:

Name Type Description Default
system CoupledAeroelastic

The coupled aeroelastic system.

 required
case AeroelasticCase

Converged static solution around which to linearise.

 required
omega Array

Frequencies to sample in rad/s, (n_freq,).

 required
objective FRFObjective

Function (full_states, design_variables, H_gust) -> scalar where H_gust is the complex gust FRF, (n_freq, n_outputs).

 required
frf_flowfield FrequencyFlowField | None

Frequency-domain turbulence spectrum. If None, the raw transfer function is passed to the objective.

 None
grads_to_compute AeroelasticGradsToCompute | None

Which design variable gradients to request.

 None
batch_size int

Batch size for Jacobian materialisation.

 32

Returns:

Type Description
tuple[Array, AeroelasticDesignVariables]

Primal objective value and its gradient w.r.t. design variables.
Source code in src/flapjax/coupled/linear_gradients/frf.py 
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
def gust_frf_adjoint(
    system: CoupledAeroelastic,
    case: AeroelasticCase,
    omega: Array,
    objective: FRFObjective,
    frf_flowfield: FrequencyFlowField | None = None,
    grads_to_compute: AeroelasticGradsToCompute | None = None,
    batch_size: int = 32,
) -> tuple[Array, AeroelasticDesignVariables]:
    r"""
    Compute sensitivities of an objective that depends on the gust FRF with
    respect to design variables.
    :param system: The coupled aeroelastic system.
    :param case: Converged static solution around which to linearise.
    :param omega: Frequencies to sample in rad/s, ``(n_freq,)``.
    :param objective: Function ``(full_states, design_variables, H_gust) -> scalar``
        where ``H_gust`` is the complex gust FRF, ``(n_freq, n_outputs)``.
    :param frf_flowfield: Frequency-domain turbulence spectrum. If ``None``,
        the raw transfer function is passed to the objective.
    :param grads_to_compute: Which design variable gradients to request.
    :param batch_size: Batch size for Jacobian materialisation.
    :return: Primal objective value and its gradient w.r.t. design variables.
    """
    if grads_to_compute is None:
        grads_to_compute = AeroelasticGradsToCompute()

    varphi_eq = case.structure.varphi
    dv_ref = system.get_design_variables(case=case, grads_to_compute=grads_to_compute)
    n_dof = system.structure.n_dof
    solve_dofs = jnp.array(
        get_solve_dofs(
            n_dof=n_dof,
            prescribed_dofs=case.structure.prescribed_dofs,
        )
    )

    # compute primal gust FRF
    h_gust = _compute_gust_frf_base(
        system, dv_ref, varphi_eq, case, omega, frf_flowfield
    )

    n_out = h_gust.shape[1] // 2

    def _unpack_h(h_: Array) -> AeroelasticOutputUnflattened:
        return AeroelasticOutputUnflattened(q=h_[:, :n_out], q_dot=h_[:, n_out:])

    def _objective_of_dv_h_varphi(
        dv_: AeroelasticDesignVariables, h_: Array, varphi_flat_: Array
    ) -> Array:
        # differentiable object w.r.t. its arguments
        states_, _ = system.aeroelastic_states_res_from_dv_varphi(
            dv=dv_,
            varphi=varphi_flat_.reshape(-1, 6),
            thrust=case.structure.thrust,
            t=case.aero.t,
            i_ts=0,
            use_horseshoe=False,
        )
        return objective(states_, dv_, _unpack_h(h_))

    j_val, vjp_j = jax.vjp(_objective_of_dv_h_varphi, dv_ref, h_gust, varphi_eq.ravel())

    j_shape = j_val.shape
    n_f = max(1, int(np.prod(j_shape)))
    d_j_d_x_direct_b, d_j_d_h_b, d_j_d_varphi_direct_b = jax.vmap(vjp_j)(
        jnp.eye(n_f).reshape((n_f,) + j_shape)
    )

    def _gust_frf_fn(dv_: AeroelasticDesignVariables, varphi_: Array) -> Array:
        # differentiate through FRF computation
        return _compute_gust_frf_base(system, dv_, varphi_, case, omega, frf_flowfield)

    _, vjp_gust = jax.vjp(_gust_frf_fn, dv_ref, varphi_eq)
    dv_bar_gust_b, varphi_bar_gust_b = jax.vmap(vjp_gust)(d_j_d_h_b)

    def _residual_of_varphi(varphi_vec: Array) -> Array:
        # differentiate full states w.r.t. static deformation
        return system.aeroelastic_states_res_from_dv_varphi(
            dv=dv_ref,
            varphi=varphi_vec.reshape(-1, 6),
            thrust=case.structure.thrust,
            t=case.aero.t,
            i_ts=0,
            use_horseshoe=False,
        )[1]

    _, vjp_res_v = jax.vjp(_residual_of_varphi, varphi_eq.ravel())
    p_res_p_varphi = jax.lax.map(
        lambda cot: vjp_res_v(cot)[0], jnp.eye(n_dof), batch_size=batch_size
    )

    varphi_bar_b = jnp.real(varphi_bar_gust_b.reshape(n_f, -1)) + jnp.real(
        d_j_d_varphi_direct_b.reshape(n_f, -1)
    )
    varphi_bar_free_b = varphi_bar_b[:, solve_dofs]

    j_res_free = p_res_p_varphi[jnp.ix_(solve_dofs, solve_dofs)]
    mu_free_b = jnp.linalg.solve(j_res_free.T, varphi_bar_free_b.T).T
    mu_full_b = (
        jnp.zeros((n_f, n_dof), dtype=mu_free_b.dtype).at[:, solve_dofs].set(mu_free_b)
    )

    _, vjp_res_dv = jax.vjp(
        lambda dv_: system.aeroelastic_states_res_from_dv_varphi(
            dv=dv_,
            varphi=varphi_eq,
            thrust=case.structure.thrust,
            t=case.aero.t,
            i_ts=0,
            use_horseshoe=False,
        )[1],
        dv_ref,
    )
    (dv_bar_via_eq_pos_b,) = jax.vmap(vjp_res_dv)(mu_full_b)

    def _neg_if_float(x):
        if hasattr(x, "dtype") and jnp.issubdtype(x.dtype, jnp.floating):
            return -x
        return x

    dv_bar_via_eq_b = jax.tree.map(_neg_if_float, dv_bar_via_eq_pos_b)

    def _sum_real(direct, via_gust, via_eq):
        if hasattr(direct, "dtype") and jnp.issubdtype(direct.dtype, jnp.floating):
            return direct + jnp.real(via_gust) + via_eq
        return direct

    total_b = jax.tree.map(_sum_real, d_j_d_x_direct_b, dv_bar_gust_b, dv_bar_via_eq_b)

    def _to_j_shape(leaf):
        if not hasattr(leaf, "shape") or leaf.ndim == 0:
            return leaf
        return leaf.reshape(j_shape + leaf.shape[1:])

    total = jax.tree.map(_to_j_shape, total_b)

    return j_val, total

stability_adjoint 

stability_adjoint(
    system: CoupledAeroelastic,
    case: AeroelasticCase,
    objective: StabilityObjective,
    grads_to_compute: AeroelasticGradsToCompute
    | None = None,
    batch_size: int = 32,
) -> tuple[Array, AeroelasticDesignVariables]

Compute the sensitivities of some objective that refers to the aeroelastic continuous-time eigenvalues, full system states and design variables, with respect to the design variables.

Parameters:

Name Type Description Default
system CoupledAeroelastic

The coupled aeroelastic system.

 required
case AeroelasticCase

Converged static solution around which to linearise.

 required
objective StabilityObjective

Function which takes the full aeroelastic states, aeroelastic design variables, and the full continuous-time eigenvalue vector, and returns a real-valued objective.

 required
grads_to_compute AeroelasticGradsToCompute | None

Which design variable gradients to request.

 None
batch_size int

Batch size for Jacobian materialisation.

 32

Returns:

Type Description
tuple[Array, AeroelasticDesignVariables]

Primal value of the objective, and its gradient with respect to the design variables.
Source code in src/flapjax/coupled/linear_gradients/stability.py 
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
def stability_adjoint(
    system: CoupledAeroelastic,
    case: AeroelasticCase,
    objective: StabilityObjective,
    grads_to_compute: AeroelasticGradsToCompute | None = None,
    batch_size: int = 32,
) -> tuple[Array, AeroelasticDesignVariables]:
    r"""
    Compute the sensitivities of some objective that refers to the aeroelastic continuous-time eigenvalues, full system
    states and design variables, with respect to the design variables.
    :param system: The coupled aeroelastic system.
    :param case: Converged static solution around which to linearise.
    :param objective: Function which takes the full aeroelastic states, aeroelastic design variables, and the full
    continuous-time eigenvalue vector, and returns a real-valued objective.
    :param grads_to_compute: Which design variable gradients to request.
    :param batch_size: Batch size for Jacobian materialisation.
    :return: Primal value of the objective, and its gradient with respect to the design variables.
    """
    n_struct_modes = None  # model structure not implemented

    if grads_to_compute is None:
        grads_to_compute = AeroelasticGradsToCompute()

    # extract base parameters
    varphi_eq = case.structure.varphi
    dv_ref = system.get_design_variables(case=case, grads_to_compute=grads_to_compute)
    dt = system.aero.dt
    n_dof = system.structure.n_dof
    solve_dofs = jnp.array(
        get_solve_dofs(
            n_dof=n_dof,
            prescribed_dofs=case.structure.prescribed_dofs,
        )
    )

    # assemble system matrix and compute eigendecomposition
    a_matrix = assemble_a(
        system,
        dv_ref,
        varphi_eq,
        case,
        n_struct_modes=n_struct_modes,
    )
    lam_d, phi_r, phi_l = eig_left_right(a_matrix)
    lam_c = jnp.log(lam_d) / dt  # continuous time eigenvalues

    def _objective_of_dv_lam_varphi(
        dv_: AeroelasticDesignVariables, lam_c_: Array, varphi_flat_: Array
    ) -> Array:
        states_, _ = system.aeroelastic_states_res_from_dv_varphi(
            dv=dv_,
            varphi=varphi_flat_.reshape(-1, 6),
            thrust=case.structure.thrust,
            t=case.aero.t,
            i_ts=0,
            use_horseshoe=False,
        )
        return objective(states_, dv_, lam_c_)

    # VJP of the objective against it's arguments
    j_val, vjp_j = jax.vjp(
        _objective_of_dv_lam_varphi, dv_ref, lam_c, varphi_eq.ravel()
    )

    # allow for arbitrary shape
    j_shape = j_val.shape
    n_f = max(1, int(np.prod(j_shape)))
    d_j_d_x_direct_b, d_j_d_lambda_c_b, d_j_d_varphi_direct_b = jax.vmap(vjp_j)(
        jnp.eye(n_f).reshape((n_f,) + j_shape)
    )  # sensitivities through objective direct path

    # build one A-cotangent per output row, with filter to ignore defective modes
    c_denom = jnp.einsum("ij,ij->j", phi_l, phi_r)
    denom_ok = jnp.abs(c_denom) > 1e-12

    def _make_a_bar_row(dj_row: Array) -> Array:
        w_d = jnp.conj(dj_row) / (dt * lam_d)
        coeff = jnp.where(denom_ok, w_d / c_denom, 0.0)
        return jnp.real(phi_l @ jnp.diag(coeff) @ phi_r.T)

    a_bar_b = jax.vmap(_make_a_bar_row)(d_j_d_lambda_c_b)  # (n_f, N, N)

    # create a JVP for the system matrix construction
    def _a_fn(dv_, varphi_):
        return assemble_a(
            system,
            dv_,
            varphi_,
            case,
            n_struct_modes=n_struct_modes,
        )

    _, vjp_a = jax.vjp(_a_fn, dv_ref, varphi_eq)
    dv_bar_a_b, varphi_bar_a_b = jax.vmap(vjp_a)(a_bar_b)  # leading (n_f,)

    # static residual Jacobian
    def _residual_of_varphi(varphi_vec: Array) -> Array:
        return system.aeroelastic_states_res_from_dv_varphi(
            dv=dv_ref,
            varphi=varphi_vec.reshape(-1, 6),
            thrust=case.structure.thrust,
            t=case.aero.t,
            i_ts=0,
            use_horseshoe=False,
        )[1]

    _, vjp_res_v = jax.vjp(_residual_of_varphi, varphi_eq.ravel())
    p_res_p_varphi = jax.lax.map(
        lambda cot: vjp_res_v(cot)[0], jnp.eye(n_dof), batch_size=batch_size
    )

    # combine the sensitivities with respect to varphi from the full states in the objective and the path through the
    # system matrix
    varphi_bar_b = jnp.real(varphi_bar_a_b.reshape(n_f, -1)) + jnp.real(
        d_j_d_varphi_direct_b.reshape(n_f, -1)
    )  # (n_f, n_dof)
    varphi_bar_free_b = varphi_bar_b[:, solve_dofs]  # (n_f, n_free)

    j_res_free = p_res_p_varphi[jnp.ix_(solve_dofs, solve_dofs)]
    mu_free_b = jnp.linalg.solve(j_res_free.T, varphi_bar_free_b.T).T  # (n_f, n_free)
    mu_full_b = (
        jnp.zeros((n_f, n_dof), dtype=mu_free_b.dtype).at[:, solve_dofs].set(mu_free_b)
    )

    _, vjp_res_dv = jax.vjp(
        lambda dv_: system.aeroelastic_states_res_from_dv_varphi(
            dv=dv_,
            varphi=varphi_eq,
            thrust=case.structure.thrust,
            t=case.aero.t,
            i_ts=0,
            use_horseshoe=False,
        )[1],
        dv_ref,
    )
    (dv_bar_via_eq_pos_b,) = jax.vmap(vjp_res_dv)(mu_full_b)

    # handles sign flip for floating point types
    def _neg_if_float(x):
        if hasattr(x, "dtype") and jnp.issubdtype(x.dtype, jnp.floating):
            return -x
        return x

    dv_bar_via_eq_b = jax.tree.map(_neg_if_float, dv_bar_via_eq_pos_b)

    # combine contributions
    def _sum_real(direct, via_a, via_eq):
        if hasattr(direct, "dtype") and jnp.issubdtype(direct.dtype, jnp.floating):
            return direct + jnp.real(via_a) + via_eq
        return direct

    total_b = jax.tree.map(_sum_real, d_j_d_x_direct_b, dv_bar_a_b, dv_bar_via_eq_b)

    # reshape back to original shape
    def _to_j_shape(leaf):
        if not hasattr(leaf, "shape") or leaf.ndim == 0:
            return leaf
        return leaf.reshape(j_shape + leaf.shape[1:])

    total = jax.tree.map(_to_j_shape, total_b)

    return j_val, total

coupled

BaseCoupledAeroelastic 

BaseCoupledAeroelastic(
    structure: BeamStructure,
    aero: UVLM,
    fsi_convergence_settings: ConvergenceSettings = DEFAULT_FSI_CONVERGENCE_SETTINGS,
)
Source code in src/flapjax/coupled/coupled.py 
53
54
55
56
57
58
59
60
61
def __init__(
    self,
    structure: BeamStructure,
    aero: UVLM,
    fsi_convergence_settings: ConvergenceSettings = DEFAULT_FSI_CONVERGENCE_SETTINGS,
):
    self.structure: BeamStructure = structure
    self.aero: UVLM = aero
    self.fsi_convergence_settings: ConvergenceSettings = fsi_convergence_settings
get_design_variables 
get_design_variables(
    case: AeroelasticCase,
    grads_to_compute: AeroelasticGradsToCompute | None,
) -> AeroelasticDesignVariables

Obtain the design variables describing the wing.

Parameters:

Name Type Description Default
case AeroelasticCase
required
grads_to_compute AeroelasticGradsToCompute | None

Data structure which describes which design variables should be obtained. If none, all variables are obtained.

 required

Returns:

Type Description
AeroelasticDesignVariables

AeroelasticDesignVariables object.
Source code in src/flapjax/coupled/coupled.py 
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
def get_design_variables(
    self,
    case: AeroelasticCase,
    grads_to_compute: AeroelasticGradsToCompute | None,
) -> AeroelasticDesignVariables:
    r"""
    Obtain the design variables describing the wing.
    :param case:
    :param grads_to_compute: Data structure which describes which design variables should be obtained. If none, all
    variables are obtained.
    :return: AeroelasticDesignVariables object.
    """
    return AeroelasticDesignVariables(
        structure_dv=self.structure.get_design_variables(
            struct_case=case.structure,
            thrust_t=case.structure.thrust,
            grads_to_compute=grads_to_compute.structure
            if grads_to_compute is not None
            else None,
        ),
        aero_dv=self.aero.get_design_variables(
            cs_ang_t=case.aero.cs_ang,
            cs_vel_t=case.aero.cs_vel,
            grads_to_compute=grads_to_compute.aero
            if grads_to_compute is not None
            else None,
        ),
    )
reference_configuration 
reference_configuration(
    prescribed_dofs: Sequence[int]
    | Array
    | slice
    | int = (),
    horseshoe: bool = False,
    use_f_ext_follower: bool = False,
    use_f_ext_dead: bool = False,
    t_init: float | Array = 0.0,
) -> AeroelasticCase

Obtain the static aeroelastic object describing the undeformed wing.

Parameters:

Name Type Description Default
prescribed_dofs Sequence[int] | Array | slice | int

Prescribed dofs for the structure. Defaults to no prescribed dofs.

 ()
horseshoe bool

Horseshoe flag.

 False
use_f_ext_follower bool

If true, allocate an array for follower forces.

 False
use_f_ext_dead bool

If true, allocate an array for dead forces.

 False
t_init float | Array

Initial time

 0.0

Returns:

Type Description
AeroelasticCase

Static aeroelastic object for undeformed wing
Source code in src/flapjax/coupled/coupled.py 
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
def reference_configuration(
    self,
    prescribed_dofs: Sequence[int] | Array | slice | int = (),
    horseshoe: bool = False,
    use_f_ext_follower: bool = False,
    use_f_ext_dead: bool = False,
    t_init: float | Array = 0.0,
) -> AeroelasticCase:
    r"""
    Obtain the static aeroelastic object describing the undeformed wing.
    :param prescribed_dofs: Prescribed dofs for the structure. Defaults to no prescribed dofs.
    :param horseshoe: Horseshoe flag.
    :param use_f_ext_follower: If true, allocate an array for follower forces.
    :param use_f_ext_dead: If true, allocate an array for dead forces.

    :param t_init: Initial time
    :return: Static aeroelastic object for undeformed wing
    """
    prescribed_dofs = self.structure.make_prescribed_dofs_tuple(prescribed_dofs)
    return AeroelasticCase(
        structure=self.structure.reference_configuration(
            use_f_grav=self.structure.use_gravity,
            use_f_ext_dead=use_f_ext_dead,
            use_f_ext_follower=use_f_ext_follower,
            use_f_aero=True,
            prescribed_dofs=prescribed_dofs,
        ),
        aero=self.aero.static_solve(
            t=t_init, hg=self.structure.hg0, horseshoe=horseshoe
        ),
    )
initialise_dynamic 
initialise_dynamic(
    static_case: AeroelasticCase,
    prescribed_dofs: Sequence[int] | Array | slice | int,
) -> AeroelasticCase

Initialise a dynamic aeroelastic snapshot from a static aeroelastic case. This takes a static aeroelastic case obtained under clamped conditions (i.e. relative_motion is True for the free stream), and sets the structural velocity to be that of the freestream.

Parameters:

Name Type Description Default
static_case AeroelasticCase

Static aeroelastic case.

 required
prescribed_dofs Sequence[int] | Array | slice | int

Prescribed dofs. This is often useful for updating from a clamped trim to a free-flying dynamic case.

 required

Returns:

Type Description
AeroelasticCase

Dynamic aeroelastic snapshot.
Source code in src/flapjax/coupled/coupled.py 
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
def initialise_dynamic(
    self,
    static_case: AeroelasticCase,
    prescribed_dofs: Sequence[int] | Array | slice | int,
) -> AeroelasticCase:
    r"""
    Initialise a dynamic aeroelastic snapshot from a static aeroelastic case. This takes a static aeroelastic case
    obtained under clamped conditions (i.e. `relative_motion` is True for the free stream), and sets the structural
    velocity to be that of the freestream.
    :param static_case: Static aeroelastic case.
    :param prescribed_dofs: Prescribed dofs. This is often useful for updating from a clamped trim to a free-flying
    dynamic case.
    :return: Dynamic aeroelastic snapshot.
    """
    u_inf = self.aero.flowfield.u_inf  # flowfield velocity to set
    self.aero.flowfield.relative_motion = (
        False  # the output will have relative motion disabled
    )

    rmat_struct = static_case.structure.hg[:, :3, :3]  # deformed rotations
    v_local = jnp.einsum(
        "ijk,j->ik", rmat_struct, -u_inf
    )  # local frame velocity, (n_nodes, 3)

    dynamic_case = static_case.to_dynamic(t=None)
    dynamic_case.structure.v = dynamic_case.structure.v.at[:, :3].set(v_local)
    dynamic_case.structure.prescribed_dofs = (
        self.structure.make_prescribed_dofs_tuple(prescribed_dofs)
    )
    dynamic_case.structure.free_dofs = get_solve_dofs(
        n_dof=self.structure.n_dof,
        prescribed_dofs=dynamic_case.structure.prescribed_dofs,
    )

    return dynamic_case

data_structures

AeroelasticCase 

AeroelasticCase(structure: StructureCase, aero: AeroCase)

Coupled aeroelastic case, which is a wrapper around a StructureCase and a AeroCase pair.

A single instance may represent any of three flavours (derived from the wrapped structure):

  • Static: static snapshot structure and snapshot aero.
  • Dynamic snapshot: dynamic snapshotstructure and snapshot aero.
  • Dynamic trajectory: batched structure and batched aero.

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

Source code in src/flapjax/coupled/data_structures.py 
46
47
48
def __init__(self, structure: StructureCase, aero: AeroCase):
    self.structure: StructureCase = structure
    self.aero: AeroCase = aero
to_dynamic 
to_dynamic() -> AeroelasticCase

to_dynamic(t: None) -> AeroelasticCase

to_dynamic(t: Array) -> AeroelasticCase

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

Convert a static snapshot to a dynamic snapshot (t=None) or a batched trajectory (t provided). Calling on an already-dynamic case returns self.

Source code in src/flapjax/coupled/data_structures.py 
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
def to_dynamic(self, t: Array | None = None) -> AeroelasticCase:
    """Convert a static snapshot to a dynamic snapshot (``t=None``) or a batched
    trajectory (``t`` provided). Calling on an already-dynamic case returns
    ``self``.
    """
    if self.is_dynamic:
        return self
    if t is None:
        return AeroelasticCase(
            structure=self.structure.to_dynamic(), aero=self.aero
        )
    return AeroelasticCase(
        structure=self.structure.to_dynamic(t),
        aero=self.aero.to_dynamic(i_ts=0, n_tstep=len(t)),
    )
to_static 
to_static() -> AeroelasticCase

Return a static AeroelasticCase, dropping the structure's velocity/acceleration fields.

Source code in src/flapjax/coupled/data_structures.py 
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
def to_static(self) -> AeroelasticCase:
    """Return a static AeroelasticCase, dropping the structure's
    velocity/acceleration fields.
    """
    if not self.is_dynamic:
        return self
    if self.is_batched:
        raise ValueError(
            "to_static() on a batched AeroelasticCase is ambiguous; index a "
            "single time step first (e.g. `case[i_ts].to_static()`)."
        )
    return AeroelasticCase(
        structure=self.structure.to_static(),
        aero=self.aero,
    )
initialise classmethod 
initialise(
    initial_snapshot: AeroelasticCase,
    t: Array,
    use_f_ext_follower: bool,
    use_f_ext_dead: bool,
    structure: BeamStructure,
    x0_aero: ArrayList,
) -> AeroelasticCase

Build a batched AeroelasticCase from any single-timestep case.

Source code in src/flapjax/coupled/data_structures.py 
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
@classmethod
def initialise(
    cls,
    initial_snapshot: AeroelasticCase,
    t: Array,
    use_f_ext_follower: bool,
    use_f_ext_dead: bool,
    structure: BeamStructure,
    x0_aero: ArrayList,
) -> AeroelasticCase:
    """Build a batched AeroelasticCase from any single-timestep case."""
    if initial_snapshot.is_batched:
        if initial_snapshot.structure.n_tstep != 1:
            raise ValueError("initial_snapshot.structure.n_tstep != 1")
        if initial_snapshot.aero.n_tstep != 1:
            raise ValueError("initial_snapshot.aero.n_tstep != 1")
        init_struct: StructureCase = initial_snapshot.structure[0]
        init_aero: AeroCase = initial_snapshot.aero[0]
    elif not initial_snapshot.is_dynamic:
        init_struct = initial_snapshot.structure.to_dynamic(t=None)
        init_aero = initial_snapshot.aero
    else:
        init_struct = initial_snapshot.structure
        init_aero = initial_snapshot.aero

    struct_case = StructureCase.initialise(
        initial_snapshot=init_struct,
        t=t,
        use_f_ext_aero=True,
        use_f_ext_follower=use_f_ext_follower,
        use_f_ext_dead=use_f_ext_dead,
    )
    aero_case = AeroCase.initialise(initial_snapshot=init_aero, n_tstep=len(t))

    # compute aerodynamic forcing at timestep 0
    f_aero_init = aero_case.project_forcing_to_beam(
        i_ts=0,
        rmat=struct_case.hg[0, :, :3, :3],
        x0_aero=x0_aero,
        include_unsteady=False,
    )
    f_aero_local = structure.make_f_dead_ext(
        f_ext=f_aero_init, rmat=struct_case.hg[0, :, :3, :3]
    )

    if struct_case.f_ext_aero is None:
        raise ValueError("f_ext_aero cannot be None")

    struct_case.f_ext_aero = struct_case.f_ext_aero.at[0, ...].set(f_aero_local)

    return AeroelasticCase(structure=struct_case, aero=aero_case)
get_full_states 
get_full_states(
    i_ts: int | Array | None = None,
) -> AeroelasticFullStates

Get the full aeroelastic states. For a batched case, i_ts selects the timestep; for a snapshot, i_ts is ignored.

Source code in src/flapjax/coupled/data_structures.py 
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
def get_full_states(self, i_ts: int | Array | None = None) -> AeroelasticFullStates:
    r"""
    Get the full aeroelastic states. For a batched case, ``i_ts`` selects the
    timestep; for a snapshot, ``i_ts`` is ignored.
    """
    if self.is_batched:
        if i_ts is None:
            raise ValueError("i_ts must be provided for batched AeroelasticCase")
        return AeroelasticFullStates(
            structure=self.structure.get_full_states(i_ts=i_ts),
            aero=self.aero.get_states(i_ts=i_ts),
        )
    if self.structure.f_ext_aero is None:
        raise ValueError("f_ext_aero is None")
    return AeroelasticFullStates(
        structure=self.structure.get_full_states(),
        aero=self.aero.get_states(i_ts=0 if self.is_dynamic else None),
    )
get_minimal_states 
get_minimal_states(
    i_ts: int | Array,
) -> AeroelasticMinimalStates

Get minimal aeroelastic states at the given timestep (batched only).

Source code in src/flapjax/coupled/data_structures.py 
176
177
178
179
180
181
182
183
184
185
def get_minimal_states(self, i_ts: int | Array) -> AeroelasticMinimalStates:
    r"""Get minimal aeroelastic states at the given timestep (batched only)."""
    if not self.is_batched:
        raise TypeError(
            "get_minimal_states only supported for batched AeroelasticCase"
        )
    return AeroelasticMinimalStates(
        structure=self.structure.get_minimal_states(i_ts=i_ts),
        aero=self.aero.get_states(i_ts=i_ts),
    )
plot 
plot(
    directory: PathLike | str,
    index: int
    | Sequence[int]
    | Array
    | slice
    | None = None,
    n_interp: int = 0,
    plot_bound: bool = True,
    plot_wake: bool = True,
) -> tuple[Path, Sequence[Path]]

Plot the aeroelastic case.

Parameters:

Name Type Description Default
directory PathLike | str

Directory to save the plots to.

 required
index int | Sequence[int] | Array | slice | None

For batched cases, timestep indices to plot; ignored for snapshots.

 None
n_interp int

Number of interpolation points for plotting the structure.

 0
plot_bound bool

Whether to plot the bound aerodynamic panels.

 True
plot_wake bool

Whether to plot the wake aerodynamic panels.

 True

Returns:

Type Description
tuple[Path, Sequence[Path]]

(structure_path, aero_paths). For batched cases the structure path is a PVD; for snapshots it is a single VTU.
Source code in src/flapjax/coupled/data_structures.py 
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
def plot(
    self,
    directory: os.PathLike | str,
    index: int | Sequence[int] | Array | slice | None = None,
    n_interp: int = 0,
    plot_bound: bool = True,
    plot_wake: bool = True,
) -> tuple[Path, Sequence[Path]]:
    r"""Plot the aeroelastic case.
    :param directory: Directory to save the plots to.
    :param index: For batched cases, timestep indices to plot; ignored for
        snapshots.
    :param n_interp: Number of interpolation points for plotting the structure.
    :param plot_bound: Whether to plot the bound aerodynamic panels.
    :param plot_wake: Whether to plot the wake aerodynamic panels.
    :return: ``(structure_path, aero_paths)``. For batched cases the
        structure path is a PVD; for snapshots it is a single VTU.
    """
    if self.is_batched:
        struct_out: Path = self.structure.plot(
            directory=directory, n_interp=n_interp, index=index
        )
    else:
        struct_out = self.structure.plot(directory=directory, n_interp=n_interp)
    aero_out: Sequence[Path] = self.aero.plot(
        directory=directory,  # type: ignore
        plot_bound=plot_bound,
        plot_wake=plot_wake,
        index=index,
    )
    return struct_out, aero_out

gradients

coupled

CoupledAeroelastic 
CoupledAeroelastic(
    structure: BeamStructure,
    aero: UVLM,
    fsi_convergence_settings: ConvergenceSettings = DEFAULT_FSI_CONVERGENCE_SETTINGS,
)

Bases: BaseCoupledAeroelastic

Source code in src/flapjax/coupled/coupled.py 
53
54
55
56
57
58
59
60
61
def __init__(
    self,
    structure: BeamStructure,
    aero: UVLM,
    fsi_convergence_settings: ConvergenceSettings = DEFAULT_FSI_CONVERGENCE_SETTINGS,
):
    self.structure: BeamStructure = structure
    self.aero: UVLM = aero
    self.fsi_convergence_settings: ConvergenceSettings = fsi_convergence_settings
aeroelastic_states_res_from_dv_varphi 
aeroelastic_states_res_from_dv_varphi(
    dv: AeroelasticDesignVariables,
    varphi: Array,
    thrust: dict[str, Array],
    i_ts: int,
    t: Array,
    use_horseshoe: bool,
) -> tuple[AeroelasticFullStates, Array]

Obtain useful states and forcing residual from design variables and a minimal configuration vector.

Source code in src/flapjax/coupled/gradients/coupled.py 
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
def aeroelastic_states_res_from_dv_varphi(
    self,
    dv: AeroelasticDesignVariables,
    varphi: Array,
    thrust: dict[str, Array],
    i_ts: int,
    t: Array,
    use_horseshoe: bool,
) -> tuple[AeroelasticFullStates, Array]:
    r"""
    Obtain useful states and forcing residual from design variables and a minimal configuration vector.
    """

    # make a copy of the structure object to prevent modifying the original states
    inner_case = pytree_clone(self)

    struct_dv = dv.structure
    aero_dv = dv.aero
    struct = self.structure
    aero = self.aero
    flowfield = (
        aero.flowfield.from_design_variables(design_variables=aero_dv.flowfield)
        if aero_dv.flowfield is not None
        else aero.flowfield
    )
    inner_case.set_design_variables(
        coords=dv_or(struct_dv.x0, struct.x0),
        k_cs=dv_or(struct_dv.k_cs, struct.k_cs),
        m_cs=dv_or(struct_dv.m_cs, struct.m_cs),
        m_lumped=dv_or(struct_dv.m_lumped, struct.m_lumped)
        if struct.use_lumped_mass
        else None,
        thrust_reference=dv_or(struct_dv.thrust_t, struct.thrust_reference),
        flowfield=flowfield,
        delta_w=aero.delta_w,
        dt=aero.dt,
        x0_aero=dv_or(aero_dv.zeta_b0, aero.zeta_b0),
        orientation_euler=dv_or(
            struct_dv.orientation_euler, struct.orientation_euler
        ),
        cs_angles_reference=dv_or(aero_dv.cs_ang_t, aero.cs_ang0),
        remove_checks=True,
    )

    exp_varphi = vmap(exp_se3)(varphi.reshape(-1, 6))  # (n_nodes, 4, 4)
    hg = jnp.einsum(
        "ijk,ikl->ijl", inner_case.structure.hg0, exp_varphi
    )  # (n_nodes, 4, 4)

    # evaluate aero forcing and project to beam nodes
    aero_sol = inner_case.aero.static_solve(hg=hg, t=t, horseshoe=use_horseshoe)
    f_ext_aero_global = aero_sol.project_forcing_to_beam(
        i_ts=0,
        rmat=hg[:, :3, :3],
        x0_aero=self.aero.zeta_b0,
        include_unsteady=False,
    )

    d = inner_case.structure.make_d(hg)
    p_d = inner_case.structure.make_p_d(d)
    eps = inner_case.structure.make_eps(d)
    f_elem = inner_case.structure.make_f_elem(eps=eps)

    if inner_case.structure.use_gravity:
        m_t = inner_case.structure.make_m_t(d)
    else:
        m_t = None

    if dv.structure.f_ext_dead is not None:
        f_ext_dead = inner_case.structure.make_f_dead_ext(
            dv.structure.f_ext_dead, hg[:, :3, :3]
        )
    else:
        f_ext_dead = None

    f_dead_total = inner_case.structure.make_f_ext_dead_tot(
        f_ext_dead, f_ext_aero_global, i_load_step=None
    )

    f_res = inner_case.structure.make_f_res(
        solve_dofs=None,
        p_d=p_d,
        eps=eps,
        hg=hg,
        f_ext_follower_n=dv.structure.f_ext_follower,
        f_ext_dead_n=f_dead_total,
        thrust_n=dv.structure.thrust_t
        if dv.structure.thrust_t is not None
        else thrust,
        dynamic=False,
        m_t=m_t,
        c_l=None,
        c_l_lumped=None,
        v=None,
        v_dot=None,
    )[0]

    struct_states = StructureFullStates(
        hg=hg,
        varphi=varphi,
        eps=eps,
        f_elem=f_elem,
        f_res=f_res.reshape(-1, 6),
        v=None,
        v_dot=None,
    )

    aero_states = aero_sol.get_states(i_ts=i_ts)

    return AeroelasticFullStates(structure=struct_states, aero=aero_states), f_res
static_adjoint 
static_adjoint(
    case: AeroelasticCase,
    objective: AeroelasticObjectiveFunction,
    grads_to_compute: AeroelasticGradsToCompute = DEFAULT_GRADS_TO_COMPUTE,
    optional_jacobians: OptionalJacobians
    | None = DEFAULT_OPTIONAL_JACOBIANS,
    ad_mode: ADMode = "forward",
    batch_size: int | None = 32,
) -> tuple[AeroelasticDesignVariables, 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
case AeroelasticCase

AeroelasticCase containing the current state of the aeroelastic system.

 required
objective AeroelasticObjectiveFunction

Objective function that takes the structure and design variables and returns an array

 required
grads_to_compute AeroelasticGradsToCompute

Data structure which specifies which gradients to compute. This is used to speed up the adjoint solve by only computing the necessary Jacobian blocks.

 DEFAULT_GRADS_TO_COMPUTE
optional_jacobians OptionalJacobians | None

OptionalJacobians object specifying which Jacobians to compute.

 DEFAULT_OPTIONAL_JACOBIANS
ad_mode ADMode

Optional use of either forward or reverse adjoint. For passing the initial state sensitivities to a dynamic solve, only forward mode can be used to give the required adjoint.

 'forward'
batch_size int | None

Batch size for computing p_res_p_varphi to reduce memory usage. Ignored when matrix_free is True.

 32

Returns:

Type Description
tuple[AeroelasticDesignVariables, Array]

Gradient of objective function output with respect to design variables.
Source code in src/flapjax/coupled/gradients/coupled.py 
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
@jax.jit(static_argnums=(0, 1, 2, 3, 4, 5, 6))
def static_adjoint(
    self,
    case: AeroelasticCase,
    objective: AeroelasticObjectiveFunction,
    grads_to_compute: AeroelasticGradsToCompute = DEFAULT_GRADS_TO_COMPUTE,
    optional_jacobians: OptionalJacobians | None = DEFAULT_OPTIONAL_JACOBIANS,
    ad_mode: ADMode = "forward",
    batch_size: int | None = 32,
) -> tuple[AeroelasticDesignVariables, 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 case: AeroelasticCase containing the current state of the aeroelastic system.
    :param objective: Objective function that takes the structure and design variables and returns an array
    :param grads_to_compute: Data structure which specifies which gradients to compute. This is used to speed up the
    adjoint solve by only computing the necessary Jacobian blocks.
    :param optional_jacobians: OptionalJacobians object specifying which Jacobians to compute.
    :param ad_mode: Optional use of either forward or reverse adjoint. For passing the initial state sensitivities
    to a dynamic solve, only forward mode can be used to give the required adjoint.
    :param batch_size: Batch size for computing p_res_p_varphi to reduce memory usage. Ignored when
    ``matrix_free`` is True.
    :return: Gradient of objective function output with respect to design variables.
    """

    if ad_mode not in ("forward", "reverse"):
        raise ValueError("ad_mode must be either 'forward' or 'reverse'")

    jax_print("Computing static adjoint", verbose_level="normal")

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

    if optional_jacobians is not None:
        self.structure.optional_jacobians = optional_jacobians

    dv = self.get_design_variables(case=case, grads_to_compute=grads_to_compute)
    states = case.get_full_states()

    # find shape of objective function output without evaluating function
    f_properties = jax.eval_shape(lambda: objective(states, dv, None))
    f_shape = f_properties.shape
    j0_shape = f_shape if len(f_shape) > 0 else (1,)
    n_f = f_properties.size
    n_x = dv.structure.n_x + dv.aero.n_x
    n_u_full = self.structure.n_dof

    varphi = case.structure.varphi

    if case.aero.static_horseshoe is None:
        raise ValueError("static_horseshoe not defined")
    static_horseshoe: bool = case.aero.static_horseshoe

    # function for computing sensitivity of objective to design variables and degrees of freedom
    # to obtain the actual Jacobian we must pull back the identity through it
    vjp_fn = self.compute_p_j0_p_x(
        case=case,
        objective=objective,
        grads_to_compute=grads_to_compute,
        horseshoe=static_horseshoe,
    )

    cot_j0 = jnp.eye(n_f).reshape(n_f, *j0_shape)  # seed for backpropogation
    p_j_p_varphi_raw, p_j_p_x_raw = jax.vmap(vjp_fn)(cot_j0)  # sensitivities

    p_j_p_varphi_flat = p_j_p_varphi_raw.reshape(n_f, -1)  # (n_f, n_dof)
    p_j_p_x_flat = AeroelasticDesignVariables(
        structure_dv=StructureDesignVariables(
            **{
                k: getattr(p_j_p_x_raw.structure, k) for k in dv.structure.to_dict()
            },
            f_shape=(n_f,),
        ),
        aero_dv=AeroDesignVariables(
            **{k: getattr(p_j_p_x_raw.aero, k) for k in dv.aero.to_dict()},
            f_shape=(n_f,),
        ),
    ).ravel_jacobian(f_size=n_f, x_size=n_x)

    def _residual(varphi_vec: Array, dv_: AeroelasticDesignVariables) -> Array:
        r"""
        Helper function to give the static aeroelastic residual for a given deformation and design variables.
        """
        return self.aeroelastic_states_res_from_dv_varphi(
            dv=dv_,
            varphi=varphi_vec.reshape(self.structure.n_nodes, 6),
            thrust=case.structure.thrust,
            t=case.aero.t,
            i_ts=0,
            use_horseshoe=static_horseshoe,
        )[1]

    if ad_mode == "forward":
        # single joint VJP shared between p_res_p_varphi and p_res_p_x
        _, vjp_res_both = jax.vjp(_residual, varphi.ravel(), dv)
        p_res_p_varphi, p_res_p_x = jax.lax.map(
            vjp_res_both,
            jnp.eye(n_u_full),
            batch_size=batch_size,
        )
        # solve for adjoint
        adj = jnp.linalg.solve(
            p_res_p_varphi[jnp.ix_(solve_dofs, solve_dofs)],
            p_res_p_x.ravel_jacobian(f_size=n_u_full, x_size=n_x)[solve_dofs, :],
        )

        d_f_d_x_dict = dv.from_adjoint(
            f_shape,
            p_j_p_x_flat - p_j_p_varphi_flat[:, solve_dofs] @ adj,
        )
    else:
        # construct residual Jacobian
        _, vjp_res_varphi = jax.vjp(lambda v: _residual(v, dv), varphi.ravel())
        p_res_p_varphi = jax.lax.map(
            lambda cot: vjp_res_varphi(cot)[0],
            jnp.eye(n_u_full),
            batch_size=batch_size,
        )
        adj = jnp.linalg.solve(
            p_res_p_varphi[jnp.ix_(solve_dofs, solve_dofs)].T,
            p_j_p_varphi_flat[:, solve_dofs].T,
        ).T  # (n_f, n_solve_dofs)

        adj_full = (
            jnp.zeros((n_f, n_u_full), dtype=adj.dtype).at[:, solve_dofs].set(adj)
        )

        # sensitivity of residual w.r.t. design variables
        _, vjp_res_dv = jax.vjp(lambda dv_: _residual(varphi.ravel(), dv_), dv)
        (adj_p_res_p_x_raw,) = jax.vmap(vjp_res_dv)(adj_full)

        adj_p_res_p_x_flat = AeroelasticDesignVariables(
            structure_dv=StructureDesignVariables(
                **{
                    k: getattr(adj_p_res_p_x_raw.structure, k)
                    for k in dv.structure.to_dict()
                },
                f_shape=(n_f,),
            ),
            aero_dv=AeroDesignVariables(
                **{
                    k: getattr(adj_p_res_p_x_raw.aero, k) for k in dv.aero.to_dict()
                },
                f_shape=(n_f,),
            ),
        ).ravel_jacobian(f_size=n_f, x_size=n_x)

        d_f_d_x_dict = dv.from_adjoint(
            f_shape,
            p_j_p_x_flat - adj_p_res_p_x_flat,
        )

    return dv.split_adjoint(d_f_d_x=d_f_d_x_dict, f_shape=f_shape), adj
compute_p_j0_p_x 
compute_p_j0_p_x(
    case: AeroelasticCase,
    objective: AeroelasticObjectiveFunction,
    grads_to_compute: AeroelasticGradsToCompute | None,
    horseshoe: bool = False,
    include_q0: bool = False,
) -> Callable[
    ..., tuple[Array, AeroelasticDesignVariables]
]

Build the VJP of the initial-timestep objective for pertubations in the design variables, and optionally the initial states.

Parameters:

Name Type Description Default
case AeroelasticCase

AeroelasticCase solution for the initial timestep.

 required
objective AeroelasticObjectiveFunction

Objective function which takes the system full states, design variables and timestep index.

 required
grads_to_compute AeroelasticGradsToCompute | None

Grads to compute when computing design gradient.

 required
horseshoe bool

Flag for using horseshoe wake.

 False
include_q0 bool

If True, the returned VJP also propagates a cotangent through q0.

 False

Returns:

Type Description
Callable[..., tuple[Array, AeroelasticDesignVariables]]

VJP for cotangents of the design variables, and optionally the initial states.
Source code in src/flapjax/coupled/gradients/coupled.py 
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
def compute_p_j0_p_x(
    self,
    case: AeroelasticCase,
    objective: AeroelasticObjectiveFunction,
    grads_to_compute: AeroelasticGradsToCompute | None,
    horseshoe: bool = False,
    include_q0: bool = False,
) -> Callable[
    ...,
    tuple[Array, AeroelasticDesignVariables],
]:
    r"""
    Build the VJP of the initial-timestep objective for pertubations in the design variables, and optionally the
    initial states.
    :param case: AeroelasticCase solution for the initial timestep.
    :param objective: Objective function which takes the system full states, design variables and timestep index.
    :param grads_to_compute: Grads to compute when computing design gradient.
    :param horseshoe: Flag for using horseshoe wake.
    :param include_q0: If True, the returned VJP also propagates a cotangent through ``q0``.
    :return: VJP for cotangents of the design variables, and optionally the initial states.
    """

    # design variables with variables that we don't require gradients omitted to speed up computations.
    dv = self.get_design_variables(case=case, grads_to_compute=grads_to_compute)

    # design variables with no omissions
    dv_full = self.get_design_variables(case=case, grads_to_compute=None)

    varphi = case.structure.varphi
    n_dof = self.structure.n_dof

    @jax.checkpoint
    def objective_from_varphi(
        varphi_: Array,
        dv_: AeroelasticDesignVariables,
    ) -> Array | tuple[Array, Array]:
        inner_case = self.case_from_dv(dv=dv_)

        assert (
            dv_full.aero.cs_ang_t is not None and dv_full.aero.cs_vel_t is not None
        )

        # solve aero problem
        hg = inner_case.structure.compute_hg_from_varphi(varphi=varphi_)
        _, _, gamma_b, gamma_w, _, _, zeta_w, _, f_steady, _, _, _, _, _ = (
            inner_case.aero.base_solve(
                q_nm1=None,
                t_n=case.aero.t,
                hg_n=hg,
                hg_nm1=None,
                hg_dot_n=None,
                static=True,
                horseshoe=horseshoe,
                cs_ang_n={
                    k: jnp.atleast_1d(v)[0]
                    for k, v in (
                        dv_.aero.cs_ang_t
                        if dv_.aero.cs_ang_t is not None
                        else dv_full.aero.cs_ang_t
                    ).items()
                },
                cs_ang_nm1=None,
                cs_vel_n={
                    k: jnp.atleast_1d(v)[0]
                    for k, v in (
                        dv_.aero.cs_vel_t
                        if dv_.aero.cs_vel_t is not None
                        else dv_full.aero.cs_vel_t
                    ).items()
                },
            )
        )

        f_aero_beam_global = project_forcing_to_beam(
            f_total=f_steady,
            rmat=hg[:, :3, :3],
            dof_mapping=inner_case.aero.dof_mapping,
            x0_aero=inner_case.aero.zeta_b0,
            mirror_edge_low=inner_case.aero.mirror_edge_low,
            mirror_edge_high=inner_case.aero.mirror_edge_high,
        )

        f_aero_beam_local = transform_nodal_vect(
            vect=f_aero_beam_global, rmat=jnp.transpose(hg[:, :3, :3], (0, 2, 1))
        )

        q_aero = AeroFullStates(
            gamma_b=gamma_b,
            gamma_w=gamma_w,
            zeta_w=zeta_w,
            gamma_b_dot=ArrayList.zeros_like(gamma_b),
        )

        # assume initial velocities and accelerations are zero
        q_structure = StructureMinimalStates(
            varphi=varphi_,
            v=jnp.zeros_like(varphi_),
            v_dot=jnp.zeros_like(varphi_),
            a=jnp.zeros_like(varphi_),
            f_ext_aero=f_aero_beam_local,
        )

        q0 = AeroelasticMinimalStates(structure=q_structure, aero=q_aero).ravel()
        q0_full = self.minimal_states_to_full_states(
            i_ts=0,
            q=AeroelasticMinimalStates.from_vector(
                vect=q0,
                n_dof=n_dof,
                aero_shapes=q_aero.shapes(),
            ),
            dv=dv_,
            dv_full=dv_full,
        )
        j0 = jnp.atleast_1d(objective(q0_full, dv_, 0))
        if include_q0:
            return j0, q0
        return j0

    _, vjp_fn = jax.vjp(objective_from_varphi, varphi, dv)
    return vjp_fn
timestep_residual 
timestep_residual(
    i_ts: int | Array,
    t: Array,
    q_nm1: AeroelasticMinimalStates,
    q_n: AeroelasticMinimalStates,
    dv_: AeroelasticDesignVariables,
    dv_full: AeroelasticDesignVariables,
    thrust_t: dict[str, Array],
    solve_dofs: tuple[int, ...],
    approx_grads: bool,
) -> Array

Compute the coupled aeroelastic residual vector. This is used for the matrix-free time domain case by applying VJP to this function.

Parameters:

Name Type Description Default
i_ts int | Array

Time step index.

 required
t Array

Time at step n.

 required
q_nm1 AeroelasticMinimalStates

Minimal states at step n-1.

 required
q_n AeroelasticMinimalStates

Minimal states at step n.

 required
dv_ AeroelasticDesignVariables

Aeroelastic design variables (may have some fields omitted).

 required
dv_full AeroelasticDesignVariables

Aeroelastic design variables without omissions.

 required
thrust_t dict[str, Array]

Thrust time history, {key: [n_tstep]}.

 required
solve_dofs tuple[int, ...]

Structural degrees of freedom which are solved for.

 required
approx_grads bool

If True, remove some gradient terms which are generally small.

 required

Returns:

Type Description
Array

Coupled residual (n_adj_dof, ).
Source code in src/flapjax/coupled/gradients/coupled.py 
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
def timestep_residual(
    self,
    i_ts: int | Array,
    t: Array,
    q_nm1: AeroelasticMinimalStates,
    q_n: AeroelasticMinimalStates,
    dv_: AeroelasticDesignVariables,
    dv_full: AeroelasticDesignVariables,
    thrust_t: dict[str, Array],
    solve_dofs: tuple[int, ...],
    approx_grads: bool,
) -> Array:
    r"""
    Compute the coupled aeroelastic residual vector. This is used for the matrix-free time domain case by applying
    VJP to this function.
    :param i_ts: Time step index.
    :param t: Time at step n.
    :param q_nm1: Minimal states at step n-1.
    :param q_n: Minimal states at step n.
    :param dv_: Aeroelastic design variables (may have some fields omitted).
    :param dv_full: Aeroelastic design variables without omissions.
    :param thrust_t: Thrust time history, {key: [n_tstep]}.
    :param solve_dofs: Structural degrees of freedom which are solved for.
    :param approx_grads: If True, remove some gradient terms which are generally small.
    :return: Coupled residual ``(n_adj_dof, )``.
    """
    assert (
        q_n.structure.f_ext_aero is not None
        and q_nm1.structure.f_ext_aero is not None
    )

    struct_res = self.structure.timestep_residual(
        i_ts=i_ts,
        q_nm1=q_nm1.structure,
        q_n=q_n.structure,
        dv_=dv_.structure,
        thrust_t=thrust_t,
        solve_dofs=solve_dofs,
        approx_grads=approx_grads,
    )

    # Rematerialise the aero pass to reduce memory usage
    @jax.checkpoint
    def _aero_forward(
        varphi_nm1_: Array,
        varphi_n_: Array,
        v_n_: Array,
        t_n_: Array,
        q_nm1_aero_: AeroFullStates,
        q_n_aero_: AeroFullStates,
        dv__: AeroelasticDesignVariables,
        dv_full_: AeroelasticDesignVariables,
        f_aero_beam_n_: Array,
    ) -> Array:
        return self.aero.timestep_residual(
            i_ts=i_ts,
            varphi_nm1=varphi_nm1_,
            varphi_n=varphi_n_,
            v_n=v_n_,
            t_n=t_n_,
            q_n=q_n_aero_,
            q_nm1=q_nm1_aero_,
            dv=dv__,
            dv_full=dv_full_,
            f_aero_beam_n=f_aero_beam_n_,
            struct_obj=self.structure,
            approx_grads=approx_grads,
        )

    # evaluate checkpointed function
    aero_res = _aero_forward(
        varphi_nm1_=q_nm1.structure.varphi,
        varphi_n_=q_n.structure.varphi,
        v_n_=q_n.structure.v,
        t_n_=t,
        q_nm1_aero_=q_nm1.aero,
        q_n_aero_=q_n.aero,
        dv__=dv_,
        dv_full_=dv_full,
        f_aero_beam_n_=q_n.structure.f_ext_aero,
    )

    # remove forces from degrees of freedom which are not solved for
    n_aero_states: int = q_n.aero.n_states
    solve_dofs_arr = jnp.array(solve_dofs)
    aero_res_solve = jnp.concatenate(
        (aero_res[:n_aero_states], aero_res[n_aero_states:][solve_dofs_arr])
    )

    return jnp.concatenate((struct_res, aero_res_solve))
timestep_residual_jacobians 
timestep_residual_jacobians(
    i_ts: int | Array,
    t: Array,
    q_nm1: AeroelasticMinimalStates,
    q_n: AeroelasticMinimalStates,
    dv_: AeroelasticDesignVariables,
    dv_full: AeroelasticDesignVariables,
    thrust_t: dict[str, Array],
    solve_dofs: tuple[int, ...],
    approx_grads: bool,
    n_profile_loops: int | None,
    jac_options: dict,
    mode: ADMode = "reverse",
    map_batch_size: int | None = None,
) -> tuple[
    Array,
    Array,
    StructureDesignVariables,
    AeroelasticDesignVariables,
    dict[str, dict[str, float]] | None,
    dict[str, dict[str, float]] | None,
]

Compute the required time-domain residual Jacobians for the adjoint solution.

Parameters:

Name Type Description Default
i_ts int | Array

Time step index.

 required
t Array

Time.

 required
q_nm1 AeroelasticMinimalStates

Minimal degrees of freedom at timestep n-1.

 required
q_n AeroelasticMinimalStates

Minimal degrees of freedom at timestep n.

 required
dv_ AeroelasticDesignVariables

Design variables for which to obtain gradients.

 required
dv_full AeroelasticDesignVariables

All design variables.

 required
thrust_t dict[str, Array]

Thrust at each time step.

 required
solve_dofs tuple[int, ...]

Degrees of freedom to solve for.

 required
approx_grads bool

Whether to use approximate gradients for the structural dynamic subproblem.

 required
n_profile_loops int | None

Number of profile loops. Used for profiling routines only.

 required
jac_options dict

Options for Jacobian computation, allowing for approximations to be introduced.

 required
mode ADMode

Mode for automatic differentiation, either forward or reverse.

 'reverse'
map_batch_size int | None

Batch size used for vectorising Jacobian construction.

 None

Returns:

Type Description
tuple[Array, Array, StructureDesignVariables, AeroelasticDesignVariables, dict[str, dict[str, float]] | None, dict[str, dict[str, float]] | None]

Jacobian of residual with respect to previous degrees of freedom, Jacobian of residual with respect to current degrees of freedom. Jacobian of v_dot residual with respect to structural design variables, Jacobian of aero residual with respect to design variables. Can also include compile time and run time when profiling is used.
Source code in src/flapjax/coupled/gradients/coupled.py 
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
def timestep_residual_jacobians(
    self,
    i_ts: int | Array,
    t: Array,
    q_nm1: AeroelasticMinimalStates,
    q_n: AeroelasticMinimalStates,
    dv_: AeroelasticDesignVariables,
    dv_full: AeroelasticDesignVariables,
    thrust_t: dict[str, Array],
    solve_dofs: tuple[int, ...],
    approx_grads: bool,
    n_profile_loops: int | None,
    jac_options: dict,
    mode: ADMode = "reverse",
    map_batch_size: int | None = None,
) -> tuple[
    Array,
    Array,
    StructureDesignVariables,
    AeroelasticDesignVariables,
    dict[str, dict[str, float]] | None,
    dict[str, dict[str, float]] | None,
]:
    r"""
    Compute the required time-domain residual Jacobians for the adjoint solution.
    :param i_ts: Time step index.
    :param t: Time.
    :param q_nm1: Minimal degrees of freedom at timestep n-1.
    :param q_n: Minimal degrees of freedom at timestep n.
    :param dv_: Design variables for which to obtain gradients.
    :param dv_full: All design variables.
    :param thrust_t: Thrust at each time step.
    :param solve_dofs: Degrees of freedom to solve for.
    :param approx_grads: Whether to use approximate gradients for the structural dynamic subproblem.
    :param n_profile_loops: Number of profile loops. Used for profiling routines only.
    :param jac_options: Options for Jacobian computation, allowing for approximations to be introduced.
    :param mode: Mode for automatic differentiation, either ``forward`` or ``reverse``.
    :param map_batch_size: Batch size used for vectorising Jacobian construction.
    :return: Jacobian of residual with respect to previous degrees of freedom, Jacobian of residual with respect to
    current degrees of freedom. Jacobian of v_dot residual with respect to structural design variables, Jacobian of
    aero residual with respect to design variables. Can also include compile time and run time when profiling is
    used.
    """

    assert (
        q_n.structure.f_ext_aero is not None
        and q_nm1.structure.f_ext_aero is not None
    )

    (
        p_aero_res_p_q_aero_nm1,
        p_aero_res_p_q_aero_n,
        p_aero_res_d_dv,
        p_aero_res_p_q_struct_nm1,
        p_aero_res_p_q_struct_n,
        aero_compile_time,
        aero_run_time,
    ) = self.aero.timestep_residual_jacobians(
        i_ts=i_ts,
        varphi_nm1=q_nm1.structure.varphi,
        varphi_n=q_n.structure.varphi,
        v_n=q_n.structure.v,
        t_n=t,
        q_n=q_n.aero,
        q_nm1=q_nm1.aero,
        dv=dv_,
        dv_full=dv_full,
        f_aero_beam_n=q_n.structure.f_ext_aero,
        struct_obj=self.structure,
        approx_grads=approx_grads,
        solve_dofs=solve_dofs,
        n_profile_loops=n_profile_loops,
        jac_options=jac_options,
        mode=mode,
        map_batch_size=map_batch_size,
    )

    n_aero_dof, _ = p_aero_res_p_q_struct_nm1.shape

    (
        p_struct_res_p_q_struct_nm1,
        p_struct_res_p_q_struct_n,
        p_v_dot_res_p_struct_dv,
        p_v_dot_res_p_f_ext_nm1,
        p_v_dot_res_p_f_ext_n,
        struct_compile_time,
        struct_run_time,
    ) = self.structure.timestep_residual_jacobians(
        i_ts=i_ts,
        q_nm1=q_nm1.structure,
        q_n=q_n.structure,
        f_ext_aero_n=q_n.structure.f_ext_aero,
        f_ext_aero_nm1=q_nm1.structure.f_ext_aero,
        thrust_t=thrust_t,
        dv=dv_.structure,
        solve_dofs=solve_dofs,
        approx_grads=approx_grads,
        n_profile_loops=n_profile_loops,
        jac_options=jac_options,
        mode=mode,
    )

    assert p_v_dot_res_p_f_ext_nm1 is not None and p_v_dot_res_p_f_ext_n is not None

    # reduce aero-to-struct Jacobian columns to solve_dofs
    n_solve = len(solve_dofs)
    n_struct_res = p_struct_res_p_q_struct_nm1.shape[0]
    solve_dofs_arr = jnp.array(solve_dofs)
    struct_col_ix = jnp.concatenate(
        [solve_dofs_arr + i * self.structure.n_dof for i in range(4)]
    )
    p_aero_res_p_q_struct_nm1 = p_aero_res_p_q_struct_nm1[:, struct_col_ix]
    p_aero_res_p_q_struct_n = p_aero_res_p_q_struct_n[:, struct_col_ix]

    # create struct-to-aero cross-coupling (only v_dot residual depends on f_ext_aero)
    # f_aero block is n_solve-wide (last n_solve cols of aero state), so slice Jacobian to solve_dofs cols
    p_struct_res_p_q_aero_nm1 = jnp.zeros((n_struct_res, n_aero_dof))
    p_struct_res_p_q_aero_nm1 = p_struct_res_p_q_aero_nm1.at[
        jnp.arange(n_solve) + 2 * n_solve, -n_solve:
    ].set(p_v_dot_res_p_f_ext_nm1[:, solve_dofs_arr])

    p_struct_res_p_q_aero_n = jnp.zeros((n_struct_res, n_aero_dof))
    p_struct_res_p_q_aero_n = p_struct_res_p_q_aero_n.at[
        jnp.arange(n_solve) + 2 * n_solve, -n_solve:
    ].set(p_v_dot_res_p_f_ext_n[:, solve_dofs_arr])

    p_res_p_q_nm1 = jnp.block(
        [
            [p_struct_res_p_q_struct_nm1, p_struct_res_p_q_aero_nm1],
            [p_aero_res_p_q_struct_nm1, p_aero_res_p_q_aero_nm1],
        ]
    )

    p_res_p_q_n = jnp.block(
        [
            [p_struct_res_p_q_struct_n, p_struct_res_p_q_aero_n],
            [p_aero_res_p_q_struct_n, p_aero_res_p_q_aero_n],
        ]
    )

    if n_profile_loops is not None:
        assert (
            aero_compile_time is not None
            and aero_run_time is not None
            and struct_compile_time is not None
            and struct_run_time is not None
        )
        compile_time = aero_compile_time | struct_compile_time
        run_time = aero_run_time | struct_run_time
    else:
        compile_time = None
        run_time = None
    return (
        p_res_p_q_nm1,
        p_res_p_q_n,
        p_v_dot_res_p_struct_dv,
        p_aero_res_d_dv,
        compile_time,
        run_time,
    )
construct_approximate_jacobians 
construct_approximate_jacobians(
    sol: AeroelasticCase,
    jacobian_approximations: AeroelasticJacobianApproximations,
) -> dict[str, dict[str, Callable[..., Any] | None]]

Compute approximations for Jacobians which are specified in the jacobian_approximations data structure. The aerodynamic residual approximations are delegated to UVLM.construct_approximate_jacobians, and the structural residual approximations to BeamStructure.construct_approximate_jacobians. The two dictionaries are merged into a single result.

Parameters:

Name Type Description Default
sol AeroelasticCase

Solution for which approximations will be created for the initial time step.

 required
jacobian_approximations AeroelasticJacobianApproximations

Data structure which defines which approximations to create.

 required

Returns:

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

Dictionary of approximations keyed by residual name.
Source code in src/flapjax/coupled/gradients/coupled.py 
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
def construct_approximate_jacobians(
    self,
    sol: AeroelasticCase,
    jacobian_approximations: AeroelasticJacobianApproximations,
) -> dict[str, dict[str, Callable[..., Any] | None]]:
    r"""
    Compute approximations for Jacobians which are specified in the jacobian_approximations data structure. The
    aerodynamic residual approximations are delegated to ``UVLM.construct_approximate_jacobians``, and the
    structural residual approximations to ``BeamStructure.construct_approximate_jacobians``. The two
    dictionaries are merged into a single result.
    :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 keyed by residual name.
    """
    dv = self.get_design_variables(case=sol, grads_to_compute=None)
    solve_dofs = get_solve_dofs(
        n_dof=self.structure.n_dof,
        prescribed_dofs=sol.structure.prescribed_dofs,
    )

    aero_options = self.aero.construct_approximate_jacobians(
        aero_sol=sol.aero,
        structure_sol=sol.structure,
        struct_obj=self.structure,
        dv=dv,
        dv_full=dv,
        solve_dofs=solve_dofs,
        jacobian_approximations=jacobian_approximations.aero,
    )

    struct_options = self.structure.construct_approximate_jacobians(
        sol=sol.structure,
        jacobian_approximations=jacobian_approximations.structure,
    )

    return aero_options | struct_options
evaluate_dynamic_objective 
evaluate_dynamic_objective(
    case: AeroelasticCase,
    objective: AeroelasticObjectiveFunction,
) -> Array

Evaluate the dynamic objective for a given case.

Parameters:

Name Type Description Default
case AeroelasticCase

Dynamic aeroelastic case object.

 required
objective AeroelasticObjectiveFunction

Objective function to be evaluated.

 required

Returns:

Type Description
Array

Value of subobjective at every time step, [n_tstep].
Source code in src/flapjax/coupled/gradients/coupled.py 
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
def evaluate_dynamic_objective(
    self, case: AeroelasticCase, objective: AeroelasticObjectiveFunction
) -> Array:
    r"""
    Evaluate the dynamic objective for a given case.
    :param case: Dynamic aeroelastic case object.
    :param objective: Objective function to be evaluated.
    :return: Value of subobjective at every time step, [n_tstep].
    """
    n_tstep = case.structure.n_tstep

    dv = self.get_design_variables(case=case, grads_to_compute=None)

    return jax.vmap(
        lambda i_ts: jnp.atleast_1d(
            objective(case.get_full_states(i_ts=i_ts), dv, i_ts)
        )
    )(jnp.arange(n_tstep)).reshape(n_tstep, -1)
make_frozen_wake_preconditioner 
make_frozen_wake_preconditioner(
    case: AeroelasticCase,
    dv_full: AeroelasticDesignVariables,
    solve_dofs: tuple[int, ...],
    approx_grads: bool = False,
    precond_i_ts: int = 0,
    batch_size: int | None = 32,
) -> Callable[[Array], Array]

Build a preconditioner for the coupled aeroelastic system which skips the wake grid and circulation. When applied to the matrix-free system GMRES, this found a good reduction in the number of iterations required whilst avoiding the large memory overhead involved in computing the full Jacobian due to the large number of wake states.

Parameters:

Name Type Description Default
case AeroelasticCase

Dynamic aeroelastic case object.

 required
dv_full AeroelasticDesignVariables

Full dynamic aeroelastic design variables.

 required
solve_dofs tuple[int, ...]

Solve degree of freedom index.

 required
approx_grads bool

Approximate gradient of the coupled aeroelastic system, removing some negligible terms.

 False
precond_i_ts int

Time step index for which to create the preconditioner. Defaults to 0.

 0
batch_size int | None

Batch size for mapping the Jacobian construction on the aerodynamic system.

 32

Returns:

Type Description
Callable[[Array], Array]

Preconditioner function.
Source code in src/flapjax/coupled/gradients/coupled.py 
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
def make_frozen_wake_preconditioner(
    self,
    case: AeroelasticCase,
    dv_full: AeroelasticDesignVariables,
    solve_dofs: tuple[int, ...],
    approx_grads: bool = False,
    precond_i_ts: int = 0,
    batch_size: int | None = 32,
) -> Callable[[Array], Array]:
    r"""
    Build a preconditioner for the coupled aeroelastic system which skips the wake grid and circulation. When
    applied to the matrix-free system GMRES, this found a good reduction in the number of iterations required whilst
    avoiding the large memory overhead involved in computing the full Jacobian due to the large number of wake
    states.
    :param case: Dynamic aeroelastic case object.
    :param dv_full: Full dynamic aeroelastic design variables.
    :param solve_dofs: Solve degree of freedom index.
    :param approx_grads: Approximate gradient of the coupled aeroelastic system, removing some negligible terms.
    :param precond_i_ts: Time step index for which to create the preconditioner. Defaults to 0.
    :param batch_size: Batch size for mapping the Jacobian construction on the aerodynamic system.
    :return: Preconditioner function.
    """
    precond_q_nm1 = case.get_minimal_states(i_ts=max(precond_i_ts - 1, 0))
    precond_q_n = case.get_minimal_states(i_ts=precond_i_ts)
    precond_t_n = case.structure.t[precond_i_ts]

    dv_precond = self.get_design_variables(case=case, grads_to_compute=None)

    assert precond_q_n.structure.f_ext_aero is not None
    assert precond_q_nm1.structure.f_ext_aero is not None

    # populate jac_options with the expected residual/argname keys, all mapped to None so AD is used everywhere
    jac_options = self.construct_approximate_jacobians(
        sol=case,
        jacobian_approximations=AeroelasticJacobianApproximations(),
    )

    # compute structural Jacobians
    (
        _,
        p_struct_res_p_q_struct_n,
        _,
        _,
        p_v_dot_res_p_f_ext_n,
        *_,
    ) = self.structure.timestep_residual_jacobians(
        i_ts=precond_i_ts,
        q_nm1=precond_q_nm1.structure,
        q_n=precond_q_n.structure,
        f_ext_aero_n=precond_q_n.structure.f_ext_aero,
        f_ext_aero_nm1=precond_q_nm1.structure.f_ext_aero,
        thrust_t=case.structure.thrust,
        dv=dv_precond.structure,
        solve_dofs=solve_dofs,
        approx_grads=approx_grads,
        n_profile_loops=None,
        jac_options=jac_options,
    )

    # compute aero Jacobians
    (
        _,
        p_aero_res_p_q_aero_n,
        _,
        _,
        p_aero_res_p_q_struct_n,
        *_,
    ) = self.aero.timestep_residual_jacobians(
        i_ts=precond_i_ts,
        varphi_nm1=precond_q_nm1.structure.varphi,
        varphi_n=precond_q_n.structure.varphi,
        v_n=precond_q_n.structure.v,
        t_n=precond_t_n,
        q_n=precond_q_n.aero,
        q_nm1=precond_q_nm1.aero,
        dv=dv_precond,
        dv_full=dv_full,
        f_aero_beam_n=precond_q_n.structure.f_ext_aero,
        struct_obj=self.structure,
        approx_grads=approx_grads,
        solve_dofs=solve_dofs,
        n_profile_loops=None,
        jac_options=jac_options,
        compute_wake_gradients=False,
        map_batch_size=batch_size,
    )

    assert p_v_dot_res_p_f_ext_n is not None

    solve_dofs_arr = jnp.array(solve_dofs)
    struct_col_ix = jnp.concatenate(
        [solve_dofs_arr + i * self.structure.n_dof for i in range(4)]
    )
    p_aero_res_p_q_struct_n = p_aero_res_p_q_struct_n[:, struct_col_ix]

    # assemble Jacobians
    n_solve = len(solve_dofs)
    n_struct_res = p_struct_res_p_q_struct_n.shape[0]
    n_aero_reduced = p_aero_res_p_q_aero_n.shape[0]
    p_struct_res_p_q_aero_n = jnp.zeros((n_struct_res, n_aero_reduced))
    p_struct_res_p_q_aero_n = p_struct_res_p_q_aero_n.at[
        jnp.arange(n_solve) + 2 * n_solve, -n_solve:
    ].set(p_v_dot_res_p_f_ext_n[:, solve_dofs_arr])

    p_res_p_q_n_reduced = jnp.block(
        [
            [p_struct_res_p_q_struct_n, p_struct_res_p_q_aero_n],
            [p_aero_res_p_q_struct_n, p_aero_res_p_q_aero_n],
        ]
    )

    # compute the LU decomposition for fast reuse when solving
    precond_lu = jax.scipy.linalg.lu_factor(p_res_p_q_n_reduced.T)

    # adjoint state counts and placement
    n_struct = 4 * n_solve
    n_gamma_b = int(precond_q_n.aero.gamma_b.ravel().size)
    n_gamma_w = int(precond_q_n.aero.gamma_w.ravel().size)
    n_gamma_b_dot = int(precond_q_n.aero.gamma_b_dot.ravel().size)
    n_zeta_w = int(precond_q_n.aero.zeta_w.ravel().size)

    gamma_w_start = n_struct + n_gamma_b
    gamma_w_end = gamma_w_start + n_gamma_w
    zeta_w_start = gamma_w_end + n_gamma_b_dot
    zeta_w_end = zeta_w_start + n_zeta_w

    def apply_precond(vec: Array) -> Array:
        # split the system to remove gamma_w and zeta_w. The removed blocks use the negative identity as preconditioner.
        vec_pre_gw = vec[:gamma_w_start]  # struct + gamma_b
        vec_gw = vec[gamma_w_start:gamma_w_end]  # gamma_w (identity)
        vec_mid = vec[gamma_w_end:zeta_w_start]  # gamma_b_dot
        vec_zw = vec[zeta_w_start:zeta_w_end]  # zeta_w (identity)
        vec_post_zw = vec[zeta_w_end:]  # f_ext_aero
        vec_reduced = jnp.concatenate([vec_pre_gw, vec_mid, vec_post_zw])
        x_reduced = jax.scipy.linalg.lu_solve(precond_lu, vec_reduced)

        x_pre_gw = x_reduced[:gamma_w_start]
        x_mid = x_reduced[gamma_w_start : gamma_w_start + n_gamma_b_dot]
        x_post_zw = x_reduced[gamma_w_start + n_gamma_b_dot :]

        return jnp.concatenate([x_pre_gw, -vec_gw, x_mid, -vec_zw, x_post_zw])

    return apply_precond
dynamic_adjoint 
dynamic_adjoint(
    case: AeroelasticCase,
    objective: AeroelasticObjectiveFunction,
    matrix_free: bool = True,
    jacobian_approximations: AeroelasticJacobianApproximations
    | None = None,
    grads_to_compute: AeroelasticGradsToCompute
    | None = DEFAULT_GRADS_TO_COMPUTE,
    p_varphi_p_x: Array | None = None,
    save_adjoint: bool = False,
    approx_grads: bool = True,
    i_ts_adjoint_range: tuple[int | None, int | None] = (
        None,
        None,
    ),
    include_initial_state_grad: bool = True,
    gmres_mode: Literal[
        "batched", "incremental"
    ] = "incremental",
    gmres_warm_start: bool = True,
    gmres_precond: bool = True,
    gmres_restart: int = 50,
    i_ts_preconditioner: int = 0,
    preconditioner_batch_size: int | None = 16,
    preconditioner: Callable[[Array], Array] | None = None,
) -> tuple[AeroelasticDesignVariables, Array, Array | None]

Compute the adjoint of a coupled dynamic aeroelastic system.

Parameters:

Name Type Description Default
case AeroelasticCase

Dynamic aeroelastic case

 required
objective AeroelasticObjectiveFunction

Objective function that takes the system full states, design variables and timestep index, and returns an array.

 required
matrix_free bool

If true, do not explicitly compute the residual Jacobians and instead use the VJP and GMRES to solve.

 True
jacobian_approximations AeroelasticJacobianApproximations | None

Data structure which specifies Jacobian approximations to use for each part of the problem.

 None
grads_to_compute AeroelasticGradsToCompute | None

Specify which design variables for which to compute gradients for. If None, all available gradients are computed.

 DEFAULT_GRADS_TO_COMPUTE
p_varphi_p_x Array | None

Gradient of initial twists with respect to design variables. In practice, this is found from the static solve.

 None
save_adjoint bool

Whether to save the adjoint of the dynamic aeroelastic system.

 False
approx_grads bool

Whether to use gradient approximation or not. This removes some negligible contributions in the structural dynamic system.

 True
i_ts_adjoint_range tuple[int | None, int | None]

Optional (start, end) window of time steps for which to compute the adjoint. Either entry may be None to leave that side untruncated. When start > 1 the initial-state gradient contribution is automatically skipped.

 (None, None)
include_initial_state_grad bool

If False, skip the _initial_timestep_grad_contribution call that solves the static adjoint at t = 0 and propagates p_varphi_p_x. Intended for profiling only.

 True
gmres_mode Literal['batched', 'incremental']

If using matrix free, sets the mode for GMRES. Batched is preferred for GPU, whereas incremental may be preferred on CPU.

 'incremental'
gmres_warm_start bool

If True, use the previous timestep adjoint vector as the first guess for the current value. Otherwise, initialise with the zero vector.

 True
gmres_precond bool

If True and preconditioner is None, build the frozen-wake preconditioner internally. Ignored when preconditioner is supplied.

 True
gmres_restart int

Number of times to restart the GMRES algorithm.

 50
i_ts_preconditioner int

Timestep at which to build the preconditioner if requested.

 0
preconditioner_batch_size int | None

Batch size for creating the Jacobians for the preconditioner. Ignored if preconditioner is supplied.

 16
preconditioner Callable[[Array], Array] | None

Optional pass a prebuild preconditioner. Useful for profiling.

 None

Returns:

Type Description
tuple[AeroelasticDesignVariables, Array, Array | None]

Gradient of sum of objective across timesteps with respect to design variables, objective at each time step, and optional adjoint states.
Source code in src/flapjax/coupled/gradients/coupled.py 
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
@jax.jit(static_argnums=(0, 1, 2, 3, 4, 5, 7, 8, 9, 10, 11, 12, 13, 14, 15, 17))
def dynamic_adjoint(
    self,
    case: AeroelasticCase,
    objective: AeroelasticObjectiveFunction,
    matrix_free: bool = True,
    jacobian_approximations: AeroelasticJacobianApproximations | None = None,
    grads_to_compute: AeroelasticGradsToCompute | None = DEFAULT_GRADS_TO_COMPUTE,
    p_varphi_p_x: Array | None = None,
    save_adjoint: bool = False,
    approx_grads: bool = True,
    i_ts_adjoint_range: tuple[int | None, int | None] = (None, None),
    include_initial_state_grad: bool = True,
    gmres_mode: Literal["batched", "incremental"] = "incremental",
    gmres_warm_start: bool = True,
    gmres_precond: bool = True,
    gmres_restart: int = 50,
    i_ts_preconditioner: int = 0,
    preconditioner_batch_size: int | None = 16,
    preconditioner: Callable[[Array], Array] | None = None,
) -> tuple[AeroelasticDesignVariables, Array, Array | None]:
    r"""
    Compute the adjoint of a coupled dynamic aeroelastic system.
    :param case: Dynamic aeroelastic case
    :param objective: Objective function that takes the system full states, design variables and timestep index,
    and returns an array.
    :param matrix_free: If true, do not explicitly compute the residual Jacobians and instead use the VJP and GMRES
    to solve.
    :param jacobian_approximations: Data structure which specifies Jacobian approximations to use for each part of
    the problem.
    :param grads_to_compute: Specify which design variables for which to compute gradients for. If None, all
    available gradients are computed.
    :param p_varphi_p_x: Gradient of initial twists with respect to design variables. In practice, this is found
    from the static solve.
    :param save_adjoint: Whether to save the adjoint of the dynamic aeroelastic system.
    :param approx_grads: Whether to use gradient approximation or not. This removes some negligible contributions
    in the structural dynamic system.
    :param i_ts_adjoint_range: Optional ``(start, end)`` window of time steps for which to compute the adjoint.
     Either entry may be ``None`` to leave that side untruncated. When ``start > 1`` the initial-state gradient
     contribution is automatically skipped.
    :param include_initial_state_grad: If False, skip the ``_initial_timestep_grad_contribution`` call that solves
    the static adjoint at ``t = 0`` and propagates ``p_varphi_p_x``. Intended for profiling only.
    :param gmres_mode: If using matrix free, sets the mode for GMRES. Batched is preferred for GPU, whereas
    incremental may be preferred on CPU.
    :param gmres_warm_start: If True, use the previous timestep adjoint vector as the first guess for the current
    value. Otherwise, initialise with the zero vector.
    :param gmres_precond: If True and ``preconditioner`` is None, build the frozen-wake preconditioner internally.
    Ignored when ``preconditioner`` is supplied.
    :param gmres_restart: Number of times to restart the GMRES algorithm.
    :param i_ts_preconditioner: Timestep at which to build the preconditioner if requested.
    :param preconditioner_batch_size: Batch size for creating the Jacobians for the preconditioner. Ignored if
    ``preconditioner`` is supplied.
    :param preconditioner: Optional pass a prebuild preconditioner. Useful for profiling.
    :return: Gradient of sum of objective across timesteps with respect to design variables, objective at each time
    step, and optional adjoint states.
    """

    # make copies to prevent contaminating input object with tracer
    case = deepcopy(case)
    p_varphi_p_x = deepcopy(p_varphi_p_x)

    solve_dofs: tuple[int, ...] = get_solve_dofs(
        n_dof=self.structure.n_dof,
        prescribed_dofs=case.structure.prescribed_dofs,
    )

    solve_dofs_arr: Array = jnp.array(solve_dofs)

    n_tstep = case.structure.n_tstep

    dv = self.get_design_variables(case=case, grads_to_compute=grads_to_compute)

    dv_full = self.get_design_variables(case=case, grads_to_compute=None)

    assert case.aero.static_horseshoe is not None
    static_horseshoe: bool = case.aero.static_horseshoe

    full_states_init = case.get_full_states(i_ts=0)
    minimal_states_init = case.get_minimal_states(i_ts=0)

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

    n_solve = len(solve_dofs)

    j_eval = jax.vmap(
        lambda i_ts: jnp.atleast_1d(
            objective(case.get_full_states(i_ts=i_ts), dv, i_ts)
        )
    )(jnp.arange(n_tstep)).reshape(n_tstep, n_j)

    jac_options = self.construct_approximate_jacobians(
        sol=case,
        jacobian_approximations=jacobian_approximations
        if jacobian_approximations is not None
        else AeroelasticJacobianApproximations(),
    )

    # define 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 = 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_ > n_tstep - 1:
        raise ValueError(
            f"i_ts_adjoint_range end must be <= n_tstep - 1 = {n_tstep - 1}, got "
            f"{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

    @jax.jit
    def objective_jacobians(
        i_ts: int, q_n: AeroelasticMinimalStates
    ) -> tuple[Array, AeroelasticDesignVariables]:
        # function to obtain the Jacobians of the objective w.r.t. the minimal states and the design variables
        p_j_n_p_q_n, p_j_n_p_x = jax.jacrev(
            lambda q_free, dv__: jnp.atleast_1d(
                objective(
                    self.minimal_states_to_full_states(
                        i_ts=i_ts,
                        q=AeroelasticMinimalStates.from_vector(
                            vect=q_n.ravel().at[free_state_ix].set(q_free),
                            n_dof=self.structure.n_dof,
                            aero_shapes=minimal_states_init.aero.shapes(),
                        ),
                        dv=dv__,
                        dv_full=dv_full,
                    ),
                    dv__,
                    i_ts,
                )
            ),
            argnums=(0, 1),
            allow_int=True,
        )(q_n.ravel()[free_state_ix], dv)
        return p_j_n_p_q_n, p_j_n_p_x

    assert dv_full.aero.cs_ang_t is not None and dv_full.aero.cs_vel_t is not None

    # create the initial d_j_d_x sensitivites which are accumulated though the solve process
    dv_grad_init = AeroelasticDesignVariables.zeros(
        system=self, case=case, grads_to_compute=grads_to_compute, j_shape=j_shape
    )

    n_dof: int = self.structure.n_dof
    n_aero_states: int = minimal_states_init.aero.n_states
    free_state_ix: Array = jnp.concatenate(
        [solve_dofs_arr + i * n_dof for i in range(4)]
        + [jnp.arange(5 * n_dof, 5 * n_dof + n_aero_states)]
        + [solve_dofs_arr + 4 * n_dof]
    )
    n_adj_dof = 4 * n_solve + n_aero_states + n_solve

    adj_full_init: Array | None = (
        jnp.zeros((case.structure.n_tstep + 1, n_j, n_adj_dof))
        if save_adjoint
        else None
    )

    d_j_d_x: AeroelasticDesignVariables
    if matrix_free:
        if preconditioner is not None:
            apply_precond = preconditioner
        elif gmres_precond:
            apply_precond = self.make_frozen_wake_preconditioner(
                case=case,
                dv_full=dv_full,
                solve_dofs=solve_dofs,
                approx_grads=approx_grads,
                precond_i_ts=i_ts_preconditioner,
                batch_size=preconditioner_batch_size,
            )
            jax_print(
                "Built frozen-wake preconditioner",
                verbose_level="normal",
            )
        else:
            apply_precond = None

        def matrix_free_body(
            rev_i_ts_: int,
            carry: tuple[AeroelasticDesignVariables, Array, Array, Array],
        ) -> tuple[AeroelasticDesignVariables, Array, Array, Array]:
            d_j_d_x_, adj_np1, adj_t_p_r_np1_p_q_n, adj_full_ = carry

            i_ts = i_ts_end_adj_ - rev_i_ts_
            i_ts_nm1 = jnp.maximum(i_ts - 1, 0)
            q_nm1 = case.get_minimal_states(i_ts=i_ts_nm1)
            q_n = case.get_minimal_states(i_ts=i_ts)
            t_n = case.structure.t[i_ts]

            p_j_n_p_q_n, p_j_n_p_x = objective_jacobians(i_ts=i_ts, q_n=q_n)

            def _residual_all(
                q_n_: AeroelasticMinimalStates,
                q_nm1_: AeroelasticMinimalStates,
                dv_: AeroelasticDesignVariables,
            ) -> Array:
                return self.timestep_residual(
                    i_ts=i_ts,
                    t=t_n,
                    q_nm1=q_nm1_,
                    q_n=q_n_,
                    dv_=dv_,
                    dv_full=dv_full,
                    thrust_t=case.structure.thrust,
                    solve_dofs=solve_dofs,
                    approx_grads=approx_grads,
                )

            # single VJP shared between the GMRES matvec, the coupling term and the design-variable pull
            _, pull_all = jax.vjp(_residual_all, q_n, q_nm1, dv)

            def matvec_qn_t(v: Array) -> Array:
                if map_verbosity_level(get_verbosity()) >= map_verbosity_level(
                    "normal"
                ):
                    # print a dot for every GMRES iteration. Due to the jax GMRES function not returning the number
                    # of iterations, this at least allows us to count the dots!
                    def _print_gmres_dot() -> None:
                        sys.stdout.write(".")
                        sys.stdout.flush()

                    jax.debug.callback(_print_gmres_dot, ordered=True)

                d_q: AeroelasticMinimalStates = pull_all(v)[0]
                return d_q.to_free_dofs(solve_dofs_arr=solve_dofs_arr)

            b_rhs = -(p_j_n_p_q_n.reshape(n_j, -1) + adj_t_p_r_np1_p_q_n)

            def _solve_row(b_row: Array, x0_row: Array) -> tuple[Array, Array]:
                # noinspection PyTypeChecker
                x, info = jax.scipy.sparse.linalg.gmres(
                    matvec_qn_t,
                    b_row,
                    x0=x0_row if gmres_warm_start else None,
                    tol=1e-6,
                    atol=1e-6,
                    restart=gmres_restart,
                    maxiter=50,
                    M=apply_precond,
                    solve_method=gmres_mode,
                )
                return x, info

            adj_n, gmres_info = jax.vmap(_solve_row)(b_rhs, adj_np1)

            def _pull_row(
                a: Array,
            ) -> tuple[Array, AeroelasticDesignVariables]:
                q_nm1_cot: AeroelasticMinimalStates
                _, q_nm1_cot, dv_cot = pull_all(a)

                return q_nm1_cot.to_free_dofs(solve_dofs_arr=solve_dofs_arr), dv_cot

            adj_t_p_r_n_p_q_nm1, dv_grads = jax.vmap(_pull_row)(adj_n)

            d_j_d_x_ += dv_grads
            d_j_d_x_ += p_j_n_p_x

            jax_print(
                "\nSolved adjoint for timestep {i_ts} (GMRES converged={converged}, max|adj|={ma:.2e})",
                i_ts=i_ts,
                converged=jnp.max(gmres_info) == 0,
                ma=jnp.max(jnp.abs(adj_n)),
                verbose_level="normal",
            )

            if save_adjoint:
                adj_full_ = adj_full_.at[i_ts].set(adj_n)

            return d_j_d_x_, adj_n, adj_t_p_r_n_p_q_nm1, adj_full_

        d_j_d_x, _, future_row, adj_full = jax.lax.fori_loop(
            lower=0,
            upper=n_adj_iters,
            body_fun=matrix_free_body,
            init_val=(
                dv_grad_init,
                jnp.zeros((n_j, n_adj_dof)),
                jnp.zeros((n_j, n_adj_dof)),
                adj_full_init,
            ),
        )
    else:

        def step_body(
            rev_i_ts_: int,
            carry: tuple[AeroelasticDesignVariables, Array, Array, Array],
        ) -> tuple[AeroelasticDesignVariables, Array, Array, Array]:
            d_j_d_x_, adj_np1, p_r_np1_p_q_n, adj_full_ = carry

            i_ts = i_ts_end_adj_ - rev_i_ts_
            i_ts_nm1 = jnp.maximum(i_ts - 1, 0)
            q_nm1 = case.get_minimal_states(i_ts=i_ts_nm1)
            q_n = case.get_minimal_states(i_ts=i_ts)
            (
                p_res_p_q_nm1,
                p_res_p_q_n,
                p_v_dot_res_p_struct_dv_,
                p_aero_res_d_dv_,
                *_,
            ) = self.timestep_residual_jacobians(
                i_ts=i_ts,
                t=case.structure.t[i_ts],
                q_nm1=q_nm1,
                q_n=q_n,
                dv_=dv,
                dv_full=dv_full,
                thrust_t=case.structure.thrust,
                solve_dofs=solve_dofs,
                approx_grads=approx_grads,
                n_profile_loops=None,
                jac_options=jac_options,
            )
            p_j_n_p_q_n_, p_j_n_p_x_ = objective_jacobians(i_ts=i_ts, q_n=q_n)

            # solve adjoint step
            b = -(p_j_n_p_q_n_.reshape(n_j, -1) + adj_np1 @ p_r_np1_p_q_n).T
            adj_n = jnp.linalg.solve(p_res_p_q_n.T, b).T

            jax_print(
                "Solved adjoint for timestep {i_ts}",
                i_ts=i_ts,
                verbose_level="normal",
            )

            # add sentitivity of aerodynamic problem through full aero residual
            d_j_d_x_ += p_aero_res_d_dv_.premultiply_adj(adj_n[:, 4 * n_solve :])

            # add sensitivity of structural problem through v_dot residual
            d_j_d_x_.structure += p_v_dot_res_p_struct_dv_.premultiply_adj(
                adj_n[:, 2 * n_solve : 3 * n_solve]
            )
            d_j_d_x_ += p_j_n_p_x_

            if save_adjoint:
                adj_full_ = adj_full_.at[i_ts].set(adj_n)

            return d_j_d_x_, adj_n, p_res_p_q_nm1, adj_full_

        d_j_d_x, adj_last, p_r1_p_q0, adj_full = jax.lax.fori_loop(
            lower=0,
            upper=n_adj_iters,
            body_fun=step_body,
            init_val=(
                dv_grad_init,
                jnp.zeros((n_j, n_adj_dof)),
                jnp.zeros((n_adj_dof, n_adj_dof)),
                adj_full_init,
            ),
        )
        future_row = adj_last @ p_r1_p_q0

    # solve initial timestep adjoint, as there is no r0. Skipped when the adjoint window truncates early time steps
    if include_initial_state_grad and i_ts_start_adj_ <= 1:
        future_cot_q0_full = jnp.zeros((n_j, minimal_states_init.n_states))
        if case.structure.n_tstep > 1:
            future_cot_q0_full = future_cot_q0_full.at[:, free_state_ix].set(
                future_row
            )

        d_j_d_x += self._initial_timestep_grad_contribution(
            case=case[0].to_static(),
            objective=objective,
            grads_to_compute=grads_to_compute,
            p_varphi_p_x=p_varphi_p_x,
            solve_dofs=solve_dofs_arr,
            adj_t_p_r1_p_q0=future_cot_q0_full,
            horseshoe=static_horseshoe,
        )

    # restore original shape of j, and cut off zeros for past-end timestep and initial timestep which are always 0
    adj = (
        adj_full.reshape(adj_full.shape[0], *j_shape, *adj_full.shape[2:])[1:-1]
        if save_adjoint
        else None
    )

    d_j_d_x.mapping = dv.mapping

    return d_j_d_x, j_eval, adj
dynamic_adjoint_profile 
dynamic_adjoint_profile(
    case: AeroelasticCase,
    approx_grads: bool,
    jacobian_approximations: AeroelasticJacobianApproximations
    | None = None,
    grads_to_compute: AeroelasticGradsToCompute
    | None = None,
    i_ts: int = 1,
    n_profile_loops: int = 10,
) -> tuple[
    dict[str, dict[str, float]], dict[str, dict[str, float]]
]

Function to time evaluation of the Jacobians used for the coupled aeroelastic adjoint solution for the case where the full Jacobian is computed.

Parameters:

Name Type Description Default
case AeroelasticCase

Dynamic aeroelastic case from which to extract states.

 required
approx_grads bool

If True, neglect small gradient terms.

 required
jacobian_approximations AeroelasticJacobianApproximations | None

Define which blocks of the adjoint Jacobians will be substituted for approximations.

 None
grads_to_compute AeroelasticGradsToCompute | None

AeroelasticGradsToCompute object describing which design gradients to compute. If None, all gradients will be computed.

 None
i_ts int

Time step index where to evaluate residual Jacobians.

 1
n_profile_loops int

Number of times to loop the Jacobian evaluation time for averaging the runtime.

 10

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/coupled/gradients/coupled.py 
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
def dynamic_adjoint_profile(
    self,
    case: AeroelasticCase,
    approx_grads: bool,
    jacobian_approximations: AeroelasticJacobianApproximations | None = None,
    grads_to_compute: AeroelasticGradsToCompute | None = None,
    i_ts: int = 1,
    n_profile_loops: int = 10,
) -> tuple[dict[str, dict[str, float]], dict[str, dict[str, float]]]:
    r"""
    Function to time evaluation of the Jacobians used for the coupled aeroelastic adjoint solution for the case
    where the full Jacobian is computed.
    :param case: Dynamic aeroelastic case from which to extract states.
    :param approx_grads: If True, neglect small gradient terms.
    :param jacobian_approximations: Define which blocks of the adjoint Jacobians will be substituted for
    approximations.
    :param grads_to_compute: AeroelasticGradsToCompute object describing which design gradients to compute. If
    None, all gradients will be computed.
    :param i_ts: Time step index where to evaluate residual Jacobians.
    :param n_profile_loops: Number of times to loop the Jacobian evaluation time for averaging the runtime.
    :return: Dictionary of {residual_name: {gradient_argument: val}} for compile time and run time respectively.
    """

    print_table_title(inner_width=95, title="Aeroelastic Adjoint Profile")

    jac_options = self.construct_approximate_jacobians(
        sol=case,
        jacobian_approximations=jacobian_approximations
        if jacobian_approximations is not None
        else AeroelasticJacobianApproximations(),
    )

    *_, compile_time, run_time = self.timestep_residual_jacobians(
        i_ts=i_ts,
        t=case.aero.t[i_ts],
        q_nm1=case.get_minimal_states(i_ts=i_ts - 1),
        q_n=case.get_minimal_states(i_ts=i_ts),
        dv_=self.get_design_variables(case=case, grads_to_compute=grads_to_compute),
        dv_full=self.get_design_variables(case=case, grads_to_compute=None),
        thrust_t=case.structure.thrust,
        solve_dofs=get_solve_dofs(
            n_dof=self.structure.n_dof,
            prescribed_dofs=case.structure.prescribed_dofs,
        ),
        approx_grads=approx_grads,
        n_profile_loops=n_profile_loops,
        jac_options=jac_options,
    )

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

    print_table_line(inner_width=95)

    return compile_time, run_time
trim 
trim(
    prescribed_dofs: Sequence[int] | Array | slice | int,
    zero_force_dofs: Sequence[int] | Array | slice | int,
    trim_cs: Sequence[str | Sequence[str]] | str | None,
    thrust_nodes: Sequence[str | Sequence[str]]
    | str
    | None,
    trim_orientation: str | Sequence[str] | None = "x",
    trim_hinges: Sequence[str | Sequence[str]]
    | str
    | None = None,
    trim_f_abs_tolerance: float = 0.01,
    f_ext_follower: Array | None = None,
    f_ext_dead: Array | None = None,
    t: float | Array = 0.0,
    load_steps: int = 1,
    trim_relaxation: float = 0.9,
    horseshoe: bool = False,
    method: Literal[
        "adjoint", "finite_difference"
    ] = "finite_difference",
    broyden_fd_step: float = 0.001,
    max_iter: int = 100,
) -> tuple[AeroelasticCase, TrimVariables]

Trim an aircraft such that the resulting sum of forces on the aircraft is zero without any supports.

Parameters:

Name Type Description Default
prescribed_dofs Sequence[int] | Array | slice | int

Degrees of freedom which are clamped for the trim process.

 required
zero_force_dofs Sequence[int] | Array | slice | int

Degrees freedom where we wish to drive the clamping force to zero. This is not necessarily the same as prescribed_dofs, as in some cases there are degrees of freedom we will allow to have a non-zero clamping force. For example in the case of a clamped cantilever wing where we wish to find the angle of attack that gives lift equal to the weight, there would be a nonzero pitching moment.

 required
trim_cs Sequence[str | Sequence[str]] | str | None

Keys of control surfaces which are to be used to trim the aircraft. Each element may either be a single control-surface key (that surface gets its own independent deflection) or a sequence of keys (those surfaces are tied together and share a single deflection).

 required
thrust_nodes Sequence[str | Sequence[str]] | str | None

Keys of thrust nodes which are to be used to trim the aircraft. Each element may either be a single node key (that node gets its own independent thrust value) or a sequence of node keys (those nodes are tied together and share a single thrust value).

 required
trim_orientation str | Sequence[str] | None

Inertial axis (or axes if a sequence is provided) around which the aircraft is rotated about at the clamp to achieve trim.

 'x'
trim_hinges Sequence[str | Sequence[str]] | str | None

Names of MultibodyHinge constraints (keys into the structure's named constraints) whose rotation should be solved as a trim variable.

 None
trim_f_abs_tolerance float

Absolute maximum force residual at the clamped nodes for convergence to be achieved.

 0.01
f_ext_follower Array | None

External follower forces, [n_nodes, 6].

 None
f_ext_dead Array | None

external dead forces, [n_nodes, 6].

 None
t float | Array

Time at which to trim the aircraft, default zero.

 0.0
load_steps int

Number of load steps used for the static solution.

 1
trim_relaxation float

Relaxation factor for updates to degrees of freedom used to achieve trim.

 0.9
horseshoe bool

If true, use a horseshoe wake formulation.

 False
method Literal['adjoint', 'finite_difference']

"adjoint" rebuilds the trim Jacobian each iteration via the adjoint method. "finite_difference" approximates the Jacobian with a forward finite-difference sweep. Usually the latter is faster assuming a small number of trim variables.

 'finite_difference'
broyden_fd_step float

Step size used for the finite-difference bootstrap of the Broyden Jacobian.

 0.001
max_iter int

Maximum number of trim iterations. A warning is emitted if the routine fails to converge within this limit.

 100

Returns:

Type Description
tuple[AeroelasticCase, TrimVariables]

Aeroelastic solution object for the trimmed aircraft.
Source code in src/flapjax/coupled/gradients/coupled.py 
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
def trim(
    self,
    prescribed_dofs: Sequence[int] | Array | slice | int,
    zero_force_dofs: Sequence[int] | Array | slice | int,
    trim_cs: Sequence[str | Sequence[str]] | str | None,
    thrust_nodes: Sequence[str | Sequence[str]] | str | None,
    trim_orientation: str | Sequence[str] | None = "x",
    trim_hinges: Sequence[str | Sequence[str]] | str | None = None,
    trim_f_abs_tolerance: float = 1e-2,
    f_ext_follower: Array | None = None,
    f_ext_dead: Array | None = None,
    t: float | Array = 0.0,
    load_steps: int = 1,
    trim_relaxation: float = 0.9,
    horseshoe: bool = False,
    method: Literal["adjoint", "finite_difference"] = "finite_difference",
    broyden_fd_step: float = 1e-3,
    max_iter: int = 100,
) -> tuple[AeroelasticCase, TrimVariables]:
    r"""
    Trim an aircraft such that the resulting sum of forces on the aircraft is zero without any supports.
    :param prescribed_dofs: Degrees of freedom which are clamped for the trim process.
    :param zero_force_dofs: Degrees freedom where we wish to drive the clamping force to zero. This is not
    necessarily the same as prescribed_dofs, as in some cases there are degrees of freedom we will allow to have a
    non-zero clamping force. For example in the case of a clamped cantilever wing where we wish to find the angle
    of attack that gives lift equal to the weight, there would be a nonzero pitching moment.
    :param trim_cs: Keys of control surfaces which are to be used to trim the aircraft. Each element may either be
    a single control-surface key (that surface gets its own independent deflection) or a sequence of keys (those
    surfaces are tied together and share a single deflection).
    :param thrust_nodes: Keys of thrust nodes which are to be used to trim the aircraft. Each element may either be
    a single node key (that node gets its own independent thrust value) or a sequence of node keys (those nodes are
    tied together and share a single thrust value).
    :param trim_orientation: Inertial axis (or axes if a sequence is provided) around which the aircraft is
    rotated about at the clamp to achieve trim.
    :param trim_hinges: Names of ``MultibodyHinge`` constraints (keys into the structure's named constraints)
    whose rotation should be solved as a trim variable.
    :param trim_f_abs_tolerance: Absolute maximum force residual at the clamped nodes for convergence to be achieved.
    :param f_ext_follower: External follower forces, [n_nodes, 6].
    :param f_ext_dead: external dead forces, [n_nodes, 6].
    :param t: Time at which to trim the aircraft, default zero.
    :param load_steps: Number of load steps used for the static solution.
    :param trim_relaxation: Relaxation factor for updates to degrees of freedom used to achieve trim.
    :param horseshoe: If true, use a horseshoe wake formulation.
    :param method: "adjoint" rebuilds the trim Jacobian each iteration via the adjoint method. "finite_difference"
    approximates the Jacobian with a forward finite-difference sweep. Usually the latter is faster assuming a small
    number of trim variables.
    :param broyden_fd_step: Step size used for the finite-difference bootstrap of the Broyden Jacobian.
    :param max_iter: Maximum number of trim iterations. A warning is emitted if the routine fails to converge within
    this limit.
    :return: Aeroelastic solution object for the trimmed aircraft.
    """

    # parse groups for paired thrust nodes/control surfaces/hinges
    cs_groups: list[tuple[str, ...]] = parse_groups(trim_cs, "trim_cs")
    thrust_groups: list[tuple[str, ...]] = parse_groups(
        thrust_nodes, "thrust_nodes"
    )
    hinge_groups: list[tuple[str, ...]] = parse_groups(trim_hinges, "trim_hinges")

    check_unique_members(cs_groups, "trim_cs")
    check_unique_members(thrust_groups, "thrust_nodes")
    check_unique_members(hinge_groups, "trim_hinges")

    if hinge_groups and method == "adjoint":
        raise ValueError(
            "trim_hinges is only supported with method='finite_difference'."
        )

    trim_orientation_: Sequence[str] = (
        [trim_orientation]
        if isinstance(trim_orientation, str)
        else trim_orientation
        if trim_orientation is not None
        else []
    )

    zero_force_dofs_: tuple[int, ...] = self.structure.make_prescribed_dofs_tuple(
        zero_force_dofs
    )

    prescribed_dofs_: tuple[int, ...] = self.structure.make_prescribed_dofs_tuple(
        prescribed_dofs
    )

    if not self.structure.use_gravity:
        warn("Gravity is not enabled. Trim may result in unexpected behaviour.")

    # initial set of variables. Tied members share a single value, initialised from the first member.
    trim_variables_init: TrimVariables = TrimVariables(
        cs_ang={group_key(g): self.aero.cs_ang0[g[0]] for g in cs_groups},
        thrust={
            group_key(g): self.structure.thrust_reference[g[0]]
            for g in thrust_groups
        },
        trim_angles={
            k: self.structure.orientation_euler[ORIENTATION_DICT[k]]
            for k in trim_orientation_
        },
        hinge_angle={group_key(g): jnp.array(0.0) for g in hinge_groups},
    )

    ae_sol_init = self.reference_configuration(
        horseshoe=horseshoe,
        prescribed_dofs=prescribed_dofs_,
        use_f_ext_dead=f_ext_dead is not None,
        use_f_ext_follower=f_ext_follower is not None,
    )

    inner_case = deepcopy(self)

    if method == "adjoint":

        def trim_body(
            i_iter: int,
            trim_variables_: TrimVariables,
            sol_: AeroelasticCase,
            f_clamp_: Array,
        ) -> tuple[int, TrimVariables, AeroelasticCase, Array]:
            i_iter, _, tv, sol, fc = self.trim_iter(
                i_iter,
                inner_case,
                trim_variables_,
                sol_,
                f_clamp_,
                prescribed_dofs=prescribed_dofs_,
                zero_force_dofs=zero_force_dofs_,
                f_ext_dead=f_ext_dead,
                f_ext_follower=f_ext_follower,
                t=jnp.array(t),
                load_steps=load_steps,
                horseshoe=horseshoe,
                cs_groups=cs_groups,
                thrust_groups=thrust_groups,
                trim_orientation=trim_orientation_,
                trim_relaxation=trim_relaxation,
            )
            return i_iter, tv, sol, fc

        f_clamp_init = jnp.full((len(zero_force_dofs_)), 1e10)
        print_table_title(title="Trim (Adjoint)", inner_width=104)
        trim_variables_init.print_header(f_clamp=f_clamp_init)
        _, trim_variables, ae_sol, f_clamp_final = jax.lax.while_loop(
            lambda args_: jnp.logical_and(
                jnp.any(jnp.abs(args_[3]) >= trim_f_abs_tolerance),
                args_[0] < max_iter,
            ),
            body_fun=lambda args_: trim_body(*args_),
            init_val=(
                0,
                trim_variables_init,
                ae_sol_init,
                f_clamp_init,
            ),
        )
        print_table_line(inner_width=104)
    elif method == "finite_difference":
        print_table_title(title="Trim (Finite Difference)", inner_width=104)

        b_approx_init, f_clamp_init, ae_sol_bootstrap = self._trim_fd_jacobian(
            inner_case=inner_case,
            trim_variables=trim_variables_init,
            fd_step=broyden_fd_step,
            prescribed_dofs=prescribed_dofs_,
            zero_force_dofs=zero_force_dofs_,
            f_ext_follower=f_ext_follower,
            f_ext_dead=f_ext_dead,
            t=jnp.array(t),
            load_steps=load_steps,
            horseshoe=horseshoe,
            cs_groups=cs_groups,
            thrust_groups=thrust_groups,
            hinge_groups=hinge_groups,
            trim_orientation=trim_orientation_,
        )

        def trim_body_broyden(
            i_iter: int,
            trim_variables_: TrimVariables,
            sol_: AeroelasticCase,
            f_clamp_: Array,
            b_approx_: Array,
        ) -> tuple[int, TrimVariables, AeroelasticCase, Array, Array]:
            i_iter, _, tv, sol, fc, b = self._trim_iter_fd(
                i_iter,
                inner_case,
                trim_variables_,
                sol_,
                f_clamp_,
                b_approx_,
                prescribed_dofs=prescribed_dofs_,
                zero_force_dofs=zero_force_dofs_,
                f_ext_dead=f_ext_dead,
                f_ext_follower=f_ext_follower,
                t=jnp.array(t),
                load_steps=load_steps,
                horseshoe=horseshoe,
                cs_groups=cs_groups,
                thrust_groups=thrust_groups,
                hinge_groups=hinge_groups,
                trim_orientation=trim_orientation_,
                trim_relaxation=trim_relaxation,
            )
            return i_iter, tv, sol, fc, b

        trim_variables_init.print_header(f_clamp=f_clamp_init)
        _, trim_variables, ae_sol, f_clamp_final, _ = jax.lax.while_loop(
            lambda args_: jnp.logical_and(
                jnp.any(jnp.abs(args_[3]) >= trim_f_abs_tolerance),
                args_[0] < max_iter,
            ),
            body_fun=lambda args_: trim_body_broyden(*args_),
            init_val=(
                0,
                trim_variables_init,
                ae_sol_bootstrap,
                f_clamp_init,
                b_approx_init,
            ),
        )
        print_table_line(inner_width=104)
    else:
        raise ValueError(f"Unknown trim method: {method!r}.")

    if bool(jnp.any(jnp.isnan(f_clamp_final))):
        warn("Trim residual is NaN - solution diverged")
    elif bool(jnp.any(jnp.abs(f_clamp_final) >= trim_f_abs_tolerance)):
        warn(f"Trim did not converge within max_iter={max_iter} iterations ")

    new_orientation: Array = self.structure.orientation_euler
    for k, v in trim_variables.trim_angles.items():
        new_orientation = new_orientation.at[ORIENTATION_DICT[k]].set(v)

    # set solutions into case object
    self.set_design_variables(
        coords=self.structure.x0_reference,
        k_cs=self.structure.k_cs,
        m_cs=self.structure.m_cs,
        m_lumped=self.structure.m_lumped
        if self.structure.use_lumped_mass
        else None,
        dt=self.aero.dt,
        flowfield=self.aero.flowfield,
        delta_w=self.aero.delta_w,
        x0_aero=self.aero.zeta_b0,
        thrust_reference=self.structure.thrust_reference
        | expand_groups(trim_variables.thrust, thrust_groups),
        orientation_euler=new_orientation,
        cs_angles_reference=self.aero.cs_ang0
        | expand_groups(trim_variables.cs_ang, cs_groups),
        remove_checks=True,
    )

    if hinge_groups:
        # release the angle-pinning constraint used to condition the trim solve
        self._revert_hinge_trim(hinge_groups)

    return ae_sol, trim_variables
trim_angles_to_euler 
trim_angles_to_euler(
    trim_angles: dict[str, Array],
) -> Array

Find the 3 Euler angles describing the aircraft orientation. This allows for any combination to be set by the trim routine, with values not passed using the fixed values provided in the reference orientation.

Parameters:

Name Type Description Default
trim_angles dict[str, Array]

Dictionary of axis-angle pairs.

 required

Returns:

Type Description
Array

Euler angles describing the aircraft orientation, (3, ).
Source code in src/flapjax/coupled/gradients/coupled.py 
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
def trim_angles_to_euler(self, trim_angles: dict[str, Array]) -> Array:
    r"""
    Find the 3 Euler angles describing the aircraft orientation. This allows for any combination to be set by
    the trim routine, with values not passed using the fixed values provided in the reference orientation.
    :param trim_angles: Dictionary of axis-angle pairs.
    :return: Euler angles describing the aircraft orientation, ``(3, )``.
    """

    orientation_euler = self.structure.orientation_euler
    for k, v in trim_angles.items():
        orientation_euler = orientation_euler.at[ORIENTATION_DICT[k]].set(v)
    return orientation_euler
get_design_variables 
get_design_variables(
    case: AeroelasticCase,
    grads_to_compute: AeroelasticGradsToCompute | None,
) -> AeroelasticDesignVariables

Obtain the design variables describing the wing.

Parameters:

Name Type Description Default
case AeroelasticCase
required
grads_to_compute AeroelasticGradsToCompute | None

Data structure which describes which design variables should be obtained. If none, all variables are obtained.

 required

Returns:

Type Description
AeroelasticDesignVariables

AeroelasticDesignVariables object.
Source code in src/flapjax/coupled/coupled.py 
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
def get_design_variables(
    self,
    case: AeroelasticCase,
    grads_to_compute: AeroelasticGradsToCompute | None,
) -> AeroelasticDesignVariables:
    r"""
    Obtain the design variables describing the wing.
    :param case:
    :param grads_to_compute: Data structure which describes which design variables should be obtained. If none, all
    variables are obtained.
    :return: AeroelasticDesignVariables object.
    """
    return AeroelasticDesignVariables(
        structure_dv=self.structure.get_design_variables(
            struct_case=case.structure,
            thrust_t=case.structure.thrust,
            grads_to_compute=grads_to_compute.structure
            if grads_to_compute is not None
            else None,
        ),
        aero_dv=self.aero.get_design_variables(
            cs_ang_t=case.aero.cs_ang,
            cs_vel_t=case.aero.cs_vel,
            grads_to_compute=grads_to_compute.aero
            if grads_to_compute is not None
            else None,
        ),
    )
reference_configuration 
reference_configuration(
    prescribed_dofs: Sequence[int]
    | Array
    | slice
    | int = (),
    horseshoe: bool = False,
    use_f_ext_follower: bool = False,
    use_f_ext_dead: bool = False,
    t_init: float | Array = 0.0,
) -> AeroelasticCase

Obtain the static aeroelastic object describing the undeformed wing.

Parameters:

Name Type Description Default
prescribed_dofs Sequence[int] | Array | slice | int

Prescribed dofs for the structure. Defaults to no prescribed dofs.

 ()
horseshoe bool

Horseshoe flag.

 False
use_f_ext_follower bool

If true, allocate an array for follower forces.

 False
use_f_ext_dead bool

If true, allocate an array for dead forces.

 False
t_init float | Array

Initial time

 0.0

Returns:

Type Description
AeroelasticCase

Static aeroelastic object for undeformed wing
Source code in src/flapjax/coupled/coupled.py 
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
def reference_configuration(
    self,
    prescribed_dofs: Sequence[int] | Array | slice | int = (),
    horseshoe: bool = False,
    use_f_ext_follower: bool = False,
    use_f_ext_dead: bool = False,
    t_init: float | Array = 0.0,
) -> AeroelasticCase:
    r"""
    Obtain the static aeroelastic object describing the undeformed wing.
    :param prescribed_dofs: Prescribed dofs for the structure. Defaults to no prescribed dofs.
    :param horseshoe: Horseshoe flag.
    :param use_f_ext_follower: If true, allocate an array for follower forces.
    :param use_f_ext_dead: If true, allocate an array for dead forces.

    :param t_init: Initial time
    :return: Static aeroelastic object for undeformed wing
    """
    prescribed_dofs = self.structure.make_prescribed_dofs_tuple(prescribed_dofs)
    return AeroelasticCase(
        structure=self.structure.reference_configuration(
            use_f_grav=self.structure.use_gravity,
            use_f_ext_dead=use_f_ext_dead,
            use_f_ext_follower=use_f_ext_follower,
            use_f_aero=True,
            prescribed_dofs=prescribed_dofs,
        ),
        aero=self.aero.static_solve(
            t=t_init, hg=self.structure.hg0, horseshoe=horseshoe
        ),
    )
initialise_dynamic 
initialise_dynamic(
    static_case: AeroelasticCase,
    prescribed_dofs: Sequence[int] | Array | slice | int,
) -> AeroelasticCase

Initialise a dynamic aeroelastic snapshot from a static aeroelastic case. This takes a static aeroelastic case obtained under clamped conditions (i.e. relative_motion is True for the free stream), and sets the structural velocity to be that of the freestream.

Parameters:

Name Type Description Default
static_case AeroelasticCase

Static aeroelastic case.

 required
prescribed_dofs Sequence[int] | Array | slice | int

Prescribed dofs. This is often useful for updating from a clamped trim to a free-flying dynamic case.

 required

Returns:

Type Description
AeroelasticCase

Dynamic aeroelastic snapshot.
Source code in src/flapjax/coupled/coupled.py 
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
def initialise_dynamic(
    self,
    static_case: AeroelasticCase,
    prescribed_dofs: Sequence[int] | Array | slice | int,
) -> AeroelasticCase:
    r"""
    Initialise a dynamic aeroelastic snapshot from a static aeroelastic case. This takes a static aeroelastic case
    obtained under clamped conditions (i.e. `relative_motion` is True for the free stream), and sets the structural
    velocity to be that of the freestream.
    :param static_case: Static aeroelastic case.
    :param prescribed_dofs: Prescribed dofs. This is often useful for updating from a clamped trim to a free-flying
    dynamic case.
    :return: Dynamic aeroelastic snapshot.
    """
    u_inf = self.aero.flowfield.u_inf  # flowfield velocity to set
    self.aero.flowfield.relative_motion = (
        False  # the output will have relative motion disabled
    )

    rmat_struct = static_case.structure.hg[:, :3, :3]  # deformed rotations
    v_local = jnp.einsum(
        "ijk,j->ik", rmat_struct, -u_inf
    )  # local frame velocity, (n_nodes, 3)

    dynamic_case = static_case.to_dynamic(t=None)
    dynamic_case.structure.v = dynamic_case.structure.v.at[:, :3].set(v_local)
    dynamic_case.structure.prescribed_dofs = (
        self.structure.make_prescribed_dofs_tuple(prescribed_dofs)
    )
    dynamic_case.structure.free_dofs = get_solve_dofs(
        n_dof=self.structure.n_dof,
        prescribed_dofs=dynamic_case.structure.prescribed_dofs,
    )

    return dynamic_case

data_structures

TrimVariables 
TrimVariables(
    cs_ang: dict[str, Array],
    thrust: dict[str, Array],
    trim_angles: dict[str, Array],
    hinge_angle: dict[str, Array] | None = None,
)
Source code in src/flapjax/coupled/gradients/data_structures.py 
41
42
43
44
45
46
47
48
49
50
51
def __init__(
    self,
    cs_ang: dict[str, Array],
    thrust: dict[str, Array],
    trim_angles: dict[str, Array],
    hinge_angle: dict[str, Array] | None = None,
):
    self.cs_ang: dict[str, Array] = cs_ang
    self.thrust: dict[str, Array] = thrust
    self.trim_angles: dict[str, Array] = trim_angles
    self.hinge_angle: dict[str, Array] = {} if hinge_angle is None else hinge_angle
print_header 
print_header(f_clamp: Array | None) -> None

Print the column-header row. Call once before the iteration loop.

Source code in src/flapjax/coupled/gradients/data_structures.py 
71
72
73
74
75
76
77
78
79
80
81
82
83
84
def print_header(self, f_clamp: Array | None) -> None:
    """Print the column-header row. Call once before the iteration loop."""
    specs = self._column_specs(f_clamp)
    col_widths = [max(len(self._header_label(k, u)), vw) for k, u, vw, _ in specs]

    cells = ["iter".rjust(self._ITER_W)]
    cells += [
        self._header_label(k, u).rjust(cw)
        for (k, u, _, _), cw in zip(specs, col_widths)
    ]
    inner = " | ".join(cells)
    padding = self._INNER_WIDTH - len(inner) - 2
    jax_print("| " + inner + " " * padding + " |", verbose_level="normal")
    print_table_line(inner_width=self._INNER_WIDTH)
print_values 
print_values(i_iter: int, f_clamp: Array | None) -> None

Print one row of numeric values, aligned with the header.

Source code in src/flapjax/coupled/gradients/data_structures.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
def print_values(self, i_iter: int, f_clamp: Array | None) -> None:
    """Print one row of numeric values, aligned with the header."""
    specs = self._column_specs(f_clamp)
    col_widths = [max(len(self._header_label(k, u)), vw) for k, u, vw, _ in specs]

    def _scalar(x: Array) -> Array:
        return jnp.ravel(x)[0]

    values: list[Array] = []
    values.extend(jnp.rad2deg(_scalar(v)) for v in self.cs_ang.values())
    values.extend(_scalar(v) for v in self.thrust.values())
    values.extend(jnp.rad2deg(_scalar(v)) for v in self.trim_angles.values())
    values.extend(jnp.rad2deg(_scalar(v)) for v in self.hinge_angle.values())
    if f_clamp is not None:
        values.extend(_scalar(f_clamp[i]) for i in range(f_clamp.shape[0]))

    placeholders = [f"c{i}" for i in range(len(specs))]
    cells = [f"{{i_iter:>{self._ITER_W}}}"]
    cells += [
        f"{{{ph}:>{cw}{fmt}}}"
        for ph, (_, _, _, fmt), cw in zip(placeholders, specs, col_widths)
    ]
    inner = " | ".join(cells)
    rendered_len = self._ITER_W + sum(col_widths) + 3 * len(specs)
    padding = self._INNER_WIDTH - rendered_len - 2

    kwargs = {"i_iter": i_iter, **dict(zip(placeholders, values))}
    jax_print("| " + inner + " " * padding + " |", **kwargs, verbose_level="normal")

utils

group_key 
group_key(group: Sequence[str]) -> str

create key for a group of names (thrust nodes or control surfaces).

Source code in src/flapjax/coupled/gradients/utils.py 
11
12
13
def group_key(group: Sequence[str]) -> str:
    """create key for a group of names (thrust nodes or control surfaces)."""
    return "+".join(group)
expand_groups 
expand_groups(
    group_values: dict[str, Array],
    groups: Sequence[Sequence[str]],
) -> dict[str, Array]

Expand a group-keyed dict into a per-member dict so tied members share one value.

Source code in src/flapjax/coupled/gradients/utils.py 
16
17
18
19
20
21
22
23
24
25
26
def expand_groups(
    group_values: dict[str, Array],
    groups: Sequence[Sequence[str]],
) -> dict[str, Array]:
    """Expand a group-keyed dict into a per-member dict so tied members share one value."""
    per_member: dict[str, Array] = {}
    for group in groups:
        v = group_values[group_key(group)]
        for member in group:
            per_member[member] = v
    return per_member
parse_groups 
parse_groups(
    value: Sequence[str | Sequence[str]] | str | None,
    arg_name: str,
) -> list[tuple[str, ...]]

Convert a user-specified group of control surfaces or thrust nodes into a list of tuples, where each tuple contains the names of the members in that group.

Parameters:

Name Type Description Default
value Sequence[str | Sequence[str]] | str | None

Input value for groupings.

 required
arg_name str

Name of the argument being parsed.

 required

Returns:

Type Description
list[tuple[str, ...]]

List of tuples containing the names of the members in each group.
Source code in src/flapjax/coupled/gradients/utils.py 
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
def parse_groups(
    value: Sequence[str | Sequence[str]] | str | None,
    arg_name: str,
) -> list[tuple[str, ...]]:
    r"""
    Convert a user-specified group of control surfaces or thrust nodes into a list of tuples, where each tuple contains
    the names of the members in that group.
    :param value: Input value for groupings.
    :param arg_name: Name of the argument being parsed.
    :return: List of tuples containing the names of the members in each group.
    """
    if value is None:
        return []
    if isinstance(value, str):
        return [(value,)]
    if isinstance(value, Sequence):
        groups: list[tuple[str, ...]] = []
        for item in value:
            if isinstance(item, str):
                groups.append((item,))
            elif isinstance(item, Sequence):
                if not item:
                    raise ValueError(
                        f"Each {arg_name} group must contain at least one key."
                    )
                groups.append(tuple(item))
            else:
                raise ValueError(
                    f"{arg_name} entries must be a string or sequence of strings. Got {type(item)}."
                )
        return groups
    raise ValueError(
        f"{arg_name} must be a string, a sequence, or None. Got {type(value)}."
    )
check_unique_members 
check_unique_members(
    groups: Sequence[Sequence[str]], arg_name: str
) -> None

Ensure that each entry only appears in one group. Raises a ValueError if any entry appears in more than one group.

Parameters:

Name Type Description Default
groups Sequence[Sequence[str]]

List of groups to check.

 required
arg_name str

Name of the argument being parsed.

 required
Source code in src/flapjax/coupled/gradients/utils.py 
65
66
67
68
69
70
71
72
73
74
75
76
77
78
def check_unique_members(groups: Sequence[Sequence[str]], arg_name: str) -> None:
    r"""
    Ensure that each entry only appears in one group. Raises a ValueError if any entry appears in more than one group.
    :param groups: List of groups to check.
    :param arg_name: Name of the argument being parsed.
    """
    seen: set[str] = set()
    for group in groups:
        for member in group:
            if member in seen:
                raise ValueError(
                    f"{arg_name} entry {member!r} appears in more than one trim group."
                )
            seen.add(member)

linear

linear_coupled

LinearCoupled 
LinearCoupled(
    case: BaseCoupledAeroelastic,
    reference: AeroelasticCase,
    batch_size: int | None,
    n_struct_modes: int | None,
    wake_type: LinearWakeType = "frozen",
    bound_upwash: bool = True,
    wake_upwash: bool = False,
    unsteady_force: bool = True,
    int_order: Literal[3, 4, 5] = BASE_LOBATTO_ORDER,
    *,
    skip_checks: bool = False,
    prescribed_dofs: Sequence[int]
    | Array
    | slice
    | int
    | None = None,
)

Bases: LinearModel[AeroelasticCase, AeroelasticInputUnflattened, AeroelasticStateUnflattened, AeroelasticOutputUnflattened, AeroelasticLinearResult]

Source code in src/flapjax/coupled/linear/linear_coupled.py 
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
def __init__(
    self,
    case: BaseCoupledAeroelastic,
    reference: AeroelasticCase,
    batch_size: int | None,
    n_struct_modes: int | None,
    wake_type: LinearWakeType = "frozen",
    bound_upwash: bool = True,
    wake_upwash: bool = False,
    unsteady_force: bool = True,
    int_order: Literal[3, 4, 5] = BASE_LOBATTO_ORDER,
    *,
    skip_checks: bool = False,
    prescribed_dofs: Sequence[int] | Array | slice | int | None = None,
):
    if prescribed_dofs is not None:
        prescribed_dofs = case.structure.make_prescribed_dofs_tuple(prescribed_dofs)

    self.aero = LinearUVLM(
        case=case.aero,
        reference=reference.aero,
        wake_type=wake_type,
        bound_upwash=bound_upwash,
        wake_upwash=wake_upwash,
        unsteady_force=unsteady_force,
        skip_linearisation=True,
        skip_checks=skip_checks,
    )
    self.structure = LinearBeam(
        beam=case.structure,
        reference=reference.structure,
        dt=case.aero.dt,
        n_modes=n_struct_modes,
        int_order=int_order,
        prescribed_dofs=prescribed_dofs,
    )

    effective_prescribed = (
        prescribed_dofs
        if prescribed_dofs is not None
        else reference.structure.prescribed_dofs
    )
    self.n_beam_nodal_dof: int = case.structure.n_dof - len(effective_prescribed)
    self.n_beam_input_dof: int = (
        self.structure.n_modes
        if self.structure.modal_inputs
        else self.n_beam_nodal_dof
    )
    self.n_beam_state_dof: int = (
        self.structure.n_modes
        if self.structure.modal_states
        else self.n_beam_nodal_dof
    )
    self.n_beam_output_dof: int = (
        self.structure.n_modes
        if self.structure.modal_outputs
        else self.n_beam_nodal_dof
    )

    self.n_nodes: int = case.structure.n_nodes
    self.free_dofs: Array = jnp.array(
        get_solve_dofs(
            n_dof=case.structure.n_dof,
            prescribed_dofs=effective_prescribed,
        )
    )

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

    self._case: BaseCoupledAeroelastic = case
    self.unsteady_force: bool = unsteady_force
    if batch_size is not False:
        self.sys = self.linearise(batch_size=batch_size)
step 
step(
    gamma_b_vec: Array | None = None,
    gamma_w_vec: Array | None = None,
    gamma_b_nm1_vec: Array | None = None,
    zeta_w_vec: Array | None = None,
    nu_b_vec: Array | None = None,
    nu_w_vec: Array | None = None,
    f_ext: Array | None = None,
    q_nodal: Array | None = None,
    q_dot_nodal: Array | None = None,
) -> tuple[
    AeroelasticStateUnflattened,
    AeroelasticOutputUnflattened,
]

Step solution from states at timestep n and inputs at timestep n+1 to give states at timestep n+1 and outputs at timestep n.

Source code in src/flapjax/coupled/linear/linear_coupled.py 
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
def step(
    self,
    gamma_b_vec: Array | None = None,
    gamma_w_vec: Array | None = None,
    gamma_b_nm1_vec: Array | None = None,
    zeta_w_vec: Array | None = None,
    nu_b_vec: Array | None = None,
    nu_w_vec: Array | None = None,
    f_ext: Array | None = None,
    q_nodal: Array | None = None,
    q_dot_nodal: Array | None = None,
) -> tuple[AeroelasticStateUnflattened, AeroelasticOutputUnflattened]:
    r"""
    Step solution from states at timestep n and inputs at timestep n+1 to give states at timestep n+1 and outputs
    at timestep n.
    """
    ref = self.reference

    # unravel vector inputs, falling back to reference values when None
    gamma_b = (
        ArrayList.from_vector(vect=gamma_b_vec, shapes=ref.aero.gamma_b.shape)
        if gamma_b_vec is not None
        else ref.aero.gamma_b
    )
    gamma_w = (
        ArrayList.from_vector(vect=gamma_w_vec, shapes=ref.aero.gamma_w.shape)
        if gamma_w_vec is not None
        else ref.aero.gamma_w
    )
    gamma_b_nm1 = (
        (
            ArrayList.from_vector(
                vect=gamma_b_nm1_vec, shapes=ref.aero.gamma_b.shape
            )
            if gamma_b_nm1_vec is not None
            else ref.aero.gamma_b
        )
        if self.aero.unsteady_force
        else None
    )
    zeta_w = (
        ArrayList.from_vector(vect=zeta_w_vec, shapes=ref.aero.zeta_w.shape)
        if zeta_w_vec is not None
        else ref.aero.zeta_w
    )
    nu_b = (
        (
            ArrayList.from_vector(vect=nu_b_vec, shapes=ref.aero.zeta_b.shape)
            if nu_b_vec is not None
            else ArrayList.zeros_like(ref.aero.zeta_b)
        )
        if self.aero.bound_upwash
        else None
    )
    nu_w = (
        (
            ArrayList.from_vector(vect=nu_w_vec, shapes=ref.aero.zeta_w.shape)
            if nu_w_vec is not None
            else ArrayList.zeros_like(ref.aero.zeta_w)
        )
        if self.aero.wake_upwash
        else None
    )
    q_nodal = (
        q_nodal if q_nodal is not None else jnp.zeros(len(self.structure.free_dofs))
    )
    q_dot_nodal = (
        q_dot_nodal
        if q_dot_nodal is not None
        else jnp.zeros(len(self.structure.free_dofs))
    )

    # fill in prescribed dofs with zeros
    q_full = (
        jnp.zeros(self.n_nodes * 6)
        .at[self.free_dofs]
        .set(q_nodal)
        .reshape(self.n_nodes, 6)
    )
    q_dot_full = (
        jnp.zeros(self.n_nodes * 6)
        .at[self.free_dofs]
        .set(q_dot_nodal)
        .reshape(self.n_nodes, 6)
    )

    # total perturbed coordinates and time derivative
    hg = jnp.einsum("ijk,ikl->ijl", ref.structure.hg, vmap(exp_se3)(q_full))
    hg_dot = jnp.einsum(
        "ijk,ikl->ijl",
        ref.structure.hg,
        vmap(ha_to_ha_tilde)(q_dot_full),
    )

    # aerodynamic grid
    zeta_b = self.aero.case.hg_to_zeta_b(hg_n=hg, cs_ang_n=self.aero.case.cs_ang0)
    zeta_b_dot = self.aero.case.hg_dot_to_zeta_b_dot(
        hg_n=hg,
        hg_dot_n=hg_dot,
        cs_ang_n=self.aero.case.cs_ang0,
        cs_vel_n=self.aero.case.cs_vel0,
    )

    # pass through aerodynamic system
    u_n_aero = AeroInputUnflattened(
        zeta_b=zeta_b, zeta_b_dot=zeta_b_dot, nu_b=nu_b, nu_w=nu_w
    )

    x_n_aero = AeroStateUnflattened(
        gamma_b=gamma_b,
        gamma_w=gamma_w,
        gamma_b_nm1=gamma_b_nm1,
        zeta_w=zeta_w if self.aero.prescribed_wake else None,
        zeta_b=zeta_b if self.aero.prescribed_wake else None,
    )

    u_n_aero_vec = self.aero.pack_input_vector(u_n_aero)
    x_n_aero_vec = self.aero.pack_state_vector(x_n_aero)

    x_np1_aero_vec, y_np1_aero_vec = self.aero.step_vec(
        x_vec=x_n_aero_vec, u_vec=u_n_aero_vec
    )

    x_np1_aero = self.aero.unpack_state_vector(x=x_np1_aero_vec)
    y_np1_aero = self.aero.unpack_output_vector(y=y_np1_aero_vec)
    assert isinstance(x_np1_aero, AeroStateUnflattened) and isinstance(
        y_np1_aero, AeroOutputUnflattened
    ), (
        "Unpacked aero state and output must be of type AeroStateUnflattened and AeroOutputUnflattened."
    )

    # total aero forces on the grid, from the aero step (returns totals)
    f_aero_np1 = y_np1_aero.f_steady
    if self.unsteady_force:
        assert y_np1_aero.f_unsteady is not None
        # don't want to mutate in place
        # noinspection augment-assignment
        f_aero_np1 = f_aero_np1 + y_np1_aero.f_unsteady

    # project total aero forces onto the beam under the current (perturbed) rotation
    rmat = hg[:, :3, :3]
    f_aero_beam_total = project_forcing_to_beam(
        f_total=f_aero_np1,
        rmat=rmat,
        dof_mapping=self.aero.case.dof_mapping,
        x0_aero=self.aero.case.zeta_b0,
        mirror_edge_low=self.aero.case.mirror_edge_low,
        mirror_edge_high=self.aero.case.mirror_edge_high,
    )

    # subtract the reference contribution so the aero forcing fed into the beam operator
    # is a pure perturbation (the beam sys.a / sys.b operate on perturbations)
    f_aero_ref_total = ref.aero.f_steady
    if self.aero.unsteady_force:
        # noinspection augment-assignment
        f_aero_ref_total = f_aero_ref_total + ref.aero.f_unsteady
    f_aero_beam_ref = project_forcing_to_beam(
        f_total=f_aero_ref_total,
        rmat=ref.structure.hg[:, :3, :3],
        dof_mapping=self.aero.case.dof_mapping,
        x0_aero=self.aero.case.zeta_b0,
        mirror_edge_low=self.aero.case.mirror_edge_low,
        mirror_edge_high=self.aero.case.mirror_edge_high,
    )
    delta_f_aero_beam = f_aero_beam_total - f_aero_beam_ref

    f_ext_: Array = f_ext if f_ext is not None else jnp.zeros(self.free_dofs.size)

    # scatter free-dof forcing to full (n_nodes * 6) vectors expected by the beam B operator, and add aero
    # forcing (global frame, perturbation) as an external force
    f_ext_full = (
        jnp.zeros(self.n_nodes * 6).at[self.free_dofs].set(f_ext_)
        + delta_f_aero_beam.ravel()
    )  # [n_nodes * 6]

    if self.structure.modal_inputs:
        f_ext_full = self.structure.nodal_to_modal(f_ext_full[self.free_dofs])

    # beam step in perturbation form (discrete-time Tustin)
    x_beam_n = jnp.concatenate(
        [
            self.structure.nodal_to_modal(q_nodal)
            if self.structure.modal_states
            else q_nodal,
            self.structure.nodal_to_modal(q_dot_nodal)
            if self.structure.modal_states
            else q_dot_nodal,
        ]
    )

    x_beam_np1 = (
        self.structure.sys.a @ x_beam_n
        + self.structure.sys.b[:, self.structure.input_slices["f_ext"].slices]
        @ f_ext_full
    )

    q_np1 = x_beam_np1[: self.n_beam_state_dof]
    q_dot_np1 = x_beam_np1[self.n_beam_state_dof :]

    state_np1 = AeroelasticStateUnflattened(
        gamma_b=x_np1_aero.gamma_b,
        gamma_w=x_np1_aero.gamma_w,
        gamma_b_nm1=x_np1_aero.gamma_b_nm1,
        zeta_w=x_np1_aero.zeta_w,
        q=q_np1,
        q_dot=q_dot_np1,
    )
    output_n = AeroelasticOutputUnflattened(q=q_np1, q_dot=q_dot_np1)

    return state_np1, output_n
gamma_b_step 
gamma_b_step(
    gamma_b_n_vec: Array,
    gamma_w_n_vec: Array,
    q_n: Array,
    q_dot_n: Array,
    zeta_w_n_vec: Array | None = None,
    nu_b_n_vec: Array | None = None,
) -> Array

Bound circulation at timestep n+1 as a function of states at n and inputs at n+1.

Source code in src/flapjax/coupled/linear/linear_coupled.py 
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
def gamma_b_step(
    self,
    gamma_b_n_vec: Array,
    gamma_w_n_vec: Array,
    q_n: Array,
    q_dot_n: Array,
    zeta_w_n_vec: Array | None = None,
    nu_b_n_vec: Array | None = None,
) -> Array:
    r"""
    Bound circulation at timestep n+1 as a function of states at n and inputs at n+1.
    """
    x_np1, _ = self.step(
        nu_b_vec=nu_b_n_vec,
        gamma_b_vec=gamma_b_n_vec,
        gamma_w_vec=gamma_w_n_vec,
        zeta_w_vec=zeta_w_n_vec,
        q_nodal=self.structure.modal_to_nodal(q_n)
        if self.structure.modal_states
        else q_n,
        q_dot_nodal=self.structure.modal_to_nodal(q_dot_n)
        if self.structure.modal_states
        else q_dot_n,
    )
    return x_np1.gamma_b.ravel()
wake_prop_step 
wake_prop_step(
    gamma_b_n_vec: Array,
    gamma_w_n_vec: Array,
    q_n: Array,
    zeta_w_n_vec: Array | None = None,
    nu_w_n_vec: Array | None = None,
) -> tuple[Array | None, Array]

Wake propagation as a function of states at n and inputs at n+1.

Source code in src/flapjax/coupled/linear/linear_coupled.py 
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
def wake_prop_step(
    self,
    gamma_b_n_vec: Array,
    gamma_w_n_vec: Array,
    q_n: Array,
    zeta_w_n_vec: Array | None = None,
    nu_w_n_vec: Array | None = None,
) -> tuple[Array | None, Array]:
    r"""
    Wake propagation as a function of states at n and inputs at n+1.
    """
    x_new, _ = self.step(
        nu_w_vec=nu_w_n_vec,
        gamma_b_vec=gamma_b_n_vec,
        gamma_w_vec=gamma_w_n_vec,
        zeta_w_vec=zeta_w_n_vec,
        q_nodal=self.structure.modal_to_nodal(q_n)
        if self.structure.modal_states
        else q_n,
    )
    assert not ((x_new.zeta_w is not None) ^ self.aero.prescribed_wake), (
        "zeta_w should be None only if prescribed_wake is False."
    )
    return (
        x_new.zeta_w.ravel() if x_new.zeta_w is not None else None,
        x_new.gamma_w.ravel(),
    )
q_step 
q_step(
    q_n: Array,
    q_dot_n: Array,
    gamma_b_n_vec: Array,
    gamma_w_n_vec: Array,
    f_ext: Array | None,
    zeta_w_n_vec: Array | None = None,
    gamma_b_nm1_vec: Array | None = None,
    nu_b_n_vec: Array | None = None,
    nu_w_n_vec: Array | None = None,
) -> Array

Beam displacement at timestep n+1 as a function of states at n and external forcing.

Source code in src/flapjax/coupled/linear/linear_coupled.py 
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
def q_step(
    self,
    q_n: Array,
    q_dot_n: Array,
    gamma_b_n_vec: Array,
    gamma_w_n_vec: Array,
    f_ext: Array | None,
    zeta_w_n_vec: Array | None = None,
    gamma_b_nm1_vec: Array | None = None,
    nu_b_n_vec: Array | None = None,
    nu_w_n_vec: Array | None = None,
) -> Array:
    r"""
    Beam displacement at timestep n+1 as a function of states at n and external forcing.
    """
    x_new, _ = self.step(
        f_ext=f_ext,
        gamma_b_vec=gamma_b_n_vec,
        gamma_w_vec=gamma_w_n_vec,
        gamma_b_nm1_vec=gamma_b_nm1_vec,
        zeta_w_vec=zeta_w_n_vec,
        nu_b_vec=nu_b_n_vec,
        nu_w_vec=nu_w_n_vec,
        q_nodal=self.structure.modal_to_nodal(q_n)
        if self.structure.modal_states
        else q_n,
        q_dot_nodal=self.structure.modal_to_nodal(q_dot_n)
        if self.structure.modal_states
        else q_dot_n,
    )
    return x_new.q.ravel()
q_dot_step 
q_dot_step(
    q_n: Array,
    q_dot_n: Array,
    gamma_b_n_vec: Array,
    gamma_w_n_vec: Array,
    f_ext: Array | None,
    zeta_w_n_vec: Array | None,
    gamma_b_nm1_vec: Array | None = None,
    nu_b_n_vec: Array | None = None,
    nu_w_n_vec: Array | None = None,
) -> Array

Beam velocity at timestep n+1 as a function of states at n and external forcing.

Source code in src/flapjax/coupled/linear/linear_coupled.py 
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
def q_dot_step(
    self,
    q_n: Array,
    q_dot_n: Array,
    gamma_b_n_vec: Array,
    gamma_w_n_vec: Array,
    f_ext: Array | None,
    zeta_w_n_vec: Array | None,
    gamma_b_nm1_vec: Array | None = None,
    nu_b_n_vec: Array | None = None,
    nu_w_n_vec: Array | None = None,
) -> Array:
    r"""
    Beam velocity at timestep n+1 as a function of states at n and external forcing.
    """
    x_new, _ = self.step(
        f_ext=f_ext,
        gamma_b_vec=gamma_b_n_vec,
        gamma_w_vec=gamma_w_n_vec,
        gamma_b_nm1_vec=gamma_b_nm1_vec,
        zeta_w_vec=zeta_w_n_vec,
        nu_b_vec=nu_b_n_vec,
        nu_w_vec=nu_w_n_vec,
        q_nodal=self.structure.modal_to_nodal(q_n)
        if self.structure.modal_states
        else q_n,
        q_dot_nodal=self.structure.modal_to_nodal(q_dot_n)
        if self.structure.modal_states
        else q_dot_n,
    )
    return x_new.q_dot.ravel()
compute_jacobians 
compute_jacobians() -> tuple[
    dict[
        str,
        tuple[
            Callable[..., Any],
            dict[str, Any],
            Sequence[str],
        ],
    ],
    dict[str, dict[str, Callable[..., Array]]],
]

Returns:

Type Description
tuple[dict[str, tuple[Callable[..., Any], dict[str, Any], Sequence[str]]], dict[str, dict[str, Callable[..., Array]]]]

Tuple. First entry is dictionaries with keys being the function name (e.g., gamma_b, gamma_w), with each entry containing the relevant stepping function, the arguments for the function, and the name of the arguments for which to obtain derivatives. Second entry is a dictionary of explicit Jacobians functions that take the same arguments as the first output.
Source code in src/flapjax/coupled/linear/linear_coupled.py 
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
def compute_jacobians(
    self,
) -> tuple[
    dict[str, tuple[Callable[..., Any], dict[str, Any], Sequence[str]]],
    dict[str, dict[str, Callable[..., Array]]],
]:
    r"""
    :return: Tuple. First entry is dictionaries with keys being the function name (e.g., gamma_b, gamma_w), with
    each entry containing the relevant stepping function, the arguments for the function, and the name of the
    arguments for which to obtain derivatives. Second entry is a dictionary of explicit Jacobians functions that
    take the same arguments as the first output.
    """
    ref = self.reference

    # bound circulation
    gamma_b_args: dict[str, Any] = {
        "gamma_b_n_vec": ref.aero.gamma_b.ravel(),
        "gamma_w_n_vec": ref.aero.gamma_w.ravel(),
        "q_n": jnp.zeros((self.n_beam_state_dof,)),
        "q_dot_n": jnp.zeros((self.n_beam_state_dof,)),
        "zeta_w_n_vec": ref.aero.zeta_w.ravel(),
        "nu_b_n_vec": jnp.zeros(ref.aero.zeta_b.size)
        if self.aero.bound_upwash
        else None,
    }
    gamma_b_diff = ["gamma_w_n_vec", "q_n", "q_dot_n"]
    if self.aero.prescribed_wake:
        gamma_b_diff.append("zeta_w_n_vec")
    if self.aero.bound_upwash:
        gamma_b_diff.append("nu_b_n_vec")

    # wake
    wake_args: dict[str, Any] = {
        "gamma_b_n_vec": ref.aero.gamma_b.ravel(),
        "gamma_w_n_vec": ref.aero.gamma_w.ravel(),
        "q_n": jnp.zeros((self.n_beam_state_dof,)),
        "zeta_w_n_vec": ref.aero.zeta_w.ravel(),
        "nu_w_n_vec": jnp.zeros(ref.aero.zeta_w.size)
        if self.aero.wake_upwash
        else None,
    }
    gamma_w_diff = ["gamma_b_n_vec", "gamma_w_n_vec"]
    zeta_w_diff = ["q_n"]
    if self.aero.prescribed_wake:
        zeta_w_diff.append("zeta_w_n_vec")
    if self.aero.wake_upwash:
        zeta_w_diff.append("nu_w_n_vec")
    if self.aero.free_wake:
        zeta_w_diff.extend(["gamma_b_n_vec", "gamma_w_n_vec"])

    q_args: dict[str, Any] = {
        "q_n": jnp.zeros((self.n_beam_state_dof,)),
        "q_dot_n": jnp.zeros((self.n_beam_state_dof,)),
        "gamma_b_n_vec": ref.aero.gamma_b.ravel(),
        "gamma_w_n_vec": ref.aero.gamma_w.ravel(),
        "f_ext": jnp.zeros((self.n_beam_input_dof,)),
        "zeta_w_n_vec": ref.aero.zeta_w.ravel(),
        "gamma_b_nm1_vec": ref.aero.gamma_b.ravel()
        if self.aero.unsteady_force
        else None,
        "nu_b_n_vec": jnp.zeros(ref.aero.zeta_b.size)
        if self.aero.bound_upwash
        else None,
        "nu_w_n_vec": jnp.zeros(ref.aero.zeta_w.size)
        if self.aero.wake_upwash
        else None,
    }
    q_diff = [
        "q_n",
        "q_dot_n",
        "gamma_b_n_vec",
        "gamma_w_n_vec",
        "f_ext",
    ]

    if self.aero.prescribed_wake:
        q_diff.append("zeta_w_n_vec")
    if self.aero.unsteady_force:
        q_diff.append("gamma_b_nm1_vec")
    if self.aero.bound_upwash:
        q_diff.append("nu_b_n_vec")
    if self.aero.wake_upwash:
        q_diff.append("nu_w_n_vec")

    linear_args: dict[
        str, tuple[Callable[..., Any], dict[str, Any], Sequence[str]]
    ] = {
        "gamma_b": (self.gamma_b_step, gamma_b_args, gamma_b_diff),
        "gamma_w": (
            lambda *args, **kwargs: self.wake_prop_step(**kwargs)[1],
            wake_args,
            gamma_w_diff,
        ),
        "gamma_b_nm1": (
            lambda *args, **kwargs: None,
            {
                "gamma_b_n_vec": ref.aero.gamma_b.ravel(),
            },
            ["gamma_b_n_vec"],
        ),
        "q": (self.q_step, q_args, q_diff),
        "q_dot": (self.q_dot_step, q_args, q_diff),
    }

    # add zeta_w for linearisation
    if self.aero.prescribed_wake:
        linear_args["zeta_w"] = (
            lambda *args, **kwargs: self.wake_prop_step(**kwargs)[0],
            wake_args,
            zeta_w_diff,
        )

    # define Jacobians we know nicely as jac_options
    jac_options = {
        "q": {
            "q_n": lambda *args, **kwargs: self.structure.sys.a[
                : self.n_beam_state_dof, : self.n_beam_state_dof
            ],
            "q_dot_n": lambda *args, **kwargs: self.structure.sys.a[
                : self.n_beam_state_dof, self.n_beam_state_dof :
            ],
        },
        "q_dot": {
            "q_n": lambda *args, **kwargs: self.structure.sys.a[
                self.n_beam_state_dof :, : self.n_beam_state_dof
            ],
            "q_dot_n": lambda *args, **kwargs: self.structure.sys.a[
                self.n_beam_state_dof :, self.n_beam_state_dof :
            ],
        },
    }
    if self.aero.unsteady_force:
        jac_options["gamma_b_nm1"] = {
            "gamma_b_n_vec": lambda *args, **kwargs: jnp.eye(ref.aero.gamma_b.size)
        }

    return linear_args, jac_options
create_jacobians 
create_jacobians(
    mode: ADMode | dict[str, ADMode] = "reverse",
    batch_size: int | None = None,
    n_profile_loops: int | None = None,
    jac_options: dict[
        str, dict[str, Callable[..., Any] | None]
    ]
    | None = None,
) -> tuple[
    dict[str, dict[str, Array]],
    dict[str, dict[str, float]] | None,
    dict[str, dict[str, float]] | None,
]

Assemble the per-residual Jacobians for the coupled linearisation. When n_profile_loops is set, all Jacobian constructions are looped locally so that they can be timed.

Source code in src/flapjax/coupled/linear/linear_coupled.py 
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
def create_jacobians(
    self,
    mode: ADMode | dict[str, ADMode] = "reverse",
    batch_size: int | None = None,
    n_profile_loops: int | None = None,
    jac_options: dict[str, dict[str, Callable[..., Any] | None]] | None = None,
) -> tuple[
    dict[str, dict[str, Array]],
    dict[str, dict[str, float]] | None,
    dict[str, dict[str, float]] | None,
]:
    """
    Assemble the per-residual Jacobians for the coupled linearisation. When
    ``n_profile_loops`` is set, all Jacobian constructions are looped locally so that they can be timed.
    """
    res_args, jac_options_exp = self.compute_jacobians()
    jac_options_total: dict[str, dict[str, Callable[..., Any] | None]] = (
        jac_options if jac_options is not None else {}
    ) | jac_options_exp

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

    aero_residual_names = {"gamma_b", "gamma_w", "gamma_b_nm1", "zeta_w"}
    delegate_aero = n_profile_loops is None

    for res_name, (res_func, args, diff_arg_names) in res_args.items():
        if delegate_aero and res_name in aero_residual_names:
            continue

        res_jac_options: dict[str, Callable[..., Any] | None] = {
            arg: None for arg in diff_arg_names
        }
        if res_name in jac_options_total:
            for arg, entry in jac_options_total[res_name].items():
                if arg in res_jac_options:
                    res_jac_options[arg] = entry

        if isinstance(mode, str):
            res_mode: ADMode = mode
        elif isinstance(mode, dict):
            try:
                res_mode = mode[res_name]
            except KeyError:
                res_mode = "reverse"
        else:
            raise NotImplementedError

        jacs, res_compile_time, res_run_time = jacrev_custom(
            func=res_func,
            jac_options=res_jac_options,
            n_profile_loops=n_profile_loops,
            func_name=res_name,
            map_batch_size=batch_size,
            mode=res_mode,
        )(**args)

        jacobians[res_name] = jacs
        if n_profile_loops is not None:
            assert res_compile_time is not None and res_run_time is not None
            compile_time[res_name] = res_compile_time
            run_time[res_name] = res_run_time

    if delegate_aero:
        # delegate aero linearisation to the aero system, with the beam kinematics wrapped as an input projection
        beam_proj = self._build_beam_projection()
        aero_mode: ADMode | dict[str, ADMode]
        if isinstance(mode, dict):
            aero_mode = {k: mode[k] for k in aero_residual_names if k in mode}
        else:
            aero_mode = mode
        needed_aero_residuals = {"gamma_b", "gamma_w"}
        if self.aero.unsteady_force:
            needed_aero_residuals.add("gamma_b_nm1")
        if self.aero.prescribed_wake:
            needed_aero_residuals.add("zeta_w")
        aero_jacs = self.aero.create_jacobians(
            mode=aero_mode,
            batch_size=batch_size,
            input_projection=beam_proj,
            residual_names=tuple(needed_aero_residuals),
        )
        jacobians.update(aero_jacs)

    return (
        jacobians,
        compile_time if n_profile_loops is not None else None,
        run_time if n_profile_loops is not None else None,
    )
linearise_profile 
linearise_profile(
    n_profile_loops: int = 3,
) -> tuple[
    dict[str, dict[str, float]], dict[str, dict[str, float]]
]

Profile forming the Jacobians required for the linearised model.

Parameters:

Name Type Description Default
n_profile_loops int

Number of times to loop Jacobian creation for averaging.

 3

Returns:

Type Description
tuple[dict[str, dict[str, float]], dict[str, dict[str, float]]]

Dictionaries of compile and run times for each sub function.
Source code in src/flapjax/coupled/linear/linear_coupled.py 
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
def linearise_profile(
    self,
    n_profile_loops: int = 3,
) -> tuple[dict[str, dict[str, float]], dict[str, dict[str, float]]]:
    r"""
    Profile forming the Jacobians required for the linearised model.
    :param n_profile_loops: Number of times to loop Jacobian creation for averaging.
    :return: Dictionaries of compile and run times for each sub function.
    """

    print_table_title(inner_width=95, title="Aeroelastic Adjoint Profile")

    _, compile_time, run_time = self.create_jacobians(
        n_profile_loops=n_profile_loops,
        jac_options=None,
        mode={"gamma_b": "forward"} if self.structure.modal_states else {},  # type: ignore
    )

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

    print_table_line(inner_width=95)

    return compile_time, run_time
modal 
modal(
    n_modes: int | None = None,
    freq_range: tuple[float | Array, float | Array] = (
        0.0,
        jnp.inf,
    ),
    damp_range: tuple[float | Array, float | Array] = (
        -jnp.inf,
        jnp.inf,
    ),
    min_struct_content: float | Array = 0.0,
    remove_complex_conjugate: bool = True,
    plot_eigvals: bool = False,
    sort: Literal["frequency", "damping"] = "frequency",
    plot_xlim: tuple[float, float] = (-500.0, 50.0),
    plot_ylim: tuple[float, float] = (-400.0, 400.0),
    n_plot_vtk: int = 0,
    vtu_directory: PathLike | str = "./modal",
    n_phase: int = 8,
    n_interp: int = 0,
    max_disp: float = 0.2,
    max_ang: float = 0.2,
    max_gamma: float = 100.0,
) -> Array

Compute stability eigenvalues of the linear system A matrix.

Parameters:

Name Type Description Default
n_modes int | None

Number of modes to be kept. If None, all eigenvalues are returned.

 None
freq_range tuple[float | Array, float | Array]

(min, max) natural frequency window in Hz. Modes outside are pushed past the truncation and dropped when n_modes is set.

 (0.0, inf)
damp_range tuple[float | Array, float | Array]

(min, max) damping-ratio window. Modes outside are pushed past the truncation and dropped when n_modes is set.

 (-inf, inf)
min_struct_content float | Array

Minimum fraction of eigenvector energy that must live in the beam (q, q_dot) states in range [0, 1]. Modes below this threshold (typically wake convection modes) are pushed past the truncation.

 0.0
remove_complex_conjugate bool

If true, one mode from each complex-conjugate pair is dropped.

 True
plot_eigvals bool

If true, plot the eigenvalues with Matplotlib.

 False
sort Literal['frequency', 'damping']

Method for sorting eigenvalues before truncation, can be either "frequency" or "damping".

 'frequency'
plot_xlim tuple[float, float]

Range of real component to be used for plotting.

 (-500.0, 50.0)
plot_ylim tuple[float, float]

Range of imaginary component to be used for plotting.

 (-400.0, 400.0)
n_plot_vtk int

Number of modes (starting from the most damped) to write to VTK for visualisation. Set to 0 to skip plotting.

 0
vtu_directory PathLike | str

Directory for plotting vtu files, defaults to "./modal".

 './modal'
n_phase int

Number of phase samples of the complex eigenvector to plot per mode.

 8
n_interp int

Number of interpolation points to add along each beam element in the beam VTU output.

 0
max_disp float

Maximum linear displacement used to normalise the plotted mode shape (in reference units).

 0.2
max_ang float

Maximum angular displacement used to normalise the plotted mode shape.

 0.2
max_gamma float

Maximum circulation used to normalise the plotted mode shape.

 100.0

Returns:

Type Description
Array

Continuous-time eigenvalues of the system A matrix, (n_states, ) or (n_states, 2) if to_components=True.
Source code in src/flapjax/coupled/linear/linear_coupled.py 
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
def modal(
    self,
    n_modes: int | None = None,
    freq_range: tuple[float | Array, float | Array] = (0.0, jnp.inf),
    damp_range: tuple[float | Array, float | Array] = (-jnp.inf, jnp.inf),
    min_struct_content: float | Array = 0.0,
    remove_complex_conjugate: bool = True,
    plot_eigvals: bool = False,
    sort: Literal["frequency", "damping"] = "frequency",
    plot_xlim: tuple[float, float] = (-500.0, 50.0),
    plot_ylim: tuple[float, float] = (-400.0, 400.0),
    n_plot_vtk: int = 0,
    vtu_directory: os.PathLike | str = "./modal",
    n_phase: int = 8,
    n_interp: int = 0,
    max_disp: float = 0.2,
    max_ang: float = 0.2,
    max_gamma: float = 100.0,
) -> Array:
    r"""
    Compute stability eigenvalues of the linear system A matrix.
    :param n_modes: Number of modes to be kept. If None, all eigenvalues are returned.
    :param freq_range: (min, max) natural frequency window in Hz. Modes outside are pushed past the
    truncation and dropped when n_modes is set.
    :param damp_range: (min, max) damping-ratio window. Modes outside are pushed past the truncation and
    dropped when n_modes is set.
    :param min_struct_content: Minimum fraction of eigenvector energy that must live in the beam
    ``(q, q_dot)`` states in range `[0, 1]`. Modes below this threshold (typically wake convection modes) are pushed
    past the truncation.
    :param remove_complex_conjugate: If true, one mode from each complex-conjugate pair is dropped.
    :param plot_eigvals: If true, plot the eigenvalues with Matplotlib.
    :param sort: Method for sorting eigenvalues before truncation, can be either "frequency" or "damping".
    :param plot_xlim: Range of real component to be used for plotting.
    :param plot_ylim: Range of imaginary component to be used for plotting.
    :param n_plot_vtk: Number of modes (starting from the most damped) to write to VTK for visualisation. Set to 0
    to skip plotting.
    :param vtu_directory: Directory for plotting vtu files, defaults to "./modal".
    :param n_phase: Number of phase samples of the complex eigenvector to plot per mode.
    :param n_interp: Number of interpolation points to add along each beam element in the beam VTU output.
    :param max_disp: Maximum linear displacement used to normalise the plotted mode shape (in reference units).
    :param max_ang: Maximum angular displacement used to normalise the plotted mode shape.
    :param max_gamma: Maximum circulation used to normalise the plotted mode shape.
    :return: Continuous-time eigenvalues of the system A matrix, ``(n_states, )`` or ``(n_states, 2)`` if ``to_components=True``.
    """

    evals_d, evecs = jnp.linalg.eig(self.sys.a)
    evals = jnp.log(evals_d) / self.dt  # convert to continuous time

    # order from most to least damped and truncate
    omega_damped = jnp.abs(evals.imag)
    damping = -evals.real / jnp.abs(evals)
    omega_natural = omega_damped / jnp.sqrt(1.0 - damping**2)

    freq_natural_hz = omega_natural / (2.0 * jnp.pi)

    match sort:
        case "frequency":
            idx = omega_natural.argsort()
        case "damping":
            idx = damping.argsort()

    # push conjugate partners past the truncation point (indexed by original position, then re-sorted so it
    # aligns with `idx`). Stable-argsort preserves the primary sort within each group.
    if remove_complex_conjugate:
        partner = conjugate_partner_mask(
            freq_hz=freq_natural_hz, damping=damping, tiebreaker=evals.real
        )
        idx = idx[jnp.argsort(partner[idx], stable=True)]

    # fraction of eigenvector energy in the structural states — used to reject
    # aero-only modes (e.g. wake convection)
    q_slice = self.state_slices["q"].slices
    q_dot_slice = self.state_slices["q_dot"].slices
    evec_sq = jnp.abs(evecs) ** 2
    struct_content = (
        evec_sq[q_slice].sum(axis=0) + evec_sq[q_dot_slice].sum(axis=0)
    ) / evec_sq.sum(axis=0)

    # push modes outside the requested natural-frequency / damping window to the back so truncation to
    # n_modes keeps only the in-range ones.
    in_range = (
        (freq_natural_hz[idx] >= freq_range[0])
        & (freq_natural_hz[idx] <= freq_range[1])
        & (damping[idx] >= damp_range[0])
        & (damping[idx] <= damp_range[1])
        & (struct_content[idx] >= min_struct_content)
    )
    idx = idx[jnp.argsort(~in_range, stable=True)]

    if n_modes is not None:
        idx = idx[:n_modes]

    freq_damped_ordered = omega_damped[idx] / (2.0 * jnp.pi)
    freq_natural_ordered = omega_natural[idx] / (2.0 * jnp.pi)
    damping_ordered = damping[idx]

    # write to console
    if n_modes is not None:
        print_table_line(inner_width=71)
        jax_print(
            "| Mode | Damped Frequency [Hz] | Natural Frequency [Hz] | Damping Ratio |",
            verbose_level="normal",
        )
        print_table_line(inner_width=71)
        for i_mode in range(n_modes):
            jax_print(
                "| {mode:>4d} | {freq_damped:>21.3f} | {freq_natural:>22.3f} | {damp:>13.6f} |",
                mode=i_mode + 1,
                freq_damped=freq_damped_ordered[i_mode],
                freq_natural=freq_natural_ordered[i_mode],
                damp=damping_ordered[i_mode],
                verbose_level="normal",
            )
        print_table_line(inner_width=71)

    if plot_eigvals:
        _, ax = plt.subplots()
        ax.scatter(
            evals.real,
            evals.imag,
        )
        ax.set_xlim(*plot_xlim)
        ax.set_ylim(*plot_ylim)
        ax.set_xlabel("Re(eig) [1/s]")
        ax.set_ylabel("Im(eig) [1/s]")
        ax.set_title("Eigenvalues")
        plt.show()

    if n_plot_vtk > 0:
        evecs_ordered = evecs[:, idx[:n_plot_vtk]].T  # (m, n_states)

        q_mode = evecs_ordered[
            :, self.state_slices["q"].slices
        ]  # (m, n_free_dof | n_modes)
        if self.structure.modal_states:
            q_mode = self.structure.modal_to_nodal(q_mode)  # (m, n_free_dof)
        q_full = (
            jnp.zeros((n_plot_vtk, self.n_nodes * 6), dtype=complex)
            .at[:, self.free_dofs]
            .set(q_mode)
        )

        def _extract_array_list(name: str) -> ArrayList | None:
            component = self.state_slices[name]
            if not component.enabled:
                return None
            return ArrayList(
                [
                    evecs_ordered[:, s].reshape((n_plot_vtk, *shape))
                    for s, shape in zip(component.slices, component.shapes)
                ]
            )

        gamma_b_full = _extract_array_list("gamma_b")
        gamma_w_full = _extract_array_list("gamma_w")
        zeta_w_full = _extract_array_list("zeta_w")

        plot_modes_vtu(
            reference=self.reference,
            directory=vtu_directory,
            q_full=q_full.reshape(n_plot_vtk, self.n_nodes, 6),
            freqs=freq_damped_ordered,
            dampings=damping_ordered,
            gamma_b_full=gamma_b_full,
            gamma_w_full=gamma_w_full,
            zeta_w_full=zeta_w_full,
            uvlm=self.aero.case,
            n_phase=n_phase,
            n_interp=n_interp,
            max_disp=max_disp,
            max_ang=max_ang,
            max_gamma=max_gamma,
        )

    return evals[idx]
modal_rescaled 
modal_rescaled(
    velocity: float | Array,
    density: float | Array,
    chord: float | Array,
    n_modes: int | None = None,
    freq_range: tuple[float | Array, float | Array] = (
        0.0,
        jnp.inf,
    ),
    damp_range: tuple[float | Array, float | Array] = (
        -jnp.inf,
        jnp.inf,
    ),
    min_struct_content: float | Array = 0.0,
    remove_complex_conjugate: bool = True,
    sort: Literal["frequency", "damping"] = "frequency",
) -> Array

Compute eigenvalues of the rescaled linear system at a new velocity, density, and chord length without re-linearising.

Parameters:

Name Type Description Default
velocity float | Array

Freestream velocity magnitude(s) at the new condition(s), scalar or (*n_points,).

 required
density float | Array

Flow density(s) at the new condition(s), scalar or broadcastable with velocity.

 required
chord float | Array

Reference chord length(s) at the new condition(s), scalar or broadcastable with velocity.

 required
n_modes int | None

Number of modes to keep. If None, all eigenvalues are returned.

 None
freq_range tuple[float | Array, float | Array]

(min, max) natural frequency window in Hz.

 (0.0, inf)
damp_range tuple[float | Array, float | Array]

(min, max) damping-ratio window.

 (-inf, inf)
min_struct_content float | Array

Minimum structural eigenvector energy fraction in [0, 1].

 0.0
remove_complex_conjugate bool

Drop one partner from each conjugate pair.

 True
sort Literal['frequency', 'damping']

Sort eigenvalues by "frequency" or "damping".

 'frequency'

Returns:

Type Description
Array

Continuous-time eigenvalues of the rescaled system(s), (n_out,) for scalar inputs or (*n_points, n_out) for batched inputs.
Source code in src/flapjax/coupled/linear/linear_coupled.py 
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
def modal_rescaled(
    self,
    velocity: float | Array,
    density: float | Array,
    chord: float | Array,
    n_modes: int | None = None,
    freq_range: tuple[float | Array, float | Array] = (0.0, jnp.inf),
    damp_range: tuple[float | Array, float | Array] = (-jnp.inf, jnp.inf),
    min_struct_content: float | Array = 0.0,
    remove_complex_conjugate: bool = True,
    sort: Literal["frequency", "damping"] = "frequency",
) -> Array:
    r"""
    Compute eigenvalues of the rescaled linear system at a new velocity, density, and chord length without
    re-linearising.
    :param velocity: Freestream velocity magnitude(s) at the new condition(s), scalar or ``(*n_points,)``.
    :param density: Flow density(s) at the new condition(s), scalar or broadcastable with ``velocity``.
    :param chord: Reference chord length(s) at the new condition(s), scalar or broadcastable with ``velocity``.
    :param n_modes: Number of modes to keep. If ``None``, all eigenvalues are returned.
    :param freq_range: (min, max) natural frequency window in Hz.
    :param damp_range: (min, max) damping-ratio window.
    :param min_struct_content: Minimum structural eigenvector energy fraction in ``[0, 1]``.
    :param remove_complex_conjugate: Drop one partner from each conjugate pair.
    :param sort: Sort eigenvalues by ``"frequency"`` or ``"damping"``.
    :return: Continuous-time eigenvalues of the rescaled system(s), ``(n_out,)`` for scalar
        inputs or ``(*n_points, n_out)`` for batched inputs.
    """
    velocity, density, chord = jnp.broadcast_arrays(
        jnp.asarray(velocity, dtype=float),
        jnp.asarray(density, dtype=float),
        jnp.asarray(chord, dtype=float),
    )
    batch_shape = velocity.shape

    def _single(velocity_: Array, density_: Array, chord_: Array) -> Array:
        r"""
        Rescale and diagonalise the linear system at a single (velocity, density, chord) point.
        """
        # reference conditions
        u_ref = self._case.aero.flowfield.u_inf_mag
        rho_ref = self._case.aero.flowfield.rho
        m_chord = self.reference.aero.gamma_b[0].shape[0]

        dt_new = chord_ / (m_chord * velocity_)

        # discretise structural system at new dt
        struct_cont = self.structure.linearise_continuous()
        n_struct = struct_cont.a.shape[0]
        eye_s = jnp.eye(n_struct)

        mat_inv_new = jnp.linalg.inv(eye_s - struct_cont.a * 0.5 * dt_new)
        a_struct_new = mat_inv_new @ (eye_s + struct_cont.a * 0.5 * dt_new)
        b_struct_new = mat_inv_new @ (struct_cont.b * dt_new)

        f_ext_slice = self.structure.input_slices["f_ext"].slices
        b_d_f_ref = self.structure.sys.b[:, f_ext_slice]
        b_d_f_new = b_struct_new[:, f_ext_slice]

        # rescaled A matrix
        a_ref = self.sys.a

        # structural and aero index ranges in the coupled state vector
        q_start = self.state_slices["q"].slices.start
        q_dot_end = self.state_slices["q_dot"].slices.stop
        struct_slice = slice(q_start, q_dot_end)

        # reference structural discrete-time A_d
        a_struct_ref = self.structure.sys.a

        # aero-induced contribution in the structural rows
        struct_rows = a_ref[struct_slice, :]
        delta = struct_rows.at[:, struct_slice].add(-a_struct_ref)

        # extract force Jacobian
        force_jac = jnp.linalg.pinv(b_d_f_ref) @ delta

        # per-column force scaling
        n_total = a_ref.shape[0]
        q_state_slice = self.state_slices["q"].slices
        base_force_scale = (density_ * velocity_) / (rho_ref * u_ref)
        force_col_scale = jnp.ones(n_total) * base_force_scale
        force_col_scale = force_col_scale.at[q_state_slice].set(
            (density_ * velocity_**2) / (rho_ref * u_ref**2)
        )
        delta_new = b_d_f_new @ (force_jac * force_col_scale[None, :])

        a_new = a_ref.at[struct_slice, :].set(delta_new)
        a_new = a_new.at[struct_slice, struct_slice].add(a_struct_new)

        vel_ratio = velocity_ / u_ref
        a_new = a_new.at[:q_start, q_state_slice].multiply(vel_ratio)

        # eigenvalue computation
        evals_d, evecs = jnp.linalg.eig(a_new)
        evals = jnp.log(evals_d) / dt_new

        omega_damped = jnp.abs(evals.imag)
        damping = -evals.real / jnp.abs(evals)
        omega_natural = omega_damped / jnp.sqrt(1.0 - damping**2)
        freq_natural_hz = omega_natural / (2.0 * jnp.pi)

        match sort:
            case "frequency":
                idx = omega_natural.argsort()
            case "damping":
                idx = damping.argsort()

        if remove_complex_conjugate:
            partner = conjugate_partner_mask(
                freq_hz=freq_natural_hz, damping=damping, tiebreaker=evals.real
            )
            idx = idx[jnp.argsort(partner[idx], stable=True)]

        q_slice = self.state_slices["q"].slices
        q_dot_slice = self.state_slices["q_dot"].slices
        evec_sq = jnp.abs(evecs) ** 2
        struct_content = (
            evec_sq[q_slice].sum(axis=0) + evec_sq[q_dot_slice].sum(axis=0)
        ) / evec_sq.sum(axis=0)

        in_range = (
            (freq_natural_hz[idx] >= freq_range[0])
            & (freq_natural_hz[idx] <= freq_range[1])
            & (damping[idx] >= damp_range[0])
            & (damping[idx] <= damp_range[1])
            & (struct_content[idx] >= min_struct_content)
        )
        idx = idx[jnp.argsort(~in_range, stable=True)]

        if n_modes is not None:
            idx = idx[:n_modes]

        return evals[idx]

    if batch_shape == ():
        return _single(velocity, density, chord)

    n_points = velocity.size

    # map scaling across multiple points
    evals_flat = vmap(_single)(
        velocity.reshape(n_points),
        density.reshape(n_points),
        chord.reshape(n_points),
    )  # (n_points, n_out)
    return evals_flat.reshape(*batch_shape, evals_flat.shape[-1])
frf 
frf(
    omega: Array, flowfield: FrequencyFlowField | None
) -> AeroelasticOutputUnflattened

Compute the frequency response function for a gust input.

Parameters:

Name Type Description Default
omega Array

Frequencies in rad/s, (n_freq,).

 required
flowfield FrequencyFlowField | None

Frequency-domain turbulence spectrum. If None, the raw transfer function is returned.

 required

Returns:

Type Description
AeroelasticOutputUnflattened

Gust FRF with q and q_dot fields, each (n_freq, n_dof).
Source code in src/flapjax/coupled/linear/linear_coupled.py 
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
def frf(
    self,
    omega: Array,
    flowfield: FrequencyFlowField | None,
) -> AeroelasticOutputUnflattened:
    r"""
    Compute the frequency response function for a gust input.
    :param omega: Frequencies in rad/s, ``(n_freq,)``.
    :param flowfield: Frequency-domain turbulence spectrum. If ``None``,
        the raw transfer function is returned.
    :return: Gust FRF with ``q`` and ``q_dot`` fields, each ``(n_freq, n_dof)``.
    """
    h = self.frf_base(omega, flowfield)
    n_out = self.n_beam_output_dof
    return AeroelasticOutputUnflattened(q=h[:, :n_out], q_dot=h[:, n_out:])
frf_base 
frf_base(
    omega: Array, flowfield: FrequencyFlowField | None
) -> Array

Compute the gust FRF as a flat vector.

Parameters:

Name Type Description Default
omega Array

Frequencies in rad/s, (n_freq,).

 required
flowfield FrequencyFlowField | None

Frequency-domain turbulence spectrum. If None, the raw transfer function is returned.

 required

Returns:

Type Description
Array

Complex FRF, (n_freq, n_outputs).
Source code in src/flapjax/coupled/linear/linear_coupled.py 
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
def frf_base(
    self,
    omega: Array,
    flowfield: FrequencyFlowField | None,
) -> Array:
    r"""
    Compute the gust FRF as a flat vector.
    :param omega: Frequencies in rad/s, ``(n_freq,)``.
    :param flowfield: Frequency-domain turbulence spectrum. If ``None``,
        the raw transfer function is returned.
    :return: Complex FRF, ``(n_freq, n_outputs)``.
    """
    from flapjax.coupled.linear_gradients.frf import gust_penetration_vector

    u_inf = (
        flowfield.u_inf
        if flowfield is not None
        else self._case.aero.flowfield.u_inf_mag
    )

    zeta_b0 = self._reference.aero.zeta_b
    vertex_x = jnp.concatenate([z[..., 0].ravel() for z in zeta_b0])
    g = gust_penetration_vector(omega=omega, vertex_x=vertex_x, u_inf=u_inf)

    a = self.sys.a
    c = self.sys.c
    n_states = a.shape[0]

    linear_upwash = LinearCoupled(
        case=self._case,
        reference=self._reference,
        batch_size=False,
        n_struct_modes=None,
        bound_upwash=True,
        skip_checks=True,
    )
    nu_b_zero = jnp.zeros(self._reference.aero.zeta_b.size)

    def _step_nu_b(nu_b_vec: Array) -> Array:
        state_np1, _ = linear_upwash.step(nu_b_vec=nu_b_vec)
        return linear_upwash.pack_state_vector(state_np1)

    def _b_g_single(g_k: Array) -> Array:
        _, bg_re = jax.jvp(_step_nu_b, (nu_b_zero,), (g_k.real,))
        _, bg_im = jax.jvp(_step_nu_b, (nu_b_zero,), (g_k.imag,))
        return bg_re + 1j * bg_im

    b_g = jax.vmap(_b_g_single)(g)

    z = jnp.exp(1j * omega * self.dt)
    eye = jnp.eye(n_states)

    def _solve_single(z_k: Array, bg_k: Array) -> Array:
        return z_k * c @ jnp.linalg.solve(z_k * eye - a, bg_k)

    h = jax.vmap(_solve_single)(z, b_g)

    if flowfield is not None:
        # scale with flowfield PSD if available
        h *= jnp.sqrt(flowfield.psd(omega))[:, None]

    return h
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 
286
287
288
289
290
291
292
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 
294
295
296
297
298
299
300
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 
362
363
364
365
366
367
368
369
370
371
372
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 
374
375
376
377
378
379
380
381
382
383
384
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 
386
387
388
389
390
391
392
393
394
395
396
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 
522
523
524
525
526
527
528
529
530
531
532
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 
534
535
536
537
538
539
540
541
542
543
544
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 
546
547
548
549
550
551
552
553
554
555
556
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 
558
559
560
561
562
563
564
565
566
567
568
569
570
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 
572
573
574
575
576
577
578
579
580
581
582
583
584
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 
586
587
588
589
590
591
592
593
594
595
596
597
598
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 
627
628
629
630
631
632
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 
634
635
636
637
638
639
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 
641
642
643
644
645
646
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)))

linear_aero_coupled

NonlinearBeamLinearAero 

NonlinearBeamLinearAero(
    structure: BeamStructure,
    aero: LinearUVLM,
    fsi_convergence_settings: ConvergenceSettings = DEFAULT_FSI_CONVERGENCE_SETTINGS,
)

Aeroelastic case coupling a nonlinear beam with a pre-built linear UVLM model. Useful for relatively cheap time-domain aeroelastic simulations where the aerodynamic nonlinearities are not significant.

Parameters:

Name Type Description Default
structure BeamStructure

Nonlinear BeamStructure with design variables set.

 required
aero LinearUVLM

Pre-built LinearUVLM linearised about some reference state.

 required
fsi_convergence_settings ConvergenceSettings

FSI iteration convergence controls.

 DEFAULT_FSI_CONVERGENCE_SETTINGS
Source code in src/flapjax/coupled/linear_aero_coupled.py 
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
def __init__(
    self,
    structure: BeamStructure,
    aero: LinearUVLM,
    fsi_convergence_settings: ConvergenceSettings = DEFAULT_FSI_CONVERGENCE_SETTINGS,
) -> None:
    r"""
    :param structure: Nonlinear ``BeamStructure`` with design variables set.
    :param aero: Pre-built ``LinearUVLM`` linearised about some reference state.
    :param fsi_convergence_settings: FSI iteration convergence controls.
    """
    self.structure: BeamStructure = structure
    self.aero: LinearUVLM = aero
    self.fsi_convergence_settings: ConvergenceSettings = fsi_convergence_settings
    self.include_unsteady_force: bool = aero.unsteady_force
get_state 
get_state(
    i_ts: int, case: AeroCase
) -> AeroStateUnflattened

Extract linear aero states at time step i_ts from the case object.

Parameters:

Name Type Description Default
i_ts int

Time step index.

 required
case AeroCase

Batched AeroCase containing the linear aero states.

 required

Returns:

Type Description
AeroStateUnflattened

Linear aero states at time step i_ts.
Source code in src/flapjax/coupled/linear_aero_coupled.py 
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
def get_state(self, i_ts: int, case: AeroCase) -> AeroStateUnflattened:
    r"""
    Extract linear aero states at time step ``i_ts`` from the case object.
    :param i_ts: Time step index.
    :param case: Batched ``AeroCase`` containing the linear aero states.
    :return: Linear aero states at time step ``i_ts``.
    """
    gamma_b = case.gamma_b.index_all(i_ts, ...)
    gamma_w = case.gamma_w.index_all(i_ts, ...)

    if self.aero.unsteady_force:
        i_prev = jnp.maximum(i_ts - 1, 0)
        gamma_b_nm1 = case.gamma_b.index_all(i_prev, ...)
    else:
        gamma_b_nm1 = None

    if self.aero.prescribed_wake:
        assert case.zeta_w is not None
        zeta_w = case.zeta_w.index_all(i_ts, ...)
        zeta_b_state = case.zeta_b.index_all(i_ts, ...)
    else:
        zeta_w = None
        zeta_b_state = None

    return AeroStateUnflattened(
        gamma_b=gamma_b,
        gamma_w=gamma_w,
        gamma_b_nm1=gamma_b_nm1,
        zeta_w=zeta_w,
        zeta_b=zeta_b_state,
    )
case_solve 
case_solve(
    case: AeroCase,
    i_ts: int,
    hg_n: Array | None,
    hg_nm1: Array | None,
    hg_dot_n: Array | None,
    static: bool,
    horseshoe: bool,
    cs_ang_n: dict[str, Array],
    cs_ang_nm1: dict[str, Array] | None,
    cs_vel_n: dict[str, Array] | None,
) -> AeroCase

Step the linear aero system one step and write the result into case. Called by BeamStructure.base_dynamic_solve inside the FSI loop. Only supports dynamic solves; hg_nm1, cs_ang_nm1 and horseshoe are part of the DynamicAeroSolver protocol but not consumed here.

Source code in src/flapjax/coupled/linear_aero_coupled.py 
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
def case_solve(
    self,
    case: AeroCase,
    i_ts: int,
    hg_n: Array | None,
    hg_nm1: Array | None,
    hg_dot_n: Array | None,
    static: bool,
    horseshoe: bool,
    cs_ang_n: dict[str, Array],
    cs_ang_nm1: dict[str, Array] | None,
    cs_vel_n: dict[str, Array] | None,
) -> AeroCase:
    r"""
    Step the linear aero system one step and write the result into ``case``. Called by
    ``BeamStructure.base_dynamic_solve`` inside the FSI loop. Only supports dynamic solves;
    ``hg_nm1``, ``cs_ang_nm1`` and ``horseshoe`` are part of the ``DynamicAeroSolver`` protocol
    but not consumed here.
    """
    del hg_nm1, cs_ang_nm1, horseshoe
    if static:
        raise NotImplementedError(
            "case_solve(static=True) is not supported for the linear aero adapter"
        )
    assert hg_n is not None and hg_dot_n is not None

    uvlm = self.aero.case
    sys = self.aero.sys
    cs_vel_n_ = cs_vel_n if cs_vel_n is not None else {}

    zeta_b_n = uvlm.hg_to_zeta_b(hg_n=hg_n, cs_ang_n=cs_ang_n)
    zeta_b_dot_n = uvlm.hg_dot_to_zeta_b_dot(
        hg_n=hg_n,
        hg_dot_n=hg_dot_n,
        cs_ang_n=cs_ang_n,
        cs_vel_n=cs_vel_n_,
    )

    # flowfield perturbation from the linearisation reference, applied as extra input upwash
    t_n = case.t[i_ts]
    t_ref = self.aero.reference.t

    if self.aero.bound_upwash:
        nu_b_n = ArrayList.zeros_like(zeta_b_n)
        nu_b_n += self.flowfield.surf_vmap_call(
            xs=self.aero.reference.zeta_b, t=t_n
        ) - self.flowfield.surf_vmap_call(xs=self.aero.reference.zeta_b, t=t_ref)
    else:
        nu_b_n = None

    if self.aero.wake_upwash:
        nu_w_n = ArrayList.zeros_like(self.aero.reference.zeta_w)
        nu_w_n += self.flowfield.surf_vmap_call(
            xs=self.aero.reference.zeta_w, t=t_n
        ) - self.flowfield.surf_vmap_call(xs=self.aero.reference.zeta_w, t=t_ref)
    else:
        nu_w_n = None

    u_n = AeroInputUnflattened(
        zeta_b=zeta_b_n,
        zeta_b_dot=zeta_b_dot_n,
        nu_b=nu_b_n,
        nu_w=nu_w_n,
    )
    u_n_vec = self.aero.pack_input_vector(u_n)
    u_ref_vec = self.aero.pack_input_vector(self.aero.reference_inputs)
    du_n = u_n_vec - u_ref_vec

    x_nm1_unflat = self.get_state(i_ts=i_ts - 1, case=case)
    x_nm1_vec = self.aero.pack_state_vector(x_nm1_unflat)
    x_ref_vec = self.aero.pack_state_vector(self.aero.reference_states)
    dx_prev = x_nm1_vec - x_ref_vec

    dx_n = sys.a @ dx_prev + sys.b @ du_n
    dy_n = sys.c @ dx_n + sys.d @ du_n

    x_n_vec = dx_n + x_ref_vec
    y_ref_vec = self.aero.pack_output_vector(self.aero.reference_outputs)
    y_n_vec = dy_n + y_ref_vec

    x_n_unflat = self.aero.unpack_state_vector(x_n_vec)
    y_n_unflat = self.aero.unpack_output_vector(y_n_vec)

    case.set_arraylist_at_ts("zeta_b", zeta_b_n, i_ts)
    case.set_arraylist_at_ts("zeta_b_dot", zeta_b_dot_n, i_ts)
    case.set_arraylist_at_ts("c", compute_c(zeta_b_n), i_ts)
    case.set_arraylist_at_ts("nc", compute_nc(zeta_b_n), i_ts)
    case.set_arraylist_at_ts("gamma_b", x_n_unflat.gamma_b, i_ts)
    case.set_arraylist_at_ts("gamma_w", x_n_unflat.gamma_w, i_ts)
    case.set_arraylist_at_ts("f_steady", y_n_unflat.f_steady, i_ts)

    if self.aero.unsteady_force:
        assert y_n_unflat.f_unsteady is not None
        gamma_b_dot_n = ArrayList(
            [
                (gb - gb_prev) / self.dt
                for gb, gb_prev in zip(x_n_unflat.gamma_b, x_nm1_unflat.gamma_b)
            ]
        )
        case.set_arraylist_at_ts("gamma_b_dot", gamma_b_dot_n, i_ts)
        case.set_arraylist_at_ts("f_unsteady", y_n_unflat.f_unsteady, i_ts)

    assert case.zeta_w is not None
    if self.aero.prescribed_wake:
        assert x_n_unflat.zeta_w is not None
        case.set_arraylist_at_ts("zeta_w", x_n_unflat.zeta_w, i_ts)
    else:
        case.set_arraylist_at_ts("zeta_w", self.aero.reference.zeta_w, i_ts)

    case.t = case.t.at[i_ts].set(t_n)

    return case
reference_configuration 
reference_configuration(
    prescribed_dofs: Sequence[int]
    | Array
    | slice
    | int = (),
    use_f_ext_follower: bool = False,
    use_f_ext_dead: bool = False,
) -> AeroelasticCase

Aeroelastic snapshot built from the beam's reference configuration and the aero linearisation reference state.

Parameters:

Name Type Description Default
prescribed_dofs Sequence[int] | Array | slice | int

Prescribed DOFs for the beam structure. Defaults to no prescribed DOFs.

 ()
use_f_ext_follower bool

Whether to include follower forces in the reference configuration.

 False
use_f_ext_dead bool

Whether to include dead forces in the reference configuration.

 False

Returns:

Type Description
AeroelasticCase

Aeroelastic snapshot at the reference configuration.
Source code in src/flapjax/coupled/linear_aero_coupled.py 
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
def reference_configuration(
    self,
    prescribed_dofs: Sequence[int] | Array | slice | int = (),
    use_f_ext_follower: bool = False,
    use_f_ext_dead: bool = False,
) -> AeroelasticCase:
    r"""
    Aeroelastic snapshot built from the beam's reference configuration and
    the aero linearisation reference state.
    :param prescribed_dofs: Prescribed DOFs for the beam structure. 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.
    :return: Aeroelastic snapshot at the reference configuration.
    """
    prescribed_dofs_tuple = self.structure.make_prescribed_dofs_tuple(
        prescribed_dofs
    )
    return AeroelasticCase(
        structure=self.structure.reference_configuration(
            use_f_grav=self.structure.use_gravity,
            use_f_ext_dead=use_f_ext_dead,
            use_f_ext_follower=use_f_ext_follower,
            use_f_aero=True,
            prescribed_dofs=prescribed_dofs_tuple,
        ),
        aero=self.aero.reference_snapshot(),
    )
dynamic_solve 
dynamic_solve(
    init_case: AeroelasticCase | None,
    prescribed_dofs: Sequence[int] | Array | slice | int,
    n_tstep: int,
    f_ext_follower: Array | None = None,
    f_ext_dead: Array | None = None,
    t_init: float = 0.0,
    load_steps: int = 1,
    thrust_t: dict[str, Array] | None = None,
    cs_ang_t: dict[str, Array] | None = None,
    cs_vel_t: dict[str, Array] | None = None,
) -> AeroelasticCase

Dynamic aeroelastic solve with nonlinear beam and linear UVLM.

Source code in src/flapjax/coupled/linear_aero_coupled.py 
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
def dynamic_solve(
    self,
    init_case: AeroelasticCase | None,
    prescribed_dofs: Sequence[int] | Array | slice | int,
    n_tstep: int,
    f_ext_follower: Array | None = None,
    f_ext_dead: Array | None = None,
    t_init: float = 0.0,
    load_steps: int = 1,
    thrust_t: dict[str, Array] | None = None,
    cs_ang_t: dict[str, Array] | None = None,
    cs_vel_t: dict[str, Array] | None = None,
) -> AeroelasticCase:
    r"""
    Dynamic aeroelastic solve with nonlinear beam and linear UVLM.
    """
    ref_cs_ang = self.aero.reference.cs_ang
    if cs_ang_t is not None:
        for key, series in cs_ang_t.items():
            if key not in ref_cs_ang:
                raise ValueError(
                    f"cs_ang_t key '{key}' not present in the linearisation reference"
                )
            if series.shape[0] != n_tstep:
                raise ValueError(
                    f"Inconsistent number of time steps for control surface input cs_ang_t['{key}']"
                )

            # TODO: implement linear control surfaces
            if not jnp.allclose(series, ref_cs_ang[key]):
                warn(f"cs_ang_t['{key}'] deviates from the linearisation reference")
    cs_ang_t_ = (
        cs_ang_t
        if cs_ang_t is not None
        else {k: jnp.full(n_tstep, v) for k, v in ref_cs_ang.items()}
    )

    if cs_vel_t is None:
        cs_vel_t_ = cs_ang_to_cs_vel(cs_ang_t=cs_ang_t_, dt=self.aero.dt)
    else:
        cs_vel_t_ = cs_vel_t

    if thrust_t is not None:
        if thrust_t.keys() != dict(self.structure.thrust_direction).keys():
            raise ValueError("Mismatch in keys for thrust")
        for k, v in thrust_t.items():
            check_arr_shape(v, (n_tstep,), name=f"thrust_t['{k}']")
        thrust_t_: dict[str, Array] = thrust_t
    else:
        thrust_t_ = {
            k: jnp.full(n_tstep, v)
            for k, v in self.structure.thrust_reference.items()
        }

    prescribed_dofs_tuple = self.structure.make_prescribed_dofs_tuple(
        prescribed_dofs
    )
    solve_dofs = get_solve_dofs(
        n_dof=self.structure.n_dof, prescribed_dofs=prescribed_dofs_tuple
    )

    t = jnp.arange(n_tstep) * self.aero.dt + t_init

    self.structure.time_integrator = TimeIntegrator(
        spectral_radius=self.structure.spectral_radius, dt=self.aero.dt
    )

    if init_case is None:
        initial_snapshot = self.reference_configuration(
            prescribed_dofs=prescribed_dofs_tuple,
            use_f_ext_follower=f_ext_follower is not None,
            use_f_ext_dead=f_ext_dead is not None,
        ).to_dynamic(t=None)
    else:
        initial_snapshot = init_case

    case = AeroelasticCase.initialise(
        initial_snapshot=initial_snapshot,
        t=t,
        use_f_ext_follower=f_ext_follower is not None,
        use_f_ext_dead=f_ext_dead is not None,
        structure=self.structure,
        x0_aero=self.aero.case.zeta_b0,
    )

    if f_ext_follower is not None and case.structure.f_ext_follower is not None:
        case.structure.f_ext_follower = case.structure.f_ext_follower.at[
            0, ...
        ].set(f_ext_follower[0, ...])
    if f_ext_dead is not None and case.structure.f_ext_dead is not None:
        case.structure.f_ext_dead = case.structure.f_ext_dead.at[0, ...].set(
            self.structure.make_f_dead_ext(
                f_ext=f_ext_dead[0, ...], rmat=case.structure.hg[0, :, :3, :3]
            )
        )

    case.structure.prescribed_dofs = prescribed_dofs_tuple

    fsi_converge_status = ConvergenceStatus(self.fsi_convergence_settings)
    fsi_converge_status.print_header(dynamic=True)

    out = self.structure.base_dynamic_solve(
        struct_case=case.structure,
        struct_convergence_status=ConvergenceStatus(
            self.structure.struct_convergence_settings
        ),
        t=t,
        solve_dofs=solve_dofs,
        load_steps=load_steps,
        f_ext_follower=f_ext_follower,
        f_ext_dead=f_ext_dead,
        aero_obj=self,
        aero_case=case.aero,
        fsi_convergence_status=fsi_converge_status,
        thrust_t=thrust_t_,
        cs_ang_t=cs_ang_t_,
        cs_vel_t=cs_vel_t_,
    )

    fsi_converge_status.print_line(dynamic=True)
    return out

linear_gradients

frf

gust_penetration_vector 
gust_penetration_vector(
    omega: Array, vertex_x: Array, u_inf: float | Array
) -> Array

Build the frequency-dependent gust-to-upwash mapping vector with delay.

Parameters:

Name Type Description Default
omega Array

Sampling frequencies in rad/s, (n_freq,).

 required
vertex_x Array

Streamwise coordinate of each grid vertex, (n_vertices,).

 required
u_inf float | Array

Freestream velocity magnitude in m/s.

 required

Returns:

Type Description
Array

Delayed vector, (n_freq, 3 * n_vertices).
Source code in src/flapjax/coupled/linear_gradients/frf.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
def gust_penetration_vector(
    omega: Array,
    vertex_x: Array,
    u_inf: float | Array,
) -> Array:
    r"""
    Build the frequency-dependent gust-to-upwash mapping vector with delay.

    :param omega: Sampling frequencies in rad/s, ``(n_freq,)``.
    :param vertex_x: Streamwise coordinate of each grid vertex, ``(n_vertices,)``.
    :param u_inf: Freestream velocity magnitude in m/s.
    :return: Delayed vector, ``(n_freq, 3 * n_vertices)``.
    """
    n_verts = vertex_x.shape[0]
    phase = jnp.exp(-1j * omega[:, None] * vertex_x[None, :] / u_inf)
    g_full = jnp.zeros((omega.shape[0], 3 * n_verts), dtype=complex)
    g_full = g_full.at[:, 2::3].set(phase)
    return g_full
compute_gust_frf 
compute_gust_frf(
    system: CoupledAeroelastic,
    dv: AeroelasticDesignVariables,
    varphi: Array,
    case: AeroelasticCase,
    omega: Array,
    frf_flowfield: FrequencyFlowField | None = None,
    batch_size: int | None = 4,
) -> AeroelasticOutputUnflattened

Compute the gust frequency response function using JVP through the coupled step.

When frf_flowfield is provided the result is PSD-weighted (each frequency row scaled by sqrt(psd(omega))). When None, the raw transfer function is returned.

Parameters:

Name Type Description Default
system CoupledAeroelastic

The coupled aeroelastic system.

 required
dv AeroelasticDesignVariables

Aeroelastic design variables.

 required
varphi Array

Structural configuration, (n_nodes, 6).

 required
case AeroelasticCase

Converged aeroelastic case around which to linearise.

 required
omega Array

Angular frequencies in rad/s, (n_freq,).

 required
frf_flowfield FrequencyFlowField | None

Frequency-domain turbulence spectrum. If None, the raw transfer function is returned.

 None
batch_size int | None

Batch size for Jacobian materialisation of the A matrix.

 4

Returns:

Type Description
AeroelasticOutputUnflattened

Gust FRF with q and q_dot fields.
Source code in src/flapjax/coupled/linear_gradients/frf.py 
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
def compute_gust_frf(
    system: CoupledAeroelastic,
    dv: AeroelasticDesignVariables,
    varphi: Array,
    case: AeroelasticCase,
    omega: Array,
    frf_flowfield: FrequencyFlowField | None = None,
    batch_size: int | None = 4,
) -> AeroelasticOutputUnflattened:
    r"""
    Compute the gust frequency response function using JVP through the coupled step.

    When ``frf_flowfield`` is provided the result is PSD-weighted (each
    frequency row scaled by ``sqrt(psd(omega))``). When ``None``, the raw
    transfer function is returned.

    :param system: The coupled aeroelastic system.
    :param dv: Aeroelastic design variables.
    :param varphi: Structural configuration, ``(n_nodes, 6)``.
    :param case: Converged aeroelastic case around which to linearise.
    :param omega: Angular frequencies in rad/s, ``(n_freq,)``.
    :param frf_flowfield: Frequency-domain turbulence spectrum. If ``None``,
        the raw transfer function is returned.
    :param batch_size: Batch size for Jacobian materialisation of the A matrix.
    :return: Gust FRF with ``q`` and ``q_dot`` fields.
    """
    ref, inner = build_reference_case(system, dv, varphi, case)
    linear = inner.linearise(
        reference=ref, skip_checks=True, batch_size=batch_size, n_struct_modes=None
    )
    return linear.frf(omega=omega, flowfield=frf_flowfield)
gust_frf_adjoint 
gust_frf_adjoint(
    system: CoupledAeroelastic,
    case: AeroelasticCase,
    omega: Array,
    objective: FRFObjective,
    frf_flowfield: FrequencyFlowField | None = None,
    grads_to_compute: AeroelasticGradsToCompute
    | None = None,
    batch_size: int = 32,
) -> tuple[Array, AeroelasticDesignVariables]

Compute sensitivities of an objective that depends on the gust FRF with respect to design variables.

Parameters:

Name Type Description Default
system CoupledAeroelastic

The coupled aeroelastic system.

 required
case AeroelasticCase

Converged static solution around which to linearise.

 required
omega Array

Frequencies to sample in rad/s, (n_freq,).

 required
objective FRFObjective

Function (full_states, design_variables, H_gust) -> scalar where H_gust is the complex gust FRF, (n_freq, n_outputs).

 required
frf_flowfield FrequencyFlowField | None

Frequency-domain turbulence spectrum. If None, the raw transfer function is passed to the objective.

 None
grads_to_compute AeroelasticGradsToCompute | None

Which design variable gradients to request.

 None
batch_size int

Batch size for Jacobian materialisation.

 32

Returns:

Type Description
tuple[Array, AeroelasticDesignVariables]

Primal objective value and its gradient w.r.t. design variables.
Source code in src/flapjax/coupled/linear_gradients/frf.py 
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
def gust_frf_adjoint(
    system: CoupledAeroelastic,
    case: AeroelasticCase,
    omega: Array,
    objective: FRFObjective,
    frf_flowfield: FrequencyFlowField | None = None,
    grads_to_compute: AeroelasticGradsToCompute | None = None,
    batch_size: int = 32,
) -> tuple[Array, AeroelasticDesignVariables]:
    r"""
    Compute sensitivities of an objective that depends on the gust FRF with
    respect to design variables.
    :param system: The coupled aeroelastic system.
    :param case: Converged static solution around which to linearise.
    :param omega: Frequencies to sample in rad/s, ``(n_freq,)``.
    :param objective: Function ``(full_states, design_variables, H_gust) -> scalar``
        where ``H_gust`` is the complex gust FRF, ``(n_freq, n_outputs)``.
    :param frf_flowfield: Frequency-domain turbulence spectrum. If ``None``,
        the raw transfer function is passed to the objective.
    :param grads_to_compute: Which design variable gradients to request.
    :param batch_size: Batch size for Jacobian materialisation.
    :return: Primal objective value and its gradient w.r.t. design variables.
    """
    if grads_to_compute is None:
        grads_to_compute = AeroelasticGradsToCompute()

    varphi_eq = case.structure.varphi
    dv_ref = system.get_design_variables(case=case, grads_to_compute=grads_to_compute)
    n_dof = system.structure.n_dof
    solve_dofs = jnp.array(
        get_solve_dofs(
            n_dof=n_dof,
            prescribed_dofs=case.structure.prescribed_dofs,
        )
    )

    # compute primal gust FRF
    h_gust = _compute_gust_frf_base(
        system, dv_ref, varphi_eq, case, omega, frf_flowfield
    )

    n_out = h_gust.shape[1] // 2

    def _unpack_h(h_: Array) -> AeroelasticOutputUnflattened:
        return AeroelasticOutputUnflattened(q=h_[:, :n_out], q_dot=h_[:, n_out:])

    def _objective_of_dv_h_varphi(
        dv_: AeroelasticDesignVariables, h_: Array, varphi_flat_: Array
    ) -> Array:
        # differentiable object w.r.t. its arguments
        states_, _ = system.aeroelastic_states_res_from_dv_varphi(
            dv=dv_,
            varphi=varphi_flat_.reshape(-1, 6),
            thrust=case.structure.thrust,
            t=case.aero.t,
            i_ts=0,
            use_horseshoe=False,
        )
        return objective(states_, dv_, _unpack_h(h_))

    j_val, vjp_j = jax.vjp(_objective_of_dv_h_varphi, dv_ref, h_gust, varphi_eq.ravel())

    j_shape = j_val.shape
    n_f = max(1, int(np.prod(j_shape)))
    d_j_d_x_direct_b, d_j_d_h_b, d_j_d_varphi_direct_b = jax.vmap(vjp_j)(
        jnp.eye(n_f).reshape((n_f,) + j_shape)
    )

    def _gust_frf_fn(dv_: AeroelasticDesignVariables, varphi_: Array) -> Array:
        # differentiate through FRF computation
        return _compute_gust_frf_base(system, dv_, varphi_, case, omega, frf_flowfield)

    _, vjp_gust = jax.vjp(_gust_frf_fn, dv_ref, varphi_eq)
    dv_bar_gust_b, varphi_bar_gust_b = jax.vmap(vjp_gust)(d_j_d_h_b)

    def _residual_of_varphi(varphi_vec: Array) -> Array:
        # differentiate full states w.r.t. static deformation
        return system.aeroelastic_states_res_from_dv_varphi(
            dv=dv_ref,
            varphi=varphi_vec.reshape(-1, 6),
            thrust=case.structure.thrust,
            t=case.aero.t,
            i_ts=0,
            use_horseshoe=False,
        )[1]

    _, vjp_res_v = jax.vjp(_residual_of_varphi, varphi_eq.ravel())
    p_res_p_varphi = jax.lax.map(
        lambda cot: vjp_res_v(cot)[0], jnp.eye(n_dof), batch_size=batch_size
    )

    varphi_bar_b = jnp.real(varphi_bar_gust_b.reshape(n_f, -1)) + jnp.real(
        d_j_d_varphi_direct_b.reshape(n_f, -1)
    )
    varphi_bar_free_b = varphi_bar_b[:, solve_dofs]

    j_res_free = p_res_p_varphi[jnp.ix_(solve_dofs, solve_dofs)]
    mu_free_b = jnp.linalg.solve(j_res_free.T, varphi_bar_free_b.T).T
    mu_full_b = (
        jnp.zeros((n_f, n_dof), dtype=mu_free_b.dtype).at[:, solve_dofs].set(mu_free_b)
    )

    _, vjp_res_dv = jax.vjp(
        lambda dv_: system.aeroelastic_states_res_from_dv_varphi(
            dv=dv_,
            varphi=varphi_eq,
            thrust=case.structure.thrust,
            t=case.aero.t,
            i_ts=0,
            use_horseshoe=False,
        )[1],
        dv_ref,
    )
    (dv_bar_via_eq_pos_b,) = jax.vmap(vjp_res_dv)(mu_full_b)

    def _neg_if_float(x):
        if hasattr(x, "dtype") and jnp.issubdtype(x.dtype, jnp.floating):
            return -x
        return x

    dv_bar_via_eq_b = jax.tree.map(_neg_if_float, dv_bar_via_eq_pos_b)

    def _sum_real(direct, via_gust, via_eq):
        if hasattr(direct, "dtype") and jnp.issubdtype(direct.dtype, jnp.floating):
            return direct + jnp.real(via_gust) + via_eq
        return direct

    total_b = jax.tree.map(_sum_real, d_j_d_x_direct_b, dv_bar_gust_b, dv_bar_via_eq_b)

    def _to_j_shape(leaf):
        if not hasattr(leaf, "shape") or leaf.ndim == 0:
            return leaf
        return leaf.reshape(j_shape + leaf.shape[1:])

    total = jax.tree.map(_to_j_shape, total_b)

    return j_val, total

stability

build_reference_case 
build_reference_case(
    system: CoupledAeroelastic,
    dv: AeroelasticDesignVariables,
    varphi: Array,
    case: AeroelasticCase,
) -> tuple[AeroelasticCase, BaseCoupledAeroelastic]

Rebuild the linearisation solution and aeroelastic object as a function of the design variables and static deformation.

Parameters:

Name Type Description Default
system CoupledAeroelastic

Coupled aeroelastic system.

 required
dv AeroelasticDesignVariables

Aeroelastic design variables.

 required
varphi Array

Structural configuration, (n_nodes, 6).

 required
case AeroelasticCase

A converged case used for linearisation.

 required

Returns:

Type Description
tuple[AeroelasticCase, BaseCoupledAeroelastic]

Solution object for reference equilibrium and the inner coupled aeroelastic object with passed design variables.
Source code in src/flapjax/coupled/linear_gradients/stability.py 
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
def build_reference_case(
    system: CoupledAeroelastic,
    dv: AeroelasticDesignVariables,
    varphi: Array,
    case: AeroelasticCase,
) -> tuple[AeroelasticCase, BaseCoupledAeroelastic]:
    r"""
    Rebuild the linearisation solution and aeroelastic object as a function of the design variables and static
    deformation.
    :param system: Coupled aeroelastic system.
    :param dv: Aeroelastic design variables.
    :param varphi: Structural configuration, ``(n_nodes, 6)``.
    :param case: A converged case used for linearisation.
    :return: Solution object for reference equilibrium and the inner coupled aeroelastic object with passed design
     variables.
    """
    inner = system.case_from_dv(dv)
    hg = inner.structure.compute_hg_from_varphi(varphi=varphi)
    aero_sol = inner.aero.static_solve(
        t=case.aero.t,
        hg=hg,
        horseshoe=False,
        cs_ang=inner.aero.cs_ang0,
    )

    # aero forcing, transformed to local frame
    f_ext_aero_global = aero_sol.project_forcing_to_beam(
        i_ts=0,
        rmat=hg[:, :3, :3],
        x0_aero=inner.aero.zeta_b0,
        include_unsteady=False,
    )
    f_ext_aero_local = transform_nodal_vect(
        f_ext_aero_global, jnp.transpose(hg[:, :3, :3], (0, 2, 1))
    )

    # find the structural states so that they propogate the derivatives correctly
    d_arr = inner.structure.make_d(hg)
    eps_arr = inner.structure.make_eps(d=d_arr)

    struct_case = pytree_clone(case.structure)
    struct_case.hg = hg
    struct_case.varphi = varphi
    struct_case.d = d_arr
    struct_case.eps = eps_arr
    struct_case.f_ext_aero = f_ext_aero_local

    return AeroelasticCase(structure=struct_case, aero=aero_sol), inner
assemble_a 
assemble_a(
    system: CoupledAeroelastic,
    dv: AeroelasticDesignVariables,
    varphi: Array,
    case: AeroelasticCase,
    n_struct_modes: int | None = None,
    batch_size: int | None = 4,
) -> Array

Assemble the discrete-time aeroelastic system matrix as a function of deformation and design variables.

Parameters:

Name Type Description Default
system CoupledAeroelastic

The coupled aeroelastic system.

 required
dv AeroelasticDesignVariables

Aeroelastic design variables.

 required
varphi Array

Structural configuration, (n_nodes, 6).

 required
case AeroelasticCase

A converged aeroelastic case around which to linearise.

 required
n_struct_modes int | None

Structural mode count. Currently not supported.

 None
batch_size int | None

Batch size for constructing the Jacobian.

 4

Returns:

Type Description
Array

Discrete-time system matrix, (n_states, n_states).
Source code in src/flapjax/coupled/linear_gradients/stability.py 
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
def assemble_a(
    system: CoupledAeroelastic,
    dv: AeroelasticDesignVariables,
    varphi: Array,
    case: AeroelasticCase,
    n_struct_modes: int | None = None,
    batch_size: int | None = 4,
) -> Array:
    r"""
    Assemble the discrete-time aeroelastic system matrix as a function of deformation and design variables.
    :param system: The coupled aeroelastic system.
    :param dv: Aeroelastic design variables.
    :param varphi: Structural configuration, ``(n_nodes, 6)``.
    :param case: A converged aeroelastic case around which to linearise.
    :param n_struct_modes: Structural mode count. Currently not supported.
    :param batch_size: Batch size for constructing the Jacobian.
    :return: Discrete-time system matrix, ``(n_states, n_states)``.
    """
    if n_struct_modes is not None:
        raise NotImplementedError(
            "Modal reduction for structural system with derivatives not supported"
        )

    # find the aeroelastic object and the reference case for linearisation
    ref, inner = build_reference_case(system, dv, varphi, case)

    # linearise system and return the system matrix
    linear = inner.linearise(
        reference=ref,
        skip_checks=True,
        batch_size=batch_size,
        n_struct_modes=n_struct_modes,
    )
    return linear.sys.a
eig_left_right 
eig_left_right(a: Array) -> tuple[Array, Array, Array]

Compute right and left eigenvectors of a square matrix, matched by nearest eigenvalue.

Parameters:

Name Type Description Default
a Array

Square matrix, (n, n).

 required

Returns:

Type Description
tuple[Array, Array, Array]

Tuple of eigenvalues, right eigenvectors, and left eigenvectors, with ordering matched by nearest eigenvalue.
Source code in src/flapjax/coupled/linear_gradients/stability.py 
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
def eig_left_right(a: Array) -> tuple[Array, Array, Array]:
    r"""
    Compute right and left eigenvectors of a square matrix, matched by nearest
    eigenvalue.
    :param a: Square matrix, ``(n, n)``.
    :return: Tuple of eigenvalues, right eigenvectors, and left eigenvectors, with ordering matched by nearest
    eigenvalue.
    """
    lam_r, phi_r = jnp.linalg.eig(a)
    lam_l, phi_l = jnp.linalg.eig(a.T)
    n = lam_r.shape[0]

    def body(i, carry):
        # function to loop to perform matching of left and right eigenvectors by nearest eigenvalue
        used, idx_ = carry
        d = jnp.where(used, jnp.inf, jnp.abs(lam_r[i] - lam_l))
        j = jnp.argmin(d)
        return used.at[j].set(True), idx_.at[i].set(j)

    used0 = jnp.zeros(n, dtype=bool)
    idx0 = jnp.zeros(n, dtype=int)
    _, idx = jax.lax.fori_loop(0, n, body, (used0, idx0))
    return lam_r, phi_r, phi_l[:, idx]
stability_adjoint 
stability_adjoint(
    system: CoupledAeroelastic,
    case: AeroelasticCase,
    objective: StabilityObjective,
    grads_to_compute: AeroelasticGradsToCompute
    | None = None,
    batch_size: int = 32,
) -> tuple[Array, AeroelasticDesignVariables]

Compute the sensitivities of some objective that refers to the aeroelastic continuous-time eigenvalues, full system states and design variables, with respect to the design variables.

Parameters:

Name Type Description Default
system CoupledAeroelastic

The coupled aeroelastic system.

 required
case AeroelasticCase

Converged static solution around which to linearise.

 required
objective StabilityObjective

Function which takes the full aeroelastic states, aeroelastic design variables, and the full continuous-time eigenvalue vector, and returns a real-valued objective.

 required
grads_to_compute AeroelasticGradsToCompute | None

Which design variable gradients to request.

 None
batch_size int

Batch size for Jacobian materialisation.

 32

Returns:

Type Description
tuple[Array, AeroelasticDesignVariables]

Primal value of the objective, and its gradient with respect to the design variables.
Source code in src/flapjax/coupled/linear_gradients/stability.py 
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
def stability_adjoint(
    system: CoupledAeroelastic,
    case: AeroelasticCase,
    objective: StabilityObjective,
    grads_to_compute: AeroelasticGradsToCompute | None = None,
    batch_size: int = 32,
) -> tuple[Array, AeroelasticDesignVariables]:
    r"""
    Compute the sensitivities of some objective that refers to the aeroelastic continuous-time eigenvalues, full system
    states and design variables, with respect to the design variables.
    :param system: The coupled aeroelastic system.
    :param case: Converged static solution around which to linearise.
    :param objective: Function which takes the full aeroelastic states, aeroelastic design variables, and the full
    continuous-time eigenvalue vector, and returns a real-valued objective.
    :param grads_to_compute: Which design variable gradients to request.
    :param batch_size: Batch size for Jacobian materialisation.
    :return: Primal value of the objective, and its gradient with respect to the design variables.
    """
    n_struct_modes = None  # model structure not implemented

    if grads_to_compute is None:
        grads_to_compute = AeroelasticGradsToCompute()

    # extract base parameters
    varphi_eq = case.structure.varphi
    dv_ref = system.get_design_variables(case=case, grads_to_compute=grads_to_compute)
    dt = system.aero.dt
    n_dof = system.structure.n_dof
    solve_dofs = jnp.array(
        get_solve_dofs(
            n_dof=n_dof,
            prescribed_dofs=case.structure.prescribed_dofs,
        )
    )

    # assemble system matrix and compute eigendecomposition
    a_matrix = assemble_a(
        system,
        dv_ref,
        varphi_eq,
        case,
        n_struct_modes=n_struct_modes,
    )
    lam_d, phi_r, phi_l = eig_left_right(a_matrix)
    lam_c = jnp.log(lam_d) / dt  # continuous time eigenvalues

    def _objective_of_dv_lam_varphi(
        dv_: AeroelasticDesignVariables, lam_c_: Array, varphi_flat_: Array
    ) -> Array:
        states_, _ = system.aeroelastic_states_res_from_dv_varphi(
            dv=dv_,
            varphi=varphi_flat_.reshape(-1, 6),
            thrust=case.structure.thrust,
            t=case.aero.t,
            i_ts=0,
            use_horseshoe=False,
        )
        return objective(states_, dv_, lam_c_)

    # VJP of the objective against it's arguments
    j_val, vjp_j = jax.vjp(
        _objective_of_dv_lam_varphi, dv_ref, lam_c, varphi_eq.ravel()
    )

    # allow for arbitrary shape
    j_shape = j_val.shape
    n_f = max(1, int(np.prod(j_shape)))
    d_j_d_x_direct_b, d_j_d_lambda_c_b, d_j_d_varphi_direct_b = jax.vmap(vjp_j)(
        jnp.eye(n_f).reshape((n_f,) + j_shape)
    )  # sensitivities through objective direct path

    # build one A-cotangent per output row, with filter to ignore defective modes
    c_denom = jnp.einsum("ij,ij->j", phi_l, phi_r)
    denom_ok = jnp.abs(c_denom) > 1e-12

    def _make_a_bar_row(dj_row: Array) -> Array:
        w_d = jnp.conj(dj_row) / (dt * lam_d)
        coeff = jnp.where(denom_ok, w_d / c_denom, 0.0)
        return jnp.real(phi_l @ jnp.diag(coeff) @ phi_r.T)

    a_bar_b = jax.vmap(_make_a_bar_row)(d_j_d_lambda_c_b)  # (n_f, N, N)

    # create a JVP for the system matrix construction
    def _a_fn(dv_, varphi_):
        return assemble_a(
            system,
            dv_,
            varphi_,
            case,
            n_struct_modes=n_struct_modes,
        )

    _, vjp_a = jax.vjp(_a_fn, dv_ref, varphi_eq)
    dv_bar_a_b, varphi_bar_a_b = jax.vmap(vjp_a)(a_bar_b)  # leading (n_f,)

    # static residual Jacobian
    def _residual_of_varphi(varphi_vec: Array) -> Array:
        return system.aeroelastic_states_res_from_dv_varphi(
            dv=dv_ref,
            varphi=varphi_vec.reshape(-1, 6),
            thrust=case.structure.thrust,
            t=case.aero.t,
            i_ts=0,
            use_horseshoe=False,
        )[1]

    _, vjp_res_v = jax.vjp(_residual_of_varphi, varphi_eq.ravel())
    p_res_p_varphi = jax.lax.map(
        lambda cot: vjp_res_v(cot)[0], jnp.eye(n_dof), batch_size=batch_size
    )

    # combine the sensitivities with respect to varphi from the full states in the objective and the path through the
    # system matrix
    varphi_bar_b = jnp.real(varphi_bar_a_b.reshape(n_f, -1)) + jnp.real(
        d_j_d_varphi_direct_b.reshape(n_f, -1)
    )  # (n_f, n_dof)
    varphi_bar_free_b = varphi_bar_b[:, solve_dofs]  # (n_f, n_free)

    j_res_free = p_res_p_varphi[jnp.ix_(solve_dofs, solve_dofs)]
    mu_free_b = jnp.linalg.solve(j_res_free.T, varphi_bar_free_b.T).T  # (n_f, n_free)
    mu_full_b = (
        jnp.zeros((n_f, n_dof), dtype=mu_free_b.dtype).at[:, solve_dofs].set(mu_free_b)
    )

    _, vjp_res_dv = jax.vjp(
        lambda dv_: system.aeroelastic_states_res_from_dv_varphi(
            dv=dv_,
            varphi=varphi_eq,
            thrust=case.structure.thrust,
            t=case.aero.t,
            i_ts=0,
            use_horseshoe=False,
        )[1],
        dv_ref,
    )
    (dv_bar_via_eq_pos_b,) = jax.vmap(vjp_res_dv)(mu_full_b)

    # handles sign flip for floating point types
    def _neg_if_float(x):
        if hasattr(x, "dtype") and jnp.issubdtype(x.dtype, jnp.floating):
            return -x
        return x

    dv_bar_via_eq_b = jax.tree.map(_neg_if_float, dv_bar_via_eq_pos_b)

    # combine contributions
    def _sum_real(direct, via_a, via_eq):
        if hasattr(direct, "dtype") and jnp.issubdtype(direct.dtype, jnp.floating):
            return direct + jnp.real(via_a) + via_eq
        return direct

    total_b = jax.tree.map(_sum_real, d_j_d_x_direct_b, dv_bar_a_b, dv_bar_via_eq_b)

    # reshape back to original shape
    def _to_j_shape(leaf):
        if not hasattr(leaf, "shape") or leaf.ndim == 0:
            return leaf
        return leaf.reshape(j_shape + leaf.shape[1:])

    total = jax.tree.map(_to_j_shape, total_b)

    return j_val, total