Aero

flapjax.aero

AeroCase

AeroCase(
    zeta_b: ArrayList,
    zeta_b_dot: ArrayList,
    zeta_w: ArrayList,
    c: ArrayList | None,
    n: ArrayList | None,
    gamma_b: ArrayList,
    gamma_b_dot: ArrayList | None,
    gamma_w: ArrayList,
    f_steady: ArrayList,
    f_unsteady: ArrayList | None,
    alpha: ArrayList | None,
    cl: ArrayList | None,
    cd: ArrayList | None,
    cm: ArrayList | None,
    cs_ang: dict[str, Array],
    cs_vel: dict[str, Array],
    kernels: Sequence[KernelFunction],
    mirror_point: Array | None,
    mirror_normal: Array | None,
    mirror_edge_low: ArrayList | None,
    mirror_edge_high: ArrayList | None,
    flowfield: FlowField,
    surf_b_names: Sequence[str],
    surf_w_names: Sequence[str],
    t: Array,
    i_ts: Array | int,
    dof_mapping: ArrayList,
    static_horseshoe: bool,
    free_wake: bool,
    gamma_dot_relaxation: float | Array,
    batch_size: int | None,
)

Contains an aerodynamic solution across one or many timesteps.

A single instance may represent either:

• Snapshot (single timestep): array leaves within have no leading time axis (e.g. zeta_b[i_surf].shape == (m+1, n+1, 3), gamma_b[i_surf].shape == (m, n)). t is a scalar and i_ts is an integer
• Batched (many timesteps): array leaves carry a leading n_tstep axis (e.g. zeta_b[i_surf].shape == (n_tstep, m+1, n+1, 3)). t is (n_tstep,) and i_ts is a (n_tstep, ) array of indices.

Use is_batched to distinguish at runtime.

Parameters:

Name	Type	Description	Default
zeta_b	ArrayList	Bound grid coordinates, batched: (n_surf, )(n_tstep, zeta_m, zeta_n, 3) / snapshot: (n_surf, )(zeta_m, zeta_n, 3).	required
zeta_b_dot	ArrayList	Bound grid velocities, same layout as zeta_b.	required
zeta_w	ArrayList	Wake grid coordinates or None.	required
c	ArrayList | None	Bound collocation points or None.	required
n	ArrayList | None	Bound grid normals or None.	required
gamma_b	ArrayList	Bound circulation strengths, batched: (n_surf, )(n_tstep, m, n) / snapshot: (n_surf, )(m, n).	required
gamma_b_dot	ArrayList | None	Bound circulation time derivatives or None.	required
gamma_w	ArrayList	Wake circulation strengths.	required
f_steady	ArrayList	Steady force contributions.	required
f_unsteady	ArrayList | None	Unsteady force contributions or None.	required
alpha	ArrayList | None	Per-strip effective angle of attack extracted from the UVLM sectional lift; batched: (n_surf, )(n_tstep, n), snapshot: (n_surf, )(n, ), or None.	required
cl	ArrayList | None	Per-strip lift coefficient sampled from the airfoil polars, or None for no polars.	required
cd	ArrayList | None	Per-strip drag coefficient sampled from the airfoil polars, or None for no polars.	required
cm	ArrayList | None	Per-strip moment coefficient sampled from the airfoil polars, or None for no polars.	required
cs_ang	dict[str, Array]	Control surface angle time history, {name: (n_tstep,)} (batched) or {name: ()} (snapshot).	required
cs_vel	dict[str, Array]	Control surface velocity time history.	required
kernels	Sequence[KernelFunction]	Kernel functions for both bound and wake source grids.	required
mirror_point	Array | None	Point on mirror plane, (3, ) or None.	required
mirror_normal	Array | None	Normal on mirror plane, (3, ) or None.	required
mirror_edge_low	ArrayList | None	Per-surface booleans marking whether that surface's n=0 edge lies on the mirror plane, (n_surf, )(), or None.	required
mirror_edge_high	ArrayList | None	As mirror_edge_low, for the n=-1 edge.	required
flowfield	FlowField	FlowField object which includes background velocity and density.	required
surf_b_names	Sequence[str]	Names of bound surfaces, (n_surf, ).	required
surf_w_names	Sequence[str]	Names of wake surfaces, (n_surf, ).	required
t	Array	Time; batched: (n_tstep, ), snapshot: scalar.	required
i_ts	Array | int	Timestep index; batched: (n_tstep, ), snapshot: int.	required
dof_mapping	ArrayList	Map from aero grid to beam DOFs, (n_surf, )(zeta_n, ).	required
static_horseshoe	bool	If true, a horseshoe formulation was used for the initial static solution.	required
free_wake	bool	Free-wake formulation flag.	required
gamma_dot_relaxation	float | Array	Circulation time derivative filter.	required
batch_size	int | None	Batch size used for AIC vectorisation.	required

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

def __init__(
    self,
    zeta_b: ArrayList,
    zeta_b_dot: ArrayList,
    zeta_w: ArrayList,
    c: ArrayList | None,
    n: ArrayList | None,
    gamma_b: ArrayList,
    gamma_b_dot: ArrayList | None,
    gamma_w: ArrayList,
    f_steady: ArrayList,
    f_unsteady: ArrayList | None,
    alpha: ArrayList | None,
    cl: ArrayList | None,
    cd: ArrayList | None,
    cm: ArrayList | None,
    cs_ang: dict[str, Array],
    cs_vel: dict[str, Array],
    kernels: Sequence[KernelFunction],
    mirror_point: Array | None,
    mirror_normal: Array | None,
    mirror_edge_low: ArrayList | None,
    mirror_edge_high: ArrayList | None,
    flowfield: FlowField,
    surf_b_names: Sequence[str],
    surf_w_names: Sequence[str],
    t: Array,
    i_ts: Array | int,
    dof_mapping: ArrayList,
    static_horseshoe: bool,
    free_wake: bool,
    gamma_dot_relaxation: float | Array,
    batch_size: int | None,
) -> None:
    r"""
    :param zeta_b: Bound grid coordinates, batched: ``(n_surf, )(n_tstep, zeta_m, zeta_n, 3)`` /
        snapshot: ``(n_surf, )(zeta_m, zeta_n, 3)``.
    :param zeta_b_dot: Bound grid velocities, same layout as ``zeta_b``.
    :param zeta_w: Wake grid coordinates or ``None``.
    :param c: Bound collocation points or ``None``.
    :param n: Bound grid normals or ``None``.
    :param gamma_b: Bound circulation strengths, batched: ``(n_surf, )(n_tstep, m, n)`` /
        snapshot: ``(n_surf, )(m, n)``.
    :param gamma_b_dot: Bound circulation time derivatives or ``None``.
    :param gamma_w: Wake circulation strengths.
    :param f_steady: Steady force contributions.
    :param f_unsteady: Unsteady force contributions or ``None``.
    :param alpha: Per-strip effective angle of attack extracted from the UVLM sectional lift; batched:
        ``(n_surf, )(n_tstep, n)``, snapshot: ``(n_surf, )(n, )``, or ``None``.
    :param cl: Per-strip lift coefficient sampled from the airfoil polars, or ``None`` for no polars.
    :param cd: Per-strip drag coefficient sampled from the airfoil polars, or ``None`` for no polars.
    :param cm: Per-strip moment coefficient sampled from the airfoil polars, or ``None`` for no polars.
    :param cs_ang: Control surface angle time history, ``{name: (n_tstep,)}`` (batched) or ``{name: ()}`` (snapshot).
    :param cs_vel: Control surface velocity time history.
    :param kernels: Kernel functions for both bound and wake source grids.
    :param mirror_point: Point on mirror plane, ``(3, )`` or None.
    :param mirror_normal: Normal on mirror plane, ``(3, )`` or None.
    :param mirror_edge_low: Per-surface booleans marking whether that surface's ``n=0`` edge lies on the
        mirror plane, ``(n_surf, )()``, or None.
    :param mirror_edge_high: As ``mirror_edge_low``, for the ``n=-1`` edge.
    :param flowfield: ``FlowField`` object which includes background velocity and density.
    :param surf_b_names: Names of bound surfaces, ``(n_surf, )``.
    :param surf_w_names: Names of wake surfaces, ``(n_surf, )``.
    :param t: Time; batched: ``(n_tstep, )``, snapshot: scalar.
    :param i_ts: Timestep index; batched: ``(n_tstep, )``, snapshot: ``int``.
    :param dof_mapping: Map from aero grid to beam DOFs, ``(n_surf, )(zeta_n, )``.
    :param static_horseshoe: If true, a horseshoe formulation was used for the initial static solution.
    :param free_wake: Free-wake formulation flag.
    :param gamma_dot_relaxation: Circulation time derivative filter.
    :param batch_size: Batch size used for AIC vectorisation.
    """
    self.zeta_b: ArrayList = zeta_b
    self.zeta_b_dot: ArrayList = zeta_b_dot
    self.zeta_w: ArrayList = zeta_w
    self.c: ArrayList | None = c
    self.nc: ArrayList | None = n
    self.gamma_b: ArrayList = gamma_b
    self.gamma_b_dot: ArrayList | None = gamma_b_dot
    self.gamma_w: ArrayList = gamma_w
    self.f_steady: ArrayList = f_steady
    self.f_unsteady: ArrayList | None = f_unsteady
    self.alpha: ArrayList | None = alpha
    self.cl: ArrayList | None = cl
    self.cd: ArrayList | None = cd
    self.cm: ArrayList | None = cm
    self.cs_ang: dict[str, Array] = cs_ang
    self.cs_vel: dict[str, Array] = cs_vel
    self.t: Array = t
    self.i_ts: Array | int = i_ts

    self.kernels: Sequence[KernelFunction] = kernels
    self.mirror_point: Array | None = mirror_point
    self.mirror_normal: Array | None = mirror_normal
    self.mirror_edge_low: ArrayList | None = mirror_edge_low
    self.mirror_edge_high: ArrayList | None = mirror_edge_high
    self.flowfield: FlowField = flowfield
    self.surf_b_names: Sequence[str] = surf_b_names
    self.surf_w_names: Sequence[str] = surf_w_names
    self.dof_mapping: ArrayList = dof_mapping

    # settings
    self.static_horseshoe: bool = static_horseshoe
    self.free_wake: bool = free_wake
    self.gamma_dot_relaxation: float | Array = gamma_dot_relaxation
    self.batch_size: int | None = batch_size

get_states

get_states(
    i_ts: int | Array | None = None,
) -> AeroFullStates

Obtain the aerodynamic state at a given timestep (used in the adjoint solution).

Parameters:

Name	Type	Description	Default
i_ts	int | Array | None	Time step index (required for batched, ignored for snapshot).	None

Returns:

Type	Description
AeroFullStates	Aero states.

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

def get_states(self, i_ts: int | Array | None = None) -> AeroFullStates:
    r"""
    Obtain the aerodynamic state at a given timestep (used in the adjoint solution).
    :param i_ts: Time step index (required for batched, ignored for snapshot).
    :return: Aero states.
    """
    if self.is_batched:
        if i_ts is None:
            raise ValueError("i_ts must be provided for batched AeroCase")
        assert self.gamma_b_dot is not None and self.zeta_w is not None
        return AeroFullStates(
            gamma_b=self.gamma_b.index_all(i_ts, ...),
            gamma_w=self.gamma_w.index_all(i_ts, ...),
            gamma_b_dot=self.gamma_b_dot.index_all(i_ts, ...),
            zeta_w=self.zeta_w.index_all(i_ts, ...),
        )
    assert self.gamma_b_dot is not None and self.zeta_w is not None
    return AeroFullStates(
        gamma_b=self.gamma_b,
        gamma_w=self.gamma_w,
        gamma_b_dot=self.gamma_b_dot,
        zeta_w=self.zeta_w,
    )

gamma_full

gamma_full(i_ts: int | None = None) -> ArrayList

Concatenate bound and wake circulation strengths.

Parameters:

Name	Type	Description	Default
i_ts	int | None	Time step index (required for batched, ignored for snapshot).	None

Returns:

Type	Description
ArrayList	Circulation strength, (2 * n_surf,)(m | m_star, n).

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

def gamma_full(self, i_ts: int | None = None) -> ArrayList:
    r"""Concatenate bound and wake circulation strengths.
    :param i_ts: Time step index (required for batched, ignored for snapshot).
    :return: Circulation strength, ``(2 * n_surf,)(m | m_star, n)``.
    """
    if self.is_batched:
        if i_ts is None:
            raise ValueError("i_ts must be provided for batched AeroCase")
        return ArrayList(
            [
                *self.gamma_b.index_all(i_ts, ...),
                *self.gamma_w.index_all(i_ts, ...),
            ]
        )
    return ArrayList([*self.gamma_b, *self.gamma_w])

zeta_full

zeta_full(i_ts: int | None = None) -> ArrayList

Concatenate bound and wake grids.

Parameters:

Name	Type	Description	Default
i_ts	int | None	Time step index (required for batched, ignored for snapshot).	None

Returns:

Type	Description
ArrayList	Grids, (2 * n_surf,)(zeta_m | zeta_m_star, zeta_n, 3).

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

def zeta_full(self, i_ts: int | None = None) -> ArrayList:
    r"""Concatenate bound and wake grids.
    :param i_ts: Time step index (required for batched, ignored for snapshot).
    :return: Grids, ``(2 * n_surf,)(zeta_m | zeta_m_star, zeta_n, 3)``.
    """
    if self.is_batched:
        if i_ts is None:
            raise ValueError("i_ts must be provided for batched AeroCase")
        assert self.zeta_w is not None
        return ArrayList(
            [
                *self.zeta_b.index_all(i_ts, ...),
                *self.zeta_w.index_all(i_ts, ...),
            ]
        )
    assert self.zeta_w is not None
    return ArrayList([*self.zeta_b, *self.zeta_w])

set_arraylist_at_ts

set_arraylist_at_ts(
    attr: str, values: ArrayList, i_ts: int
) -> None

Set an attribute at a given timestep on a batched AeroCase.

Parameters:

Name	Type	Description	Default
attr	str	Name of the attribute to set.	required
values	ArrayList	ArrayList of per-surface values (no leading time axis).	required
i_ts	int	Time step index.	required

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

def set_arraylist_at_ts(self, attr: str, values: ArrayList, i_ts: int) -> None:
    """Set an attribute at a given timestep on a batched AeroCase.
    :param attr: Name of the attribute to set.
    :param values: ArrayList of per-surface values (no leading time axis).
    :param i_ts: Time step index.
    """
    if not self.is_batched:
        raise TypeError("set_arraylist_at_ts only supported for batched AeroCase")
    arr = getattr(self, attr)
    for i_surf, val in enumerate(values):
        arr[i_surf] = arr[i_surf].at[i_ts, ...].set(val)

get_surf_snapshot

get_surf_snapshot(
    i_ts: int, i_surf: int
) -> _AeroSurfacePlot

Get single-surface plot data for a given (timestep, surface) pair on a batched AeroCase.

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

def get_surf_snapshot(self, i_ts: int, i_surf: int) -> _AeroSurfacePlot:
    r"""Get single-surface plot data for a given ``(timestep, surface)`` pair on
    a batched AeroCase.
    """
    if not self.is_batched:
        raise TypeError("get_surf_snapshot only supported for batched AeroCase")
    assert self.zeta_w is not None
    assert self.gamma_b_dot is not None
    assert self.f_unsteady is not None
    return _AeroSurfacePlot(
        zeta_b=self.zeta_b[i_surf][i_ts, ...],
        zeta_b_dot=self.zeta_b_dot[i_surf][i_ts, ...],
        zeta_w=self.zeta_w[i_surf][i_ts, ...],
        gamma_b=self.gamma_b[i_surf][i_ts, ...],
        gamma_b_dot=self.gamma_b_dot[i_surf][i_ts, ...],
        gamma_w=self.gamma_w[i_surf][i_ts, ...],
        f_steady=self.f_steady[i_surf][i_ts, ...],
        f_unsteady=self.f_unsteady[i_surf][i_ts, ...],
        alpha=self.alpha[i_surf][i_ts, ...],
        cl=self.cl[i_surf][i_ts, ...],
        cd=self.cd[i_surf][i_ts, ...],
        cm=self.cm[i_surf][i_ts, ...],
        surf_b_name=self.surf_b_names[i_surf],
        surf_w_name=self.surf_w_names[i_surf],
        i_ts=i_ts,
    )

get_surface

get_surface(idx: int) -> _AeroSurfacePlot

Get single-surface plot data for the given surface on a snapshot AeroCase.

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

def get_surface(self, idx: int) -> _AeroSurfacePlot:
    """Get single-surface plot data for the given surface on a snapshot
    AeroCase."""
    if self.is_batched:
        raise TypeError("get_surface only supported for snapshot AeroCase")
    assert self.zeta_w is not None
    assert self.gamma_b_dot is not None
    assert self.f_unsteady is not None
    return _AeroSurfacePlot(
        zeta_b=self.zeta_b[idx],
        zeta_b_dot=self.zeta_b_dot[idx],
        zeta_w=self.zeta_w[idx],
        gamma_b=self.gamma_b[idx],
        gamma_b_dot=self.gamma_b_dot[idx],
        gamma_w=self.gamma_w[idx],
        f_steady=self.f_steady[idx],
        f_unsteady=self.f_unsteady[idx],
        alpha=self.alpha[idx],
        cl=self.cl[idx],
        cd=self.cd[idx],
        cm=self.cm[idx],
        surf_b_name=self.surf_b_names[idx],
        surf_w_name=self.surf_w_names[idx],
        i_ts=int(self.i_ts),
    )

plot

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

Plot aerodynamic surfaces to VTU files (with per-surface PVD when batched).

Parameters:

Name	Type	Description	Default
directory	PathLike | str	Directory to save files.	required
plot_bound	bool	If True, plot the bound surfaces.	True
plot_wake	bool	If True, plot the wake surfaces.	True
index	int | Sequence[int] | Array | slice | None	For batched, timestep indices to plot (all if None). Ignored for snapshots.	None

Returns:

Type	Description
Sequence[Path]	Sequence of paths to the saved PVD (batched) or VTU (snapshot) files.

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

def plot(
    self,
    directory: os.PathLike | str,
    plot_bound: bool = True,
    plot_wake: bool = True,
    index: int | Sequence[int] | Array | slice | None = None,
) -> Sequence[Path]:
    r"""Plot aerodynamic surfaces to VTU files (with per-surface PVD when
    batched).
    :param directory: Directory to save files.
    :param plot_bound: If True, plot the bound surfaces.
    :param plot_wake: If True, plot the wake surfaces.
    :param index: For batched, timestep indices to plot (all if None).
        Ignored for snapshots.
    :return: Sequence of paths to the saved PVD (batched) or VTU (snapshot) files.
    """
    directory_path = Path(directory)
    directory_path.mkdir(parents=True, exist_ok=True)

    if not self.is_batched:
        paths: list[Path] = []
        for i_surf in range(self.n_surf):
            paths.extend(
                self.get_surface(idx=i_surf).plot(
                    directory, plot_bound=plot_bound, plot_wake=plot_wake
                )
            )
        return paths

    index_ = index_to_arr(index=index, n_entries=self.n_tstep)
    pvd_paths: list[Path] = []
    for i_surf in range(self.n_surf):
        per_ts_paths: list[Sequence[Path]] = []
        for i_ts in index_:
            per_ts_paths.append(
                self.get_surf_snapshot(i_ts=i_ts, i_surf=i_surf).plot(
                    directory, plot_bound=plot_bound, plot_wake=plot_wake
                )
            )

        if plot_bound:
            bound_name = f"aero_dynamic_{self.surf_b_names[i_surf]}_ts"
            pvd_paths.append(
                write_pvd(
                    directory=directory,
                    name=bound_name,
                    file_dirs=next(zip(*per_ts_paths)),
                    times=list(self.t[index_]),
                )
            )

        if plot_wake:
            wake_name = f"aero_dynamic_{self.surf_w_names[i_surf]}_ts"
            pvd_paths.append(
                write_pvd(
                    directory=directory,
                    name=wake_name,
                    file_dirs=list(zip(*per_ts_paths))[-1],
                    times=list(self.t[index_]),
                )
            )
    return pvd_paths

project_forcing_to_beam

project_forcing_to_beam(
    i_ts: int,
    rmat: Array,
    x0_aero: ArrayList,
    include_unsteady: bool,
) -> Array

Project aerodynamic forcing at i_ts onto the beam grid (global frame).

Parameters:

Name	Type	Description	Default
i_ts	int	Time step index (ignored for snapshot).	required
rmat	Array	Rotation matrix for each node relative to reference, (n_nodes, 3, 3).	required
x0_aero	ArrayList	Reference coordinates for aerodynamic grid, (n_surf, )(zeta_m, zeta_n, 3).	required
include_unsteady	bool	If true, include unsteady forcing.	required

Returns:

Type	Description
Array	Steady and unsteady forcing projected onto the beam grid, (n_nodes, 6).

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

def project_forcing_to_beam(
    self,
    i_ts: int,
    rmat: Array,
    x0_aero: ArrayList,
    include_unsteady: bool,
) -> Array:
    r"""Project aerodynamic forcing at ``i_ts`` onto the beam grid (global frame).
    :param i_ts: Time step index (ignored for snapshot).
    :param rmat: Rotation matrix for each node relative to reference, ``(n_nodes, 3, 3)``.
    :param x0_aero: Reference coordinates for aerodynamic grid, ``(n_surf, )(zeta_m, zeta_n, 3)``.
    :param include_unsteady: If true, include unsteady forcing.
    :return: Steady and unsteady forcing projected onto the beam grid, ``(n_nodes, 6)``.
    """
    if self.is_batched:
        f_total = self.f_steady.index_all(i_ts, ...)
        if include_unsteady:
            assert self.f_unsteady is not None
            f_total += self.f_unsteady.index_all(i_ts, ...)
    else:
        f_total = self.f_steady
        if include_unsteady:
            assert self.f_unsteady is not None
            f_total = ArrayList([a + b for a, b in zip(f_total, self.f_unsteady)])

    return project_forcing_to_beam(
        f_total=f_total,
        rmat=rmat,
        x0_aero=x0_aero,
        dof_mapping=self.dof_mapping,
        mirror_edge_low=self.mirror_edge_low,
        mirror_edge_high=self.mirror_edge_high,
    )

get_v_background

get_v_background(x_target: T, i_ts: int | None = None) -> T

Background velocity at specified points and time step.

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

def get_v_background[T: Array | ArrayList](
    self, x_target: T, i_ts: int | None = None
) -> T:
    r"""Background velocity at specified points and time step."""
    t_val = self._t_at(i_ts)
    if isinstance(x_target, Array):
        return self.flowfield.vmap_call(x=x_target, t=t_val)
    elif isinstance(x_target, ArrayList):
        return self.flowfield.surf_vmap_call(xs=x_target, t=t_val)  # type: ignore
    raise NotImplementedError

get_v_ind

get_v_ind(x_target: T, i_ts: int | None = None) -> T

Induced velocity at specified points and time step.

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

def get_v_ind[T: Array | ArrayList](
    self, x_target: T, i_ts: int | None = None
) -> T:
    r"""Induced velocity at specified points and time step."""
    return compute_v_ind(
        cs=x_target,
        zetas=self.zeta_full(i_ts),
        gammas=self.gamma_full(i_ts),
        kernels=self.kernels,
        mirror_normal=self.mirror_normal,
        mirror_point=self.mirror_point,
        batch_size=self.batch_size,
    )

get_v_tot

get_v_tot(x_target: T, i_ts: int | None = None) -> T

Total (induced + background) velocity at specified points and time step.

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

def get_v_tot[T: Array | ArrayList](
    self, x_target: T, i_ts: int | None = None
) -> T:
    r"""Total (induced + background) velocity at specified points and time step."""
    return self.get_v_ind(x_target=x_target, i_ts=i_ts) + self.get_v_background(
        x_target=x_target, i_ts=i_ts
    )

to_dynamic

to_dynamic(i_ts: int, n_tstep: int) -> AeroCase

Expand this snapshot into a batched AeroCase with n_tstep timesteps, placing the current snapshot at index i_ts.

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

def to_dynamic(self, i_ts: int, n_tstep: int) -> AeroCase:
    """Expand this snapshot into a batched AeroCase with ``n_tstep``
    timesteps, placing the current snapshot at index ``i_ts``.
    """
    if self.is_batched:
        raise TypeError("to_dynamic only supported for snapshot AeroCase")

    def _expand(arr_list: ArrayList) -> ArrayList:
        out = []
        for a in arr_list:
            out.append(jnp.zeros((n_tstep, *a.shape)).at[i_ts, ...].set(a))
        return ArrayList(out)

    return AeroCase(
        zeta_b=_expand(self.zeta_b),
        zeta_b_dot=_expand(self.zeta_b_dot),
        zeta_w=_expand(self.zeta_w),
        c=_expand(self.c),
        n=_expand(self.nc),
        gamma_b=_expand(self.gamma_b),
        gamma_b_dot=_expand(self.gamma_b_dot),
        gamma_w=_expand(self.gamma_w),
        f_steady=_expand(self.f_steady),
        f_unsteady=_expand(self.f_unsteady),
        alpha=_expand(self.alpha),
        cl=_expand(self.cl),
        cd=_expand(self.cd),
        cm=_expand(self.cm),
        cs_ang={k: jnp.full(n_tstep, v) for k, v in self.cs_ang.items()},
        cs_vel={k: jnp.full(n_tstep, v) for k, v in self.cs_vel.items()},
        kernels=self.kernels,
        mirror_point=self.mirror_point,
        mirror_normal=self.mirror_normal,
        mirror_edge_low=self.mirror_edge_low,
        mirror_edge_high=self.mirror_edge_high,
        flowfield=self.flowfield,
        surf_b_names=self.surf_b_names,
        surf_w_names=self.surf_w_names,
        t=jnp.zeros(n_tstep).at[i_ts].set(self.t),
        i_ts=jnp.arange(n_tstep),
        dof_mapping=self.dof_mapping,
        static_horseshoe=self.static_horseshoe,
        free_wake=self.free_wake,
        gamma_dot_relaxation=self.gamma_dot_relaxation,
        batch_size=self.batch_size,
    )

initialise classmethod

initialise(
    initial_snapshot: AeroCase, n_tstep: int
) -> AeroCase

Create a batched AeroCase from a snapshot placed at i_ts=0.

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

@classmethod
def initialise(cls, initial_snapshot: AeroCase, n_tstep: int) -> AeroCase:
    r"""Create a batched AeroCase from a snapshot placed at ``i_ts=0``."""
    return initial_snapshot.to_dynamic(i_ts=0, n_tstep=n_tstep)

GridDiscretisation dataclass

GridDiscretisation(m: int, n: int, m_star: int)

Data class to hold discretisation parameters for each aerodynamic grid.

Parameters:

Name	Type	Description	Default
m	int	Number of panels in the chordwise direction.	required
n	int	Number of panels in the spanwise direction.	required
m_star	int	Number of wake panels in the chordwise direction.	required

ConstantFlowField

ConstantFlowField(
    u_inf: Array,
    rho: float | Array,
    relative_motion: bool,
    mach: float | Array = 0.0,
)

Bases: FlowField

Constant velocity flow field.

Parameters:

Name	Type	Description	Default
u_inf	Array	Base flow velocity, (3, ).	required
rho	float | Array	Flow density.	required
relative_motion	bool	If True, the air moves, if False, the plane moves.	required
mach	float | Array	Freestream Mach number, used to apply a Prandtl-Glauert compressibility correction. Defaults to 0 (incompressible).	0.0

Source code in src/flapjax/aero/flowfields.py

def __init__(
    self,
    u_inf: Array,
    rho: float | Array,
    relative_motion: bool,
    mach: float | Array = 0.0,
):
    r"""
    :param u_inf: Base flow velocity, ``(3, )``.
    :param rho: Flow density.
    :param relative_motion: If True, the air moves, if False, the plane moves.
    :param mach: Freestream Mach number, used to apply a Prandtl-Glauert compressibility correction.
    Defaults to 0 (incompressible).
    """
    check_arr_shape(u_inf, (3,), name="u_inf")
    self.u_inf: Array = u_inf
    self.rho: Array = jnp.array(rho)
    self.u_inf_mag: Array = jnp.linalg.norm(u_inf)
    self.u_inf_dir: Array = u_inf / self.u_inf_mag
    self.q_inf: Array = 0.5 * rho * self.u_inf_mag**2  # dynamic pressure
    self.relative_motion: bool = relative_motion

    if isinstance(mach, (int, float)) and mach >= 1.0:
        warn(
            "Prandtl-Glauert compressibility correction requires subsonic flow (Mach < 1)."
        )
    self.mach: Array = jnp.array(mach)
    self.beta: Array = jnp.sqrt(
        1.0 - self.mach**2
    )  # Prandtl-Glauert compressibility factor

vmap_call

vmap_call(x: Array, t: Array) -> Array

Vectorized version of the call method. This maps over all leading dimensions of x.

Parameters:

Name	Type	Description	Default
x	Array	Spatial coordinates, (..., 3)	required
t	Array	Time, ()	required

Returns:

Type	Description
Array	Flow field values at the specified coordinates, (..., 3)

Source code in src/flapjax/aero/flowfields.py

def vmap_call(self, x: Array, t: Array) -> Array:
    """
    Vectorized version of the __call__ method. This maps over all leading dimensions of x.
    :param x: Spatial coordinates, ``(..., 3)``
    :param t: Time, ()
    :return: Flow field values at the specified coordinates, ``(..., 3)``
    """
    n_vmap = x.ndim - 1
    func = self.__call__
    for i_dim in range(n_vmap):
        func = jax.vmap(func, in_axes=(i_dim, None), out_axes=i_dim)
    return func(x, t)

surf_vmap_call

surf_vmap_call(xs: ArrayList, t: Array) -> ArrayList

Vectorized version of the call method over a list of surfaces.

Parameters:

Name	Type	Description	Default
xs	ArrayList	Spatial coordinates, (n_surf,)(..., 3)	required
t	Array	Time, ()	required

Returns:

Type	Description
ArrayList	Flow field values at the specified coordinates, (n_surf, )(..., 3)

Source code in src/flapjax/aero/flowfields.py

def surf_vmap_call(self, xs: ArrayList, t: Array) -> ArrayList:
    """
    Vectorized version of the __call__ method over a list of surfaces.
    :param xs: Spatial coordinates, ``(n_surf,)(..., 3)``
    :param t: Time, ()
    :return: Flow field values at the specified coordinates, ``(n_surf, )(..., 3)``
    """
    return ArrayList([self.vmap_call(x, t) for x in xs])

to_design_variables

to_design_variables() -> dict[str, Array]

Extract the design variables associated with this flow field.

Returns:

Type	Description
dict[str, Array]	Dictionary of design variables.

Source code in src/flapjax/aero/flowfields.py

def to_design_variables(self) -> dict[str, Array]:
    r"""
    Extract the design variables associated with this flow field.
    :return: Dictionary of design variables.
    """
    return {"u_inf": self.u_inf, "rho": self.rho, "mach": self.mach}

from_design_variables

from_design_variables(
    design_variables: dict[str, Array],
) -> FlowField

Create a new flow field from design variables as the inverse of self.to_design_variables().

Parameters:

Name	Type	Description	Default
design_variables	dict[str, Array]	Dictionary of design variables.	required

Returns:

Type	Description
FlowField	New FlowField object.

Source code in src/flapjax/aero/flowfields.py

def from_design_variables(self, design_variables: dict[str, Array]) -> FlowField:
    r"""
    Create a new flow field from design variables as the inverse of ``self.to_design_variables()``.
    :param design_variables: Dictionary of design variables.
    :return: New FlowField object.
    """
    return self.__class__(**design_variables, relative_motion=self.relative_motion)

OneMinusCosineFlowField

OneMinusCosineFlowField(
    u_inf: Array,
    rho: float | Array,
    relative_motion: bool,
    gust_length: float | Array,
    gust_amplitude: float | Array,
    gust_travel_direction: Array | None = None,
    gust_amplitude_direction: Array | None = None,
    gust_x0: Array | None = None,
    mach: float | Array = 0.0,
)

Bases: FlowField

One minus cosine gust flow field.

Parameters:

Name	Type	Description	Default
u_inf	Array	Base flow velocity, (3, ).	required
rho	float | Array	Flow density.	required
relative_motion	bool	If True, the air moves, if False, the plane moves.	required
gust_length	float | Array	Gust length.	required
gust_amplitude	float | Array	Gust amplitude.	required
gust_travel_direction	Array | None	Vector which defines the direction that the gust travels, (3, ). Defaults to the freestream direction if None.	None
gust_amplitude_direction	Array | None	Vector which defines the direction that the gust amplitude acts, (3, ). Defaults to the z-direction if None.	None
gust_x0	Array | None	Coordinate on the initial leading edge of the gust, (3, ). Defaults to 0 if None.	None
mach	float | Array	Freestream Mach number, used to apply a Prandtl-Glauert compressibility correction. Defaults to 0 (incompressible).	0.0

Source code in src/flapjax/aero/flowfields.py

def __init__(
    self,
    u_inf: Array,
    rho: float | Array,
    relative_motion: bool,
    gust_length: float | Array,
    gust_amplitude: float | Array,
    gust_travel_direction: Array | None = None,
    gust_amplitude_direction: Array | None = None,
    gust_x0: Array | None = None,
    mach: float | Array = 0.0,
):
    r"""
    :param u_inf: Base flow velocity, ``(3, )``.
    :param rho: Flow density.
    :param relative_motion: If True, the air moves, if False, the plane moves.
    :param gust_length: Gust length.
    :param gust_amplitude: Gust amplitude.
    :param gust_travel_direction: Vector which defines the direction that the gust travels, ``(3, )``. Defaults to the
    freestream direction if None.
    :param gust_amplitude_direction: Vector which defines the direction that the gust amplitude acts, ``(3, )``. Defaults
    to the z-direction if None.
    :param gust_x0: Coordinate on the initial leading edge of the gust, ``(3, )``. Defaults to 0 if None.
    :param mach: Freestream Mach number, used to apply a Prandtl-Glauert compressibility correction.
    Defaults to 0 (incompressible).
    """
    super().__init__(u_inf, rho, relative_motion, mach=mach)

    # base gust parameters
    self.gust_amplitude: Array = jnp.array(gust_amplitude)
    self.gust_length: Array = jnp.array(gust_length)

    # direction of travel for the gust - use background flow direction as default
    # even for a gust frozen in place, this defines the orientation of the ridge
    self.gust_travel_direction: Array = (
        gust_travel_direction
        if gust_travel_direction is not None
        else self.u_inf_dir
    )
    check_arr_shape(self.gust_travel_direction, (3,), "gust_travel_direction")
    self.gust_travel_direction /= jnp.linalg.norm(self.gust_travel_direction)

    # lateral direction of the gust (direction in which the gust acts), default is in Z
    self.gust_amplitude_direction: Array = (
        jnp.array((0.0, 0.0, 1.0))
        if gust_amplitude_direction is None
        else gust_amplitude_direction
    )
    check_arr_shape(self.gust_amplitude_direction, (3,), "gust_amplitude")
    self.gust_amplitude_direction /= jnp.linalg.norm(self.gust_amplitude_direction)

    # base coordinate at the start of the gust at t=0
    self.gust_x0: Array = gust_x0 if gust_x0 is not None else jnp.zeros(3)
    check_arr_shape(self.gust_x0, (3,), "gust_x0")

vmap_call

vmap_call(x: Array, t: Array) -> Array

Vectorized version of the call method. This maps over all leading dimensions of x.

Parameters:

Name	Type	Description	Default
x	Array	Spatial coordinates, (..., 3)	required
t	Array	Time, ()	required

Returns:

Type	Description
Array	Flow field values at the specified coordinates, (..., 3)

Source code in src/flapjax/aero/flowfields.py

def vmap_call(self, x: Array, t: Array) -> Array:
    """
    Vectorized version of the __call__ method. This maps over all leading dimensions of x.
    :param x: Spatial coordinates, ``(..., 3)``
    :param t: Time, ()
    :return: Flow field values at the specified coordinates, ``(..., 3)``
    """
    n_vmap = x.ndim - 1
    func = self.__call__
    for i_dim in range(n_vmap):
        func = jax.vmap(func, in_axes=(i_dim, None), out_axes=i_dim)
    return func(x, t)

surf_vmap_call

surf_vmap_call(xs: ArrayList, t: Array) -> ArrayList

Vectorized version of the call method over a list of surfaces.

Parameters:

Name	Type	Description	Default
xs	ArrayList	Spatial coordinates, (n_surf,)(..., 3)	required
t	Array	Time, ()	required

Returns:

Type	Description
ArrayList	Flow field values at the specified coordinates, (n_surf, )(..., 3)

Source code in src/flapjax/aero/flowfields.py

def surf_vmap_call(self, xs: ArrayList, t: Array) -> ArrayList:
    """
    Vectorized version of the __call__ method over a list of surfaces.
    :param xs: Spatial coordinates, ``(n_surf,)(..., 3)``
    :param t: Time, ()
    :return: Flow field values at the specified coordinates, ``(n_surf, )(..., 3)``
    """
    return ArrayList([self.vmap_call(x, t) for x in xs])

AeroFullStates

AeroFullStates(
    gamma_b: ArrayList,
    gamma_w: ArrayList,
    gamma_b_dot: ArrayList,
    zeta_w: ArrayList,
)

Aerodynamic states used for the adjoint solve.

Parameters:

Name	Type	Description	Default
gamma_b	ArrayList	Bound panel circulation strengths. (n_surf, )(m, n)	required
gamma_w	ArrayList	Wake panel circulation strengths. (n_surf, )(m_star, n)	required
gamma_b_dot	ArrayList	Bound grid circulation time derivatives. (n_surf, )(m, n)	required
zeta_w	ArrayList	Wake grid coordinates. (n_surf, )(m_star + 1, n + 1, 3)	required

Source code in src/flapjax/aero/gradients/data_structures.py

def __init__(
    self,
    gamma_b: ArrayList,
    gamma_w: ArrayList,
    gamma_b_dot: ArrayList,
    zeta_w: ArrayList,
) -> None:
    r"""
    :param gamma_b: Bound panel circulation strengths. ``(n_surf, )(m, n)``
    :param gamma_w: Wake panel circulation strengths. ``(n_surf, )(m_star, n)``
    :param gamma_b_dot: Bound grid circulation time derivatives. ``(n_surf, )(m, n)``
    :param zeta_w: Wake grid coordinates. ``(n_surf, )(m_star + 1, n + 1, 3)``
    """
    self.gamma_b: ArrayList = gamma_b
    self.gamma_w: ArrayList = gamma_w
    self.gamma_b_dot: ArrayList = gamma_b_dot
    self.zeta_w: ArrayList = zeta_w

n_states property

n_states: int

Obtain the total number of states contained within the data structure, being the size of the vector obtained from self.ravel().

Returns:

Type	Description
int	Size of vector.

shapes

shapes() -> OrderedDict[
    str, tuple[int, ...] | ArrayListShape | None
]

Obtain the shapes of all arrays within the data structure.

Returns:

Type	Description
OrderedDict[str, tuple[int, ...] | ArrayListShape | None]	Dictionary of {name: shape} pairs of all arrays or ArrayLists within the data structure.

Source code in src/flapjax/aero/gradients/data_structures.py

def shapes(self) -> OrderedDict[str, tuple[int, ...] | ArrayListShape | None]:
    r"""
    Obtain the shapes of all arrays within the data structure.
    :return: Dictionary of {name: shape} pairs of all arrays or ArrayLists within the data structure.
    """
    return OrderedDict(
        gamma_b=self.gamma_b.shape,
        gamma_w=self.gamma_w.shape,
        gamma_b_dot=self.gamma_b_dot.shape,
        zeta_w=self.zeta_w.shape,
    )

from_vector staticmethod

from_vector(
    vect: Array,
    shapes: OrderedDict[
        str, tuple[int, ...] | ArrayListShape | None
    ],
) -> AeroFullStates

Construct an AeroFullStates object from a vector of data and a corresponding dictionary of shapes, being the inverse of self.ravel().

Parameters:

Name	Type	Description	Default
vect	Array	Aerodynamic state vector.	required
shapes	OrderedDict[str, tuple[int, ...] | ArrayListShape | None]	Dictionary of {name: shape} pairs of all arrays or ArrayLists within the data structure.	required

Returns:

Type	Description
AeroFullStates	AeroFullStates object.

Source code in src/flapjax/aero/gradients/data_structures.py

@staticmethod
def from_vector(
    vect: Array,
    shapes: OrderedDict[str, tuple[int, ...] | ArrayListShape | None],
) -> AeroFullStates:
    r"""
    Construct an AeroFullStates object from a vector of data and a corresponding dictionary of shapes, being the inverse
    of ``self.ravel()``.
    :param vect: Aerodynamic state vector.
    :param shapes: Dictionary of {name: shape} pairs of all arrays or ArrayLists within the data structure.
    :return: AeroFullStates object.
    """
    return AeroFullStates(**vect_to_arrs(vect, shapes))

ravel

ravel() -> Array

Ravel the data structure to a vector, being the inverse of cls.from_vector().

Returns:

Type	Description
Array	Data vector containing all states.

Source code in src/flapjax/aero/gradients/data_structures.py

def ravel(self) -> Array:
    r"""
    Ravel the data structure to a vector, being the inverse of ``cls.from_vector()``.
    :return: Data vector containing all states.
    """

    return jnp.concatenate(
        [
            self.gamma_b.ravel(),
            self.gamma_w.ravel(),
            self.gamma_b_dot.ravel(),
            self.zeta_w.ravel(),
        ]
    )

AeroGradsToCompute dataclass

AeroGradsToCompute(
    x0_aero: bool = True,
    flowfield: bool = False,
    cs_ang_t: bool = False,
    cs_vel_t: bool = False,
)

Class which contains flags to determine which gradients are to be computed for the aerodynamic problem during the adjoint solve. Defaults to computing only the aerodynamic grid gradients.

Parameters:

Name	Type	Description	Default
x0_aero	bool	Aerodynamic grid coordinates.	True
flowfield	bool	Flow field parameters.	False
cs_ang_t	bool	Control surface deflection angle time history.	False
cs_vel_t	bool	Control surface velocity time history.	False

AeroLinearResult

AeroLinearResult(
    reference: AeroCase,
    u_t: AeroInputUnflattened,
    x_t: AeroStateUnflattened,
    y_t: AeroOutputUnflattened,
    u_t_tot: AeroInputUnflattened,
    x_t_tot: AeroStateUnflattened,
    y_t_tot: AeroOutputUnflattened,
    n_tstep: int,
    n_surf: int,
    t: Array,
    surf_b_names: list[str],
    surf_w_names: list[str],
)

Source code in src/flapjax/aero/linear/data_structures.py

def __init__(
    self,
    reference: AeroCase,
    u_t: AeroInputUnflattened,
    x_t: AeroStateUnflattened,
    y_t: AeroOutputUnflattened,
    u_t_tot: AeroInputUnflattened,
    x_t_tot: AeroStateUnflattened,
    y_t_tot: AeroOutputUnflattened,
    n_tstep: int,
    n_surf: int,
    t: Array,
    surf_b_names: list[str],
    surf_w_names: list[str],
) -> None:
    # system results, if simulated
    self.u_t: AeroInputUnflattened = u_t
    self.x_t: AeroStateUnflattened = x_t
    self.y_t: AeroOutputUnflattened = y_t
    self.u_t_tot: AeroInputUnflattened = u_t_tot
    self.x_t_tot: AeroStateUnflattened = x_t_tot
    self.y_t_tot: AeroOutputUnflattened = y_t_tot
    self.n_tstep: int = n_tstep
    self.n_surf: int = n_surf
    self.t: Array = t
    self.surf_b_names: list[str] = surf_b_names
    self.surf_w_names: list[str] = surf_w_names
    self.reference: AeroCase = reference

plot

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

Plot the aerodynamic grid at specified time steps.

Parameters:

Name	Type	Description	Default
directory	str | PathLike	Directory to save the plots to	required
index	slice | Sequence[int] | int | Array | None	Index or slice of time steps to plot. If None, plot all time steps.	None
plot_wake	bool	If True, plot the wake grid	True

Source code in src/flapjax/aero/linear/data_structures.py

def plot(
    self,
    directory: str | os.PathLike,
    index: slice | Sequence[int] | int | Array | None = None,
    plot_wake: bool = True,
) -> None:
    r"""
    Plot the aerodynamic grid at specified time steps.
    :param directory: Directory to save the plots to
    :param index: Index or slice of time steps to plot. If None, plot all time steps.
    :param plot_wake: If True, plot the wake grid
    """
    if isinstance(index, slice):
        index_ = jnp.arange(self.n_tstep)[index]
    elif isinstance(index, Sequence):
        index_ = jnp.array(index)
    elif isinstance(index, Array):
        index_ = index
    elif isinstance(index, int):
        index_ = (index,)
    elif index is None:
        index_ = jnp.arange(self.n_tstep)
    else:
        raise TypeError("index must be a slices, sequence of ints, or Array")

    directory_path = Path(directory).resolve()
    directory_path.mkdir(parents=True, exist_ok=True)

    paths: list[Sequence[Path]] = []
    for i_ts in index_:
        snapshot = self[i_ts]
        paths.append(snapshot.plot(directory, plot_wake=plot_wake))

    for i_surf in range(2 * self.n_surf):
        try:
            surf_paths = [paths[i][i_surf] for i in range(len(index_))]
            name = (self.surf_b_names + self.surf_w_names)[i_surf] + "_ts"
            write_pvd(directory, name, surf_paths, list(self.t[index_]))
        except IndexError:
            pass

LinearUVLM

LinearUVLM(
    case: UVLM,
    reference: AeroCase,
    wake_type: LinearWakeType = "frozen",
    bound_upwash: bool = True,
    wake_upwash: bool = True,
    unsteady_force: bool = True,
    *,
    skip_checks: bool = False,
    skip_linearisation: bool = False,
)

Bases: LinearModel[AeroCase, AeroInputUnflattened, AeroStateUnflattened, AeroOutputUnflattened, AeroLinearResult]

Class to represent a linearised UVLM aerodynamic system about a reference state.

Initialise linear UVLM system about a reference state.

Parameters:

Name	Type	Description	Default
case	UVLM	UVLM case object to linearise.	required
reference	AeroCase	StaticAero representing the reference state for linearisation.	required
wake_type	LinearWakeType	Instance of LinearWakeType enum to specify wake treatment.	'frozen'
bound_upwash	bool	If true, include bound surface upwash velocities as inputs.	True
wake_upwash	bool	If true, include wake surface upwash velocities as inputs.	True
unsteady_force	bool	If true, include unsteady force.	True

Source code in src/flapjax/aero/linear/linear_uvlm.py

def __init__(
    self,
    case: UVLM,
    reference: AeroCase,
    wake_type: LinearWakeType = "frozen",
    bound_upwash: bool = True,
    wake_upwash: bool = True,
    unsteady_force: bool = True,
    *,
    skip_checks: bool = False,
    skip_linearisation: bool = False,
):
    r"""
    Initialise linear UVLM system about a reference state.
    :param case: UVLM case object to linearise.
    :param reference: StaticAero representing the reference state for linearisation.
    :param wake_type: Instance of LinearWakeType enum to specify wake treatment.
    :param bound_upwash: If true, include bound surface upwash velocities as inputs.
    :param wake_upwash: If true, include wake surface upwash velocities as inputs.
    :param unsteady_force: If true, include unsteady force.
    """
    # options
    self.prescribed_wake, self.free_wake = {
        "frozen": (False, False),
        "prescribed": (True, False),
        "free": (True, True),
    }[wake_type]
    self.unsteady_force: bool = unsteady_force
    self.bound_upwash: bool = bound_upwash
    self.wake_upwash: bool = wake_upwash

    # time info
    super().__init__(reference=reference, dt=case.dt)

    # check that the reference state is steady
    # whilst linearisation can be performed about unsteady states, the current implementation omits some terms
    # required for this, however, cannot see a practical use case for such a model. Warn the user if the reference
    # state appears unsteady.
    if (
        not skip_checks
        and max([jnp.abs(zbd).max() for zbd in reference.zeta_b_dot]) > 1e-6
    ):
        warn(
            "Reference bound surface velocities are non-zero. Ensure that the reference state is steady for linearisation."
        )

    if (
        not skip_checks
        and max([jnp.abs(gbd).max() for gbd in reference.gamma_b_dot]) > 1e-6
    ):
        warn(
            "Reference bound circulation time derivative is non-zero. Ensure that the reference state is steady for linearisation."
        )

    # kernels
    self.kernels_b: Sequence[KernelFunction] = reference.n_surf * [
        biot_savart_cutoff
    ]
    self.kernels_w: Sequence[KernelFunction] = reference.n_surf * [
        biot_savart_cutoff
    ]

    # wake propagation deltas
    self.case: UVLM = case

    # linear system
    if not skip_linearisation:
        self.sys = self.linearise()

unpack_state_vector

unpack_state_vector(x: Array) -> S

Unpack a state vector into its components.

Parameters:

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

Returns:

Type	Description
S	StateUnflattened object.

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

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

unpack_output_vector

unpack_output_vector(y: Array) -> O

Unpack an output vector into its components.

Parameters:

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

Returns:

Type	Description
O	OutputUnflattened object.

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

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

pack_input_vector

pack_input_vector(u_input: InputUnflattened) -> Array

Pack an input unflattened object into a vector.

Parameters:

Name	Type	Description	Default
u_input	InputUnflattened	InputUnflattened object.	required

Returns:

Type	Description
Array	Input vector, (n_inputs, ).

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

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

pack_state_vector

pack_state_vector(x_state: StateUnflattened) -> Array

Pack a state unflattened object into a vector.

Parameters:

Name	Type	Description	Default
x_state	StateUnflattened	StateUnflattened object.	required

Returns:

Type	Description
Array	State vector, (n_states, ).

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

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

pack_output_vector

pack_output_vector(y_output: OutputUnflattened) -> Array

Pack an output unflattened object into a vector.

Parameters:

Name	Type	Description	Default
y_output	OutputUnflattened	OutputUnflattened object.	required

Returns:

Type	Description
Array	Output vector, (n_outputs, ).

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

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

get_total_input

get_total_input(u: InputUnflattened) -> InputUnflattened

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

Parameters:

Name	Type	Description	Default
u	InputUnflattened	InputUnflattened perturbation object.	required

Returns:

Type	Description
InputUnflattened	InputUnflattened total object.

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

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

get_total_state

get_total_state(x: StateUnflattened) -> StateUnflattened

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

Parameters:

Name	Type	Description	Default
x	StateUnflattened	StateUnflattened perturbation object.	required

Returns:

Type	Description
StateUnflattened	StateUnflattened total object.

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

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

get_total_output

get_total_output(y: OutputUnflattened) -> OutputUnflattened

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

Parameters:

Name	Type	Description	Default
y	OutputUnflattened	OutputUnflattened perturbation object.	required

Returns:

Type	Description
OutputUnflattened	OutputUnflattened total object.

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

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

get_total_input_t

get_total_input_t(
    u_t: InputUnflattened,
) -> InputUnflattened

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

Parameters:

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

Returns:

Type	Description
InputUnflattened	InputUnflattened total time history object.

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

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

get_total_state_t

get_total_state_t(
    x_t: StateUnflattened,
) -> StateUnflattened

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

Parameters:

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

Returns:

Type	Description
StateUnflattened	StateUnflattened total time history object.

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

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

get_total_output_t

get_total_output_t(
    y_t: OutputUnflattened,
) -> OutputUnflattened

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

Parameters:

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

Returns:

Type	Description
OutputUnflattened	OutputUnflattened total time history object.

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

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

get_zero_input

get_zero_input() -> InputUnflattened

Get a zero input unflattened object.

Returns:

Type	Description
InputUnflattened	InputUnflattened object with zero arrays.

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

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

get_zero_state

get_zero_state() -> StateUnflattened

Get a zero state unflattened object.

Returns:

Type	Description
StateUnflattened	StateUnflattened object with zero arrays.

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

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

get_zero_output

get_zero_output() -> OutputUnflattened

Get a zero output unflattened object.

Returns:

Type	Description
OutputUnflattened	OutputUnflattened object with zero arrays.

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

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

step_vec

step_vec(x_vec: Array, u_vec: Array) -> tuple[Array, Array]

Combined state/output step in vector form, operating on total (reference + perturbation) quantities and returning perturbations relative to the reference.

Parameters:

Name	Type	Description	Default
x_vec	Array	State vector, (n_states, )	required
u_vec	Array	Input vector, (n_inputs, )	required

Returns:

Type	Description
tuple[Array, Array]	Tuple of (state perturbation vector, output perturbation vector).

Source code in src/flapjax/aero/linear/linear_uvlm.py

def step_vec(self, x_vec: Array, u_vec: Array) -> tuple[Array, Array]:
    r"""
    Combined state/output step in vector form, operating on total (reference + perturbation) quantities and
    returning perturbations relative to the reference.
    :param x_vec: State vector, ``(n_states, )``
    :param u_vec: Input vector, ``(n_inputs, )``
    :return: Tuple of (state perturbation vector, output perturbation vector).
    """

    u_np1 = self._unpack_input_vector(u_vec)
    x_n = self.unpack_state_vector(x_vec)

    assert isinstance(u_np1, AeroInputUnflattened) and isinstance(
        x_n, AeroStateUnflattened
    ), (
        "Unpacked input and state must be of type AeroInputUnflattened and AeroStateUnflattened."
    )

    x_np1, y_n = self.step(u_np1=u_np1, x_n=x_n)

    return self.pack_state_vector(x_np1), self.pack_output_vector(y_n)

step

step(
    u_np1: AeroInputUnflattened, x_n: AeroStateUnflattened
) -> tuple[AeroStateUnflattened, AeroOutputUnflattened]

From total inputs u at timestep n+1 and states x at timestep n, compute the states at timestep n+1 and outputs at timestep n.

Source code in src/flapjax/aero/linear/linear_uvlm.py

def step(
    self,
    u_np1: AeroInputUnflattened,
    x_n: AeroStateUnflattened,
) -> tuple[AeroStateUnflattened, AeroOutputUnflattened]:
    r"""
    From total inputs `u` at timestep n+1 and states `x` at timestep n, compute the states at timestep n+1 and
    outputs at timestep n.
    """
    ref = self.reference

    zeta_b_np1: ArrayList = u_np1.zeta_b
    zeta_dot_b_np1: ArrayList = u_np1.zeta_b_dot
    gamma_b_n: ArrayList = x_n.gamma_b
    gamma_w_n: ArrayList = x_n.gamma_w

    if self.unsteady_force:
        assert x_n.gamma_b_nm1 is not None, "gamma_b_nm1 is None"
        gamma_b_dot_n: ArrayList = (gamma_b_n - x_n.gamma_b_nm1) / self.dt
    else:
        gamma_b_dot_n = ref.gamma_b_dot

    if self.prescribed_wake:
        assert x_n.zeta_b is not None, "zeta_b is None"
        assert x_n.zeta_w is not None, "zeta_w is None"
        zeta_b_n: ArrayList = x_n.zeta_b
        zeta_w_n: ArrayList = x_n.zeta_w
    else:
        zeta_b_n = ref.zeta_b
        zeta_w_n = ref.zeta_w

    q_n = AeroFullStates(
        gamma_b=gamma_b_n,
        gamma_w=gamma_w_n,
        gamma_b_dot=ref.gamma_b_dot,
        zeta_w=zeta_w_n,
    )
    (
        _,
        _,
        gamma_b_np1,
        gamma_w_np1,
        _,
        _,
        zeta_w_np1,
        *_,
    ) = self.case.base_solve_from_grid(
        q_nm1=q_n,
        t_n=ref.t,
        zeta_b_n=zeta_b_np1,
        zeta_b_nm1=zeta_b_n,
        zeta_b_dot_n=zeta_dot_b_np1,
        static=False,
        horseshoe=False,
        linearise_variable_wake=True,
        nu_b=u_np1.nu_b,
        nu_w=u_np1.nu_w,
    )

    # the forcing needs to be computed seperately to find its dependence on the current states
    rho = ref.flowfield.rho

    def v_out_func(x_target: Array) -> Array:
        return ref.flowfield.vmap_call(x=x_target, t=ref.t) + compute_v_ind(
            cs=x_target,
            zetas=ArrayList([*zeta_b_np1, *zeta_w_n]),
            gammas=ArrayList([*gamma_b_n, *gamma_w_n]),
            kernels=[*self.kernels_b, *self.kernels_w],
            batch_size=self.case.batch_size,
            mirror_normal=self.case.mirror_normal,
            mirror_point=self.case.mirror_point,
        )

    f_steady_n = compute_steady_forcing(
        zeta_b=zeta_b_np1,
        zeta_dot_b=zeta_dot_b_np1,
        gamma_b=gamma_b_n,
        gamma_w=gamma_w_n,
        rho=rho,
        v_func=v_out_func,
        v_inputs=u_np1.nu_b if self.bound_upwash else None,
        mirror_point=self.case.mirror_point,
        mirror_normal=self.case.mirror_normal,
        mirror_edge_low=self.case.mirror_edge_low,
        mirror_edge_high=self.case.mirror_edge_high,
    )

    normals = compute_nc(zetas=zeta_b_np1)
    if self.unsteady_force:
        f_unsteady_n = ArrayList(
            [
                split_to_vertex(
                    rho * gamma_b_dot_n[i][..., None] * normals[i], (0, 1)
                )
                for i in range(ref.n_surf)
            ]
        )
    else:
        f_unsteady_n = None

    x_np1 = AeroStateUnflattened(
        gamma_b=gamma_b_np1,
        gamma_w=gamma_w_np1,
        gamma_b_nm1=gamma_b_n if self.unsteady_force else None,
        zeta_w=zeta_w_np1 if self.prescribed_wake else None,
        zeta_b=zeta_b_np1 if self.prescribed_wake else None,
    )
    y_n = AeroOutputUnflattened(f_steady=f_steady_n, f_unsteady=f_unsteady_n)

    return x_np1, y_n

compute_jacobians

compute_jacobians(
    input_projection: LinearInputProjection | None = None,
    output_projection: LinearOutputProjection | None = None,
    residual_names: Sequence[str] | None = None,
) -> dict[
    str,
    tuple[
        Callable[..., Any], dict[str, Any], Sequence[str]
    ],
]

Build the Jacobians for the linear system.

Parameters:

Name	Type	Description	Default
input_projection	LinearInputProjection | None	If set, projects the inputs onto a different space.	None
output_projection	LinearOutputProjection | None	If set, projects the forcing outputs onto a different space.	None
residual_names	Sequence[str] | None	If set, restrict the returned residuals to this subset to avoid redundant computation.	None

Returns:

Type	Description
dict[str, tuple[Callable[..., Any], dict[str, Any], Sequence[str]]]	Mapping of residual name to Jacobian(s).

Source code in src/flapjax/aero/linear/linear_uvlm.py

def compute_jacobians(
    self,
    input_projection: LinearInputProjection | None = None,
    output_projection: LinearOutputProjection | None = None,
    residual_names: Sequence[str] | None = None,
) -> dict[str, tuple[Callable[..., Any], dict[str, Any], Sequence[str]]]:
    r"""
    Build the Jacobians for the linear system.
    :param input_projection: If set, projects the inputs onto a different space.
    :param output_projection: If set, projects the forcing outputs onto a different space.
    :param residual_names: If set, restrict the returned residuals to this subset to avoid redundant computation.
    :return: Mapping of residual name to Jacobian(s).
    """
    ref = self.reference

    # bound circulation
    gamma_b_args: dict[str, Any] = {
        "gamma_b_n_vec": ref.gamma_b.ravel(),
        "gamma_w_n_vec": ref.gamma_w.ravel(),
        "zeta_b_np1_vec": ref.zeta_b.ravel(),
        "zeta_b_dot_np1_vec": ref.zeta_b_dot.ravel(),
    }
    gamma_b_diff = [
        "gamma_b_n_vec",
        "gamma_w_n_vec",
        "zeta_b_np1_vec",
        "zeta_b_dot_np1_vec",
    ]
    if self.prescribed_wake:
        gamma_b_args["zeta_w_n_vec"] = ref.zeta_w.ravel()
        gamma_b_args["zeta_b_n_vec"] = ref.zeta_b.ravel()
        gamma_b_diff.extend(["zeta_w_n_vec", "zeta_b_n_vec"])
    if self.bound_upwash:
        gamma_b_args["nu_b_np1_vec"] = jnp.zeros(ref.zeta_b.size)
        gamma_b_diff.append("nu_b_np1_vec")
    if self.wake_upwash:
        gamma_b_args["nu_w_np1_vec"] = jnp.zeros(ref.zeta_w.size)
        gamma_b_diff.append("nu_w_np1_vec")

    # wake propagation (gamma_w and optionally zeta_w)
    wake_args: dict[str, Any] = {
        "gamma_b_n_vec": ref.gamma_b.ravel(),
        "gamma_w_n_vec": ref.gamma_w.ravel(),
    }
    gamma_w_diff = ["gamma_b_n_vec", "gamma_w_n_vec"]
    zeta_w_diff: list[str] = []
    if self.prescribed_wake:
        wake_args["zeta_w_n_vec"] = ref.zeta_w.ravel()
        wake_args["zeta_b_np1_vec"] = ref.zeta_b.ravel()
        zeta_w_diff.extend(["zeta_w_n_vec", "zeta_b_np1_vec"])
    if self.wake_upwash:
        wake_args["nu_w_np1_vec"] = jnp.zeros(ref.zeta_w.size)
        zeta_w_diff.append("nu_w_np1_vec")
    if self.free_wake:
        zeta_w_diff.extend(["gamma_b_n_vec", "gamma_w_n_vec"])

    # steady forcing
    f_steady_args: dict[str, Any] = {
        "gamma_b_n_vec": ref.gamma_b.ravel(),
        "gamma_w_n_vec": ref.gamma_w.ravel(),
        "zeta_b_np1_vec": ref.zeta_b.ravel(),
        "zeta_b_dot_np1_vec": ref.zeta_b_dot.ravel(),
    }
    f_steady_diff = [
        "gamma_b_n_vec",
        "gamma_w_n_vec",
        "zeta_b_np1_vec",
        "zeta_b_dot_np1_vec",
    ]
    if self.prescribed_wake:
        f_steady_args["zeta_w_n_vec"] = ref.zeta_w.ravel()
        f_steady_diff.append("zeta_w_n_vec")
    if self.bound_upwash:
        f_steady_args["nu_b_np1_vec"] = jnp.zeros(ref.zeta_b.size)
        f_steady_diff.append("nu_b_np1_vec")

    residuals: 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 **kw: self.wake_prop_step(**kw)[1],
            wake_args,
            gamma_w_diff,
        ),
    }
    if self.prescribed_wake:
        residuals["zeta_w"] = (
            lambda **kw: self.wake_prop_step(**kw)[0],
            wake_args,
            zeta_w_diff,
        )
    if self.unsteady_force:
        residuals["gamma_b_nm1"] = (
            lambda **kw: kw["gamma_b_n_vec"],
            {"gamma_b_n_vec": ref.gamma_b.ravel()},
            ["gamma_b_n_vec"],
        )
    if self.prescribed_wake:
        # zeta_b state at n+1 == zeta_b input at n+1
        residuals["zeta_b"] = (
            lambda **kw: kw["zeta_b_np1_vec"],
            {"zeta_b_np1_vec": ref.zeta_b.ravel()},
            ["zeta_b_np1_vec"],
        )
    residuals["f_steady"] = (self.f_steady_step, f_steady_args, f_steady_diff)
    if self.unsteady_force:
        f_unsteady_args: dict[str, Any] = {
            "gamma_b_n_vec": ref.gamma_b.ravel(),
            "gamma_b_nm1_vec": ref.gamma_b.ravel(),
            "zeta_b_np1_vec": ref.zeta_b.ravel(),
        }
        residuals["f_unsteady"] = (
            self.f_unsteady_step,
            f_unsteady_args,
            ["gamma_b_n_vec", "gamma_b_nm1_vec", "zeta_b_np1_vec"],
        )

    # apply I/O projections if provided
    if input_projection is not None:
        residuals = self._apply_input_projection(residuals, input_projection)
    if output_projection is not None:
        residuals = self._apply_output_projection(residuals, output_projection)

    if residual_names is not None:
        residuals = {k: v for k, v in residuals.items() if k in residual_names}

    return residuals

create_jacobians

create_jacobians(
    mode: ADMode | dict[str, ADMode] = "reverse",
    batch_size: int | None = None,
    input_projection: LinearInputProjection | None = None,
    output_projection: LinearOutputProjection | None = None,
    residual_names: Sequence[str] | None = None,
) -> dict[str, dict[str, Array]]

Compute the per-residual Jacobians using :func:jacrev_custom.

Source code in src/flapjax/aero/linear/linear_uvlm.py

def create_jacobians(
    self,
    mode: ADMode | dict[str, ADMode] = "reverse",
    batch_size: int | None = None,
    input_projection: LinearInputProjection | None = None,
    output_projection: LinearOutputProjection | None = None,
    residual_names: Sequence[str] | None = None,
) -> dict[str, dict[str, Array]]:
    r"""
    Compute the per-residual Jacobians using :func:`jacrev_custom`.
    """
    residuals = self.compute_jacobians(
        input_projection=input_projection,
        output_projection=output_projection,
        residual_names=residual_names,
    )

    jacobians: dict[str, dict[str, Array]] = {}
    for res_name, (res_func, args, diff_arg_names) in residuals.items():
        res_jac_options: dict[str, Callable[..., Array] | None] = {
            arg: None for arg in diff_arg_names
        }

        res_mode: ADMode = (
            mode.get(res_name, "reverse") if isinstance(mode, dict) else mode
        )

        jacs, _, _ = jacrev_custom(
            func=res_func,
            jac_options=res_jac_options,
            n_profile_loops=None,
            func_name=res_name,
            map_batch_size=batch_size,
            mode=res_mode,
        )(**args)
        jacobians[res_name] = jacs

    return jacobians

linearise

linearise(
    batch_size: int | None = None,
    *,
    input_projection: LinearInputProjection | None = None,
    output_projection: LinearOutputProjection | None = None,
) -> LinearSystem

Build the linear state-space system.

Parameters:

Name	Type	Description	Default
batch_size	int | None	If not None, batch the Jacobian passes to reduce memory.	None
input_projection	LinearInputProjection | None	Optional projection from lower-dim inputs (e.g. beam DOFs) to (zeta_b, zeta_b_dot).	None
output_projection	LinearOutputProjection | None	Optional projection from (f_steady, f_unsteady) to a smaller output (e.g. beam force); when set, C/D rows correspond to the projected output.	None

Returns:

Type	Description
LinearSystem	LinearSystem object.

Source code in src/flapjax/aero/linear/linear_uvlm.py

def linearise(
    self,
    batch_size: int | None = None,
    *,
    input_projection: LinearInputProjection | None = None,
    output_projection: LinearOutputProjection | None = None,
) -> LinearSystem:
    r"""
    Build the linear state-space system.
    :param batch_size: If not None, batch the Jacobian passes to reduce memory.
    :param input_projection: Optional projection from lower-dim inputs (e.g. beam DOFs) to
    ``(zeta_b, zeta_b_dot)``.
    :param output_projection: Optional projection from ``(f_steady, f_unsteady)`` to a
    smaller output (e.g. beam force); when set, C/D rows correspond to the projected output.
    :return: LinearSystem object.
    """
    jacobians = self.create_jacobians(
        mode="reverse",
        batch_size=batch_size,
        input_projection=input_projection,
        output_projection=output_projection,
    )

    # (row, column, size)
    ref = self.reference
    state_specs: list[tuple[str, str, int]] = [
        ("gamma_b", "gamma_b_n_vec", ref.gamma_b.size),
        ("gamma_w", "gamma_w_n_vec", ref.gamma_w.size),
    ]
    if self.unsteady_force:
        state_specs.append(("gamma_b_nm1", "gamma_b_nm1_vec", ref.gamma_b.size))
    if self.prescribed_wake:
        state_specs.append(("zeta_w", "zeta_w_n_vec", ref.zeta_w.size))
        state_specs.append(("zeta_b", "zeta_b_n_vec", ref.zeta_b.size))
    state_names = [s[0] for s in state_specs]
    state_arg_names = [s[1] for s in state_specs]
    state_sizes = [s[2] for s in state_specs]

    if input_projection is None:
        input_specs: list[tuple[str, int]] = [
            ("zeta_b_np1_vec", ref.zeta_b.size),
            ("zeta_b_dot_np1_vec", ref.zeta_b.size),
        ]
    else:
        input_specs = list(
            zip(input_projection.arg_names, input_projection.arg_sizes)
        )
    if self.bound_upwash:
        input_specs.append(("nu_b_np1_vec", ref.zeta_b.size))
    if self.wake_upwash:
        input_specs.append(("nu_w_np1_vec", ref.zeta_w.size))
    input_arg_names = [s[0] for s in input_specs]
    input_sizes = [s[1] for s in input_specs]

    if output_projection is None:
        n_fs = sum(3 * (m + 1) * (n + 1) for (m, n) in ref.gamma_b.shape)
        output_specs: list[tuple[str, int]] = [("f_steady", n_fs)]
        if self.unsteady_force:
            output_specs.append(("f_unsteady", n_fs))
    else:
        output_specs = [(output_projection.name, output_projection.size)]
    output_names = [s[0] for s in output_specs]
    output_sizes = [s[1] for s in output_specs]

    a = construct_named_block_jacobian(
        entries=tuple(jacobians[k] for k in state_names),
        keys=state_arg_names,
        widths=state_sizes,
        heights=state_sizes,
    )
    b = construct_named_block_jacobian(
        entries=tuple(jacobians[k] for k in state_names),
        keys=input_arg_names,
        widths=input_sizes,
        heights=state_sizes,
    )
    c = construct_named_block_jacobian(
        entries=tuple(jacobians[k] for k in output_names),
        keys=state_arg_names,
        widths=state_sizes,
        heights=output_sizes,
    )
    d = construct_named_block_jacobian(
        entries=tuple(jacobians[k] for k in output_names),
        keys=input_arg_names,
        widths=input_sizes,
        heights=output_sizes,
    )

    return LinearSystem(a=a, b=b, c=c, d=d, dt=self.dt)

run

run(
    u: AeroInputUnflattened,
    x0: AeroStateUnflattened | None = None,
    flowfield: FlowField | None = None,
) -> AeroLinearResult

Run the linear system.

Parameters:

Name	Type	Description	Default
u	AeroInputUnflattened	Total input over time (reference + pertubation).	required
x0	AeroStateUnflattened | None	Initial state perturbations, defaults to zero state.	None
flowfield	FlowField | None	FlowField object to provide flow velocities for bound and wake upwash, defaults to no flow.	None

Source code in src/flapjax/aero/linear/linear_uvlm.py

def run(
    self,
    u: AeroInputUnflattened,
    x0: AeroStateUnflattened | None = None,
    flowfield: FlowField | None = None,
) -> AeroLinearResult:
    r"""
    Run the linear system.
    :param u: Total input over time (reference + pertubation).
    :param x0: Initial state perturbations, defaults to zero state.
    :param flowfield: FlowField object to provide flow velocities for bound and wake upwash, defaults to no flow.
    """
    if self.prescribed_wake and self.sys.removed_u_np1:
        warn(
            "Wake perturbations coordinates at the trailing edge are zero when removing u_np1 from the system."
        )

    if x0 is None:
        x0_vec = None
    else:
        x0_vec = self._pack_state_vector_t(x0)

    n_tstep: int = u.zeta_b[0].shape[
        0
    ]  # number of time steps from first surface, first entry
    t = self.reference.t + jnp.arange(0, n_tstep) * self.dt  # time vector

    u_tot = u

    if self.bound_upwash and flowfield is None and u_tot.nu_b is None:
        warn(
            "No flowfield or bound upwash perturbations provided. Assuming zero bound upwash perturbations."
        )
        u_tot.nu_b = ArrayList(
            [jnp.zeros((n_tstep, *zb.shape)) for zb in self.reference.zeta_b]
        )

    if self.wake_upwash and flowfield is None and u_tot.nu_w is None:
        warn(
            "No flowfield or wake upwash perturbations provided. Assuming zero wake upwash perturbations."
        )
        u_tot.nu_w = ArrayList(
            [jnp.zeros((n_tstep, *zw.shape)) for zw in self.reference.zeta_w]
        )

    # add flowfield contributions to input upwash if provided
    if flowfield is not None:
        if self.bound_upwash:
            nu_b_flow = ArrayList([])
            for i_surf in range(self.reference.n_surf):
                nu_b_flow.append(
                    vmap(flowfield.vmap_call, in_axes=(None, 0), out_axes=0)(
                        self.reference.zeta_b[i_surf],
                        t,  # type: ignore
                    )
                    - flowfield.vmap_call(self.reference.zeta_b[i_surf], t[0])[
                        None, ...
                    ]
                )
            if u_tot.nu_b is None:
                u_tot.nu_b = nu_b_flow
            else:
                u_tot.nu_b += nu_b_flow
        if self.wake_upwash:
            nu_w_flow = ArrayList([])
            for i_surf in range(self.reference.n_surf):
                nu_w_flow.append(
                    vmap(flowfield.vmap_call, in_axes=(None, 0), out_axes=0)(
                        self.reference.zeta_w[i_surf],
                        t,  # type: ignore
                    )
                    - flowfield.vmap_call(self.reference.zeta_w[i_surf], t[0])[
                        None, ...
                    ]
                )
            if u_tot.nu_w is None:
                u_tot.nu_w = nu_w_flow
            else:
                u_tot.nu_w += nu_w_flow
    u_vec = self._pack_input_vector_t(u_tot)

    # run linear system
    x_t, y_t = self.sys.run(u_vec, x0_vec)

    x_t_obj = self._unpack_state_vector_t(x_t)
    y_t_obj = self._unpack_output_vector_t(y_t)

    assert isinstance(x_t_obj, AeroStateUnflattened) and isinstance(
        y_t_obj, AeroOutputUnflattened
    ), (
        "Unpacked state and output must be of type AeroStateUnflattened and AeroOutputUnflattened."
    )

    x_t_tot_obj = self.get_total_state_t(x_t_obj)
    y_t_tot_obj = self.get_total_output_t(y_t_obj)
    u_t_tot_obj = self.get_total_input_t(u_tot)

    assert (
        isinstance(u_t_tot_obj, AeroInputUnflattened)
        and isinstance(x_t_tot_obj, AeroStateUnflattened)
        and isinstance(y_t_tot_obj, AeroOutputUnflattened)
    ), (
        "Unpacked total state and output must be of type AeroStateUnflattened and AeroOutputUnflattened."
    )

    assert isinstance(self.reference, AeroCase), (
        "Reference state must be of type AeroCase."
    )

    # save results to object
    return AeroLinearResult(
        reference=self.reference,
        u_t=u,
        x_t=x_t_obj,
        y_t=y_t_obj,
        u_t_tot=u_t_tot_obj,
        x_t_tot=x_t_tot_obj,
        y_t_tot=y_t_tot_obj,
        n_tstep=n_tstep,
        t=t,
        n_surf=self.reference.n_surf,
        surf_b_names=self.case.surf_b_names,
        surf_w_names=self.case.surf_w_names,
    )

reference_snapshot

reference_snapshot() -> AeroCase

Get the reference (initial) initial_snapshot of the aerodynamic case. This will set the timestep as -1.

Returns:

Type	Description
AeroCase	StaticAero at reference state

Source code in src/flapjax/aero/linear/linear_uvlm.py

def reference_snapshot(self) -> AeroCase:
    r"""
    Get the reference (initial) initial_snapshot of the aerodynamic case. This will set the timestep as -1.
    :return: StaticAero at reference state
    """
    return AeroCase(
        zeta_b=self.reference.zeta_b,
        zeta_b_dot=self.reference.zeta_b_dot,
        zeta_w=self.reference.zeta_w,
        gamma_b=self.reference.gamma_b,
        gamma_b_dot=self.reference.gamma_b_dot,
        gamma_w=self.reference.gamma_w,
        f_steady=self.reference.f_steady,
        f_unsteady=self.reference.f_unsteady,
        cs_ang=self.reference.cs_ang,
        cs_vel=self.reference.cs_vel,
        surf_b_names=self.case.surf_b_names,
        surf_w_names=self.case.surf_w_names,
        i_ts=-1,
        t=jnp.array(0.0),
        c=self.reference.c,
        n=self.reference.nc,
        alpha=self.reference.alpha,
        cl=self.reference.cl,
        cd=self.reference.cd,
        cm=self.reference.cm,
        kernels=self.reference.kernels,
        mirror_normal=self.reference.mirror_normal,
        mirror_point=self.reference.mirror_point,
        mirror_edge_low=self.reference.mirror_edge_low,
        mirror_edge_high=self.reference.mirror_edge_high,
        flowfield=self.reference.flowfield,
        dof_mapping=self.reference.dof_mapping,
        free_wake=self.reference.free_wake,
        gamma_dot_relaxation=self.reference.gamma_dot_relaxation,
        static_horseshoe=self.reference.static_horseshoe,
        batch_size=self.case.batch_size,
    )

plot_reference

plot_reference(
    directory: PathLike, plot_wake: bool = True
) -> Sequence[Path]

Plot the reference (initial) initial_snapshot of the aerodynamic case. This will set the timestep as -1.

Parameters:

Name	Type	Description	Default
directory	PathLike	File path to save the plots to	required
plot_wake	bool	If True, plot the wake grid	True

Source code in src/flapjax/aero/linear/linear_uvlm.py

def plot_reference(
    self, directory: os.PathLike, plot_wake: bool = True
) -> Sequence[Path]:
    r"""
    Plot the reference (initial) initial_snapshot of the aerodynamic case. This will set the timestep as -1.
    :param directory: File path to save the plots to
    :param plot_wake: If True, plot the wake grid
    """
    return self.reference_snapshot().plot(
        Path(directory).resolve(), index=None, plot_wake=plot_wake
    )

UVLM

UVLM(
    grid_shapes: Sequence[
        GridDiscretisation | tuple[int, int, int]
    ],
    dof_mapping: ArrayList | Sequence[Array] | Array,
    variable_wake_disc: bool = False,
    mirror_point: Array | None = None,
    mirror_normal: Array | None = None,
    kernel: KernelFunction | None = None,
    grid_func: AeroGridFunction | None = None,
    free_wake: bool = False,
    gamma_dot_relaxation: float | Array = 0.7,
    include_unsteady_force: bool = True,
    batch_size: int | None = 64,
    polar_data: Sequence[Any | None] | None = None,
    polar_function: Sequence[PolarFunction | None]
    | None = None,
    polar_circulation_scale: float = 0.0,
)

Class to define an unsteady vortex lattice method aerodynamic case with arbitrary number of aerodynamic surfaces.

Initialise UVLM class with all non-design parameters.

Parameters:

Name	Type	Description	Default
grid_shapes	Sequence[GridDiscretisation | tuple[int, int, int]]	Discretisations for the number of chordwise, spanwise and wake-wise panels for each surface. May be passed as a sequence of either the GridDiscretisation class or a tuple of integers ordered as (m, n, m_star).	required
dof_mapping	ArrayList | Sequence[Array] | Array	Mapping from aerodynamic grid points to structure grid points for each surface.	required
variable_wake_disc	bool	If True, allow for variable wake discretisations.	False
mirror_point	Array | None	Optional point in mirror plane, (3, ). If provided, this will apply mirroring of the aerodynamic geometry and flow about the plane defined by this point and the mirror normal.	None
mirror_normal	Array | None	Optional normal vector for mirror plane, (3, ).	None
kernel	KernelFunction | None	Input for custom kernel function to use for induced velocity calculations.	None
grid_func	AeroGridFunction | None	Input functions used for defining surfaces with control surfaces. This function should take the reference local grid coordinates zeta_b0, as well as control inputs as keyword arguments, and return the deflected local grid coordinates as an ArrayList of equal dimensionality to zeta_b0.	None
free_wake	bool	If True, include the velocity induced from the aerodynamic elements for wake propagation.	False
gamma_dot_relaxation	float | Array	Filtering parameter used for obtaining the time derivative of the circulation strengths.	0.7
include_unsteady_force	bool	If True, include forces due to apparent mass for simulation.	True
batch_size	int | None	Batch size for vectorising AIC computations. Larger values may result in faster computations, at the expense of increased memory usage. Setting to None is equivelant to a vmap.	64
polar_data	Sequence[Any | None] | None	Optional per-surface polar database used to correct the UVLM sectional forcing, which can be an arbitrary data type.	None
polar_function	Sequence[PolarFunction | None] | None	Per-surface function mapping (alpha, database) -> (cl, cd, cm) about the quarter-chord. Entries of None in the sequence indicate no polar correction for that surface, or set the input value to None to apply no polar correction for any surface.	None
polar_circulation_scale	float	Factor in [0, 1] controlling how much of the per-strip lift correction factor cl_polar / cl_uvlm is applied to the bound circulation before it is stored and convected into the wake. A valud of 0 does not correct the circulation, whereas 1 fully rescales it so the shed vortex strength matches the polar-corrected lift.	0.0

Source code in src/flapjax/aero/uvlm.py

def __init__(
    self,
    grid_shapes: Sequence[GridDiscretisation | tuple[int, int, int]],
    dof_mapping: ArrayList | Sequence[Array] | Array,
    variable_wake_disc: bool = False,
    mirror_point: Array | None = None,
    mirror_normal: Array | None = None,
    kernel: KernelFunction | None = None,
    grid_func: AeroGridFunction | None = None,
    free_wake: bool = False,
    gamma_dot_relaxation: float | Array = 0.7,
    include_unsteady_force: bool = True,
    batch_size: int | None = 64,
    polar_data: Sequence[Any | None] | None = None,
    polar_function: Sequence[PolarFunction | None] | None = None,
    polar_circulation_scale: float = 0.0,
) -> None:
    r"""
    Initialise UVLM class with all non-design parameters.
    :param grid_shapes: Discretisations for the number of chordwise, spanwise and wake-wise panels for each surface.
    May be passed as a sequence of either the GridDiscretisation class or a tuple of integers ordered as ``(m, n, m_star)``.
    :param dof_mapping: Mapping from aerodynamic grid points to structure grid points for each surface.
    :param variable_wake_disc: If True, allow for variable wake discretisations.
    :param mirror_point: Optional point in mirror plane, ``(3, )``. If provided, this will apply mirroring of the aerodynamic
    geometry and flow about the plane defined by this point and the mirror normal.
    :param mirror_normal: Optional normal vector for mirror plane, ``(3, )``.
    :param kernel: Input for custom kernel function to use for induced velocity calculations.
    :param grid_func: Input functions used for defining surfaces with control surfaces. This function should take
    the reference local grid coordinates ``zeta_b0``, as well as control inputs as keyword arguments, and return the
    deflected local grid coordinates as an ArrayList of equal dimensionality to ``zeta_b0``.
    :param free_wake: If True, include the velocity induced from the aerodynamic elements for wake propagation.
    :param gamma_dot_relaxation: Filtering parameter used for obtaining the time derivative of the circulation
    strengths.
    :param include_unsteady_force: If True, include forces due to apparent mass for simulation.
    :param batch_size: Batch size for vectorising AIC computations. Larger values may result in faster computations,
    at the expense of increased memory usage. Setting to None is equivelant to a vmap.
    :param polar_data: Optional per-surface polar database used to correct the UVLM sectional forcing, which can
    be an arbitrary data type.
    :param polar_function: Per-surface function mapping ``(alpha, database) -> (cl, cd, cm)`` about the
    quarter-chord. Entries of ``None`` in the sequence indicate no polar correction for that surface, or set the
    input value to ``None`` to apply no polar correction for any surface.
    :param polar_circulation_scale: Factor in ``[0, 1]`` controlling how much of the per-strip lift
    correction factor ``cl_polar / cl_uvlm`` is applied to the bound circulation before it is stored and convected
    into the wake. A valud of 0 does not correct the circulation, whereas 1 fully rescales it so the shed vortex
    strength matches the polar-corrected lift.
    """

    # case for single inputs
    if isinstance(dof_mapping, Array):
        dof_mapping_arrlist: ArrayList = ArrayList([dof_mapping])
    elif isinstance(dof_mapping, Sequence):
        dof_mapping_arrlist = ArrayList(dof_mapping)
    elif isinstance(dof_mapping, ArrayList):
        dof_mapping_arrlist = dof_mapping
    else:
        raise TypeError("Invalid dof mapping type")
    self.dof_mapping: ArrayList = dof_mapping_arrlist

    # number of aerodynamic surfaces
    self.n_surf: int = len(grid_shapes)

    # set grid discretisations parameters for number of panels
    grid_disc = []

    for grid in grid_shapes:
        if isinstance(grid, Sequence):
            if len(grid) != 3:
                raise ValueError(
                    "Grid shape tuple must have exactly three elements (m, n, m_star)"
                )
            grid_disc.append(GridDiscretisation(*grid))
        elif isinstance(grid, GridDiscretisation):
            grid_disc.append(grid)
        else:
            raise TypeError(
                "Grid shape must be either a Sequence of three integers or a GridDiscretisation instance"
            )
    self.grid_disc: tuple[GridDiscretisation] = tuple(grid_disc)

    # count of number of panels
    self.n_bound_panels: tuple[int, ...] = tuple(
        [gd.m * gd.n for gd in self.grid_disc]
    )
    self.n_wake_panels: tuple[int, ...] = tuple(
        [gd.m_star * gd.n for gd in self.grid_disc]
    )
    self.n_panels_tot: int = sum(self.n_bound_panels) + sum(self.n_wake_panels)

    # placeholder for aerodynamic local grid coordinates, and global coordinates for wing and wake
    self.hg_ref = None
    self.zeta_b0 = None
    self.zeta_b_ref = None
    self.zeta_w_ref = None

    # placeholder for which surfaces have spanwise edges that lie on the mirror plane
    # needed for force corrections
    self.mirror_edge_low: ArrayList | None = None
    self.mirror_edge_high: ArrayList | None = None

    self.gamma_b_slice, self.gamma_w_slice = self._make_gamma_slices()

    # store DOF mapping
    if len(self.dof_mapping) != self.n_surf:
        raise ValueError(
            f"Expected {self.n_surf} DOF mapping arrays, got {len(self.dof_mapping)}"
        )
    for i_surf, map_ in enumerate(self.dof_mapping):
        check_arr_dtype(map_, int, "dof_mapping")
        check_arr_shape(map_, (self.grid_disc[i_surf].n + 1,), "grid_disc")

    # this must be optional as it is set as a design variable later
    self.flowfield = None

    # time step length
    self._dt: Array | None = None

    # wake discretisation parameters
    self.variable_wake_disc: bool = variable_wake_disc
    self.delta_w = None

    # kernel definitions per surface (separate for wing and wake)
    self.kernels_b: Sequence[KernelFunction] = self.n_surf * [
        kernel if kernel is not None else biot_savart_epsilon
    ]
    self.kernels_w: Sequence[KernelFunction] = self.n_surf * [
        kernel if kernel is not None else biot_savart_epsilon
    ]

    # settings for solvers
    self.free_wake: bool = free_wake
    self.gamma_dot_relaxation: float | Array = gamma_dot_relaxation
    self.include_unsteady_force: bool = include_unsteady_force
    self.batch_size: int | None = batch_size

    # optional per-surface polar database and evaluation function for sectional forcing correction
    polar_data_: list[Any | None] = (
        list(polar_data) if polar_data is not None else [None] * self.n_surf
    )
    if len(polar_data_) != self.n_surf:
        raise ValueError(
            f"Expected {self.n_surf} polar_data entries, got {len(polar_data_)}"
        )

    polar_function_: list[PolarFunction | None] = (
        list(polar_function) if polar_function is not None else [None] * self.n_surf
    )
    if len(polar_function_) != self.n_surf:
        raise ValueError(
            f"Expected {self.n_surf} polar_function entries, got {len(polar_function_)}"
        )

    for i_surf, (database, func) in enumerate(zip(polar_data_, polar_function_)):
        if (database is None) != (func is None):
            raise ValueError(
                f"Surface {i_surf}: polar_data and polar_function must either both be None or both be set"
            )

    self.polar_data: tuple[Any | None, ...] = tuple(polar_data_)
    self.polar_function: tuple[PolarFunction | None, ...] = tuple(polar_function_)

    if not 0.0 <= polar_circulation_scale <= 1.0:
        raise ValueError(
            f"polar_circulation_scale must be in [0, 1], got {polar_circulation_scale}."
        )
    self.polar_circulation_scale: float = float(polar_circulation_scale)

    # mirror definitions
    if (mirror_point is None) != (mirror_normal is None):
        raise ValueError(
            "Both mirror_point and mirror_normal must be provided to apply mirroring, or both must be None to "
            "apply no mirroring."
        )

    if mirror_point is None or mirror_normal is None:
        self.mirror_point: Array | None = None
        self.mirror_normal: Array | None = None
    else:
        self.mirror_point = mirror_point
        self.mirror_normal = mirror_normal / jnp.linalg.norm(
            mirror_normal
        )  # normalise

    # surface names used for plotting
    self.surf_b_names: list[str] = [f"surf_{i}_bound" for i in range(self.n_surf)]
    self.surf_w_names: list[str] = [f"surf_{i}_wake" for i in range(self.n_surf)]

    # control surface variables
    self.grid_func: AeroGridFunction = (
        grid_func if grid_func is not None else _identity_grid_func
    )

    self.cs_ang0: dict[str, Array] = {}
    self.cs_vel0: dict[str, Array] = {}

dt property writable

dt: Array

Get the time step length.

Returns:

Type	Description
Array	Time step length.

linearise

linearise(
    reference: AeroCase,
    wake_type: LinearWakeType,
    bound_upwash: bool = True,
    wake_upwash: bool = True,
    unsteady_force: bool = True,
) -> LinearUVLM

Create linearised aerodynamic model.

Parameters:

Name	Type	Description	Default
reference	AeroCase	Reference StaticAero around which to linearise.	required
wake_type	LinearWakeType	Type of wake model to use in linearisation, with options given from the LinearWakeType class (frozen, prescribed, or free). Value of None defaults to prescribed.	required
bound_upwash	bool	If true, linearise for flowfield perturbations at the bound vortex vertex.	True
wake_upwash	bool	If true, linearise for flowfield perturbations at the wake vortex vertex.	True
unsteady_force	bool	If true, include unsteady force terms in linearisation.	True

Returns:

Type	Description
LinearUVLM	LinearUVLM model, linearised at specified time step.

Source code in src/flapjax/aero/uvlm.py

def linearise(
    self,
    reference: AeroCase,
    wake_type: LinearWakeType,
    bound_upwash: bool = True,
    wake_upwash: bool = True,
    unsteady_force: bool = True,
) -> LinearUVLM:
    r"""
    Create linearised aerodynamic model.
    :param reference: Reference StaticAero around which to linearise.
    :param wake_type: Type of wake model to use in linearisation, with options given from the LinearWakeType class
     (frozen, prescribed, or free). Value of None defaults to prescribed.
    :param bound_upwash: If true, linearise for flowfield perturbations at the bound vortex vertex.
    :param wake_upwash: If true, linearise for flowfield perturbations at the wake vortex vertex.
    :param unsteady_force: If true, include unsteady force terms in linearisation.
    :return: LinearUVLM model, linearised at specified time step.
    """

    # local import used to prevent circular import issues
    from flapjax.aero.linear.linear_uvlm import LinearUVLM

    return LinearUVLM(
        self,
        reference=reference,
        wake_type=wake_type,
        bound_upwash=bound_upwash,
        wake_upwash=wake_upwash,
        unsteady_force=unsteady_force,
    )

set_design_variables

set_design_variables(
    dt: float | Array,
    flowfield: FlowField,
    zeta_b0: ArrayList | Sequence[Array] | Array,
    hg0: Array,
    delta_w: Sequence[Array | None] | Array | None = None,
    reference_cs_angles: dict[str, Array] | None = None,
) -> None

Set aerodynamic design variables for solution.

Parameters:

Name	Type	Description	Default
dt	float | Array	Time step length	required
flowfield	FlowField	FlowField object defining the background flow in space and time	required
delta_w	Sequence[Array | None] | Array | None	Vector to define segment lengths of a variable wake discretisation per surface. If None, this will use a uniform discretisation, as in the canonical UVLM.	None
zeta_b0	ArrayList | Sequence[Array] | Array	Aerodynamic local grid coordinates, (n_surf, )(zeta_m, zeta_n, 3).	required
hg0	Array	Beam reference global grid coordinates, (n_nodes, 4, 4). dictionary input for deflections, and returns an ArrayList of the local deflected grid.	required
reference_cs_angles	dict[str, Array] | None	Dictionary of {name: angle} for each control surface at the reference. If None, defaults to no control surfaces.	None

Source code in src/flapjax/aero/uvlm.py

def set_design_variables(
    self,
    dt: float | Array,
    flowfield: FlowField,
    zeta_b0: ArrayList | Sequence[Array] | Array,
    hg0: Array,
    delta_w: Sequence[Array | None] | Array | None = None,
    reference_cs_angles: dict[str, Array] | None = None,
) -> None:
    r"""
    Set aerodynamic design variables for solution.
    :param dt: Time step length
    :param flowfield: FlowField object defining the background flow in space and time
    :param delta_w: Vector to define segment lengths of a variable wake discretisation per surface. If None, this
    will use a uniform discretisation, as in the canonical UVLM.
    :param zeta_b0: Aerodynamic local grid coordinates, ``(n_surf, )(zeta_m, zeta_n, 3)``.
    :param hg0: Beam reference global grid coordinates, ``(n_nodes, 4, 4)``.
    dictionary input for deflections, and returns an ArrayList of the local deflected grid.
    :param reference_cs_angles: Dictionary of {name: angle} for each control surface at the reference. If None,
    defaults to no control surfaces.
    """

    if isinstance(delta_w, Array):
        delta_w_seq: Sequence[Array | None] = [
            delta_w if gd.m_star > 0 else None for gd in self.grid_disc
        ]
    elif delta_w is None:
        delta_w_seq = self.n_surf * [None]
    elif isinstance(delta_w, Sequence):
        if len(delta_w) != self.n_surf:
            raise ValueError(
                "Number of delta_w entries must match number of surfaces if passed as a Sequence"
            )
        delta_w_seq = delta_w
    else:
        raise TypeError("Invalid delta_w type")

    if isinstance(zeta_b0, Array):
        x0_aero_arraylist = ArrayList([zeta_b0])
    elif isinstance(zeta_b0, Sequence):
        x0_aero_arraylist = ArrayList(zeta_b0)
    elif isinstance(zeta_b0, ArrayList):
        x0_aero_arraylist = zeta_b0
    else:
        raise TypeError("Invalid zeta_b0 type")

    # set aerodynamic local coordinates
    if len(x0_aero_arraylist) != self.n_surf:
        raise ValueError(
            f"Expected {self.n_surf} aerodynamic grid coordinate arrays, got {len(zeta_b0)}"
        )

    for i_surf in range(self.n_surf):
        check_arr_shape(
            x0_aero_arraylist[i_surf],
            (self.grid_disc[i_surf].m + 1, self.grid_disc[i_surf].n + 1, 3),
            "zeta_b0",
        )
    self.zeta_b0 = x0_aero_arraylist

    if reference_cs_angles is not None:
        self.cs_ang0 = reference_cs_angles
        self.cs_vel0 = {
            k: jnp.zeros_like(v) for k, v in reference_cs_angles.items()
        }

    # set global grid coordinates for bound and wake
    check_arr_shape(hg0, (None, 4, 4), "hg0")
    self.hg_ref = hg0
    self.zeta_b_ref = self.hg_to_zeta_b(hg_n=hg0, cs_ang_n=self.cs_ang0)

    # surface spanwise edges that sit on the mirror plane, derived from reference configuration
    self.mirror_edge_low, self.mirror_edge_high = compute_mirror_edges(
        self.zeta_b_ref, self.mirror_point, self.mirror_normal
    )

    # set flowfield
    self.flowfield = flowfield

    # set timestep
    if isinstance(dt, float):
        self._dt = jnp.array(dt)
    elif isinstance(dt, Array):
        check_arr_shape(dt, (), "dt")
        self._dt = dt
    else:
        raise TypeError("dt must be either a float or an Array scalar")

    # set wake displacement
    self.delta_w = []
    for i_surf, dw_ in enumerate(delta_w_seq):
        if dw_ is None:
            self.delta_w.append(None)
        else:
            check_arr_shape(dw_, (self.grid_disc[i_surf].m_star,), "delta_w")
            self.delta_w.append(dw_)
    self.zeta_w_ref = self.initialise_wake()

case_from_dv

case_from_dv(dv: AeroDesignVariables) -> UVLM

Create a new UVLM instance as a function of design variables, allowing it to have defined gradients w.r.t. design variables.

Parameters:

Name	Type	Description	Default
dv	AeroDesignVariables	Design variables.	required

Returns:

Type	Description
UVLM	UVLM object with the same functionality as self.

Source code in src/flapjax/aero/uvlm.py

def case_from_dv(self, dv: AeroDesignVariables) -> UVLM:
    r"""
    Create a new UVLM instance as a function of design variables, allowing it to have defined gradients w.r.t.
    design variables.
    :param dv: Design variables.
    :return: UVLM object with the same functionality as ``self``.
    """
    inner_case = pytree_clone(self)
    flowfield = (
        inner_case.flowfield.from_design_variables(dv.flowfield)
        if dv.flowfield is not None
        else self.flowfield
    )
    cs_angles = (
        {k: jnp.atleast_1d(v)[0] for k, v in dv.cs_ang_t.items()}
        if dv.cs_ang_t is not None
        else self.cs_ang0
    )
    inner_case.set_design_variables(
        dt=self.dt,
        flowfield=flowfield,
        delta_w=self.delta_w,
        zeta_b0=dv_or(dv.zeta_b0, self.zeta_b0),
        hg0=self.hg_ref,
        reference_cs_angles=cs_angles,
    )

    return inner_case

get_design_variables

get_design_variables(
    cs_ang_t: dict[str, Array],
    cs_vel_t: dict[str, Array],
    grads_to_compute: AeroGradsToCompute | None,
) -> AeroDesignVariables

Extract design variables from the aerodynamic case. As the control input time histories are defined when initialising the simulation, they are not included in self and so are passed by argument.

Parameters:

Name	Type	Description	Default
cs_ang_t	dict[str, Array]	Time history of control surface angles, {keys: (n_tstep, )}.	required
cs_vel_t	dict[str, Array]	Time history of control surface velocities, {keys: (n_tstep, )}.	required
grads_to_compute	AeroGradsToCompute | None	Data structure which describes which design variables should be obtained. If None, all variables are obtained.	required

Returns:

Type	Description
AeroDesignVariables	Aerodynamic design variables.

Source code in src/flapjax/aero/uvlm.py

def get_design_variables(
    self,
    cs_ang_t: dict[str, Array],
    cs_vel_t: dict[str, Array],
    grads_to_compute: AeroGradsToCompute | None,
) -> AeroDesignVariables:
    r"""
    Extract design variables from the aerodynamic case. As the control input time histories are defined when
    initialising the simulation, they are not included in ``self`` and so are passed by argument.
    :param cs_ang_t: Time history of control surface angles, ``{keys: (n_tstep, )}``.
    :param cs_vel_t: Time history of control surface velocities, ``{keys: (n_tstep, )}``.
    :param grads_to_compute: Data structure which describes which design variables should be obtained. If None, all
    variables are obtained.
    :return: Aerodynamic design variables.
    """
    if isinstance(grads_to_compute, AeroGradsToCompute):
        return AeroDesignVariables(
            zeta_b0=self.zeta_b0 if grads_to_compute.x0_aero else None,
            flowfield=self.flowfield.to_design_variables()
            if grads_to_compute.flowfield
            else None,
            cs_ang_t=cs_ang_t if grads_to_compute.cs_ang_t else None,
            cs_vel_t=cs_vel_t if grads_to_compute.cs_vel_t else None,
            f_shape=(),
        )
    else:  # grads_to_compute is None
        return AeroDesignVariables(
            zeta_b0=self.zeta_b0,
            flowfield=self.flowfield.to_design_variables(),
            cs_ang_t=cs_ang_t,
            cs_vel_t=cs_vel_t,
            f_shape=(),
        )

hg_to_zeta_b

hg_to_zeta_b(
    hg_n: Array, cs_ang_n: dict[str, Array]
) -> ArrayList

Convert beam global grid coordinates to aerodynamic global grid coordinates.

Parameters:

Name	Type	Description	Default
hg_n	Array	Beam global grid coordinates at time step n, (n_nodes, 4, 4).	required
cs_ang_n	dict[str, Array]	Control surface angles as {surface_name: angle} pairs at time step n.	required

Returns:

Type	Description
ArrayList	Full aerodynamic global grid coordinates for each surface, (n_surf, )(zeta_m, zeta_n, 3).

Source code in src/flapjax/aero/uvlm.py

def hg_to_zeta_b(self, hg_n: Array, cs_ang_n: dict[str, Array]) -> ArrayList:
    r"""
    Convert beam global grid coordinates to aerodynamic global grid coordinates.
    :param hg_n: Beam global grid coordinates at time step n, ``(n_nodes, 4, 4)``.
    :param cs_ang_n: Control surface angles as {surface_name: angle} pairs at time step n.
    :return: Full aerodynamic global grid coordinates for each surface, ``(n_surf, )(zeta_m, zeta_n, 3)``.
    """

    zeta_b0_cs = self.grid_func(
        self.zeta_b0, **cs_ang_n
    )  # get local aerodynamic grid for control surface deflections.

    zetas = ArrayList([])
    for i_surf in range(self.n_surf):
        this_hg = jnp.take(
            hg_n, self.dof_mapping[i_surf], axis=0
        )  # (n_nodes, 4, 4)

        zetas.append(
            vmap(vmap(se3_vect_product, (None, 0), 0), (0, 1), 1)(
                this_hg, zeta_b0_cs[i_surf]
            )
        )
    return zetas

hg_dot_to_zeta_b_dot

hg_dot_to_zeta_b_dot(
    hg_n: Array,
    hg_dot_n: Array,
    cs_ang_n: dict[str, Array],
    cs_vel_n: dict[str, Array],
) -> ArrayList

Convert beam global grid velocities to aerodynamic global grid velocities.

Parameters:

Name	Type	Description	Default
hg_n	Array	Beam global grid coordinates, (n_nodes, 4, 4).	required
hg_dot_n	Array	Beam global grid velocities, (n_nodes, 4, 4).	required
cs_ang_n	dict[str, Array]	Control surface angles as {surface_name: angle} pairs.	required
cs_vel_n	dict[str, Array]	Control surface velocities as {surface_name: velocities} pairs.	required

Returns:

Type	Description
ArrayList	Full aerodynamic global grid velocities for each surface, (n_surf, )(zeta_m, zeta_n, 3).

Source code in src/flapjax/aero/uvlm.py

def hg_dot_to_zeta_b_dot(
    self,
    hg_n: Array,
    hg_dot_n: Array,
    cs_ang_n: dict[str, Array],
    cs_vel_n: dict[str, Array],
) -> ArrayList:
    r"""
    Convert beam global grid velocities to aerodynamic global grid velocities.
    :param hg_n: Beam global grid coordinates, ``(n_nodes, 4, 4)``.
    :param hg_dot_n: Beam global grid velocities, ``(n_nodes, 4, 4)``.
    :param cs_ang_n: Control surface angles as {surface_name: angle} pairs.
    :param cs_vel_n: Control surface velocities as {surface_name: velocities} pairs.
    :return: Full aerodynamic global grid velocities for each surface, ``(n_surf, )(zeta_m, zeta_n, 3)``.
    """
    zeta_b0_cs = self.grid_func(
        self.zeta_b0, **cs_ang_n
    )  # deflected local aerodynamic grid

    # as this is where the control velocities are used, they are checked here
    for key in set(cs_ang_n.keys()) | set(cs_vel_n.keys()):
        if key not in cs_vel_n or key not in cs_ang_n:
            raise ValueError(
                f"Missing pair of control angles and velocities for control surface key '{key}'"
            )

        if cs_vel_n[key].shape != cs_vel_n[key].shape:
            raise ValueError(
                f"Mismatched shape for control surface angles {cs_vel_n[key].shape} and velocities {cs_vel_n[key].shape}"
            )

    # use jvp to find the velocity of the aerodynamic grid due to control velocity.
    _, zeta_b0_dot_cs = jax.jvp(
        lambda angs: self.grid_func(
            self.zeta_b0, **angs
        ),  # local grid velocities due to control surface
        primals=(cs_ang_n,),
        tangents=(cs_vel_n,),
    )

    zeta_dots = ArrayList([])
    for i_surf in range(self.n_surf):
        this_hg = jnp.take(
            hg_n, self.dof_mapping[i_surf], axis=0
        )  # (n_nodes, 4, 4)
        this_hg_dot = jnp.take(
            hg_dot_n, self.dof_mapping[i_surf], axis=0
        )  # (n_nodes, 4, 4)
        this_rmat = this_hg[:, :3, :3]  # (n_span, 3, 3)
        zeta_dots.append(
            vmap(vmap(se3_vect_product, (None, 0), 0), (0, 1), 1)(
                this_hg_dot, zeta_b0_cs[i_surf]
            )
            + jnp.einsum(
                "njk,mnk->mnj",
                this_rmat,
                zeta_b0_dot_cs[i_surf],
            )
        )
    return zeta_dots

initialise_wake

initialise_wake(
    zeta_b: ArrayList | None = None,
) -> ArrayList

Generate initial wake grid coordinates, based on the bound grid coordinates and the freestream conditions.

Parameters:

Name	Type	Description	Default
zeta_b	ArrayList | None	Initial wake grid coordinates, (n_surf, )(zeta_m, zeta_n, 3). If None, this will use the initialised bound grid coordinates based on hg0.	None

Returns:

Type	Description
ArrayList	Initial wake grid coordinates, (zeta_m_star, zeta_n, 3)

Source code in src/flapjax/aero/uvlm.py

def initialise_wake(self, zeta_b: ArrayList | None = None) -> ArrayList:
    r"""
    Generate initial wake grid coordinates, based on the bound grid coordinates and the freestream conditions.
    :param zeta_b: Initial wake grid coordinates, ``(n_surf, )(zeta_m, zeta_n, 3)``. If None, this will use the
    initialised bound grid coordinates based on hg0.
    :return: Initial wake grid coordinates, ``(zeta_m_star, zeta_n, 3)``
    """
    zeta_b: ArrayList = zeta_b if zeta_b is not None else self.zeta_b_ref

    zeta0_w = ArrayList([])
    for i_surf, this_delta_w in enumerate(self.delta_w):
        # get bound grid coordinates
        zeta_te = zeta_b[i_surf][-1, :, :]  # (n+1, 3)

        # set wake grid coordinates as trailing edge + displacement
        if this_delta_w is None:
            this_delta_w = (
                jnp.ones(self.grid_disc[i_surf].m_star)
                * self.dt
                * self.flowfield.u_inf_mag
            )
        grid_s = jnp.concatenate((jnp.zeros(1), jnp.cumsum(this_delta_w)))

        zeta0_w.append(
            zeta_te[None, :, :]
            + jnp.outer(grid_s, self.flowfield.u_inf_dir)[:, None, :]
        )
    return zeta0_w

compute_gamma_dot staticmethod

compute_gamma_dot(
    gamma_b_n: ArrayList,
    gamma_b_nm1: ArrayList,
    gamma_b_dot_nm1: ArrayList,
    dt: Array,
    gamma_dot_relaxation: float | Array,
) -> ArrayList

Calculate time derivative of bound circulation strengths at specified time step using finite difference.

Parameters:

Name	Type	Description	Default
gamma_b_n	ArrayList	Bound circulation strengths at timestep n, (n_surf, )(m, n).	required
gamma_b_nm1	ArrayList	Bound circulation strengths at timestep n-1, (n_surf, )(m, n).	required
gamma_b_dot_nm1	ArrayList	Filtered bound circulation strengths time derivative at timestep n-1, (n_surf, )(m, n).	required
dt	Array	Time step length.	required
gamma_dot_relaxation	float | Array	Relaxation factor which filters the time derivative.	required

Source code in src/flapjax/aero/uvlm.py

@staticmethod
def compute_gamma_dot(
    gamma_b_n: ArrayList,
    gamma_b_nm1: ArrayList,
    gamma_b_dot_nm1: ArrayList,
    dt: Array,
    gamma_dot_relaxation: float | Array,
) -> ArrayList:
    r"""
    Calculate time derivative of bound circulation strengths at specified time step using finite difference.
    :param gamma_b_n: Bound circulation strengths at timestep n, ``(n_surf, )(m, n)``.
    :param gamma_b_nm1: Bound circulation strengths at timestep n-1, ``(n_surf, )(m, n)``.
    :param gamma_b_dot_nm1: Filtered bound circulation strengths time derivative at timestep n-1, ``(n_surf, )(m, n)``.
    :param dt: Time step length.
    :param gamma_dot_relaxation: Relaxation factor which filters the time derivative.
    """

    # first obtain the current unfiltered, and previous filtered values for gamma_dot
    gamma_b_dot_curr = (gamma_b_n - gamma_b_nm1) / dt

    # blend with relaxation parameter
    return gamma_b_dot_curr * gamma_dot_relaxation + gamma_b_dot_nm1 * (
        1.0 - gamma_dot_relaxation
    )

set_gamma_w

set_gamma_w(
    gamma_vec: Array, case: AeroCase, i_ts: int
) -> None

Set wake circulation strengths from total circulation strengths at specified time step. Can be passed either a full vector of strengths, or a sequence of strengths per surface.

Parameters:

Name	Type	Description	Default
case	AeroCase	AeroCase object.	required
gamma_vec	Array	Total circulation strengths vector, (gamma_w_tot, ).	required
i_ts	int	Timestep index.	required

Source code in src/flapjax/aero/uvlm.py

@singledispatchmethod
def set_gamma_w(self, gamma_vec: Array, case: AeroCase, i_ts: int) -> None:
    r"""
    Set wake circulation strengths from total circulation strengths at specified time step. Can be passed either a
    full vector of strengths, or a sequence of strengths per surface.
    :param case: AeroCase object.
    :param gamma_vec: Total circulation strengths vector, ``(gamma_w_tot, )``.
    :param i_ts: Timestep index.
    """
    for i_surf in range(self.n_surf):
        case.gamma_w[i_surf] = (
            case.gamma_w[i_surf]
            .at[i_ts, ...]
            .set(
                gamma_vec[self.gamma_w_slice[i_surf]].reshape(
                    self.grid_disc[i_surf].m_star, self.grid_disc[i_surf].n
                )
            )
        )

base_solve

base_solve(
    q_nm1: AeroFullStates | None,
    t_n: Array,
    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,
) -> tuple[
    ArrayList,
    ArrayList,
    ArrayList,
    ArrayList,
    ArrayList | None,
    ArrayList,
    ArrayList,
    ArrayList | None,
    ArrayList,
    ArrayList | None,
    ArrayList,
    ArrayList,
    ArrayList,
    ArrayList,
]

Solve the UVLM equations for a single time step from beam coordinate inputs.

Parameters:

Name	Type	Description	Default
q_nm1	AeroFullStates | None	Minimal aerodynamic states from timestep n-1.	required
t_n	Array	Time at timestep n.	required
hg_n	Array | None	Beam global grid coordinates at time step n, (n_nodes, 4, 4).	required
hg_nm1	Array | None	Beam global grid coordinates at time step n - 1, (n_nodes, 4, 4). Required for consistent free wake modelling.	required
hg_dot_n	Array | None	Beam global grid velocities, (n_nodes, 4, 4).	required
static	bool	If True, perform a static solve.	required
horseshoe	bool	If True, replace the wake with a horseshoe wake in static solve which extends a fixed distance.	required
cs_ang_n	dict[str, Array]	Control surface angle at timestep n, {name, ()}.	required
cs_ang_nm1	dict[str, Array] | None	Control surface angle at timestep n - 1, {name, ()}.	required
cs_vel_n	dict[str, Array] | None	Control surface velocity at timestep n, {name, ()}.	required

Returns:

Type	Description
tuple[ArrayList, ArrayList, ArrayList, ArrayList, ArrayList | None, ArrayList, ArrayList, ArrayList | None, ArrayList, ArrayList | None, ArrayList, ArrayList, ArrayList, ArrayList]	Collocation points, bound normals, bound circulation, wake circulation, bound circulation time derivative, bound grid, wake grid, bound grid time derivative, steady forcing, unsteady forcing, per-strip effective angle of attack, and per-strip lift, drag and moment coefficients.

Source code in src/flapjax/aero/uvlm.py

def base_solve(
    self,
    q_nm1: AeroFullStates | None,
    t_n: Array,
    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,
) -> tuple[
    ArrayList,
    ArrayList,
    ArrayList,
    ArrayList,
    ArrayList | None,
    ArrayList,
    ArrayList,
    ArrayList | None,
    ArrayList,
    ArrayList | None,
    ArrayList,
    ArrayList,
    ArrayList,
    ArrayList,
]:
    r"""
    Solve the UVLM equations for a single time step from beam coordinate inputs.
    :param q_nm1: Minimal aerodynamic states from timestep n-1.
    :param t_n: Time at timestep n.
    :param hg_n: Beam global grid coordinates at time step n, ``(n_nodes, 4, 4)``.
    :param hg_nm1: Beam global grid coordinates at time step n - 1, ``(n_nodes, 4, 4)``. Required for consistent free
    wake modelling.
    :param hg_dot_n: Beam global grid velocities, ``(n_nodes, 4, 4)``.
    :param static: If True, perform a static solve.
    :param horseshoe: If True, replace the wake with a horseshoe wake in static solve which extends a fixed
    distance.
    :param cs_ang_n: Control surface angle at timestep n, {name, ()}.
    :param cs_ang_nm1: Control surface angle at timestep n - 1, {name, ()}.
    :param cs_vel_n: Control surface velocity at timestep n, {name, ()}.
    :return: Collocation points, bound normals, bound circulation, wake circulation, bound circulation time
    derivative, bound grid, wake grid, bound grid time derivative, steady forcing, unsteady forcing, per-strip
    effective angle of attack, and per-strip lift, drag and moment coefficients.
    """

    zeta_b_n = self.hg_to_zeta_b(
        hg_n=hg_n if hg_n is not None else self.hg_ref, cs_ang_n=cs_ang_n
    )

    if hg_dot_n is None:
        zeta_b_dot_n: ArrayList | None = None
        zeta_b_nm1: ArrayList | None = None
    else:
        assert (
            hg_n is not None
            and cs_vel_n is not None
            and hg_nm1 is not None
            and cs_ang_nm1 is not None
        )
        zeta_b_dot_n = self.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
        )
        zeta_b_nm1 = self.hg_to_zeta_b(hg_n=hg_nm1, cs_ang_n=cs_ang_nm1)

    return self.base_solve_from_grid(
        q_nm1=q_nm1,
        t_n=t_n,
        zeta_b_n=zeta_b_n,
        zeta_b_nm1=zeta_b_nm1,
        zeta_b_dot_n=zeta_b_dot_n,
        static=static,
        horseshoe=horseshoe,
    )

base_solve_from_grid

base_solve_from_grid(
    q_nm1: AeroFullStates | None,
    t_n: Array,
    zeta_b_n: ArrayList,
    zeta_b_nm1: ArrayList | None,
    zeta_b_dot_n: ArrayList | None,
    static: bool,
    horseshoe: bool,
    *,
    linearise_variable_wake: bool = False,
    nu_b: ArrayList | None = None,
    nu_w: ArrayList | None = None,
) -> tuple[
    ArrayList,
    ArrayList,
    ArrayList,
    ArrayList,
    ArrayList | None,
    ArrayList,
    ArrayList,
    ArrayList | None,
    ArrayList,
    ArrayList | None,
    ArrayList,
    ArrayList,
    ArrayList,
    ArrayList,
]

Solve the UVLM equations for a single time step from aerodynamic grid inputs.

Parameters:

Name	Type	Description	Default
q_nm1	AeroFullStates | None	Aerodynamic states carried from timestep n-1. Required for dynamic (static=False) solves.	required
t_n	Array	Time at timestep n.	required
zeta_b_n	ArrayList	Bound aerodynamic grid at timestep n, (n_surf, )(zeta_m, zeta_n, 3).	required
zeta_b_nm1	ArrayList | None	Bound aerodynamic grid at timestep n-1, used to seed the wake-convection velocity for the free-wake case. Required for dynamic solves.	required
zeta_b_dot_n	ArrayList | None	Bound grid velocity at timestep n, or None for a static solve.	required
static	bool	If True, perform a static solve (initialise wake, skip wake propagation and unsteady forcing).	required
horseshoe	bool	If True, replace the wake with a horseshoe wake in the static solve.	required
linearise_variable_wake	bool	If True, block gradients through the arc-length discretisation in wake propagation so it acts as a linear operator when differentiated. Used by the linear system; default False.	False
nu_b	ArrayList | None	Optional additive bound upwash velocity at bound vertices, (n_surf, )(zeta_m, zeta_n, 3). Used by the linear system; ignored (equivalent to zero) if not supplied.	None
nu_w	ArrayList | None	Optional additive wake-convection velocity at wake vertices, (n_surf, )(zeta_m_star, zeta_n, 3). Used by the linear system; ignored if not supplied.	None

Returns:

Type	Description
tuple[ArrayList, ArrayList, ArrayList, ArrayList, ArrayList | None, ArrayList, ArrayList, ArrayList | None, ArrayList, ArrayList | None, ArrayList, ArrayList, ArrayList, ArrayList]	Collocation points, bound normals, bound circulation, wake circulation, bound circulation time derivative, bound grid, wake grid, bound grid velocity, steady forcing, unsteady forcing, per-strip effective angle of attack, and per-strip lift, drag and moment coefficients sampled from the polars.

Source code in src/flapjax/aero/uvlm.py

def base_solve_from_grid(
    self,
    q_nm1: AeroFullStates | None,
    t_n: Array,
    zeta_b_n: ArrayList,
    zeta_b_nm1: ArrayList | None,
    zeta_b_dot_n: ArrayList | None,
    static: bool,
    horseshoe: bool,
    *,
    linearise_variable_wake: bool = False,
    nu_b: ArrayList | None = None,
    nu_w: ArrayList | None = None,
) -> tuple[
    ArrayList,
    ArrayList,
    ArrayList,
    ArrayList,
    ArrayList | None,
    ArrayList,
    ArrayList,
    ArrayList | None,
    ArrayList,
    ArrayList | None,
    ArrayList,
    ArrayList,
    ArrayList,
    ArrayList,
]:
    r"""
    Solve the UVLM equations for a single time step from aerodynamic grid inputs.
    :param q_nm1: Aerodynamic states carried from timestep n-1. Required for dynamic (``static=False``) solves.
    :param t_n: Time at timestep n.
    :param zeta_b_n: Bound aerodynamic grid at timestep n, ``(n_surf, )(zeta_m, zeta_n, 3)``.
    :param zeta_b_nm1: Bound aerodynamic grid at timestep n-1, used to seed the wake-convection velocity for the
    free-wake case. Required for dynamic solves.
    :param zeta_b_dot_n: Bound grid velocity at timestep n, or ``None`` for a static solve.
    :param static: If True, perform a static solve (initialise wake, skip wake propagation and unsteady forcing).
    :param horseshoe: If True, replace the wake with a horseshoe wake in the static solve.
    :param linearise_variable_wake: If True, block gradients through the arc-length discretisation in wake
    propagation so it acts as a linear operator when differentiated. Used by the linear system; default False.
    :param nu_b: Optional additive bound upwash velocity at bound vertices, ``(n_surf, )(zeta_m, zeta_n, 3)``. Used by
    the linear system; ignored (equivalent to zero) if not supplied.
    :param nu_w: Optional additive wake-convection velocity at wake vertices, ``(n_surf, )(zeta_m_star, zeta_n, 3)``.
    Used by the linear system; ignored if not supplied.
    :return: Collocation points, bound normals, bound circulation, wake circulation, bound circulation time
    derivative, bound grid, wake grid, bound grid velocity, steady forcing, unsteady forcing, per-strip effective
    angle of attack, and per-strip lift, drag and moment coefficients sampled from the polars.
    """
    if isinstance(
        self.gamma_dot_relaxation,
        (int, float),  # prevents evaluating if value is traced
    ) and not (0.0 < self.gamma_dot_relaxation <= 1.0):
        warn("Gamma_dot relaxation factor not in (0, 1]")

    if not static and horseshoe:
        warn(
            "Horseshoe wake not compatible with non-static solve. Overriding horseshoe to False."
        )
        horseshoe = False

    if not static and q_nm1 is None:
        raise ValueError("q_nm1 needs to be specified for dynamic solve")

    c_n = compute_c(zetas=zeta_b_n)
    nc_n = compute_nc(zetas=zeta_b_n)

    # Prandtl-Glauert compressibility transform: components parallel to the freestream are unchanged, components
    # perpendicular to the freestream are scaled into new coordinates denoted with bar
    # physical circulation is recovered afterwardss as gamma = gamma_bar / beta**2
    x_hat = self.flowfield.u_inf_dir
    beta = self.flowfield.beta

    def _pg(z: Array) -> Array:
        return prandtl_glauert_transform(z, x_hat, beta)

    zeta_b_bar_n = ArrayList([_pg(z) for z in zeta_b_n])
    c_n_bar = compute_c(zetas=zeta_b_bar_n)
    nc_n_bar = compute_nc(zetas=zeta_b_bar_n)
    mirror_point_bar = (
        _pg(self.mirror_point) if self.mirror_point is not None else None
    )

    if zeta_b_dot_n is None:
        c_dot_n_bar: ArrayList | None = None
    else:
        c_dot_n = ArrayList(
            [neighbour_average(zeta_dot, axes=(0, 1)) for zeta_dot in zeta_b_dot_n]
        )
        c_dot_n_bar = ArrayList([_pg(cd) for cd in c_dot_n])

    if static:
        if horseshoe:
            zeta_w_n = ArrayList(
                [
                    self._make_surf_horseshoe_wake(
                        zeta_b=zeta_b_n[i_surf],
                        i_surf=i_surf,
                        horseshoe_length=HORSESHOE_LENGTH,
                    )
                    for i_surf in range(self.n_surf)
                ]
            )
        else:
            zeta_w_n = self.initialise_wake(zeta_b_n)

        gamma_w_n = None  # allocate later from gamma_b
        gamma_w_bar_n: ArrayList | None = None
        zeta_w_bar_n: ArrayList | None = ArrayList([_pg(z) for z in zeta_w_n])
    else:
        assert q_nm1 is not None and zeta_b_nm1 is not None

        zeta_full = ArrayList([*zeta_b_nm1, *q_nm1.zeta_w])
        gamma_full = ArrayList([*q_nm1.gamma_b, *q_nm1.gamma_w])

        def v_wake_prop(x_: Array) -> Array:
            v = self.flowfield.vmap_call(x=x_, t=t_n)
            if self.free_wake:
                v += compute_v_ind(
                    cs=x_,
                    zetas=zeta_full,
                    gammas=gamma_full,
                    kernels=[*self.kernels_b, *self.kernels_w],
                    batch_size=self.batch_size,
                    mirror_normal=self.mirror_normal,
                    mirror_point=self.mirror_point,
                )
            return v

        zeta_w_n, gamma_w_n = propagate_wake(
            gamma_b_nm1=q_nm1.gamma_b,
            gamma_w_nm1=q_nm1.gamma_w,
            zeta_b_n=zeta_b_n,
            zeta_w_nm1=q_nm1.zeta_w,
            delta_w=self.delta_w,
            v_func=v_wake_prop,
            dt=self.dt,
            frozen_wake=False,
            linearise_variable_wake=linearise_variable_wake,
        )

        # for the linearised case, add wake upwash from input
        if nu_w is not None:
            zeta_w_n = ArrayList(
                [zw + nub * self.dt for zw, nub in zip(zeta_w_n, nu_w)]
            )

        zeta_w_bar_n = ArrayList([_pg(z) for z in zeta_w_n])
        gamma_w_bar_n = ArrayList([beta**2 * gw for gw in gamma_w_n])

    aic_solve = compute_aic_solve(
        cs=c_n_bar,
        ns=nc_n_bar,
        zetas_b=zeta_b_bar_n,
        zetas_w=zeta_w_bar_n if static else None,
        kernels_b=self.kernels_b,
        kernels_w=self.kernels_w if static else None,
        batch_size=self.batch_size,
        mirror_normal=self.mirror_normal,
        mirror_point=mirror_point_bar,
    )

    # sampled at the physical collocation points, not transformed coordinates
    v_bc_n = self.flowfield.surf_vmap_call(xs=c_n, t=t_n)  # (n_surf, )(m, n, 3)

    if not static:
        assert c_dot_n_bar is not None
        v_bc_n -= c_dot_n_bar

        if zeta_w_bar_n is None or gamma_w_bar_n is None:
            raise ValueError("zeta_w_nm1 and gamma_w_nm1 are None")

        v_bc_n += compute_v_ind(
            cs=c_n_bar,
            zetas=zeta_w_bar_n,
            gammas=gamma_w_bar_n,
            kernels=self.kernels_w,
            batch_size=self.batch_size,
            mirror_normal=self.mirror_normal,
            mirror_point=mirror_point_bar,
        )

    # for linearised case, add bound grid upwash
    if nu_b is not None:
        v_bc_n += compute_c(ArrayList([_pg(nb) for nb in nu_b]))

    v_bc_n = ArrayList.einsum("ijk,ijk->ij", v_bc_n, nc_n_bar)  # (c_tot, )

    gamma_b_bar_vec_n = jnp.linalg.solve(aic_solve, -v_bc_n.ravel())
    gamma_b_n = ArrayList(
        [g / beta**2 for g in self._vec_to_gamma_b_list(gamma_b_bar_vec_n)]
    )

    def _static_wake_from_gamma_b(gamma_b: ArrayList) -> ArrayList:
        return ArrayList(
            [
                jnp.broadcast_to(
                    gb[[-1], ...],
                    shape=(
                        1
                        if (horseshoe and gd.m_star != 0)
                        else gd.m_star,  # m_star of 0 will override horseshoe
                        gd.n,
                    ),
                )
                for gb, gd in zip(gamma_b, self.grid_disc)
            ]
        )

    if static:
        gamma_w_n = _static_wake_from_gamma_b(gamma_b_n)

    assert gamma_w_n is not None
    assert zeta_w_n is not None

    zeta_b_dot_for_forces = (
        zeta_b_dot_n
        if zeta_b_dot_n is not None
        else ArrayList([jnp.zeros_like(zb) for zb in zeta_b_n])
    )

    def v_total_func(x_: Array) -> Array:
        assert gamma_w_n is not None
        return self.flowfield.vmap_call(x=x_, t=t_n) + compute_v_ind(
            cs=x_,
            zetas=ArrayList([*zeta_b_n, *zeta_w_n]),
            gammas=ArrayList([*gamma_b_n, *gamma_w_n]),
            kernels=[*self.kernels_b, *self.kernels_w],
            batch_size=self.batch_size,
            mirror_normal=self.mirror_normal,
            mirror_point=self.mirror_point,
        )

    f_steady = compute_steady_forcing(
        zeta_b=zeta_b_n,
        zeta_dot_b=zeta_b_dot_for_forces,
        gamma_b=gamma_b_n,
        gamma_w=gamma_w_n,
        rho=self.flowfield.rho,
        v_func=v_total_func,
        v_inputs=nu_b,
        mirror_point=self.mirror_point,
        mirror_normal=self.mirror_normal,
        mirror_edge_low=self.mirror_edge_low,
        mirror_edge_high=self.mirror_edge_high,
    )

    def v_freestream_func(x: Array) -> Array:
        return self.flowfield.vmap_call(x=x, t=t_n)

    alpha_n = strip_alpha(
        zeta_b=zeta_b_n,
        f_steady=f_steady,
        v_func=v_freestream_func,
        rho=self.flowfield.rho,
        beta=beta,
    )

    # update forcing with any polar corrections; surfaces with a ``None`` database report the Prandtl-Glauert
    # corrected flat-plate (2 pi alpha / beta) values implied by the UVLM itself
    f_steady, lift_scale, cl_n, cd_n, cm_n = apply_polar_correction(
        zeta_b=zeta_b_n,
        f_steady=f_steady,
        v_func=v_freestream_func,
        rho=self.flowfield.rho,
        beta=beta,
        polar_data=self.polar_data,
        polar_function=self.polar_function,
        alpha=alpha_n,
    )

    if self.polar_circulation_scale != 0.0:
        # blend bound (and, in the static case, wake) circulation towards the polar-corrected lift
        # scale=0 leaves gamma unchanged; scale=1 fully rescales it to lift_scale
        s = self.polar_circulation_scale
        gamma_b_n = ArrayList(
            [
                gb * (1.0 + s * (ls[None, :] - 1.0))
                for gb, ls in zip(gamma_b_n, lift_scale)
            ]
        )
        if static:
            gamma_w_n = _static_wake_from_gamma_b(gamma_b_n)

    if static:
        gamma_b_dot_n: ArrayList | None = None
        f_unsteady: ArrayList | None = None
    else:
        if q_nm1 is None:
            raise ValueError("q_nm1 needs to be specified for dynamic solve")

        gamma_b_dot_n = self.compute_gamma_dot(
            gamma_b_n=gamma_b_n,
            gamma_b_nm1=q_nm1.gamma_b,
            gamma_b_dot_nm1=q_nm1.gamma_b_dot,
            dt=self.dt,
            gamma_dot_relaxation=self.gamma_dot_relaxation,
        )
        f_unsteady = ArrayList(
            [
                split_to_vertex(
                    self.flowfield.rho
                    * gamma_b_dot_n[i_surf][..., None]
                    * nc_n[i_surf],
                    (0, 1),
                )
                for i_surf in range(self.n_surf)
            ]
        )

    return (
        c_n,
        nc_n,
        gamma_b_n,
        gamma_w_n,
        gamma_b_dot_n,
        zeta_b_n,
        zeta_w_n,
        zeta_b_dot_n,
        f_steady,
        f_unsteady,
        alpha_n,
        cl_n,
        cd_n,
        cm_n,
    )

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

Solve the UVLM equations for a single time step. Can be used for both static and dynamic solves. The solution is updated in-place in the case object.

Parameters:

Name	Type	Description	Default
case	AeroCase	Solution object.	required
i_ts	int	Timestep index to solve for.	required
hg_n	Array | None	Beam global grid coordinates at time step n, (zeta_n, 4, 4).	required
hg_nm1	Array | None	Beam global grid coordinates at time step n-1, (zeta_n, 4, 4).	required
hg_dot_n	Array | None	Beam global grid velocities at time step n, (zeta_n, 4, 4).	required
static	bool	If true, perform a static solve.	required
horseshoe	bool	If true, replace the wake with a static_horseshoe wake in static solve which extends a fixed distance.	required
cs_ang_n	dict[str, Array]	Control surface angle at timestep n, {name, ()}.	required
cs_ang_nm1	dict[str, Array] | None	Control surface angle at timestep n - 1, {name, ()}.	required
cs_vel_n	dict[str, Array] | None	Control surface velocity at timestep n, {name, ()}.	required

Returns:

Type	Description
AeroCase	Solution object with data for current time step added.

Source code in src/flapjax/aero/uvlm.py

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"""
    Solve the UVLM equations for a single time step. Can be used for both static and dynamic solves. The solution
    is updated in-place in the case object.
    :param case: Solution object.
    :param i_ts: Timestep index to solve for.
    :param hg_n: Beam global grid coordinates at time step n, ``(zeta_n, 4, 4)``.
    :param hg_nm1: Beam global grid coordinates at time step n-1, ``(zeta_n, 4, 4)``.
    :param hg_dot_n: Beam global grid velocities at time step n, ``(zeta_n, 4, 4)``.
    :param static: If true, perform a static solve.
    :param horseshoe: If true, replace the wake with a static_horseshoe wake in static solve which extends a fixed
    distance.
    :param cs_ang_n: Control surface angle at timestep n, {name, ()}.
    :param cs_ang_nm1: Control surface angle at timestep n - 1, {name, ()}.
    :param cs_vel_n: Control surface velocity at timestep n, {name, ()}.
    :return: Solution object with data for current time step added.
    """

    assert case.gamma_b_dot is not None and case.zeta_w is not None

    q_nm1 = AeroFullStates(
        gamma_b=case.gamma_b.index_all(i_ts - 1, ...),
        gamma_w=case.gamma_w.index_all(i_ts - 1, ...),
        gamma_b_dot=case.gamma_b_dot.index_all(i_ts - 1, ...),
        zeta_w=case.zeta_w.index_all(i_ts - 1, ...),
    )

    if not static:
        case.t = case.t.at[i_ts].set(
            jax.lax.select(i_ts, case.t[i_ts - 1] + self.dt, 0.0)
        )

    (
        c_n,
        nc_n,
        gamma_b_n,
        gamma_w_n,
        gamma_b_dot_n,
        zeta_b_n,
        zeta_w_n,
        zeta_b_dot_n,
        f_steady,
        f_unsteady,
        alpha_n,
        cl_n,
        cd_n,
        cm_n,
    ) = self.base_solve(
        q_nm1=q_nm1,
        t_n=case.t[i_ts, ...],
        hg_n=hg_n,
        hg_nm1=hg_nm1,
        hg_dot_n=hg_dot_n,
        static=static,
        horseshoe=horseshoe,
        cs_ang_n=cs_ang_n,
        cs_ang_nm1=cs_ang_nm1,
        cs_vel_n=cs_vel_n,
    )

    case.set_arraylist_at_ts("c", values=c_n, i_ts=i_ts)
    case.set_arraylist_at_ts("nc", values=nc_n, i_ts=i_ts)
    case.set_arraylist_at_ts("gamma_b", values=gamma_b_n, i_ts=i_ts)
    case.set_arraylist_at_ts("gamma_w", values=gamma_w_n, i_ts=i_ts)
    case.set_arraylist_at_ts("zeta_b", values=zeta_b_n, i_ts=i_ts)
    case.set_arraylist_at_ts("f_steady", values=f_steady, i_ts=i_ts)
    case.set_arraylist_at_ts("alpha", values=alpha_n, i_ts=i_ts)
    case.set_arraylist_at_ts("cl", values=cl_n, i_ts=i_ts)
    case.set_arraylist_at_ts("cd", values=cd_n, i_ts=i_ts)
    case.set_arraylist_at_ts("cm", values=cm_n, i_ts=i_ts)

    if not static:
        if gamma_b_dot_n is None:
            raise ValueError("gamma_b_dot_n is None")
        if zeta_b_dot_n is None:
            raise ValueError("zeta_b_dot_n is None")
        if f_unsteady is None:
            raise ValueError("f_unsteady is None")
        case.set_arraylist_at_ts("gamma_b_dot", values=gamma_b_dot_n, i_ts=i_ts)
        case.set_arraylist_at_ts("zeta_b_dot", values=zeta_b_dot_n, i_ts=i_ts)
        case.set_arraylist_at_ts("f_unsteady", values=f_unsteady, i_ts=i_ts)

    # set wake grid coordinates. If using static_horseshoe, it will still create a regular wake for plotting
    if horseshoe:
        case.set_arraylist_at_ts(
            "zeta_w", values=self.initialise_wake(zeta_w_n), i_ts=i_ts
        )
    else:
        case.set_arraylist_at_ts("zeta_w", values=zeta_w_n, i_ts=i_ts)

    return case

initialise_case_object

initialise_case_object(
    n_tstep: int,
    static_horseshoe: bool,
    free_wake: bool,
    gamma_dot_relaxation: float | Array,
    cs_ang_t: dict[str, Array],
    cs_vel_t: dict[str, Array],
) -> AeroCase

Initialise an AeroCase object to store the solution of the aerodynamic case for a given number of time steps. All solution data is initialised to zero.

Parameters:

Name	Type	Description	Default
n_tstep	int	Number of time steps to solve for in dynamic solution.	required
static_horseshoe	bool	Whether a horseshoe formulation was used for the static case.	required
free_wake	bool	Whether to use a free wake formulation.	required
gamma_dot_relaxation	float | Array	Relaxation factor for damping gamma_dot.	required
cs_ang_t	dict[str, Array]	Control surface angle time history, {name: (n_tstep, )}.	required
cs_vel_t	dict[str, Array]	Control surface velocity time history, {name: (n_tstep, )}.	required

Returns:

Type	Description
AeroCase	AeroCase object.

Source code in src/flapjax/aero/uvlm.py

def initialise_case_object(
    self,
    n_tstep: int,
    static_horseshoe: bool,
    free_wake: bool,
    gamma_dot_relaxation: float | Array,
    cs_ang_t: dict[str, Array],
    cs_vel_t: dict[str, Array],
) -> AeroCase:
    r"""
    Initialise an AeroCase object to store the solution of the aerodynamic case for a given number of time
    steps. All solution data is initialised to zero.
    :param n_tstep: Number of time steps to solve for in dynamic solution.
    :param static_horseshoe: Whether a horseshoe formulation was used for the static case.
    :param free_wake: Whether to use a free wake formulation.
    :param gamma_dot_relaxation: Relaxation factor for damping gamma_dot.
    :param cs_ang_t: Control surface angle time history, ``{name: (n_tstep, )}``.
    :param cs_vel_t: Control surface velocity time history, ``{name: (n_tstep, )}``.
    :return: AeroCase object.
    """
    # zero initialise
    return AeroCase(
        zeta_b=ArrayList(
            [jnp.zeros((n_tstep, gd.m + 1, gd.n + 1, 3)) for gd in self.grid_disc]
        ),
        zeta_b_dot=ArrayList(
            [jnp.zeros((n_tstep, gd.m + 1, gd.n + 1, 3)) for gd in self.grid_disc]
        ),
        zeta_w=ArrayList(
            [
                jnp.zeros((n_tstep, gd.m_star + 1, gd.n + 1, 3))
                for gd in self.grid_disc
            ]
        ),
        gamma_b=ArrayList(
            [jnp.zeros((n_tstep, gd.m, gd.n)) for gd in self.grid_disc]
        ),
        gamma_b_dot=ArrayList(
            [jnp.zeros((n_tstep, gd.m, gd.n)) for gd in self.grid_disc]
        ),
        gamma_w=ArrayList(
            [jnp.zeros((n_tstep, gd.m_star, gd.n)) for gd in self.grid_disc]
        ),
        f_steady=ArrayList(
            [jnp.zeros((n_tstep, gd.m + 1, gd.n + 1, 3)) for gd in self.grid_disc]
        ),
        f_unsteady=ArrayList(
            [jnp.zeros((n_tstep, gd.m + 1, gd.n + 1, 3)) for gd in self.grid_disc]
        ),
        alpha=ArrayList([jnp.zeros((n_tstep, gd.n)) for gd in self.grid_disc]),
        cl=ArrayList([jnp.zeros((n_tstep, gd.n)) for gd in self.grid_disc]),
        cd=ArrayList([jnp.zeros((n_tstep, gd.n)) for gd in self.grid_disc]),
        cm=ArrayList([jnp.zeros((n_tstep, gd.n)) for gd in self.grid_disc]),
        c=ArrayList([jnp.zeros((n_tstep, gd.m, gd.n, 3)) for gd in self.grid_disc]),
        n=ArrayList([jnp.zeros((n_tstep, gd.m, gd.n, 3)) for gd in self.grid_disc]),
        kernels=[*self.kernels_b, *self.kernels_w],
        mirror_point=self.mirror_point,
        mirror_normal=self.mirror_normal,
        mirror_edge_low=self.mirror_edge_low,
        mirror_edge_high=self.mirror_edge_high,
        flowfield=self.flowfield,
        surf_b_names=self.surf_b_names,
        surf_w_names=self.surf_w_names,
        i_ts=jnp.arange(n_tstep),
        t=jnp.zeros(n_tstep),
        dof_mapping=self.dof_mapping,
        static_horseshoe=static_horseshoe,
        free_wake=free_wake,
        gamma_dot_relaxation=gamma_dot_relaxation,
        cs_ang=cs_ang_t,
        cs_vel=cs_vel_t,
        batch_size=self.batch_size,
    )

static_solve

static_solve(
    hg: Array | None = None,
    t: Array | float = 0.0,
    horseshoe: bool = False,
    cs_ang: dict[str, Array] | None = None,
) -> AeroCase

Solve the VLM for given static beam coordinates. TODO: add free wake

Parameters:

Name	Type	Description	Default
hg	Array | None	Beam coordinates, (zeta_n, 4, 4).	None
t	Array | float	Time at which to solve static solution, used for background flowfield evaluation.	0.0
horseshoe	bool	If true, replace the wake with a horseshoe wake which extends a fixed distance.	False
cs_ang	dict[str, Array] | None	Control surface angles, {name: ()}.	None

Returns:

Type	Description
AeroCase	AeroCase solution object.

Source code in src/flapjax/aero/uvlm.py

def static_solve(
    self,
    hg: Array | None = None,
    t: Array | float = 0.0,
    horseshoe: bool = False,
    cs_ang: dict[str, Array] | None = None,
) -> AeroCase:
    r"""
    Solve the VLM for given static beam coordinates.
    TODO: add free wake
    :param hg: Beam coordinates, ``(zeta_n, 4, 4)``.
    :param t: Time at which to solve static solution, used for background flowfield evaluation.
    :param horseshoe: If true, replace the wake with a horseshoe wake which extends a fixed distance.
    :param cs_ang: Control surface angles, ``{name: ()}``.
    :return: AeroCase solution object.
    """

    case = self.initialise_case_object(
        1,
        static_horseshoe=horseshoe,
        gamma_dot_relaxation=0.7,
        free_wake=False,
        cs_ang_t=cs_ang if cs_ang is not None else self.cs_ang0,
        cs_vel_t={
            k: jnp.zeros_like(v)
            for k, v in (cs_ang if cs_ang is not None else self.cs_ang0).items()
        },
    )
    case.t = case.t.at[0].set(t)

    out_case = self.case_solve(
        case=case,
        i_ts=0,
        hg_n=hg,
        hg_nm1=None,
        hg_dot_n=None,
        static=True,
        horseshoe=horseshoe,
        cs_ang_n=cs_ang if cs_ang is not None else self.cs_ang0,
        cs_ang_nm1=None,
        cs_vel_n=None,
    )[0]

    if horseshoe:
        # if using a horseshoe wake, return the normal wake to prevent continuity errors
        out_case.zeta_w = self.initialise_wake(zeta_b=out_case.zeta_b)

    return out_case

prescribed_dynamic_solve

prescribed_dynamic_solve(
    init_case: AeroCase,
    hg_t: Array,
    hg_dot_t: Array,
    cs_ang_t: dict[str, Array] | None = None,
    cs_vel_t: dict[str, Array] | None = None,
) -> AeroCase

Solve the UVLM for prescribed grid motions.

Parameters:

Name	Type	Description	Default
init_case	AeroCase	StaticAero object containing initial conditions for the solution at time step 0.	required
hg_t	Array	Beam coordinates over time, (n_tstep, n_nodes, 4, 4).	required
hg_dot_t	Array	Beam coordinate time derivative over time, (n_tstep, n_nodes, 4, 4).	required
cs_ang_t	dict[str, Array] | None	Control surface angle time history, {name, (n_tstep, )}.	None
cs_vel_t	dict[str, Array] | None	Control surface velocity time history, {name, (n_tstep, )}.	None

Returns:

Type	Description
AeroCase	AeroCase solution object.

Source code in src/flapjax/aero/uvlm.py

def prescribed_dynamic_solve(
    self,
    init_case: AeroCase,
    hg_t: Array,
    hg_dot_t: Array,
    cs_ang_t: dict[str, Array] | None = None,
    cs_vel_t: dict[str, Array] | None = None,
) -> AeroCase:
    r"""
    Solve the UVLM for prescribed grid motions.
    :param init_case: StaticAero object containing initial conditions for the solution at time step 0.
    :param hg_t: Beam coordinates over time, ``(n_tstep, n_nodes, 4, 4)``.
    :param hg_dot_t: Beam coordinate time derivative over time, ``(n_tstep, n_nodes, 4, 4)``.
    :param cs_ang_t: Control surface angle time history, {name, ``(n_tstep, )``}.
    :param cs_vel_t: Control surface velocity time history, {name, ``(n_tstep, )``}.
    :return: AeroCase solution object.
    """
    check_arr_shape(hg_t, (None, None, 4, 4), "hg_n")
    check_if_all_se3_g(hg_t, True)

    if hg_t.shape != hg_dot_t.shape:
        raise ValueError(
            f"hg_dot_n must have the same shape as hg_n, got {hg_dot_t.shape} vs {hg_t.shape}"
        )

    check_if_all_se3_a(hg_dot_t, True)

    n_tstep = hg_t.shape[0]

    case = init_case.to_dynamic(i_ts=0, n_tstep=n_tstep)

    def _step_func(i_ts_: int, case_: AeroCase) -> AeroCase:
        cs_angle_nm1, cs_angle_n = (
            (
                {k: v[i_ts__] for k, v in cs_ang_t.items()}
                if cs_ang_t is not None
                else {}
            )
            for i_ts__ in (i_ts_ - 1, i_ts_)
        )
        cs_velocity_n = (
            {k: v[i_ts_] for k, v in cs_vel_t.items()}
            if cs_vel_t is not None
            else {}
        )

        case_ = self.case_solve(
            case=case_,
            i_ts=i_ts_,
            hg_n=hg_t[i_ts_, ...],
            hg_nm1=hg_t[i_ts_ - 1, ...],
            hg_dot_n=hg_dot_t[i_ts_, ...],
            static=False,
            horseshoe=False,
            cs_ang_n=cs_angle_n,
            cs_ang_nm1=cs_angle_nm1,
            cs_vel_n=cs_velocity_n,
        )
        jax_print(
            "UVLM timestep {i_ts_}",
            i_ts_=i_ts_,
            verbose_level="normal",
        )
        return case_

    case = fori_loop(
        1,
        n_tstep,
        _step_func,
        init_val=case,
    )
    return case

reference_configuration

reference_configuration() -> AeroCase

Get the reference (initial) snapshot of the aerodynamic case. This will set the timestep as -1.

Returns:

Type	Description
AeroCase	StaticAero object at initial time step.

Source code in src/flapjax/aero/uvlm.py

def reference_configuration(self) -> AeroCase:
    r"""
    Get the reference (initial) snapshot of the aerodynamic case. This will set the timestep as -1.
    :return: StaticAero object at initial time step.
    """
    return AeroCase(
        zeta_b=self.zeta_b_ref,
        zeta_b_dot=ArrayList(
            [jnp.zeros((gd.m + 1, gd.n + 1, 3)) for gd in self.grid_disc]
        ),
        zeta_w=self.zeta_w_ref,
        c=compute_c(self.zeta_b_ref),
        n=compute_nc(self.zeta_b_ref),
        gamma_b=ArrayList([jnp.zeros((gd.m, gd.n)) for gd in self.grid_disc]),
        gamma_b_dot=ArrayList([jnp.zeros((gd.m, gd.n)) for gd in self.grid_disc]),
        gamma_w=ArrayList([jnp.zeros((gd.m_star, gd.n)) for gd in self.grid_disc]),
        f_steady=ArrayList(
            [jnp.zeros((gd.m + 1, gd.n + 1, 3)) for gd in self.grid_disc]
        ),
        f_unsteady=ArrayList(
            [jnp.zeros((gd.m + 1, gd.n + 1, 3)) for gd in self.grid_disc]
        ),
        alpha=ArrayList([jnp.zeros((gd.n,)) for gd in self.grid_disc]),
        cl=ArrayList([jnp.zeros((gd.n,)) for gd in self.grid_disc]),
        cd=ArrayList([jnp.zeros((gd.n,)) for gd in self.grid_disc]),
        cm=ArrayList([jnp.zeros((gd.n,)) for gd in self.grid_disc]),
        surf_b_names=self.surf_b_names,
        surf_w_names=self.surf_w_names,
        i_ts=-1,
        t=jnp.array(0.0),
        dof_mapping=self.dof_mapping,
        flowfield=self.flowfield,
        mirror_point=self.mirror_point,
        mirror_normal=self.mirror_normal,
        mirror_edge_low=self.mirror_edge_low,
        mirror_edge_high=self.mirror_edge_high,
        kernels=[*self.kernels_b, *self.kernels_w],
        static_horseshoe=False,
        gamma_dot_relaxation=0.0,
        free_wake=False,
        cs_ang=self.cs_ang0,
        cs_vel=self.cs_vel0,
        batch_size=self.batch_size,
    )

plot_reference

plot_reference(
    directory: PathLike | str, plot_wake: bool = True
) -> Sequence[Path]

Plot the reference (initial) snapshot of the aerodynamic case. This will set the timestep as -1.

Parameters:

Name	Type	Description	Default
directory	PathLike | str	Path to write files to.	required
plot_wake	bool	If True, plot the wake grid.	True

Source code in src/flapjax/aero/uvlm.py

def plot_reference(
    self, directory: os.PathLike | str, plot_wake: bool = True
) -> Sequence[Path]:
    r"""
    Plot the reference (initial) snapshot of the aerodynamic case. This will set the timestep as -1.
    :param directory: Path to write files to.
    :param plot_wake: If True, plot the wake grid.
    """
    return self.reference_configuration().plot(
        Path(directory).resolve(), plot_wake=plot_wake
    )

gamma_b_res_func

gamma_b_res_func(
    i_ts: int | Array,
    t_n: Array,
    varphi_n: Array,
    v_n: Array,
    gamma_b_n: Array,
    gamma_w_n: Array,
    zeta_w_n: Array,
    dv: AeroelasticDesignVariables,
    dv_full: AeroelasticDesignVariables,
    struct_obj: BeamStructure,
) -> Array

Bound circulation residual used for adjoint computations.

:math:\mathbf{r}_{\Gamma_b} = \left(\boldsymbol{\mathcal{A}}_{b, n} \cdot \mathbf{n}_n\right)^{-1} \left[\left( \boldsymbol{\mathcal{A}}_{w, n} \boldsymbol{\Gamma}_{w, n} +\mathbf{v}_{bc, n} - \dot{\boldsymbol{\zeta}}_c\right) \cdot \mathbf{n}_n\right] + \boldsymbol{\Gamma}_{b, n}

Parameters:

Name	Type	Description	Default
i_ts	int | Array	Time step index.	required
t_n	Array	Time at step n.	required
varphi_n	Array	varphi vector at timestep n.	required
v_n	Array	Beam velocity vector at timestep n.	required
gamma_b_n	Array	Bound circulation vector at timestep n.	required
gamma_w_n	Array	Wake circulation vector at timestep n.	required
zeta_w_n	Array	Wake grid vector at timestep n.	required
dv	AeroelasticDesignVariables	Aeroelastic design variables.	required
dv_full	AeroelasticDesignVariables	Aeroelastic design variables without omissions for the variables where gradients aren't requested.	required
struct_obj	BeamStructure	Beam structure.	required

Returns:

Type	Description
Array	Bound circulation residual.

Source code in src/flapjax/aero/uvlm.py

def gamma_b_res_func(
    self,
    i_ts: int | Array,
    t_n: Array,
    varphi_n: Array,
    v_n: Array,
    gamma_b_n: Array,
    gamma_w_n: Array,
    zeta_w_n: Array,
    dv: AeroelasticDesignVariables,
    dv_full: AeroelasticDesignVariables,
    struct_obj: BeamStructure,
) -> Array:
    r"""
    Bound circulation residual used for adjoint computations.

    :math:`\mathbf{r}_{\Gamma_b} = \left(\boldsymbol{\mathcal{A}}_{b, n} \cdot \mathbf{n}_n\right)^{-1} \left[\left(
    \boldsymbol{\mathcal{A}}_{w, n} \boldsymbol{\Gamma}_{w, n} +\mathbf{v}_{bc, n} - \dot{\boldsymbol{\zeta}}_c\right)
    \cdot \mathbf{n}_n\right] + \boldsymbol{\Gamma}_{b, n}`

    :param i_ts: Time step index.
    :param t_n: Time at step n.
    :param varphi_n: varphi vector at timestep n.
    :param v_n: Beam velocity vector at timestep n.
    :param gamma_b_n: Bound circulation vector at timestep n.
    :param gamma_w_n: Wake circulation vector at timestep n.
    :param zeta_w_n: Wake grid vector at timestep n.
    :param dv: Aeroelastic design variables.
    :param dv_full: Aeroelastic design variables without omissions for the variables where gradients aren't
    requested.
    :param struct_obj: Beam structure.
    :return: Bound circulation residual.
    """

    varphi_n = varphi_n.reshape(-1, 6)
    v_n = v_n.reshape(-1, 6)
    gamma_b_n = ArrayList.from_vector(
        vect=gamma_b_n,
        shapes=ArrayListShape([(gd.m, gd.n) for gd in self.grid_disc]),
    )
    gamma_w_n = ArrayList.from_vector(
        vect=gamma_w_n,
        shapes=ArrayListShape([(gd.m_star, gd.n) for gd in self.grid_disc]),
    )
    zeta_w_n = ArrayList.from_vector(
        vect=zeta_w_n,
        shapes=ArrayListShape(
            [(gd.m_star + 1, gd.n + 1, 3) for gd in self.grid_disc]
        ),
    )

    inner_struct = struct_obj.case_from_dv(dv=dv.structure)
    hg_n = inner_struct.compute_hg_from_varphi(varphi=varphi_n)
    hg_dot_n = inner_struct.make_hg_dot(hg=hg_n, v=v_n)

    inner_case = self.case_from_dv(dv=dv.aero)

    # get control surface deflections from design variables
    cs_ang_n, cs_vel_n = dv.aero.get_cs_n(i_ts=i_ts, dv_full=dv_full.aero)

    zeta_b_n = inner_case.hg_to_zeta_b(hg_n=hg_n, cs_ang_n=cs_ang_n)
    c_n = compute_c(
        zetas=zeta_b_n
    )  # physical collocation points, used for freestream sampling only

    # Prandtl-Glauert compressibility transformforward solve.
    x_hat = inner_case.flowfield.u_inf_dir
    beta = inner_case.flowfield.beta

    def _pg(z: Array) -> Array:
        return prandtl_glauert_transform(z, x_hat, beta)

    zeta_b_bar_n = ArrayList([_pg(z) for z in zeta_b_n])
    c_n_bar = compute_c(zetas=zeta_b_bar_n)
    nc_n_bar = compute_nc(zetas=zeta_b_bar_n)
    mirror_point_bar = (
        _pg(inner_case.mirror_point)
        if inner_case.mirror_point is not None
        else None
    )

    zeta_b_dot_n = inner_case.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
    )

    c_dot_n = ArrayList(
        [neighbour_average(zeta_dot, axes=(0, 1)) for zeta_dot in zeta_b_dot_n]
    )
    c_dot_n_bar = ArrayList([_pg(cd) for cd in c_dot_n])

    zeta_w_bar_n = ArrayList([_pg(z) for z in zeta_w_n])
    gamma_w_bar_n = ArrayList([beta**2 * gw for gw in gamma_w_n])

    aic_solve = compute_aic_solve(
        cs=c_n_bar,
        ns=nc_n_bar,
        zetas_b=zeta_b_bar_n,
        zetas_w=None,
        kernels_b=inner_case.kernels_b,
        kernels_w=None,
        batch_size=self.batch_size,
        mirror_normal=inner_case.mirror_normal,
        mirror_point=mirror_point_bar,
    )

    v_bc_n = inner_case.flowfield.surf_vmap_call(
        xs=c_n, t=t_n
    )  # (n_surf, )(m, n, 3)

    # structural component
    v_bc_n -= c_dot_n_bar

    # find wake component
    v_bc_n += compute_v_ind(
        cs=c_n_bar,
        zetas=zeta_w_bar_n,
        gammas=gamma_w_bar_n,
        kernels=inner_case.kernels_w,
        batch_size=self.batch_size,
        mirror_normal=inner_case.mirror_normal,
        mirror_point=mirror_point_bar,
    )

    v_bc_n = ArrayList.einsum("ijk,ijk->ij", v_bc_n, nc_n_bar)  # (c_tot, )

    gamma_b_bar_vec_n = jnp.linalg.solve(aic_solve, -v_bc_n.ravel())
    gamma_b_nm1_update = ArrayList(
        [g / beta**2 for g in self._vec_to_gamma_b_list(gamma_b_bar_vec_n)]
    )

    return (gamma_b_nm1_update - gamma_b_n).ravel()

wake_prop_res_func

wake_prop_res_func(
    i_ts: int | Array,
    t_n: Array,
    varphi_nm1: Array,
    varphi_n: Array,
    gamma_b_nm1: Array,
    gamma_w_nm1: Array,
    gamma_w_n: Array,
    zeta_w_nm1: Array,
    zeta_w_n: Array,
    dv: AeroelasticDesignVariables,
    dv_full: AeroelasticDesignVariables,
    struct_obj: BeamStructure,
) -> tuple[Array, Array]

Wake propagation residual for both grid coordinates and circulation strengths.

Parameters:

Name	Type	Description	Default
i_ts	int | Array	Time step index.	required
t_n	Array	Time at timestep n.	required
varphi_nm1	Array	Beam minimal coordinates vector at timestep n-1.	required
varphi_n	Array	Beam maximal coordinates vector at timestep n.	required
gamma_b_nm1	Array	Bound circulation vector at timestep n-1.	required
gamma_w_nm1	Array	Wake circulation vector at timestep n-1.	required
gamma_w_n	Array	Wake circulation vector at timestep n.	required
zeta_w_nm1	Array	Wake grid vector at timestep n-1.	required
zeta_w_n	Array	Wake grid vector at timestep n.	required
dv	AeroelasticDesignVariables	Aeroelastic design variables.	required
dv_full	AeroelasticDesignVariables	Aeroelastic design variables without omissions.	required
struct_obj	BeamStructure	Beam structure.	required

Returns:

Type	Description
tuple[Array, Array]	Wake grid and circulation residuals.

Source code in src/flapjax/aero/uvlm.py

def wake_prop_res_func(
    self,
    i_ts: int | Array,
    t_n: Array,
    varphi_nm1: Array,
    varphi_n: Array,
    gamma_b_nm1: Array,
    gamma_w_nm1: Array,
    gamma_w_n: Array,
    zeta_w_nm1: Array,
    zeta_w_n: Array,
    dv: AeroelasticDesignVariables,
    dv_full: AeroelasticDesignVariables,
    struct_obj: BeamStructure,
) -> tuple[Array, Array]:
    r"""
    Wake propagation residual for both grid coordinates and circulation strengths.

    :param i_ts: Time step index.
    :param t_n: Time at timestep n.
    :param varphi_nm1: Beam minimal coordinates vector at timestep n-1.
    :param varphi_n: Beam maximal coordinates vector at timestep n.
    :param gamma_b_nm1: Bound circulation vector at timestep n-1.
    :param gamma_w_nm1: Wake circulation vector at timestep n-1.
    :param gamma_w_n: Wake circulation vector at timestep n.
    :param zeta_w_nm1: Wake grid vector at timestep n-1.
    :param zeta_w_n: Wake grid vector at timestep n.
    :param dv: Aeroelastic design variables.
    :param dv_full: Aeroelastic design variables without omissions.
    :param struct_obj: Beam structure.
    :return: Wake grid and circulation residuals.
    """

    varphi_nm1 = varphi_nm1.reshape(-1, 6)
    varphi_n = varphi_n.reshape(-1, 6)

    gamma_b_nm1 = ArrayList.from_vector(
        vect=gamma_b_nm1,
        shapes=ArrayListShape([(gd.m, gd.n) for gd in self.grid_disc]),
    )

    gamma_w_nm1 = ArrayList.from_vector(
        vect=gamma_w_nm1,
        shapes=ArrayListShape([(gd.m_star, gd.n) for gd in self.grid_disc]),
    )

    gamma_w_n = ArrayList.from_vector(
        vect=gamma_w_n,
        shapes=ArrayListShape([(gd.m_star, gd.n) for gd in self.grid_disc]),
    )
    zeta_w_nm1 = ArrayList.from_vector(
        vect=zeta_w_nm1,
        shapes=ArrayListShape(
            [(gd.m_star + 1, gd.n + 1, 3) for gd in self.grid_disc]
        ),
    )

    zeta_w_n = ArrayList.from_vector(
        vect=zeta_w_n,
        shapes=ArrayListShape(
            [(gd.m_star + 1, gd.n + 1, 3) for gd in self.grid_disc]
        ),
    )

    inner_struct = struct_obj.case_from_dv(dv=dv.structure)
    hg_nm1 = inner_struct.compute_hg_from_varphi(varphi=varphi_nm1)
    hg_n = inner_struct.compute_hg_from_varphi(varphi=varphi_n)

    inner_case = self.case_from_dv(dv=dv.aero)

    # get control surface deflections from design variables
    cs_ang_nm1, _ = dv.aero.get_cs_n(i_ts=i_ts - 1, dv_full=dv_full.aero)
    cs_ang_n, _ = dv.aero.get_cs_n(i_ts=i_ts, dv_full=dv_full.aero)
    zeta_b_nm1 = inner_case.hg_to_zeta_b(hg_n=hg_nm1, cs_ang_n=cs_ang_nm1)
    zeta_b_n = inner_case.hg_to_zeta_b(hg_n=hg_n, cs_ang_n=cs_ang_n)

    def v_wake_prop(x_: Array) -> Array:
        v = inner_case.flowfield.vmap_call(x=x_, t=t_n)
        if self.free_wake:
            v += compute_v_ind(
                cs=x_,
                zetas=ArrayList([*zeta_b_nm1, *zeta_w_nm1]),
                gammas=ArrayList([*gamma_b_nm1, *gamma_w_nm1]),
                kernels=[*inner_case.kernels_b, *inner_case.kernels_w],
                batch_size=self.batch_size,
                mirror_normal=inner_case.mirror_normal,
                mirror_point=inner_case.mirror_point,
            )
        return v

    zeta_w_nm1_update, gamma_w_nm1_update = propagate_wake(
        gamma_b_nm1=gamma_b_nm1,
        gamma_w_nm1=gamma_w_nm1,
        zeta_b_n=zeta_b_n,
        zeta_w_nm1=zeta_w_nm1,
        delta_w=inner_case.delta_w,
        v_func=v_wake_prop,
        dt=inner_case.dt,
        frozen_wake=False,
        linearise_variable_wake=False,
    )

    return (zeta_w_nm1_update - zeta_w_n).ravel(), (
        gamma_w_nm1_update - gamma_w_n
    ).ravel()

gamma_b_dot_res_func

gamma_b_dot_res_func(
    gamma_b_nm1: Array,
    gamma_b_n: Array,
    gamma_b_dot_nm1: Array,
    gamma_b_dot_n: Array,
    dv: AeroelasticDesignVariables,
) -> Array

Bound circulation time derivative residual.

:math:\frac{g}{h} \left[\mathbf{\Gamma}_{b, n} - \mathbf{\Gamma}_{b, n-1}\right] + (1-g) \dot{\mathbf{\Gamma}}_{b, n-1} - \dot{\mathbf{\Gamma}}_{b, n}

Parameters:

Name	Type	Description	Default
gamma_b_nm1	Array	Bound circulation vector at timestep n-1.	required
gamma_b_n	Array	Bound circulation vector at timestep n.	required
gamma_b_dot_nm1	Array	Bound circulation time derivative vector at timestep n-1.	required
gamma_b_dot_n	Array	Bound circulation time derivative vector at timestep n.	required
dv	AeroelasticDesignVariables	Design variables. Whilst this function does not depend upon it, including it simplified obtaining the residual design gradient.	required

Returns:

Type	Description
Array	Bound circulation time derivative residual.

Source code in src/flapjax/aero/uvlm.py

def gamma_b_dot_res_func(
    self,
    gamma_b_nm1: Array,
    gamma_b_n: Array,
    gamma_b_dot_nm1: Array,
    gamma_b_dot_n: Array,
    dv: AeroelasticDesignVariables,
) -> Array:
    r"""
    Bound circulation time derivative residual.

    :math:`\frac{g}{h} \left[\mathbf{\Gamma}_{b, n} - \mathbf{\Gamma}_{b, n-1}\right] + (1-g)
    \dot{\mathbf{\Gamma}}_{b, n-1} - \dot{\mathbf{\Gamma}}_{b, n}`

    :param gamma_b_nm1: Bound circulation vector at timestep n-1.
    :param gamma_b_n: Bound circulation vector at timestep n.
    :param gamma_b_dot_nm1: Bound circulation time derivative vector at timestep n-1.
    :param gamma_b_dot_n: Bound circulation time derivative vector at timestep n.
    :param dv: Design variables. Whilst this function does not depend upon it, including it simplified obtaining
    the residual design gradient.
    :return: Bound circulation time derivative residual.
    """
    del dv  # intentionally unused

    gamma_b_nm1 = ArrayList.from_vector(
        vect=gamma_b_nm1,
        shapes=ArrayListShape([(gd.m, gd.n) for gd in self.grid_disc]),
    )

    gamma_b_n = ArrayList.from_vector(
        vect=gamma_b_n,
        shapes=ArrayListShape([(gd.m, gd.n) for gd in self.grid_disc]),
    )

    gamma_b_dot_nm1 = ArrayList.from_vector(
        vect=gamma_b_dot_nm1,
        shapes=ArrayListShape([(gd.m, gd.n) for gd in self.grid_disc]),
    )

    gamma_b_dot_n = ArrayList.from_vector(
        vect=gamma_b_dot_n,
        shapes=ArrayListShape([(gd.m, gd.n) for gd in self.grid_disc]),
    )

    return (
        self.gamma_dot_relaxation / self.dt * (gamma_b_n - gamma_b_nm1)
        + (1.0 - self.gamma_dot_relaxation) * gamma_b_dot_nm1
        - gamma_b_dot_n
    ).ravel()

f_aero_res_func

f_aero_res_func(
    i_ts: int | Array,
    t_n: Array,
    varphi_n: Array,
    v_n: Array,
    gamma_b_n: Array,
    gamma_w_n: Array,
    gamma_b_dot_n: Array,
    zeta_w_n: Array,
    dv: AeroelasticDesignVariables,
    dv_full: AeroelasticDesignVariables,
    struct_obj: BeamStructure,
    f_aero_beam_n: Array,
    block_grid_gradients: bool,
    solve_dofs: tuple[int, ...],
) -> Array

Aerodynamic forcing residual, compared in the local frame for compatibility with the structure.

:math:\boldsymbol{\mathcal{F}}_{\text{aero}, n}(\mathbf{\Gamma}_{b, n}, \mathbf{\Gamma}_{w, n}, \dot{\mathbf{\Gamma}}_{b, n}, \boldsymbol{\zeta}_{b, n}, \dot{\boldsymbol{\zeta}}_{b, n}, \boldsymbol{\zeta}_{w, n}) - \mathbf{f}_{\text{aero}, n}

Parameters:

Name	Type	Description	Default
i_ts	int | Array	Time step index.	required
t_n	Array	Time at timestep n.	required
varphi_n	Array	Beam minimal coordinates vector at timestep n.	required
v_n	Array	Beam velocity vector at timestep n.	required
gamma_b_n	Array	Bound circulation vector at timestep n.	required
gamma_w_n	Array	Wake circulation vector at timestep n.	required
gamma_b_dot_n	Array	Bound circulation vector time derivative at timestep n.	required
zeta_w_n	Array	Wake grid vector at timestep n.	required
dv	AeroelasticDesignVariables	Aeroelastic design variables.	required
dv_full	AeroelasticDesignVariables	Aeroelastic design variables without omissions.	required
struct_obj	BeamStructure	Beam structure.	required
f_aero_beam_n	Array	Local aerodynamic forcing projected onto beam.	required
block_grid_gradients	bool	If true, blocks the gradient path for the dependency of the steady aerodynamic forcing on the bound grid.	required
solve_dofs	tuple[int, ...]	Index of forces to keep for solution.	required

Source code in src/flapjax/aero/uvlm.py

def f_aero_res_func(
    self,
    i_ts: int | Array,
    t_n: Array,
    varphi_n: Array,
    v_n: Array,
    gamma_b_n: Array,
    gamma_w_n: Array,
    gamma_b_dot_n: Array,
    zeta_w_n: Array,
    dv: AeroelasticDesignVariables,
    dv_full: AeroelasticDesignVariables,
    struct_obj: BeamStructure,
    f_aero_beam_n: Array,
    block_grid_gradients: bool,
    solve_dofs: tuple[int, ...],
) -> Array:
    r"""
    Aerodynamic forcing residual, compared in the local frame for compatibility with the structure.

    :math:`\boldsymbol{\mathcal{F}}_{\text{aero}, n}(\mathbf{\Gamma}_{b, n}, \mathbf{\Gamma}_{w, n},
    \dot{\mathbf{\Gamma}}_{b, n}, \boldsymbol{\zeta}_{b, n}, \dot{\boldsymbol{\zeta}}_{b, n},
    \boldsymbol{\zeta}_{w, n}) - \mathbf{f}_{\text{aero}, n}`

    :param i_ts: Time step index.
    :param t_n: Time at timestep n.
    :param varphi_n: Beam minimal coordinates vector at timestep n.
    :param v_n: Beam velocity vector at timestep n.
    :param gamma_b_n: Bound circulation vector at timestep n.
    :param gamma_w_n: Wake circulation vector at timestep n.
    :param gamma_b_dot_n: Bound circulation vector time derivative at timestep n.
    :param zeta_w_n: Wake grid vector at timestep n.
    :param dv: Aeroelastic design variables.
    :param dv_full: Aeroelastic design variables without omissions.
    :param struct_obj: Beam structure.
    :param f_aero_beam_n: Local aerodynamic forcing projected onto beam.
    :param block_grid_gradients: If true, blocks the gradient path for the dependency of the steady aerodynamic
    forcing on the bound grid.
    :param solve_dofs: Index of forces to keep for solution.
    """

    varphi_n = varphi_n.reshape(-1, 6)
    v_n = v_n.reshape(-1, 6)

    gamma_b_n = ArrayList.from_vector(
        vect=gamma_b_n,
        shapes=ArrayListShape([(gd.m, gd.n) for gd in self.grid_disc]),
    )

    gamma_w_n = ArrayList.from_vector(
        vect=gamma_w_n,
        shapes=ArrayListShape([(gd.m_star, gd.n) for gd in self.grid_disc]),
    )

    gamma_b_dot_n = ArrayList.from_vector(
        vect=gamma_b_dot_n,
        shapes=ArrayListShape([(gd.m, gd.n) for gd in self.grid_disc]),
    )

    zeta_w_n = ArrayList.from_vector(
        vect=zeta_w_n,
        shapes=ArrayListShape(
            [(gd.m_star + 1, gd.n + 1, 3) for gd in self.grid_disc]
        ),
    )

    f_aero_beam_n = f_aero_beam_n.reshape(-1, 6)

    inner_struct = struct_obj.case_from_dv(dv=dv.structure)
    hg_n = inner_struct.compute_hg_from_varphi(varphi=varphi_n)
    hg_dot_n = inner_struct.make_hg_dot(hg=hg_n, v=v_n)

    inner_case = self.case_from_dv(dv=dv.aero)

    # create grid
    cs_ang_n, cs_vel_n = dv.aero.get_cs_n(i_ts=i_ts, dv_full=dv_full.aero)

    zeta_b_n = inner_case.hg_to_zeta_b(hg_n=hg_n, cs_ang_n=cs_ang_n)
    if block_grid_gradients:
        zeta_b_n = jax.lax.stop_gradient(zeta_b_n)
        zeta_w_n = jax.lax.stop_gradient(zeta_w_n)

    zeta_b_dot_n = inner_case.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
    )
    nc_n = compute_nc(zetas=zeta_b_n)

    def v_total_func(x_: Array) -> Array:
        return inner_case.flowfield.vmap_call(x=x_, t=t_n) + compute_v_ind(
            cs=x_,
            zetas=ArrayList([*zeta_b_n, *zeta_w_n]),
            gammas=ArrayList([*gamma_b_n, *gamma_w_n]),
            kernels=[*inner_case.kernels_b, *inner_case.kernels_w],
            batch_size=self.batch_size,
            mirror_normal=inner_case.mirror_normal,
            mirror_point=inner_case.mirror_point,
        )

    f_steady = compute_steady_forcing(
        zeta_b=zeta_b_n,
        zeta_dot_b=zeta_b_dot_n,
        gamma_b=gamma_b_n,
        gamma_w=gamma_w_n,
        rho=inner_case.flowfield.rho,
        v_func=v_total_func,
        v_inputs=None,
        mirror_point=inner_case.mirror_point,
        mirror_normal=inner_case.mirror_normal,
        mirror_edge_low=inner_case.mirror_edge_low,
        mirror_edge_high=inner_case.mirror_edge_high,
    )

    if self.include_unsteady_force:
        f_unsteady: ArrayList = ArrayList(
            [
                split_to_vertex(
                    inner_case.flowfield.rho
                    * gamma_b_dot_n[i_surf][..., None]
                    * nc_n[i_surf],
                    (0, 1),
                )
                for i_surf in range(inner_case.n_surf)
            ]
        )
        f_tot = f_steady + f_unsteady
    else:
        f_tot = f_steady

    # project forcing to beam (global frame)
    f_tot_beam_global = project_forcing_to_beam(
        f_total=f_tot,
        rmat=hg_n[:, :3, :3],
        dof_mapping=inner_case.dof_mapping,
        x0_aero=inner_case.zeta_b0,
        mirror_edge_low=inner_case.mirror_edge_low,
        mirror_edge_high=inner_case.mirror_edge_high,
    )

    # transform to local frame to match f_aero_beam_n
    f_tot_beam_local = transform_nodal_vect(
        vect=f_tot_beam_global, rmat=jnp.swapaxes(hg_n[:, :3, :3], -2, -1)
    )

    return (f_tot_beam_local - f_aero_beam_n).ravel()[jnp.array(solve_dofs)]

timestep_residual

timestep_residual(
    i_ts: int | Array,
    varphi_nm1: Array,
    varphi_n: Array,
    v_n: Array,
    t_n: Array,
    q_n: AeroFullStates,
    q_nm1: AeroFullStates,
    dv: AeroelasticDesignVariables,
    dv_full: AeroelasticDesignVariables,
    f_aero_beam_n: Array,
    struct_obj: BeamStructure,
    approx_grads: bool,
) -> Array

Compute the residual vector to the UVLM equations. These are given as:

:math:\left(\boldsymbol{\mathcal{A}}_{b, n} \cdot \mathbf{n}_n\right)^{-1} \left[\left( \boldsymbol{\mathcal{A}}_{w, n} \boldsymbol{\Gamma}_{w, n} +\mathbf{v}_{bc, n} - \dot{\mathbf{c}}\right) \cdot \mathbf{n}_n\right] + \boldsymbol{\Gamma}_{b, n}

:math:\boldsymbol{\mathcal{W}}_{\Gamma}(\mathbf{\Gamma}_{b, {n-1}}, \mathbf{\Gamma}_{w, {n-1}}) - \mathbf{\Gamma}_{w, n}

:math:\frac{g}{h} \left[\mathbf{\Gamma}_{b, n} - \mathbf{\Gamma}_{b, n-1}\right] + (1-g) \dot{\mathbf{\Gamma}}_{b, n-1} - \dot{\mathbf{\Gamma}}_{b, n}

:math:\boldsymbol{\mathcal{W}}_{\zeta}(\boldsymbol{\zeta}_{b, n}, \boldsymbol{\zeta}_{w, n-1}) - \boldsymbol{\zeta}_{w, n}

:math:\boldsymbol{\mathcal{F}}_{\text{aero}, n}(\mathbf{\Gamma}_{b, n}, \mathbf{\Gamma}_{w, n}, \dot{\mathbf{\Gamma}}_{b, n}, \boldsymbol{\zeta}_{b, n}, \dot{\boldsymbol{\zeta}}_{b, n}, \boldsymbol{\zeta}_{w, n}) - \mathbf{f}_{\text{aero}, n}

Parameters:

Name	Type	Description	Default
i_ts	int | Array	Time step index.	required
varphi_nm1	Array	Beam minimal coordinates at timestep n-1, (n_nodes, 6).	required
varphi_n	Array	Beam maximal coordinates at timestep n, (n_nodes, 6).	required
v_n	Array	Beam velocity at timestep n, (n_nodes, 6).	required
t_n	Array	Time at step n.	required
q_n	AeroFullStates	Aero minimal states at timestep n.	required
q_nm1	AeroFullStates	Aero minimal states at timestep n-1.	required
dv	AeroelasticDesignVariables	Aeroelastic design variables.	required
dv_full	AeroelasticDesignVariables	Aeroelastic design variables without omissions.	required
f_aero_beam_n	Array	Aerodynamic forcing for the beam at timestep n, in the local frame.	required
struct_obj	BeamStructure	Beam structure.	required
approx_grads	bool	If true, eliminate grid gradients from the aerodynamic force residual.	required

Returns:

Type	Description
Array	Residual vector.

Source code in src/flapjax/aero/uvlm.py

def timestep_residual(
    self,
    i_ts: int | Array,
    varphi_nm1: Array,
    varphi_n: Array,
    v_n: Array,
    t_n: Array,
    q_n: AeroFullStates,
    q_nm1: AeroFullStates,
    dv: AeroelasticDesignVariables,
    dv_full: AeroelasticDesignVariables,
    f_aero_beam_n: Array,
    struct_obj: BeamStructure,
    approx_grads: bool,
) -> Array:
    r"""
    Compute the residual vector to the UVLM equations. These are given as:

    :math:`\left(\boldsymbol{\mathcal{A}}_{b, n} \cdot \mathbf{n}_n\right)^{-1} \left[\left(
    \boldsymbol{\mathcal{A}}_{w, n} \boldsymbol{\Gamma}_{w, n} +\mathbf{v}_{bc, n} - \dot{\mathbf{c}}\right)
    \cdot \mathbf{n}_n\right] + \boldsymbol{\Gamma}_{b, n}`

    :math:`\boldsymbol{\mathcal{W}}_{\Gamma}(\mathbf{\Gamma}_{b, {n-1}}, \mathbf{\Gamma}_{w, {n-1}})
    - \mathbf{\Gamma}_{w, n}`

    :math:`\frac{g}{h} \left[\mathbf{\Gamma}_{b, n} - \mathbf{\Gamma}_{b, n-1}\right] + (1-g)
    \dot{\mathbf{\Gamma}}_{b, n-1} - \dot{\mathbf{\Gamma}}_{b, n}`

    :math:`\boldsymbol{\mathcal{W}}_{\zeta}(\boldsymbol{\zeta}_{b, n}, \boldsymbol{\zeta}_{w, n-1})
    - \boldsymbol{\zeta}_{w, n}`

    :math:`\boldsymbol{\mathcal{F}}_{\text{aero}, n}(\mathbf{\Gamma}_{b, n}, \mathbf{\Gamma}_{w, n},
    \dot{\mathbf{\Gamma}}_{b, n}, \boldsymbol{\zeta}_{b, n}, \dot{\boldsymbol{\zeta}}_{b, n},
    \boldsymbol{\zeta}_{w, n}) - \mathbf{f}_{\text{aero}, n}`

    :param i_ts: Time step index.
    :param varphi_nm1: Beam minimal coordinates at timestep n-1, ``(n_nodes, 6)``.
    :param varphi_n: Beam maximal coordinates at timestep n, ``(n_nodes, 6)``.
    :param v_n: Beam velocity at timestep n, ``(n_nodes, 6)``.
    :param t_n: Time at step n.
    :param q_n: Aero minimal states at timestep n.
    :param q_nm1: Aero minimal states at timestep n-1.
    :param dv: Aeroelastic design variables.
    :param dv_full: Aeroelastic design variables without omissions.
    :param f_aero_beam_n: Aerodynamic forcing for the beam at timestep n, in the local frame.
    :param struct_obj: Beam structure.
    :param approx_grads: If true, eliminate grid gradients from the aerodynamic force residual.
    :return: Residual vector.
    """

    zeta_w_res, gamma_w_res = self.wake_prop_res_func(
        i_ts=i_ts,
        varphi_n=varphi_n.ravel(),
        varphi_nm1=varphi_nm1.ravel(),
        t_n=t_n,
        dv=dv,
        dv_full=dv_full,
        gamma_b_nm1=q_nm1.gamma_b.ravel(),
        gamma_w_nm1=q_nm1.gamma_w.ravel(),
        gamma_w_n=q_n.gamma_w.ravel(),
        zeta_w_nm1=q_nm1.zeta_w.ravel(),
        zeta_w_n=q_n.zeta_w.ravel(),
        struct_obj=struct_obj,
    )

    return jnp.concatenate(
        (
            self.gamma_b_res_func(
                i_ts=i_ts,
                varphi_n=varphi_n.ravel(),
                v_n=v_n.ravel(),
                t_n=t_n,
                dv=dv,
                dv_full=dv_full,
                gamma_b_n=q_n.gamma_b.ravel(),
                gamma_w_n=q_n.gamma_w.ravel(),
                zeta_w_n=q_n.zeta_w.ravel(),
                struct_obj=struct_obj,
            ),
            gamma_w_res,
            self.gamma_b_dot_res_func(
                gamma_b_nm1=q_nm1.gamma_b.ravel(),
                gamma_b_n=q_n.gamma_b.ravel(),
                gamma_b_dot_nm1=q_nm1.gamma_b_dot.ravel(),
                gamma_b_dot_n=q_n.gamma_b_dot.ravel(),
                dv=dv,
            ),
            zeta_w_res,
            self.f_aero_res_func(
                i_ts=i_ts,
                varphi_n=varphi_n.ravel(),
                v_n=v_n.ravel(),
                t_n=t_n,
                dv=dv,
                dv_full=dv_full,
                gamma_b_n=q_n.gamma_b.ravel(),
                gamma_w_n=q_n.gamma_w.ravel(),
                gamma_b_dot_n=q_n.gamma_b_dot.ravel(),
                zeta_w_n=q_n.zeta_w.ravel(),
                f_aero_beam_n=f_aero_beam_n.ravel(),
                struct_obj=struct_obj,
                block_grid_gradients=approx_grads,
                solve_dofs=tuple(range(struct_obj.n_dof)),
            ),
        )
    )

construct_approximate_jacobians

construct_approximate_jacobians(
    aero_sol: AeroCase,
    structure_sol: StructureCase,
    struct_obj: BeamStructure,
    dv: AeroelasticDesignVariables,
    dv_full: AeroelasticDesignVariables,
    solve_dofs: tuple[int, ...],
    jacobian_approximations: AeroJacobianApproximations,
) -> dict[str, dict[str, Callable[..., Any] | None]]

Compute approximations for the aerodynamic residual Jacobians specified in the jacobian_approximations data structure. The residuals covered are the bound circulation, wake propagation, bound circulation rate, and aerodynamic forcing.

Parameters:

Name	Type	Description	Default
aero_sol	AeroCase	Aerodynamic solution from which to extract states for the initial time step.	required
structure_sol	StructureCase	Structural solution from which to extract states for the initial time step.	required
struct_obj	BeamStructure	Beam structure used by the residual functions for the kinematic transformations.	required
dv	AeroelasticDesignVariables	Aeroelastic design variables.	required
dv_full	AeroelasticDesignVariables	Aeroelastic design variables without omissions for the variables where gradients aren't requested.	required
solve_dofs	tuple[int, ...]	Active structural degrees of freedom.	required
jacobian_approximations	AeroJacobianApproximations	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/aero/uvlm.py

def construct_approximate_jacobians(
    self,
    aero_sol: AeroCase,
    structure_sol: StructureCase,
    struct_obj: BeamStructure,
    dv: AeroelasticDesignVariables,
    dv_full: AeroelasticDesignVariables,
    solve_dofs: tuple[int, ...],
    jacobian_approximations: AeroJacobianApproximations,
) -> dict[str, dict[str, Callable[..., Any] | None]]:
    r"""
    Compute approximations for the aerodynamic residual Jacobians specified in the ``jacobian_approximations`` data
    structure. The residuals covered are the bound circulation, wake propagation, bound circulation rate, and
    aerodynamic forcing.
    :param aero_sol: Aerodynamic solution from which to extract states for the initial time step.
    :param structure_sol: Structural solution from which to extract states for the initial time step.
    :param struct_obj: Beam structure used by the residual functions for the kinematic transformations.
    :param dv: Aeroelastic design variables.
    :param dv_full: Aeroelastic design variables without omissions for the variables where gradients aren't
    requested.
    :param solve_dofs: Active structural degrees of freedom.
    :param jacobian_approximations: Data structure which defines which approximations to create.
    :return: Dictionary of approximations keyed by residual name.
    """
    q_nm1 = aero_sol.get_states(0)
    q_n = aero_sol.get_states(1)
    struct_nm1 = structure_sol.get_minimal_states(0)
    struct_n = structure_sol.get_minimal_states(1)

    varphi_nm1 = struct_nm1.varphi.ravel()
    varphi_n = struct_n.varphi.ravel()
    v_n = struct_n.v.ravel()
    gamma_b_nm1 = q_nm1.gamma_b.ravel()
    gamma_b_n = q_n.gamma_b.ravel()
    gamma_w_nm1 = q_nm1.gamma_w.ravel()
    gamma_w_n = q_n.gamma_w.ravel()
    gamma_b_dot_nm1 = q_nm1.gamma_b_dot.ravel()
    gamma_b_dot_n = q_n.gamma_b_dot.ravel()
    zeta_w_nm1 = q_nm1.zeta_w.ravel()
    zeta_w_n = q_n.zeta_w.ravel()

    if struct_n.f_ext_aero is None:
        raise ValueError("Missing aerodynamic forcing states")
    f_aero_beam_n = struct_n.f_ext_aero.ravel()

    t_n = aero_sol.t[1]

    gamma_b_args = {
        "i_ts": 1,
        "t_n": t_n,
        "varphi_n": varphi_n,
        "v_n": v_n,
        "gamma_b_n": gamma_b_n,
        "gamma_w_n": gamma_w_n,
        "zeta_w_n": zeta_w_n,
        "dv": dv,
        "dv_full": dv_full,
        "struct_obj": struct_obj,
    }

    wake_args = {
        "i_ts": 1,
        "t_n": t_n,
        "varphi_nm1": varphi_nm1,
        "varphi_n": varphi_n,
        "gamma_b_nm1": gamma_b_nm1,
        "gamma_w_nm1": gamma_w_nm1,
        "gamma_w_n": gamma_w_n,
        "zeta_w_nm1": zeta_w_nm1,
        "zeta_w_n": zeta_w_n,
        "dv": dv,
        "dv_full": dv_full,
        "struct_obj": struct_obj,
    }

    gamma_b_dot_args = {
        "gamma_b_nm1": gamma_b_nm1,
        "gamma_b_n": gamma_b_n,
        "gamma_b_dot_nm1": gamma_b_dot_nm1,
        "gamma_b_dot_n": gamma_b_dot_n,
        "dv": dv,
    }

    f_aero_args = {
        "i_ts": 1,
        "t_n": t_n,
        "varphi_n": varphi_n,
        "v_n": v_n,
        "gamma_b_n": gamma_b_n,
        "gamma_w_n": gamma_w_n,
        "gamma_b_dot_n": gamma_b_dot_n,
        "zeta_w_n": zeta_w_n,
        "dv": dv,
        "dv_full": dv_full,
        "struct_obj": struct_obj,
        "f_aero_beam_n": f_aero_beam_n,
        "block_grid_gradients": True,
        "solve_dofs": solve_dofs,
    }

    res_args: dict[
        str, tuple[Callable[..., Array], dict[str, Any], Sequence[str]]
    ] = {
        "gamma_b": (
            self.gamma_b_res_func,
            gamma_b_args,
            [f.name for f in fields(GammaBApprox)],
        ),
        "gamma_w": (
            lambda **kwargs: self.wake_prop_res_func(**kwargs)[1],
            wake_args,
            [f.name for f in fields(GammaWApprox)],
        ),
        "zeta_w": (
            lambda **kwargs: self.wake_prop_res_func(**kwargs)[0],
            wake_args,
            [f.name for f in fields(ZetaWApprox)],
        ),
        "gamma_b_dot": (
            self.gamma_b_dot_res_func,
            gamma_b_dot_args,
            [f.name for f in fields(GammaBDotApprox)],
        ),
        "f_aero": (
            self.f_aero_res_func,
            f_aero_args,
            [f.name for f in fields(FAeroApprox)],
        ),
    }

    return construct_approximation(
        res_args=res_args, jacobian_approximations=jacobian_approximations
    )

timestep_residual_jacobians

timestep_residual_jacobians(
    i_ts: int | Array,
    varphi_nm1: Array,
    varphi_n: Array,
    v_n: Array,
    t_n: Array,
    q_n: AeroFullStates,
    q_nm1: AeroFullStates,
    dv: AeroelasticDesignVariables,
    dv_full: AeroelasticDesignVariables,
    f_aero_beam_n: Array,
    struct_obj: BeamStructure,
    approx_grads: bool,
    solve_dofs: tuple[int, ...],
    n_profile_loops: int | None,
    jac_options: dict[str, dict[str, Any]],
    compute_wake_gradients: bool = True,
    mode: ADMode = "reverse",
    map_batch_size: int | None = None,
) -> tuple[
    Array,
    Array,
    AeroelasticDesignVariables,
    Array,
    Array,
    dict[str, dict[str, float]] | None,
    dict[str, dict[str, float]] | None,
]

Compute the Jacobians of the aerodynamic problem.

Parameters:

Name	Type	Description	Default
i_ts	int | Array	Time step index.	required
varphi_nm1	Array	Minimal structural coordinates at timestep n-1, (n_nodes, 6).	required
varphi_n	Array	Minimal structural coordinates at timestep n, (n_nodes, 6).	required
v_n	Array	Structural velocity at timestep n, (n_nodes, 6).	required
t_n	Array	Time at timestep n.	required
q_n	AeroFullStates	Aerodynamic minimal states at timestep n.	required
q_nm1	AeroFullStates	Aerodynamic minimal states at timestep n-1.	required
dv	AeroelasticDesignVariables	Aeroelastic design variables.	required
dv_full	AeroelasticDesignVariables	Aeroelastic design variables without omissions.	required
f_aero_beam_n	Array	Aerodynamic forcing in local frame of reference at timestep n, (n_nodes, 6).	required
struct_obj	BeamStructure	Structural object.	required
approx_grads	bool	If true, eliminate grid gradients from force computation.	required
solve_dofs	tuple[int, ...]	Degrees of freedom to solve for. This removes non-active forcing entries.	required
n_profile_loops	int | None	Optional number of loops for profiling function.	required
jac_options	dict[str, dict[str, Any]]	Dictionary of optional functions which can be used to substitute Jacobian evaluations from AD.	required
compute_wake_gradients	bool	If False, skip Jacobians involving the wake for zeta_w and gamma_w.	True
mode	ADMode	AD mode for obtaining gradients, either forward or reverse.	'reverse'
map_batch_size	int | None	Batch size for vectorising Jacobian construction.	None

Returns:

Type	Description
tuple[Array, Array, AeroelasticDesignVariables, Array, Array, dict[str, dict[str, float]] | None, dict[str, dict[str, float]] | None]	Gradients of aerodynamic residual with respect to previous states, current states, and design variables. Additionally, includes Jacobians of the aerodynamic residual with respect to the structural displacement and velocity, and profiling times for compilation and run time.

Source code in src/flapjax/aero/uvlm.py

def timestep_residual_jacobians(
    self,
    i_ts: int | Array,
    varphi_nm1: Array,
    varphi_n: Array,
    v_n: Array,
    t_n: Array,
    q_n: AeroFullStates,
    q_nm1: AeroFullStates,
    dv: AeroelasticDesignVariables,
    dv_full: AeroelasticDesignVariables,
    f_aero_beam_n: Array,
    struct_obj: BeamStructure,
    approx_grads: bool,
    solve_dofs: tuple[int, ...],
    n_profile_loops: int | None,
    jac_options: dict[str, dict[str, Any]],
    compute_wake_gradients: bool = True,
    mode: ADMode = "reverse",
    map_batch_size: int | None = None,
) -> tuple[
    Array,
    Array,
    AeroelasticDesignVariables,
    Array,
    Array,
    dict[str, dict[str, float]] | None,
    dict[str, dict[str, float]] | None,
]:
    r"""
    Compute the Jacobians of the aerodynamic problem.
    :param i_ts: Time step index.
    :param varphi_nm1: Minimal structural coordinates at timestep n-1, ``(n_nodes, 6)``.
    :param varphi_n: Minimal structural coordinates at timestep n, ``(n_nodes, 6)``.
    :param v_n: Structural velocity at timestep n, ``(n_nodes, 6)``.
    :param t_n: Time at timestep n.
    :param q_n: Aerodynamic minimal states at timestep n.
    :param q_nm1: Aerodynamic minimal states at timestep n-1.
    :param dv: Aeroelastic design variables.
    :param dv_full: Aeroelastic design variables without omissions.
    :param f_aero_beam_n: Aerodynamic forcing in local frame of reference at timestep n, ``(n_nodes, 6)``.
    :param struct_obj: Structural object.
    :param approx_grads: If true, eliminate grid gradients from force computation.
    :param solve_dofs: Degrees of freedom to solve for. This removes non-active forcing entries.
    :param n_profile_loops: Optional number of loops for profiling function.
    :param jac_options: Dictionary of optional functions which can be used to substitute Jacobian evaluations from
    AD.
    :param compute_wake_gradients: If False, skip Jacobians involving the wake for ``zeta_w`` and ``gamma_w``.
    :param mode: AD mode for obtaining gradients, either ``forward`` or ``reverse``.
    :param map_batch_size: Batch size for vectorising Jacobian construction.
    :return: Gradients of aerodynamic residual with respect to previous states, current states, and design
    variables. Additionally, includes Jacobians of the aerodynamic residual with respect to the structural
    displacement and velocity, and profiling times for compilation and run time.
    """

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

    varphi_nm1 = varphi_nm1.ravel()
    varphi_n = varphi_n.ravel()
    v_n = v_n.ravel()
    gamma_b_nm1 = q_nm1.gamma_b.ravel()
    gamma_b_n = q_n.gamma_b.ravel()
    gamma_w_nm1 = q_nm1.gamma_w.ravel()
    gamma_w_n = q_n.gamma_w.ravel()
    gamma_b_dot_nm1 = q_nm1.gamma_b_dot.ravel()
    gamma_b_dot_n = q_n.gamma_b_dot.ravel()
    zeta_w_nm1 = q_nm1.zeta_w.ravel()
    zeta_w_n = q_n.zeta_w.ravel()

    if not compute_wake_gradients:
        # optionally omit wake computations
        jac_options: dict[str, dict[str, Callable[..., Array] | None]] = {
            k: dict(v) for k, v in jac_options.items()
        }

        # remove entries that involve the wake
        for res_name, wake_keys in (
            ("gamma_b", ("zeta_w_n", "gamma_w_n")),
            ("gamma_w", ("zeta_w_n", "zeta_w_nm1")),
            ("f_aero", ("zeta_w_n", "gamma_w_n")),
        ):
            for key in wake_keys:
                jac_options[res_name].pop(key, None)

    # stop AD tape where required
    gamma_b_gamma_w_n_input = (
        jax.lax.stop_gradient(gamma_w_n)
        if not compute_wake_gradients
        else gamma_w_n
    )
    gamma_b_zeta_w_n_input = (
        jax.lax.stop_gradient(zeta_w_n) if not compute_wake_gradients else zeta_w_n
    )

    d_gamma_b, compile_time["gamma_b"], run_time["gamma_b"] = jacrev_custom(
        func=self.gamma_b_res_func,
        jac_options=jac_options["gamma_b"],
        n_profile_loops=n_profile_loops,
        func_name="gamma_b",
        mode=mode,
        map_batch_size=map_batch_size,
    )(
        i_ts=i_ts,
        t_n=t_n,
        varphi_n=varphi_n,
        v_n=v_n,
        gamma_b_n=gamma_b_n,
        gamma_w_n=gamma_b_gamma_w_n_input,
        zeta_w_n=gamma_b_zeta_w_n_input,
        dv=dv,
        dv_full=dv_full,
        struct_obj=struct_obj,
    )

    if compute_wake_gradients:
        d_gamma_w, compile_time["gamma_w"], run_time["gamma_w"] = jacrev_custom(
            func=lambda **kwargs: self.wake_prop_res_func(**kwargs)[1],
            jac_options=jac_options["gamma_w"],
            n_profile_loops=n_profile_loops,
            func_name="gamma_w",
            mode=mode,
            map_batch_size=map_batch_size,
        )(
            i_ts=i_ts,
            t_n=t_n,
            varphi_nm1=varphi_nm1,
            varphi_n=varphi_n,
            gamma_b_nm1=gamma_b_nm1,
            gamma_w_nm1=gamma_w_nm1,
            gamma_w_n=gamma_w_n,
            zeta_w_nm1=zeta_w_nm1,
            zeta_w_n=zeta_w_n,
            dv=dv,
            dv_full=dv_full,
            struct_obj=struct_obj,
        )

        d_zeta_w, compile_time["zeta_w"], run_time["zeta_w"] = jacrev_custom(
            func=lambda **kwargs: self.wake_prop_res_func(**kwargs)[0],
            jac_options=jac_options["zeta_w"],
            n_profile_loops=n_profile_loops,
            func_name="zeta_w",
            mode=mode,
            map_batch_size=map_batch_size,
        )(
            i_ts=i_ts,
            t_n=t_n,
            varphi_nm1=varphi_nm1,
            varphi_n=varphi_n,
            gamma_b_nm1=gamma_b_nm1,
            gamma_w_nm1=gamma_w_nm1,
            gamma_w_n=gamma_w_n,
            zeta_w_nm1=zeta_w_nm1,
            zeta_w_n=zeta_w_n,
            dv=dv,
            dv_full=dv_full,
            struct_obj=struct_obj,
        )
    else:
        # skip computations
        d_gamma_w = {}
        d_zeta_w = {}

    d_gamma_b_dot, compile_time["gamma_b_dot"], run_time["gamma_b_dot"] = (
        jacrev_custom(
            func=self.gamma_b_dot_res_func,
            jac_options=jac_options["gamma_b_dot"],
            n_profile_loops=n_profile_loops,
            func_name="gamma_b_dot",
            mode=mode,
            map_batch_size=map_batch_size,
        )(
            gamma_b_nm1=gamma_b_nm1,
            gamma_b_n=gamma_b_n,
            gamma_b_dot_nm1=gamma_b_dot_nm1,
            gamma_b_dot_n=gamma_b_dot_n,
            dv=dv,
        )
    )

    f_aero_gamma_w_n_input = (
        jax.lax.stop_gradient(gamma_w_n)
        if not compute_wake_gradients
        else gamma_w_n
    )
    f_aero_zeta_w_n_input = (
        jax.lax.stop_gradient(zeta_w_n) if not compute_wake_gradients else zeta_w_n
    )

    d_f_aero, compile_time["f_aero"], run_time["f_aero"] = jacrev_custom(
        func=self.f_aero_res_func,
        jac_options=jac_options["f_aero"],
        n_profile_loops=n_profile_loops,
        func_name="f_aero",
        static_argnames=("block_grid_gradients", "solve_dofs"),
        mode=mode,
        map_batch_size=map_batch_size,
    )(
        i_ts=i_ts,
        t_n=t_n,
        varphi_n=varphi_n,
        v_n=v_n,
        gamma_b_n=gamma_b_n,
        gamma_w_n=f_aero_gamma_w_n_input,
        gamma_b_dot_n=gamma_b_dot_n,
        zeta_w_n=f_aero_zeta_w_n_input,
        dv=dv,
        dv_full=dv_full,
        struct_obj=struct_obj,
        f_aero_beam_n=f_aero_beam_n.ravel(),
        block_grid_gradients=approx_grads,
        solve_dofs=solve_dofs,
    )
    # slice f_aero Jacobian to get forcing only on active degrees of freedom
    d_f_aero["f_aero_beam_n"] = d_f_aero["f_aero_beam_n"][:, jnp.array(solve_dofs)]

    # Jacobians block widths and heights, assembled in degree of freedom order
    n_solve_dof = len(solve_dofs)
    aero_entries_list: list[dict[str, Any]] = [d_gamma_b]
    aero_heights_list: list[int] = [gamma_b_n.size]
    aero_n_keys_list: list[str] = ["gamma_b_n"]
    aero_nm1_keys_list: list[str] = ["gamma_b_nm1"]
    if compute_wake_gradients:
        aero_entries_list.append(d_gamma_w)
        aero_heights_list.append(gamma_w_n.size)
        aero_n_keys_list.append("gamma_w_n")
        aero_nm1_keys_list.append("gamma_w_nm1")
    aero_entries_list.append(d_gamma_b_dot)
    aero_heights_list.append(gamma_b_n.size)
    aero_n_keys_list.append("gamma_b_dot_n")
    aero_nm1_keys_list.append("gamma_b_dot_nm1")

    if compute_wake_gradients:
        aero_entries_list.append(d_zeta_w)
        aero_heights_list.append(zeta_w_n.size)
        aero_n_keys_list.append("zeta_w_n")
        aero_nm1_keys_list.append("zeta_w_nm1")
    aero_entries_list.append(d_f_aero)
    aero_heights_list.append(n_solve_dof)
    aero_n_keys_list.append("f_aero_beam_n")
    aero_nm1_keys_list.append("f_aero_beam_nm1")

    aero_entries = tuple(aero_entries_list)
    aero_heights: tuple[int, ...] = tuple(aero_heights_list)
    aero_n_keys = tuple(aero_n_keys_list)
    aero_nm1_keys = tuple(aero_nm1_keys_list)
    aero_widths = aero_heights

    struct_sizes = (
        struct_obj.n_dof,
        struct_obj.n_dof,
        struct_obj.n_dof,
        struct_obj.n_dof,
    )

    d_aero_res_d_q_nm1 = construct_named_block_jacobian(
        entries=aero_entries,
        keys=aero_nm1_keys,
        widths=aero_widths,
        heights=aero_heights,
    )

    d_aero_res_d_q_n = construct_named_block_jacobian(
        entries=aero_entries,
        keys=aero_n_keys,
        widths=aero_widths,
        heights=aero_heights,
    )

    # residual of aero problem w.r.t. structural states
    d_struct_res_d_q_nm1 = construct_named_block_jacobian(
        entries=aero_entries,
        keys=("varphi_nm1", "v_nm1", "v_dot_nm1", "a_nm1"),
        widths=struct_sizes,
        heights=aero_heights,
    )

    d_struct_res_d_q_n = construct_named_block_jacobian(
        entries=aero_entries,
        keys=("varphi_n", "v_n", "v_dot_n", "a_n"),
        widths=struct_sizes,
        heights=aero_heights,
    )

    # handle design gradients: include a row per included residual, mirroring
    # the aero block layout above.
    dv_rows: list[Any] = [d_gamma_b["dv"]]
    if compute_wake_gradients:
        dv_rows.append(d_gamma_w["dv"])
    dv_rows.append(d_gamma_b_dot["dv"])
    if compute_wake_gradients:
        dv_rows.append(d_zeta_w["dv"])
    dv_rows.append(d_f_aero["dv"])

    from flapjax.coupled.data_structures import AeroelasticDesignVariables

    d_res_d_dv = AeroelasticDesignVariables.concatenate(*dv_rows)

    return (
        d_aero_res_d_q_nm1,
        d_aero_res_d_q_n,
        d_res_d_dv,
        d_struct_res_d_q_nm1,
        d_struct_res_d_q_n,
        compile_time if n_profile_loops is not None else None,
        run_time if n_profile_loops is not None else None,
    )

add_control_surface

add_control_surface(
    grid: Array,
    angle: Array,
    m_slice: Array | Sequence[int] | slice,
    n_slice: Array | Sequence[int] | slice,
    hinge_axis: Array = HINGE_AXIS_DEFAULT,
) -> Array

Add a control surface to a panel grid.

Parameters:

Name	Type	Description	Default
grid	Array	Grid without deflection of this surface, (zeta_m, zeta_n, 3).	required
angle	Array	Angle in radians through which the control surface will be deflected.	required
m_slice	Array | Sequence[int] | slice	Slice of chordwise strips to include in the control surface.	required
n_slice	Array | Sequence[int] | slice	Slice of spanwise strips to include in the control surface.	required
hinge_axis	Array	Axis of the hinge surface in the local frame, (3, ).	HINGE_AXIS_DEFAULT

Returns:

Type	Description
Array	Deflected aerodynamic grid, (zeta_m, zeta_n, 3).

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

def add_control_surface(
    grid: Array,
    angle: Array,
    m_slice: Array | Sequence[int] | slice,
    n_slice: Array | Sequence[int] | slice,
    hinge_axis: Array = HINGE_AXIS_DEFAULT,
) -> Array:
    r"""
    Add a control surface to a panel grid.
    :param grid: Grid without deflection of this surface, ``(zeta_m, zeta_n, 3)``.
    :param angle: Angle in radians through which the control surface will be deflected.
    :param m_slice: Slice of chordwise strips to include in the control surface.
    :param n_slice: Slice of spanwise strips to include in the control surface.
    :param hinge_axis: Axis of the hinge surface in the local frame, ``(3, )``.
    :return: Deflected aerodynamic grid, ``(zeta_m, zeta_n, 3)``.
    """

    m_slice_arr: Array = index_to_arr(index=m_slice, n_entries=grid.shape[0])
    n_slice_arr: Array = index_to_arr(index=n_slice, n_entries=grid.shape[1])

    # grid for deflected surfaces
    grid_out = grid

    def inner_func(n_idx: Array) -> Array:
        hinge_point = grid[m_slice_arr[0], n_idx, :]  # (3, )

        crv = hinge_axis * angle  # cartesian rotation vector for surface, (3, ).
        rmat = exp_so3(crv)  # rotation matrix for rotating surface

        # transform coordinates to rotate control surface
        return (
            jnp.einsum(
                "ij,hj->hi",
                rmat,
                (grid[m_slice_arr, n_idx, :] - hinge_point[None, :]),
            )
            + hinge_point[None, :]
        )

    # update grid
    return grid_out.at[jnp.ix_(m_slice_arr, n_slice_arr, jnp.arange(3))].set(
        vmap(inner_func, in_axes=0, out_axes=1)(n_slice_arr)
    )

make_rectangular_grid

make_rectangular_grid(
    m: int,
    n: int,
    chord: Array | float,
    ea: Array | float,
    camber_line: tuple[Array, Array] | None = None,
    twist: Array | float = 0.0,
) -> Array

Create a rectangular aerodynamic grid.

Parameters:

Name	Type	Description	Default
m	int	Number of panels in the chordwise direction.	required
n	int	Number of panels in the spanwise direction.	required
chord	Array | float	Surface chord length.	required
ea	Array | float	Elastic axis location as fraction of chord.	required
camber_line	tuple[Array, Array] | None	Optional mean camber line as a pair (x/c, z/c) of equal-length vectors, giving the camber-line height at a set of chordwise stations, with x/c running from 0 (leading edge) to 1 (trailing edge). If None (default), the section is a flat plate.	None
twist	Array | float	Built-in geometric twist angle in radians, uniform over the whole grid, applied by rotating the local chord/camber section about the local spanwise axis (HINGE_AXIS_DEFAULT). This is a purely aerodynamic incidence offset: unlike a beam's y_vector-defined twist (which only reorients the structural cross-section's stiffness/mass axes), it actually changes the panels' angle of attack, since a uniform twist produces no curvature for the structural solver to pick up on its own. Default 0 (no twist).	0.0

Returns:

Type	Description
Array	Local grid points for planar wing, (zeta_m, zeta_n, 3).

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

def make_rectangular_grid(
    m: int,
    n: int,
    chord: Array | float,
    ea: Array | float,
    camber_line: tuple[Array, Array] | None = None,
    twist: Array | float = 0.0,
) -> Array:
    r"""
    Create a rectangular aerodynamic grid.
    :param m: Number of panels in the chordwise direction.
    :param n: Number of panels in the spanwise direction.
    :param chord: Surface chord length.
    :param ea: Elastic axis location as fraction of chord.
    :param camber_line: Optional mean camber line as a pair ``(x/c, z/c)`` of equal-length vectors, giving
    the camber-line height at a set of chordwise stations, with ``x/c`` running from 0 (leading edge)
    to 1 (trailing edge). If None (default), the section is a flat plate.
    :param twist: Built-in geometric twist angle in radians, uniform over the whole grid, applied by rotating the
    local chord/camber section about the local spanwise axis (``HINGE_AXIS_DEFAULT``). This is a purely
    aerodynamic incidence offset: unlike a beam's ``y_vector``-defined twist (which only reorients the
    structural cross-section's stiffness/mass axes), it actually changes the panels' angle of attack, since a
    uniform twist produces no curvature for the structural solver to pick up on its own. Default 0 (no twist).
    :return: Local grid points for planar wing, ``(zeta_m, zeta_n, 3)``.
    """

    x_over_c = jnp.linspace(0.0, 1.0, m + 1)
    grid = jnp.zeros((m + 1, n + 1, 3))
    grid = grid.at[..., 0].set((x_over_c * chord - ea * chord)[:, None])
    if camber_line is not None:
        camber_x, camber_z = camber_line
        z_over_c = jnp.interp(x_over_c, camber_x, camber_z)
        grid = grid.at[..., 2].set((z_over_c * chord)[:, None])
    rmat = exp_so3(HINGE_AXIS_DEFAULT * twist)
    grid = jnp.einsum("ij,mnj->mni", rmat, grid)
    return grid

aic

compute_aic_grid

compute_aic_grid(
    c: Array,
    n: Array | None,
    zeta: Array,
    kernel: KernelFunction,
    batch_size: int | None,
) -> Array

Compute the aerodynamic influence coefficient (AIC) across grids of points. When normal is provided, fuses the dot product inside each map step so the trailing 3-component axis is never accumulated, saving memory.

Parameters:

Name	Type	Description	Default
c	Array	Collocation points, (c_m, c_n, 3).	required
n	Array | None	Normal vectors at collocation points, (c_m, c_n, 3), or None.	required
zeta	Array	Grid vertices, (zeta_m, zeta_n, 3).	required
kernel	KernelFunction	Kernel function to compute the influence.	required
batch_size	int | None	Batch size for vectorising AIC computations.	required

Returns:

Type	Description
Array	(c_m, c_n, zeta_m, zeta_n, 3) if normal is None, else (c_m, c_n, zeta_m, zeta_n).

Source code in src/flapjax/aero/aic.py

def compute_aic_grid(
    c: Array,
    n: Array | None,
    zeta: Array,
    kernel: KernelFunction,
    batch_size: int | None,
) -> Array:
    """
    Compute the aerodynamic influence coefficient (AIC) across grids of points. When normal is provided, fuses the dot
    product inside each map step so the trailing 3-component axis is never accumulated, saving memory.
    :param c: Collocation points, ``(c_m, c_n, 3)``.
    :param n: Normal vectors at collocation points, ``(c_m, c_n, 3)``, or None.
    :param zeta: Grid vertices, ``(zeta_m, zeta_n, 3)``.
    :param kernel: Kernel function to compute the influence.
    :param batch_size: Batch size for vectorising AIC computations.
    :return: ``(c_m, c_n, zeta_m, zeta_n, 3)`` if normal is None, else ``(c_m, c_n, zeta_m, zeta_n)``.
    """
    c_m, c_n = c.shape[:2]
    m_panels, n_panels = zeta.shape[0] - 1, zeta.shape[1] - 1

    m_vect_flat = jnp.stack((zeta[:-1, :, :], zeta[1:, :, :]), axis=-2).reshape(
        -1, 2, 3
    )
    n_vect_flat = jnp.stack((zeta[:, :-1, :], zeta[:, 1:, :]), axis=-2).reshape(
        -1, 2, 3
    )

    # account for the degenerate case where there are no source panels to prevent division by zero
    if not c_m or not c_n or not m_panels or not n_panels:
        return jnp.zeros((c_m, c_n, m_panels, n_panels))

    @jax.checkpoint
    def row(args: tuple) -> Array:
        # compute the influence of all spanwise (m) and chordwise (n) filaments before combining. This prevents any
        # duplicate computations.
        ci, ni = args
        m_influence = vmap(kernel, (None, 0), 0)(ci, m_vect_flat)
        m_influence_ni = jnp.dot(m_influence, ni).reshape(
            m_panels, n_panels + 1
        )  # [m, n+1]
        n_influence = vmap(kernel, (None, 0), 0)(ci, n_vect_flat)
        n_influence_ni = jnp.dot(n_influence, ni).reshape(
            m_panels + 1, n_panels
        )  # [m+1, n]
        return -jnp.diff(m_influence_ni, axis=1) + jnp.diff(
            n_influence_ni, axis=0
        )  # [m, n]

    return jax.lax.map(
        row,
        (c.reshape(-1, 3), n.reshape(-1, 3) if n is not None else None),
        batch_size=batch_size,
    ).reshape(c_m, c_n, m_panels, n_panels)

compute_aic_sys

compute_aic_sys(
    zetas: ArrayList,
    cs: ArrayList,
    ns: ArrayList,
    kernels: Sequence[KernelFunction],
    batch_size: int | None,
    mirror_point: Array | None,
    mirror_normal: Array | None,
) -> list[list[Array]]

Compute the AIC matrix for a system of elements. Returns a list of AIC matrices, one for each element.

Parameters:

Name	Type	Description	Default
zetas	ArrayList	List of source points to compute the AIC from, (n_source,)(zeta_m, zeta_n, 3).	required
cs	ArrayList	List of target points to compute the AIC at, (n_target,)(c_m, c_n, 3).	required
ns	ArrayList	Bound normal vectors, (c_m, c_n, 3). If None, no projection will be done.	required
kernels	Sequence[KernelFunction]	List of kernel functions to use for each source surface, (n_source, ).	required
batch_size	int | None	Batch size for vectorising AIC computations.	required
mirror_normal	Array | None	Normal vector to mirror across, (3, ). If None, no mirroring will be done.	required
mirror_point	Array | None	Point on mirror plane, (3, ). If None, no mirroring will be done.	required

Returns:

Type	Description
list[list[Array]]	Nested sequences of AIC matrices, (n_target,)(n_source, c_m, c_n, zeta_m, zeta_n, 3), or (n_target,)(n_source,)(c_m, c_n, zeta_m, zeta_n) if projected onto normals.

Source code in src/flapjax/aero/aic.py

def compute_aic_sys(
    zetas: ArrayList,
    cs: ArrayList,
    ns: ArrayList,
    kernels: Sequence[KernelFunction],
    batch_size: int | None,
    mirror_point: Array | None,
    mirror_normal: Array | None,
) -> list[list[Array]]:
    """
    Compute the AIC matrix for a system of elements. Returns a list of AIC matrices, one for each element.
    :param zetas: List of source points to compute the AIC from, ``(n_source,)(zeta_m, zeta_n, 3)``.
    :param cs: List of target points to compute the AIC at, ``(n_target,)(c_m, c_n, 3)``.
    :param ns: Bound normal vectors, ``(c_m, c_n, 3)``. If None, no projection will be done.
    :param kernels: List of kernel functions to use for each source surface, ``(n_source, )``.
    :param batch_size: Batch size for vectorising AIC computations.
    :param mirror_normal: Normal vector to mirror across, ``(3, )``. If None, no mirroring will be done.
    :param mirror_point: Point on mirror plane, ``(3, )``. If None, no mirroring will be done.
    :return: Nested sequences of AIC matrices, ``(n_target,)(n_source, c_m, c_n, zeta_m, zeta_n, 3)``, or
    ``(n_target,)(n_source,)(c_m, c_n, zeta_m, zeta_n)`` if projected onto normals.
    """

    aic_mats = []
    for c, n in zip(cs, ns):
        aic_mats.append([])
        for zeta, kernel in zip(zetas, kernels):
            # compute the AIC matrix, [n_cx, n_cy, n_ex, n_ey, 3]
            aic_ = compute_aic_grid(
                c=c,
                n=n,
                zeta=zeta,
                kernel=kernel,
                batch_size=batch_size,
            )

            if mirror_point is not None and mirror_normal is not None:
                # add influence from mirrored grid, if specified
                zeta_mirror = mirror_grid(
                    zeta=zeta,
                    mirror_point=mirror_point,
                    mirror_normal=mirror_normal,
                )
                aic_ -= compute_aic_grid(
                    c=c, n=n, zeta=zeta_mirror, kernel=kernel, batch_size=batch_size
                )
            aic_mats[-1].append(aic_)
    return aic_mats

reshape_aic_sys

reshape_aic_sys(aic_mat: Array) -> Array

Reshape an AIC matrix such that the source and target dimensions are flattened.

Parameters:

Name	Type	Description	Default
aic_mat	Array	Input AIC matrix, (c_m, c_n, zeta_m, zeta_n) or (c_m, c_n, zeta_m, zeta_n, 3).	required

Returns:

Type	Description
Array	Reshaped AIC matrix, (c_m*c_n, zeta_m*zeta_n) or (c_m*c_n, zeta_m*zeta_n, 3).

Source code in src/flapjax/aero/aic.py

def reshape_aic_sys(aic_mat: Array) -> Array:
    r"""
    Reshape an AIC matrix such that the source and target dimensions are flattened.
    :param aic_mat: Input AIC matrix, ``(c_m, c_n, zeta_m, zeta_n)`` or ``(c_m, c_n, zeta_m, zeta_n, 3)``.
    :return: Reshaped AIC matrix, ``(c_m*c_n, zeta_m*zeta_n)`` or ``(c_m*c_n, zeta_m*zeta_n, 3)``.
    """
    shape = aic_mat.shape
    return aic_mat.reshape([shape[0] * shape[1], shape[2] * shape[3]])

assemble_aic_sys

assemble_aic_sys(
    aic_mats: Sequence[Sequence[Array]],
) -> Array

Assemble a nested sequence of AIC matrices into a single AIC matrix.

Parameters:

Name	Type	Description	Default
aic_mats	Sequence[Sequence[Array]]	Nested sequence of AIC matrices, (n_target,)(n_source,)(c_m, c_n, zeta_m, zeta_n) or (n_target,)(n_source,)(c_m, c_n, zeta_m, zeta_n, 3).	required

Returns:

Type	Description
Array	Assembled AIC matrix, (c_tot, zeta_tot) or (c_tot, zeta_tot, 3).

Source code in src/flapjax/aero/aic.py

def assemble_aic_sys(aic_mats: Sequence[Sequence[Array]]) -> Array:
    r"""
    Assemble a nested sequence of AIC matrices into a single AIC matrix.
    :param aic_mats: Nested sequence of AIC matrices, ``(n_target,)(n_source,)(c_m, c_n, zeta_m, zeta_n)`` or ``(n_target,)(n_source,)(c_m, c_n, zeta_m, zeta_n, 3)``.
    :return: Assembled AIC matrix, ``(c_tot, zeta_tot)`` or ``(c_tot, zeta_tot, 3)``.
    """
    aic_mats_reshaped = [
        [reshape_aic_sys(aic) for aic in aic_row] for aic_row in aic_mats
    ]
    return block_axis(aic_mats_reshaped, axes=(0, 1))

compute_aic_solve

compute_aic_solve(
    cs: ArrayList,
    ns: ArrayList,
    zetas_b: ArrayList,
    zetas_w: ArrayList | None,
    kernels_b: Sequence[KernelFunction],
    kernels_w: Sequence[KernelFunction] | None,
    batch_size: int | None,
    mirror_point: Array | None,
    mirror_normal: Array | None,
) -> Array

Compute the AIC matrix used for the UVLM solve step.

Parameters:

Name	Type	Description	Default
cs	ArrayList	List of target points to compute the AIC at, (n_target,)(c_m, c_n, 3).	required
ns	ArrayList	Bound normal vectors, (n_target,)(c_m, c_n, 3). If None, no projection will be done.	required
zetas_b	ArrayList	Bound aerodynamic grids, (n_source,)(zeta_m, zeta_n, 3).	required
zetas_w	ArrayList | None	Wake aerodynamic grids, (n_source,)(zeta_m_star, zeta_n, 3). This is only passed in the static case, as in the dynamic case the wake influence is instead included in the boundary conditions.	required
kernels_b	Sequence[KernelFunction]	Bound grid kernels.	required
kernels_w	Sequence[KernelFunction] | None	Wake grid kernels.	required
batch_size	int | None	Batch size for vectorising AIC computations.	required
mirror_normal	Array | None	Normal vector to mirror across, (3, ). If None, no mirroring will be done.	required
mirror_point	Array | None	Point on mirror plane, (3, ). If None, no mirroring will be done.	required

Returns:

Type	Description
Array	Square AIC matrix for the solve step, (c_tot, zeta_tot).

Source code in src/flapjax/aero/aic.py

def compute_aic_solve(
    cs: ArrayList,
    ns: ArrayList,
    zetas_b: ArrayList,
    zetas_w: ArrayList | None,
    kernels_b: Sequence[KernelFunction],
    kernels_w: Sequence[KernelFunction] | None,
    batch_size: int | None,
    mirror_point: Array | None,
    mirror_normal: Array | None,
) -> Array:
    r"""
    Compute the AIC matrix used for the UVLM solve step.
    :param cs: List of target points to compute the AIC at, ``(n_target,)(c_m, c_n, 3)``.
    :param ns: Bound normal vectors, ``(n_target,)(c_m, c_n, 3)``. If None, no projection will be done.
    :param zetas_b: Bound aerodynamic grids, ``(n_source,)(zeta_m, zeta_n, 3)``.
    :param zetas_w: Wake aerodynamic grids, ``(n_source,)(zeta_m_star, zeta_n, 3)``. This is only passed in the static case,
    as in the dynamic case the wake influence is instead included in the boundary conditions.
    :param kernels_b: Bound grid kernels.
    :param kernels_w: Wake grid kernels.
    :param batch_size: Batch size for vectorising AIC computations.
    :param mirror_normal: Normal vector to mirror across, ``(3, )``. If None, no mirroring will be done.
    :param mirror_point: Point on mirror plane, ``(3, )``. If None, no mirroring will be done.
    :return: Square AIC matrix for the solve step, ``(c_tot, zeta_tot)``.
    """
    aic_b_mats = compute_aic_sys(
        cs=cs,
        ns=ns,
        zetas=zetas_b,
        kernels=kernels_b,
        batch_size=batch_size,
        mirror_point=mirror_point,
        mirror_normal=mirror_normal,
    )

    if zetas_w is not None:
        if kernels_w is None:
            raise ValueError("kernels_w must not be None")
        aic_w_mats = compute_aic_sys(
            cs=cs,
            ns=ns,
            zetas=zetas_w,
            kernels=kernels_w,
            batch_size=batch_size,
            mirror_point=mirror_point,
            mirror_normal=mirror_normal,
        )

        aic_b_mats = add_wake_influence(aic_b_mats, aic_w_mats)

    return assemble_aic_sys(aic_b_mats)

add_wake_influence

add_wake_influence(
    aic_bs: list[list[Array]], aic_ws: list[list[Array]]
) -> list[list[Array]]

Lump the wake influence onto the last column of the bound AIC matrices. This captures the steady Kutta condition by ensuring that the trailing edge panels have the same strength as all wake panels along a streamline.

Parameters:

Name	Type	Description	Default
aic_bs	list[list[Array]]	Bound influence matrices, (n_target,)(n_source,)(c_m, c_n, zeta_m, zeta_n, 3).	required
aic_ws	list[list[Array]]	Wake influence matrices, (n_target,)(n_source,)(c_m, c_n, zeta_m_star, zeta_n, 3).	required

Returns:

Type	Description
list[list[Array]]	Updated bound influence matrices, (n_target,)(n_source,)(c_m, c_n, zeta_m, zeta_n, 3).

Source code in src/flapjax/aero/aic.py

def add_wake_influence(
    aic_bs: list[list[Array]], aic_ws: list[list[Array]]
) -> list[list[Array]]:
    r"""
    Lump the wake influence onto the last column of the bound AIC matrices. This captures the steady Kutta condition
    by ensuring that the trailing edge panels have the same strength as all wake panels along a streamline.
    :param aic_bs: Bound influence matrices, ``(n_target,)(n_source,)(c_m, c_n, zeta_m, zeta_n, 3)``.
    :param aic_ws: Wake influence matrices, ``(n_target,)(n_source,)(c_m, c_n, zeta_m_star, zeta_n, 3)``.
    :return: Updated bound influence matrices, ``(n_target,)(n_source,)(c_m, c_n, zeta_m, zeta_n, 3)``.
    """
    for i in range(len(aic_bs)):
        for j in range(len(aic_bs[i])):
            aic_bs[i][j] = (
                aic_bs[i][j].at[:, :, -1, :].add(jnp.sum(aic_ws[i][j], axis=2))
            )
    return aic_bs

v_ind_vmap

v_ind_vmap(
    c: Array,
    zeta: Array,
    gamma: Array,
    kernel: KernelFunction,
    batch_size: int | None,
) -> Array

Compute the induced velocity by the aerodynamic elements at some points in space for a single source-target panel system. This is done without materialising the full AIC matrix, instead directly computing its contraction with the circulation strength.

Parameters:

Name	Type	Description	Default
c	Array	Points at which to sample the velocity, (c_m, c_n, 3).	required
zeta	Array	Filament grid, (zeta_m, zeta_n, 2, 3).	required
gamma	Array	Circulation strengths, (zeta_m, zeta_n).	required
kernel	KernelFunction	Kernel function.	required
batch_size	int | None	Batch size for vectorising AIC computations.	required

Returns:

Type	Description
Array	Induced velocity, (c_m, c_n, 3).

Source code in src/flapjax/aero/aic.py

def v_ind_vmap(
    c: Array,
    zeta: Array,
    gamma: Array,
    kernel: KernelFunction,
    batch_size: int | None,
) -> Array:
    """
    Compute the induced velocity by the aerodynamic elements at some points in space for a single source-target panel
    system. This is done without materialising the full AIC matrix, instead directly computing its contraction with the
    circulation strength.
    :param c: Points at which to sample the velocity, ``(c_m, c_n, 3)``.
    :param zeta: Filament grid, ``(zeta_m, zeta_n, 2, 3)``.
    :param gamma: Circulation strengths, ``(zeta_m, zeta_n)``.
    :param kernel: Kernel function.
    :param batch_size: Batch size for vectorising AIC computations.
    :return: Induced velocity, ``(c_m, c_n, 3)``.
    """
    c_m, c_n = c.shape[:2]
    c_flat = c.reshape(-1, 3)
    zeta_flat = zeta.reshape(-1, 2, 3)
    gamma_flat = gamma.ravel()  # [zeta_m * zeta_n]

    # account for case where zeta is empty
    if zeta.size == 0:
        return jnp.zeros_like(c)

    @jax.checkpoint
    def row(ci: Array) -> Array:
        influence = vmap(kernel, (None, 0), 0)(ci, zeta_flat)  # [zeta_m * zeta_n, 3]

        if influence.shape[0] != gamma_flat.shape[0]:
            pass

        return jnp.einsum("lm,l->m", influence, gamma_flat)  # [3]

    result = jax.lax.map(row, c_flat, batch_size=batch_size)  # [c_m * c_n, 3]
    return result.reshape(c_m, c_n, 3)

compute_v_ind

compute_v_ind(
    cs: T,
    zetas: ArrayList,
    gammas: ArrayList,
    kernels: Sequence[KernelFunction],
    mirror_point: Array | None,
    mirror_normal: Array | None,
    batch_size: int | None,
) -> T

Compute the induced velocity by multiple surfaces of aerodynamic elements at one or multiple grids of points in space. This is done without materialising the full AIC matrix, instead directly computing its contraction with the circulation strength.

Parameters:

Name	Type	Description	Default
cs	T	Points at which to sample the velocity, (c_m, c_n, 3) or (n_target,)(c_m, c_n, 3).	required
zetas	ArrayList	Filament grid, (n_source,)(zeta_m, zeta_n, 3).	required
gammas	ArrayList	Circulation strengths, (n_source,)(zeta_m, zeta_n).	required
kernels	Sequence[KernelFunction]	Kernel function.	required
mirror_point	Array | None	Mirror point, (3, ). If None, no mirroring will be done.	required
mirror_normal	Array | None	Normal mirror vector, (3, ). If None, no mirroring will be done.	required
batch_size	int | None	Batch size for vectorising AIC computations.	required

Returns:

Type	Description
T	Array or ArrayList of induced velocity, (c_m, c_n, 3) or (n_target,)(c_m, c_n, 3).

Source code in src/flapjax/aero/aic.py

def compute_v_ind[T: Array | ArrayList](
    cs: T,
    zetas: ArrayList,
    gammas: ArrayList,
    kernels: Sequence[KernelFunction],
    mirror_point: Array | None,
    mirror_normal: Array | None,
    batch_size: int | None,
) -> T:
    """
    Compute the induced velocity by multiple surfaces of aerodynamic elements at one or multiple grids of points in
    space. This is done without materialising the full AIC matrix, instead directly computing its contraction with the circulation strength.
    :param cs: Points at which to sample the velocity, ``(c_m, c_n, 3)`` or ``(n_target,)(c_m, c_n, 3)``.
    :param zetas: Filament grid, ``(n_source,)(zeta_m, zeta_n, 3)``.
    :param gammas: Circulation strengths, ``(n_source,)(zeta_m, zeta_n)``.
    :param kernels: Kernel function.
    :param mirror_point: Mirror point, ``(3, )``. If None, no mirroring will be done.
    :param mirror_normal: Normal mirror vector, ``(3, )``. If None, no mirroring will be done.
    :param batch_size: Batch size for vectorising AIC computations.
    :return: Array or ArrayList of induced velocity, ``(c_m, c_n, 3)`` or ``(n_target,)(c_m, c_n, 3)``.
    """

    # convert cs to an ArrayList. If it is an Array, we will convert back before returning.
    cs_: ArrayList = ArrayList([cs]) if isinstance(cs, Array) else cs

    v = ArrayList([])
    for c in cs_:
        v.append(jnp.zeros_like(c))
        for zeta, gamma, kernel in zip(zetas, gammas, kernels):
            m_vect = jnp.stack(
                (zeta[:-1, :, :], zeta[1:, :, :]), axis=-2
            )  # [m, n+1, 2, 3]
            n_vect = jnp.stack(
                (zeta[:, :-1, :], zeta[:, 1:, :]), axis=-2
            )  # [m+1, n, 2, 3]

            gamma_eff_m = jnp.diff(jnp.pad(gamma, ((0, 0), (1, 1))), axis=1)  # [m, n+1]
            gamma_eff_n = -jnp.diff(
                jnp.pad(gamma, ((1, 1), (0, 0))), axis=0
            )  # [m+1, n]

            v[-1] += v_ind_vmap(
                c, m_vect, gamma_eff_m, kernel, batch_size
            ) + v_ind_vmap(c, n_vect, gamma_eff_n, kernel, batch_size)

            if mirror_point is not None and mirror_normal is not None:
                zeta_mirror = mirror_grid(
                    zeta=zeta,
                    mirror_point=mirror_point,
                    mirror_normal=mirror_normal,
                )
                m_vect_mirror = jnp.stack(
                    (zeta_mirror[:-1, :, :], zeta_mirror[1:, :, :]), axis=-2
                )  # [m, n+1, 2, 3]
                n_vect_mirror = jnp.stack(
                    (zeta_mirror[:, :-1, :], zeta_mirror[:, 1:, :]), axis=-2
                )  # [m+1, n, 2, 3]

                v[-1] -= v_ind_vmap(
                    c, m_vect_mirror, gamma_eff_m, kernel, batch_size
                ) + v_ind_vmap(c, n_vect_mirror, gamma_eff_n, kernel, batch_size)

    return v[0] if isinstance(cs, Array) else v

data_structures

GridDiscretisation dataclass

GridDiscretisation(m: int, n: int, m_star: int)

Data class to hold discretisation parameters for each aerodynamic grid.

Parameters:

Name	Type	Description	Default
m	int	Number of panels in the chordwise direction.	required
n	int	Number of panels in the spanwise direction.	required
m_star	int	Number of wake panels in the chordwise direction.	required

AeroCase

AeroCase(
    zeta_b: ArrayList,
    zeta_b_dot: ArrayList,
    zeta_w: ArrayList,
    c: ArrayList | None,
    n: ArrayList | None,
    gamma_b: ArrayList,
    gamma_b_dot: ArrayList | None,
    gamma_w: ArrayList,
    f_steady: ArrayList,
    f_unsteady: ArrayList | None,
    alpha: ArrayList | None,
    cl: ArrayList | None,
    cd: ArrayList | None,
    cm: ArrayList | None,
    cs_ang: dict[str, Array],
    cs_vel: dict[str, Array],
    kernels: Sequence[KernelFunction],
    mirror_point: Array | None,
    mirror_normal: Array | None,
    mirror_edge_low: ArrayList | None,
    mirror_edge_high: ArrayList | None,
    flowfield: FlowField,
    surf_b_names: Sequence[str],
    surf_w_names: Sequence[str],
    t: Array,
    i_ts: Array | int,
    dof_mapping: ArrayList,
    static_horseshoe: bool,
    free_wake: bool,
    gamma_dot_relaxation: float | Array,
    batch_size: int | None,
)

Contains an aerodynamic solution across one or many timesteps.

A single instance may represent either:

• Snapshot (single timestep): array leaves within have no leading time axis (e.g. zeta_b[i_surf].shape == (m+1, n+1, 3), gamma_b[i_surf].shape == (m, n)). t is a scalar and i_ts is an integer
• Batched (many timesteps): array leaves carry a leading n_tstep axis (e.g. zeta_b[i_surf].shape == (n_tstep, m+1, n+1, 3)). t is (n_tstep,) and i_ts is a (n_tstep, ) array of indices.

Use is_batched to distinguish at runtime.

Parameters:

Name	Type	Description	Default
zeta_b	ArrayList	Bound grid coordinates, batched: (n_surf, )(n_tstep, zeta_m, zeta_n, 3) / snapshot: (n_surf, )(zeta_m, zeta_n, 3).	required
zeta_b_dot	ArrayList	Bound grid velocities, same layout as zeta_b.	required
zeta_w	ArrayList	Wake grid coordinates or None.	required
c	ArrayList | None	Bound collocation points or None.	required
n	ArrayList | None	Bound grid normals or None.	required
gamma_b	ArrayList	Bound circulation strengths, batched: (n_surf, )(n_tstep, m, n) / snapshot: (n_surf, )(m, n).	required
gamma_b_dot	ArrayList | None	Bound circulation time derivatives or None.	required
gamma_w	ArrayList	Wake circulation strengths.	required
f_steady	ArrayList	Steady force contributions.	required
f_unsteady	ArrayList | None	Unsteady force contributions or None.	required
alpha	ArrayList | None	Per-strip effective angle of attack extracted from the UVLM sectional lift; batched: (n_surf, )(n_tstep, n), snapshot: (n_surf, )(n, ), or None.	required
cl	ArrayList | None	Per-strip lift coefficient sampled from the airfoil polars, or None for no polars.	required
cd	ArrayList | None	Per-strip drag coefficient sampled from the airfoil polars, or None for no polars.	required
cm	ArrayList | None	Per-strip moment coefficient sampled from the airfoil polars, or None for no polars.	required
cs_ang	dict[str, Array]	Control surface angle time history, {name: (n_tstep,)} (batched) or {name: ()} (snapshot).	required
cs_vel	dict[str, Array]	Control surface velocity time history.	required
kernels	Sequence[KernelFunction]	Kernel functions for both bound and wake source grids.	required
mirror_point	Array | None	Point on mirror plane, (3, ) or None.	required
mirror_normal	Array | None	Normal on mirror plane, (3, ) or None.	required
mirror_edge_low	ArrayList | None	Per-surface booleans marking whether that surface's n=0 edge lies on the mirror plane, (n_surf, )(), or None.	required
mirror_edge_high	ArrayList | None	As mirror_edge_low, for the n=-1 edge.	required
flowfield	FlowField	FlowField object which includes background velocity and density.	required
surf_b_names	Sequence[str]	Names of bound surfaces, (n_surf, ).	required
surf_w_names	Sequence[str]	Names of wake surfaces, (n_surf, ).	required
t	Array	Time; batched: (n_tstep, ), snapshot: scalar.	required
i_ts	Array | int	Timestep index; batched: (n_tstep, ), snapshot: int.	required
dof_mapping	ArrayList	Map from aero grid to beam DOFs, (n_surf, )(zeta_n, ).	required
static_horseshoe	bool	If true, a horseshoe formulation was used for the initial static solution.	required
free_wake	bool	Free-wake formulation flag.	required
gamma_dot_relaxation	float | Array	Circulation time derivative filter.	required
batch_size	int | None	Batch size used for AIC vectorisation.	required

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

def __init__(
    self,
    zeta_b: ArrayList,
    zeta_b_dot: ArrayList,
    zeta_w: ArrayList,
    c: ArrayList | None,
    n: ArrayList | None,
    gamma_b: ArrayList,
    gamma_b_dot: ArrayList | None,
    gamma_w: ArrayList,
    f_steady: ArrayList,
    f_unsteady: ArrayList | None,
    alpha: ArrayList | None,
    cl: ArrayList | None,
    cd: ArrayList | None,
    cm: ArrayList | None,
    cs_ang: dict[str, Array],
    cs_vel: dict[str, Array],
    kernels: Sequence[KernelFunction],
    mirror_point: Array | None,
    mirror_normal: Array | None,
    mirror_edge_low: ArrayList | None,
    mirror_edge_high: ArrayList | None,
    flowfield: FlowField,
    surf_b_names: Sequence[str],
    surf_w_names: Sequence[str],
    t: Array,
    i_ts: Array | int,
    dof_mapping: ArrayList,
    static_horseshoe: bool,
    free_wake: bool,
    gamma_dot_relaxation: float | Array,
    batch_size: int | None,
) -> None:
    r"""
    :param zeta_b: Bound grid coordinates, batched: ``(n_surf, )(n_tstep, zeta_m, zeta_n, 3)`` /
        snapshot: ``(n_surf, )(zeta_m, zeta_n, 3)``.
    :param zeta_b_dot: Bound grid velocities, same layout as ``zeta_b``.
    :param zeta_w: Wake grid coordinates or ``None``.
    :param c: Bound collocation points or ``None``.
    :param n: Bound grid normals or ``None``.
    :param gamma_b: Bound circulation strengths, batched: ``(n_surf, )(n_tstep, m, n)`` /
        snapshot: ``(n_surf, )(m, n)``.
    :param gamma_b_dot: Bound circulation time derivatives or ``None``.
    :param gamma_w: Wake circulation strengths.
    :param f_steady: Steady force contributions.
    :param f_unsteady: Unsteady force contributions or ``None``.
    :param alpha: Per-strip effective angle of attack extracted from the UVLM sectional lift; batched:
        ``(n_surf, )(n_tstep, n)``, snapshot: ``(n_surf, )(n, )``, or ``None``.
    :param cl: Per-strip lift coefficient sampled from the airfoil polars, or ``None`` for no polars.
    :param cd: Per-strip drag coefficient sampled from the airfoil polars, or ``None`` for no polars.
    :param cm: Per-strip moment coefficient sampled from the airfoil polars, or ``None`` for no polars.
    :param cs_ang: Control surface angle time history, ``{name: (n_tstep,)}`` (batched) or ``{name: ()}`` (snapshot).
    :param cs_vel: Control surface velocity time history.
    :param kernels: Kernel functions for both bound and wake source grids.
    :param mirror_point: Point on mirror plane, ``(3, )`` or None.
    :param mirror_normal: Normal on mirror plane, ``(3, )`` or None.
    :param mirror_edge_low: Per-surface booleans marking whether that surface's ``n=0`` edge lies on the
        mirror plane, ``(n_surf, )()``, or None.
    :param mirror_edge_high: As ``mirror_edge_low``, for the ``n=-1`` edge.
    :param flowfield: ``FlowField`` object which includes background velocity and density.
    :param surf_b_names: Names of bound surfaces, ``(n_surf, )``.
    :param surf_w_names: Names of wake surfaces, ``(n_surf, )``.
    :param t: Time; batched: ``(n_tstep, )``, snapshot: scalar.
    :param i_ts: Timestep index; batched: ``(n_tstep, )``, snapshot: ``int``.
    :param dof_mapping: Map from aero grid to beam DOFs, ``(n_surf, )(zeta_n, )``.
    :param static_horseshoe: If true, a horseshoe formulation was used for the initial static solution.
    :param free_wake: Free-wake formulation flag.
    :param gamma_dot_relaxation: Circulation time derivative filter.
    :param batch_size: Batch size used for AIC vectorisation.
    """
    self.zeta_b: ArrayList = zeta_b
    self.zeta_b_dot: ArrayList = zeta_b_dot
    self.zeta_w: ArrayList = zeta_w
    self.c: ArrayList | None = c
    self.nc: ArrayList | None = n
    self.gamma_b: ArrayList = gamma_b
    self.gamma_b_dot: ArrayList | None = gamma_b_dot
    self.gamma_w: ArrayList = gamma_w
    self.f_steady: ArrayList = f_steady
    self.f_unsteady: ArrayList | None = f_unsteady
    self.alpha: ArrayList | None = alpha
    self.cl: ArrayList | None = cl
    self.cd: ArrayList | None = cd
    self.cm: ArrayList | None = cm
    self.cs_ang: dict[str, Array] = cs_ang
    self.cs_vel: dict[str, Array] = cs_vel
    self.t: Array = t
    self.i_ts: Array | int = i_ts

    self.kernels: Sequence[KernelFunction] = kernels
    self.mirror_point: Array | None = mirror_point
    self.mirror_normal: Array | None = mirror_normal
    self.mirror_edge_low: ArrayList | None = mirror_edge_low
    self.mirror_edge_high: ArrayList | None = mirror_edge_high
    self.flowfield: FlowField = flowfield
    self.surf_b_names: Sequence[str] = surf_b_names
    self.surf_w_names: Sequence[str] = surf_w_names
    self.dof_mapping: ArrayList = dof_mapping

    # settings
    self.static_horseshoe: bool = static_horseshoe
    self.free_wake: bool = free_wake
    self.gamma_dot_relaxation: float | Array = gamma_dot_relaxation
    self.batch_size: int | None = batch_size

get_states

get_states(
    i_ts: int | Array | None = None,
) -> AeroFullStates

Obtain the aerodynamic state at a given timestep (used in the adjoint solution).

Parameters:

Name	Type	Description	Default
i_ts	int | Array | None	Time step index (required for batched, ignored for snapshot).	None

Returns:

Type	Description
AeroFullStates	Aero states.

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

def get_states(self, i_ts: int | Array | None = None) -> AeroFullStates:
    r"""
    Obtain the aerodynamic state at a given timestep (used in the adjoint solution).
    :param i_ts: Time step index (required for batched, ignored for snapshot).
    :return: Aero states.
    """
    if self.is_batched:
        if i_ts is None:
            raise ValueError("i_ts must be provided for batched AeroCase")
        assert self.gamma_b_dot is not None and self.zeta_w is not None
        return AeroFullStates(
            gamma_b=self.gamma_b.index_all(i_ts, ...),
            gamma_w=self.gamma_w.index_all(i_ts, ...),
            gamma_b_dot=self.gamma_b_dot.index_all(i_ts, ...),
            zeta_w=self.zeta_w.index_all(i_ts, ...),
        )
    assert self.gamma_b_dot is not None and self.zeta_w is not None
    return AeroFullStates(
        gamma_b=self.gamma_b,
        gamma_w=self.gamma_w,
        gamma_b_dot=self.gamma_b_dot,
        zeta_w=self.zeta_w,
    )

gamma_full

gamma_full(i_ts: int | None = None) -> ArrayList

Concatenate bound and wake circulation strengths.

Parameters:

Name	Type	Description	Default
i_ts	int | None	Time step index (required for batched, ignored for snapshot).	None

Returns:

Type	Description
ArrayList	Circulation strength, (2 * n_surf,)(m | m_star, n).

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

def gamma_full(self, i_ts: int | None = None) -> ArrayList:
    r"""Concatenate bound and wake circulation strengths.
    :param i_ts: Time step index (required for batched, ignored for snapshot).
    :return: Circulation strength, ``(2 * n_surf,)(m | m_star, n)``.
    """
    if self.is_batched:
        if i_ts is None:
            raise ValueError("i_ts must be provided for batched AeroCase")
        return ArrayList(
            [
                *self.gamma_b.index_all(i_ts, ...),
                *self.gamma_w.index_all(i_ts, ...),
            ]
        )
    return ArrayList([*self.gamma_b, *self.gamma_w])

zeta_full

zeta_full(i_ts: int | None = None) -> ArrayList

Concatenate bound and wake grids.

Parameters:

Name	Type	Description	Default
i_ts	int | None	Time step index (required for batched, ignored for snapshot).	None

Returns:

Type	Description
ArrayList	Grids, (2 * n_surf,)(zeta_m | zeta_m_star, zeta_n, 3).

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

def zeta_full(self, i_ts: int | None = None) -> ArrayList:
    r"""Concatenate bound and wake grids.
    :param i_ts: Time step index (required for batched, ignored for snapshot).
    :return: Grids, ``(2 * n_surf,)(zeta_m | zeta_m_star, zeta_n, 3)``.
    """
    if self.is_batched:
        if i_ts is None:
            raise ValueError("i_ts must be provided for batched AeroCase")
        assert self.zeta_w is not None
        return ArrayList(
            [
                *self.zeta_b.index_all(i_ts, ...),
                *self.zeta_w.index_all(i_ts, ...),
            ]
        )
    assert self.zeta_w is not None
    return ArrayList([*self.zeta_b, *self.zeta_w])

set_arraylist_at_ts

set_arraylist_at_ts(
    attr: str, values: ArrayList, i_ts: int
) -> None

Set an attribute at a given timestep on a batched AeroCase.

Parameters:

Name	Type	Description	Default
attr	str	Name of the attribute to set.	required
values	ArrayList	ArrayList of per-surface values (no leading time axis).	required
i_ts	int	Time step index.	required

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

def set_arraylist_at_ts(self, attr: str, values: ArrayList, i_ts: int) -> None:
    """Set an attribute at a given timestep on a batched AeroCase.
    :param attr: Name of the attribute to set.
    :param values: ArrayList of per-surface values (no leading time axis).
    :param i_ts: Time step index.
    """
    if not self.is_batched:
        raise TypeError("set_arraylist_at_ts only supported for batched AeroCase")
    arr = getattr(self, attr)
    for i_surf, val in enumerate(values):
        arr[i_surf] = arr[i_surf].at[i_ts, ...].set(val)

get_surf_snapshot

get_surf_snapshot(
    i_ts: int, i_surf: int
) -> _AeroSurfacePlot

Get single-surface plot data for a given (timestep, surface) pair on a batched AeroCase.

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

def get_surf_snapshot(self, i_ts: int, i_surf: int) -> _AeroSurfacePlot:
    r"""Get single-surface plot data for a given ``(timestep, surface)`` pair on
    a batched AeroCase.
    """
    if not self.is_batched:
        raise TypeError("get_surf_snapshot only supported for batched AeroCase")
    assert self.zeta_w is not None
    assert self.gamma_b_dot is not None
    assert self.f_unsteady is not None
    return _AeroSurfacePlot(
        zeta_b=self.zeta_b[i_surf][i_ts, ...],
        zeta_b_dot=self.zeta_b_dot[i_surf][i_ts, ...],
        zeta_w=self.zeta_w[i_surf][i_ts, ...],
        gamma_b=self.gamma_b[i_surf][i_ts, ...],
        gamma_b_dot=self.gamma_b_dot[i_surf][i_ts, ...],
        gamma_w=self.gamma_w[i_surf][i_ts, ...],
        f_steady=self.f_steady[i_surf][i_ts, ...],
        f_unsteady=self.f_unsteady[i_surf][i_ts, ...],
        alpha=self.alpha[i_surf][i_ts, ...],
        cl=self.cl[i_surf][i_ts, ...],
        cd=self.cd[i_surf][i_ts, ...],
        cm=self.cm[i_surf][i_ts, ...],
        surf_b_name=self.surf_b_names[i_surf],
        surf_w_name=self.surf_w_names[i_surf],
        i_ts=i_ts,
    )

get_surface

get_surface(idx: int) -> _AeroSurfacePlot

Get single-surface plot data for the given surface on a snapshot AeroCase.

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

def get_surface(self, idx: int) -> _AeroSurfacePlot:
    """Get single-surface plot data for the given surface on a snapshot
    AeroCase."""
    if self.is_batched:
        raise TypeError("get_surface only supported for snapshot AeroCase")
    assert self.zeta_w is not None
    assert self.gamma_b_dot is not None
    assert self.f_unsteady is not None
    return _AeroSurfacePlot(
        zeta_b=self.zeta_b[idx],
        zeta_b_dot=self.zeta_b_dot[idx],
        zeta_w=self.zeta_w[idx],
        gamma_b=self.gamma_b[idx],
        gamma_b_dot=self.gamma_b_dot[idx],
        gamma_w=self.gamma_w[idx],
        f_steady=self.f_steady[idx],
        f_unsteady=self.f_unsteady[idx],
        alpha=self.alpha[idx],
        cl=self.cl[idx],
        cd=self.cd[idx],
        cm=self.cm[idx],
        surf_b_name=self.surf_b_names[idx],
        surf_w_name=self.surf_w_names[idx],
        i_ts=int(self.i_ts),
    )

plot

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

Plot aerodynamic surfaces to VTU files (with per-surface PVD when batched).

Parameters:

Name	Type	Description	Default
directory	PathLike | str	Directory to save files.	required
plot_bound	bool	If True, plot the bound surfaces.	True
plot_wake	bool	If True, plot the wake surfaces.	True
index	int | Sequence[int] | Array | slice | None	For batched, timestep indices to plot (all if None). Ignored for snapshots.	None

Returns:

Type	Description
Sequence[Path]	Sequence of paths to the saved PVD (batched) or VTU (snapshot) files.

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

def plot(
    self,
    directory: os.PathLike | str,
    plot_bound: bool = True,
    plot_wake: bool = True,
    index: int | Sequence[int] | Array | slice | None = None,
) -> Sequence[Path]:
    r"""Plot aerodynamic surfaces to VTU files (with per-surface PVD when
    batched).
    :param directory: Directory to save files.
    :param plot_bound: If True, plot the bound surfaces.
    :param plot_wake: If True, plot the wake surfaces.
    :param index: For batched, timestep indices to plot (all if None).
        Ignored for snapshots.
    :return: Sequence of paths to the saved PVD (batched) or VTU (snapshot) files.
    """
    directory_path = Path(directory)
    directory_path.mkdir(parents=True, exist_ok=True)

    if not self.is_batched:
        paths: list[Path] = []
        for i_surf in range(self.n_surf):
            paths.extend(
                self.get_surface(idx=i_surf).plot(
                    directory, plot_bound=plot_bound, plot_wake=plot_wake
                )
            )
        return paths

    index_ = index_to_arr(index=index, n_entries=self.n_tstep)
    pvd_paths: list[Path] = []
    for i_surf in range(self.n_surf):
        per_ts_paths: list[Sequence[Path]] = []
        for i_ts in index_:
            per_ts_paths.append(
                self.get_surf_snapshot(i_ts=i_ts, i_surf=i_surf).plot(
                    directory, plot_bound=plot_bound, plot_wake=plot_wake
                )
            )

        if plot_bound:
            bound_name = f"aero_dynamic_{self.surf_b_names[i_surf]}_ts"
            pvd_paths.append(
                write_pvd(
                    directory=directory,
                    name=bound_name,
                    file_dirs=next(zip(*per_ts_paths)),
                    times=list(self.t[index_]),
                )
            )

        if plot_wake:
            wake_name = f"aero_dynamic_{self.surf_w_names[i_surf]}_ts"
            pvd_paths.append(
                write_pvd(
                    directory=directory,
                    name=wake_name,
                    file_dirs=list(zip(*per_ts_paths))[-1],
                    times=list(self.t[index_]),
                )
            )
    return pvd_paths

project_forcing_to_beam

project_forcing_to_beam(
    i_ts: int,
    rmat: Array,
    x0_aero: ArrayList,
    include_unsteady: bool,
) -> Array

Project aerodynamic forcing at i_ts onto the beam grid (global frame).

Parameters:

Name	Type	Description	Default
i_ts	int	Time step index (ignored for snapshot).	required
rmat	Array	Rotation matrix for each node relative to reference, (n_nodes, 3, 3).	required
x0_aero	ArrayList	Reference coordinates for aerodynamic grid, (n_surf, )(zeta_m, zeta_n, 3).	required
include_unsteady	bool	If true, include unsteady forcing.	required

Returns:

Type	Description
Array	Steady and unsteady forcing projected onto the beam grid, (n_nodes, 6).

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

def project_forcing_to_beam(
    self,
    i_ts: int,
    rmat: Array,
    x0_aero: ArrayList,
    include_unsteady: bool,
) -> Array:
    r"""Project aerodynamic forcing at ``i_ts`` onto the beam grid (global frame).
    :param i_ts: Time step index (ignored for snapshot).
    :param rmat: Rotation matrix for each node relative to reference, ``(n_nodes, 3, 3)``.
    :param x0_aero: Reference coordinates for aerodynamic grid, ``(n_surf, )(zeta_m, zeta_n, 3)``.
    :param include_unsteady: If true, include unsteady forcing.
    :return: Steady and unsteady forcing projected onto the beam grid, ``(n_nodes, 6)``.
    """
    if self.is_batched:
        f_total = self.f_steady.index_all(i_ts, ...)
        if include_unsteady:
            assert self.f_unsteady is not None
            f_total += self.f_unsteady.index_all(i_ts, ...)
    else:
        f_total = self.f_steady
        if include_unsteady:
            assert self.f_unsteady is not None
            f_total = ArrayList([a + b for a, b in zip(f_total, self.f_unsteady)])

    return project_forcing_to_beam(
        f_total=f_total,
        rmat=rmat,
        x0_aero=x0_aero,
        dof_mapping=self.dof_mapping,
        mirror_edge_low=self.mirror_edge_low,
        mirror_edge_high=self.mirror_edge_high,
    )

get_v_background

get_v_background(x_target: T, i_ts: int | None = None) -> T

Background velocity at specified points and time step.

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

def get_v_background[T: Array | ArrayList](
    self, x_target: T, i_ts: int | None = None
) -> T:
    r"""Background velocity at specified points and time step."""
    t_val = self._t_at(i_ts)
    if isinstance(x_target, Array):
        return self.flowfield.vmap_call(x=x_target, t=t_val)
    elif isinstance(x_target, ArrayList):
        return self.flowfield.surf_vmap_call(xs=x_target, t=t_val)  # type: ignore
    raise NotImplementedError

get_v_ind

get_v_ind(x_target: T, i_ts: int | None = None) -> T

Induced velocity at specified points and time step.

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

def get_v_ind[T: Array | ArrayList](
    self, x_target: T, i_ts: int | None = None
) -> T:
    r"""Induced velocity at specified points and time step."""
    return compute_v_ind(
        cs=x_target,
        zetas=self.zeta_full(i_ts),
        gammas=self.gamma_full(i_ts),
        kernels=self.kernels,
        mirror_normal=self.mirror_normal,
        mirror_point=self.mirror_point,
        batch_size=self.batch_size,
    )

get_v_tot

get_v_tot(x_target: T, i_ts: int | None = None) -> T

Total (induced + background) velocity at specified points and time step.

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

def get_v_tot[T: Array | ArrayList](
    self, x_target: T, i_ts: int | None = None
) -> T:
    r"""Total (induced + background) velocity at specified points and time step."""
    return self.get_v_ind(x_target=x_target, i_ts=i_ts) + self.get_v_background(
        x_target=x_target, i_ts=i_ts
    )

to_dynamic

to_dynamic(i_ts: int, n_tstep: int) -> AeroCase

Expand this snapshot into a batched AeroCase with n_tstep timesteps, placing the current snapshot at index i_ts.

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

def to_dynamic(self, i_ts: int, n_tstep: int) -> AeroCase:
    """Expand this snapshot into a batched AeroCase with ``n_tstep``
    timesteps, placing the current snapshot at index ``i_ts``.
    """
    if self.is_batched:
        raise TypeError("to_dynamic only supported for snapshot AeroCase")

    def _expand(arr_list: ArrayList) -> ArrayList:
        out = []
        for a in arr_list:
            out.append(jnp.zeros((n_tstep, *a.shape)).at[i_ts, ...].set(a))
        return ArrayList(out)

    return AeroCase(
        zeta_b=_expand(self.zeta_b),
        zeta_b_dot=_expand(self.zeta_b_dot),
        zeta_w=_expand(self.zeta_w),
        c=_expand(self.c),
        n=_expand(self.nc),
        gamma_b=_expand(self.gamma_b),
        gamma_b_dot=_expand(self.gamma_b_dot),
        gamma_w=_expand(self.gamma_w),
        f_steady=_expand(self.f_steady),
        f_unsteady=_expand(self.f_unsteady),
        alpha=_expand(self.alpha),
        cl=_expand(self.cl),
        cd=_expand(self.cd),
        cm=_expand(self.cm),
        cs_ang={k: jnp.full(n_tstep, v) for k, v in self.cs_ang.items()},
        cs_vel={k: jnp.full(n_tstep, v) for k, v in self.cs_vel.items()},
        kernels=self.kernels,
        mirror_point=self.mirror_point,
        mirror_normal=self.mirror_normal,
        mirror_edge_low=self.mirror_edge_low,
        mirror_edge_high=self.mirror_edge_high,
        flowfield=self.flowfield,
        surf_b_names=self.surf_b_names,
        surf_w_names=self.surf_w_names,
        t=jnp.zeros(n_tstep).at[i_ts].set(self.t),
        i_ts=jnp.arange(n_tstep),
        dof_mapping=self.dof_mapping,
        static_horseshoe=self.static_horseshoe,
        free_wake=self.free_wake,
        gamma_dot_relaxation=self.gamma_dot_relaxation,
        batch_size=self.batch_size,
    )

initialise classmethod

initialise(
    initial_snapshot: AeroCase, n_tstep: int
) -> AeroCase

Create a batched AeroCase from a snapshot placed at i_ts=0.

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

@classmethod
def initialise(cls, initial_snapshot: AeroCase, n_tstep: int) -> AeroCase:
    r"""Create a batched AeroCase from a snapshot placed at ``i_ts=0``."""
    return initial_snapshot.to_dynamic(i_ts=0, n_tstep=n_tstep)

flowfields

FlowField

FlowField(
    u_inf: Array,
    rho: float | Array,
    relative_motion: bool,
    mach: float | Array = 0.0,
)

Base class for background flow field definitions. Allows for the definition of arbitrary flow fields which are functions of space and time.

Parameters:

Name	Type	Description	Default
u_inf	Array	Base flow velocity, (3, ).	required
rho	float | Array	Flow density.	required
relative_motion	bool	If True, the air moves, if False, the plane moves.	required
mach	float | Array	Freestream Mach number, used to apply a Prandtl-Glauert compressibility correction. Defaults to 0 (incompressible).	0.0

Source code in src/flapjax/aero/flowfields.py

def __init__(
    self,
    u_inf: Array,
    rho: float | Array,
    relative_motion: bool,
    mach: float | Array = 0.0,
):
    r"""
    :param u_inf: Base flow velocity, ``(3, )``.
    :param rho: Flow density.
    :param relative_motion: If True, the air moves, if False, the plane moves.
    :param mach: Freestream Mach number, used to apply a Prandtl-Glauert compressibility correction.
    Defaults to 0 (incompressible).
    """
    check_arr_shape(u_inf, (3,), name="u_inf")
    self.u_inf: Array = u_inf
    self.rho: Array = jnp.array(rho)
    self.u_inf_mag: Array = jnp.linalg.norm(u_inf)
    self.u_inf_dir: Array = u_inf / self.u_inf_mag
    self.q_inf: Array = 0.5 * rho * self.u_inf_mag**2  # dynamic pressure
    self.relative_motion: bool = relative_motion

    if isinstance(mach, (int, float)) and mach >= 1.0:
        warn(
            "Prandtl-Glauert compressibility correction requires subsonic flow (Mach < 1)."
        )
    self.mach: Array = jnp.array(mach)
    self.beta: Array = jnp.sqrt(
        1.0 - self.mach**2
    )  # Prandtl-Glauert compressibility factor

vmap_call

vmap_call(x: Array, t: Array) -> Array

Vectorized version of the call method. This maps over all leading dimensions of x.

Parameters:

Name	Type	Description	Default
x	Array	Spatial coordinates, (..., 3)	required
t	Array	Time, ()	required

Returns:

Type	Description
Array	Flow field values at the specified coordinates, (..., 3)

Source code in src/flapjax/aero/flowfields.py

def vmap_call(self, x: Array, t: Array) -> Array:
    """
    Vectorized version of the __call__ method. This maps over all leading dimensions of x.
    :param x: Spatial coordinates, ``(..., 3)``
    :param t: Time, ()
    :return: Flow field values at the specified coordinates, ``(..., 3)``
    """
    n_vmap = x.ndim - 1
    func = self.__call__
    for i_dim in range(n_vmap):
        func = jax.vmap(func, in_axes=(i_dim, None), out_axes=i_dim)
    return func(x, t)

surf_vmap_call

surf_vmap_call(xs: ArrayList, t: Array) -> ArrayList

Vectorized version of the call method over a list of surfaces.

Parameters:

Name	Type	Description	Default
xs	ArrayList	Spatial coordinates, (n_surf,)(..., 3)	required
t	Array	Time, ()	required

Returns:

Type	Description
ArrayList	Flow field values at the specified coordinates, (n_surf, )(..., 3)

Source code in src/flapjax/aero/flowfields.py

def surf_vmap_call(self, xs: ArrayList, t: Array) -> ArrayList:
    """
    Vectorized version of the __call__ method over a list of surfaces.
    :param xs: Spatial coordinates, ``(n_surf,)(..., 3)``
    :param t: Time, ()
    :return: Flow field values at the specified coordinates, ``(n_surf, )(..., 3)``
    """
    return ArrayList([self.vmap_call(x, t) for x in xs])

to_design_variables

to_design_variables() -> dict[str, Array]

Extract the design variables associated with this flow field.

Returns:

Type	Description
dict[str, Array]	Dictionary of design variables.

Source code in src/flapjax/aero/flowfields.py

def to_design_variables(self) -> dict[str, Array]:
    r"""
    Extract the design variables associated with this flow field.
    :return: Dictionary of design variables.
    """
    return {"u_inf": self.u_inf, "rho": self.rho, "mach": self.mach}

from_design_variables

from_design_variables(
    design_variables: dict[str, Array],
) -> FlowField

Create a new flow field from design variables as the inverse of self.to_design_variables().

Parameters:

Name	Type	Description	Default
design_variables	dict[str, Array]	Dictionary of design variables.	required

Returns:

Type	Description
FlowField	New FlowField object.

Source code in src/flapjax/aero/flowfields.py

def from_design_variables(self, design_variables: dict[str, Array]) -> FlowField:
    r"""
    Create a new flow field from design variables as the inverse of ``self.to_design_variables()``.
    :param design_variables: Dictionary of design variables.
    :return: New FlowField object.
    """
    return self.__class__(**design_variables, relative_motion=self.relative_motion)

ConstantFlowField

ConstantFlowField(
    u_inf: Array,
    rho: float | Array,
    relative_motion: bool,
    mach: float | Array = 0.0,
)

Bases: FlowField

Constant velocity flow field.

Parameters:

Name	Type	Description	Default
u_inf	Array	Base flow velocity, (3, ).	required
rho	float | Array	Flow density.	required
relative_motion	bool	If True, the air moves, if False, the plane moves.	required
mach	float | Array	Freestream Mach number, used to apply a Prandtl-Glauert compressibility correction. Defaults to 0 (incompressible).	0.0

Source code in src/flapjax/aero/flowfields.py

def __init__(
    self,
    u_inf: Array,
    rho: float | Array,
    relative_motion: bool,
    mach: float | Array = 0.0,
):
    r"""
    :param u_inf: Base flow velocity, ``(3, )``.
    :param rho: Flow density.
    :param relative_motion: If True, the air moves, if False, the plane moves.
    :param mach: Freestream Mach number, used to apply a Prandtl-Glauert compressibility correction.
    Defaults to 0 (incompressible).
    """
    check_arr_shape(u_inf, (3,), name="u_inf")
    self.u_inf: Array = u_inf
    self.rho: Array = jnp.array(rho)
    self.u_inf_mag: Array = jnp.linalg.norm(u_inf)
    self.u_inf_dir: Array = u_inf / self.u_inf_mag
    self.q_inf: Array = 0.5 * rho * self.u_inf_mag**2  # dynamic pressure
    self.relative_motion: bool = relative_motion

    if isinstance(mach, (int, float)) and mach >= 1.0:
        warn(
            "Prandtl-Glauert compressibility correction requires subsonic flow (Mach < 1)."
        )
    self.mach: Array = jnp.array(mach)
    self.beta: Array = jnp.sqrt(
        1.0 - self.mach**2
    )  # Prandtl-Glauert compressibility factor

vmap_call

vmap_call(x: Array, t: Array) -> Array

Vectorized version of the call method. This maps over all leading dimensions of x.

Parameters:

Name	Type	Description	Default
x	Array	Spatial coordinates, (..., 3)	required
t	Array	Time, ()	required

Returns:

Type	Description
Array	Flow field values at the specified coordinates, (..., 3)

Source code in src/flapjax/aero/flowfields.py

def vmap_call(self, x: Array, t: Array) -> Array:
    """
    Vectorized version of the __call__ method. This maps over all leading dimensions of x.
    :param x: Spatial coordinates, ``(..., 3)``
    :param t: Time, ()
    :return: Flow field values at the specified coordinates, ``(..., 3)``
    """
    n_vmap = x.ndim - 1
    func = self.__call__
    for i_dim in range(n_vmap):
        func = jax.vmap(func, in_axes=(i_dim, None), out_axes=i_dim)
    return func(x, t)

surf_vmap_call

surf_vmap_call(xs: ArrayList, t: Array) -> ArrayList

Vectorized version of the call method over a list of surfaces.

Parameters:

Name	Type	Description	Default
xs	ArrayList	Spatial coordinates, (n_surf,)(..., 3)	required
t	Array	Time, ()	required

Returns:

Type	Description
ArrayList	Flow field values at the specified coordinates, (n_surf, )(..., 3)

Source code in src/flapjax/aero/flowfields.py

def surf_vmap_call(self, xs: ArrayList, t: Array) -> ArrayList:
    """
    Vectorized version of the __call__ method over a list of surfaces.
    :param xs: Spatial coordinates, ``(n_surf,)(..., 3)``
    :param t: Time, ()
    :return: Flow field values at the specified coordinates, ``(n_surf, )(..., 3)``
    """
    return ArrayList([self.vmap_call(x, t) for x in xs])

to_design_variables

to_design_variables() -> dict[str, Array]

Extract the design variables associated with this flow field.

Returns:

Type	Description
dict[str, Array]	Dictionary of design variables.

Source code in src/flapjax/aero/flowfields.py

def to_design_variables(self) -> dict[str, Array]:
    r"""
    Extract the design variables associated with this flow field.
    :return: Dictionary of design variables.
    """
    return {"u_inf": self.u_inf, "rho": self.rho, "mach": self.mach}

from_design_variables

from_design_variables(
    design_variables: dict[str, Array],
) -> FlowField

Create a new flow field from design variables as the inverse of self.to_design_variables().

Parameters:

Name	Type	Description	Default
design_variables	dict[str, Array]	Dictionary of design variables.	required

Returns:

Type	Description
FlowField	New FlowField object.

Source code in src/flapjax/aero/flowfields.py

def from_design_variables(self, design_variables: dict[str, Array]) -> FlowField:
    r"""
    Create a new flow field from design variables as the inverse of ``self.to_design_variables()``.
    :param design_variables: Dictionary of design variables.
    :return: New FlowField object.
    """
    return self.__class__(**design_variables, relative_motion=self.relative_motion)

OneMinusCosineFlowField

OneMinusCosineFlowField(
    u_inf: Array,
    rho: float | Array,
    relative_motion: bool,
    gust_length: float | Array,
    gust_amplitude: float | Array,
    gust_travel_direction: Array | None = None,
    gust_amplitude_direction: Array | None = None,
    gust_x0: Array | None = None,
    mach: float | Array = 0.0,
)

Bases: FlowField

One minus cosine gust flow field.

Parameters:

Name	Type	Description	Default
u_inf	Array	Base flow velocity, (3, ).	required
rho	float | Array	Flow density.	required
relative_motion	bool	If True, the air moves, if False, the plane moves.	required
gust_length	float | Array	Gust length.	required
gust_amplitude	float | Array	Gust amplitude.	required
gust_travel_direction	Array | None	Vector which defines the direction that the gust travels, (3, ). Defaults to the freestream direction if None.	None
gust_amplitude_direction	Array | None	Vector which defines the direction that the gust amplitude acts, (3, ). Defaults to the z-direction if None.	None
gust_x0	Array | None	Coordinate on the initial leading edge of the gust, (3, ). Defaults to 0 if None.	None
mach	float | Array	Freestream Mach number, used to apply a Prandtl-Glauert compressibility correction. Defaults to 0 (incompressible).	0.0

Source code in src/flapjax/aero/flowfields.py

def __init__(
    self,
    u_inf: Array,
    rho: float | Array,
    relative_motion: bool,
    gust_length: float | Array,
    gust_amplitude: float | Array,
    gust_travel_direction: Array | None = None,
    gust_amplitude_direction: Array | None = None,
    gust_x0: Array | None = None,
    mach: float | Array = 0.0,
):
    r"""
    :param u_inf: Base flow velocity, ``(3, )``.
    :param rho: Flow density.
    :param relative_motion: If True, the air moves, if False, the plane moves.
    :param gust_length: Gust length.
    :param gust_amplitude: Gust amplitude.
    :param gust_travel_direction: Vector which defines the direction that the gust travels, ``(3, )``. Defaults to the
    freestream direction if None.
    :param gust_amplitude_direction: Vector which defines the direction that the gust amplitude acts, ``(3, )``. Defaults
    to the z-direction if None.
    :param gust_x0: Coordinate on the initial leading edge of the gust, ``(3, )``. Defaults to 0 if None.
    :param mach: Freestream Mach number, used to apply a Prandtl-Glauert compressibility correction.
    Defaults to 0 (incompressible).
    """
    super().__init__(u_inf, rho, relative_motion, mach=mach)

    # base gust parameters
    self.gust_amplitude: Array = jnp.array(gust_amplitude)
    self.gust_length: Array = jnp.array(gust_length)

    # direction of travel for the gust - use background flow direction as default
    # even for a gust frozen in place, this defines the orientation of the ridge
    self.gust_travel_direction: Array = (
        gust_travel_direction
        if gust_travel_direction is not None
        else self.u_inf_dir
    )
    check_arr_shape(self.gust_travel_direction, (3,), "gust_travel_direction")
    self.gust_travel_direction /= jnp.linalg.norm(self.gust_travel_direction)

    # lateral direction of the gust (direction in which the gust acts), default is in Z
    self.gust_amplitude_direction: Array = (
        jnp.array((0.0, 0.0, 1.0))
        if gust_amplitude_direction is None
        else gust_amplitude_direction
    )
    check_arr_shape(self.gust_amplitude_direction, (3,), "gust_amplitude")
    self.gust_amplitude_direction /= jnp.linalg.norm(self.gust_amplitude_direction)

    # base coordinate at the start of the gust at t=0
    self.gust_x0: Array = gust_x0 if gust_x0 is not None else jnp.zeros(3)
    check_arr_shape(self.gust_x0, (3,), "gust_x0")

vmap_call

vmap_call(x: Array, t: Array) -> Array

Vectorized version of the call method. This maps over all leading dimensions of x.

Parameters:

Name	Type	Description	Default
x	Array	Spatial coordinates, (..., 3)	required
t	Array	Time, ()	required

Returns:

Type	Description
Array	Flow field values at the specified coordinates, (..., 3)

Source code in src/flapjax/aero/flowfields.py

def vmap_call(self, x: Array, t: Array) -> Array:
    """
    Vectorized version of the __call__ method. This maps over all leading dimensions of x.
    :param x: Spatial coordinates, ``(..., 3)``
    :param t: Time, ()
    :return: Flow field values at the specified coordinates, ``(..., 3)``
    """
    n_vmap = x.ndim - 1
    func = self.__call__
    for i_dim in range(n_vmap):
        func = jax.vmap(func, in_axes=(i_dim, None), out_axes=i_dim)
    return func(x, t)

surf_vmap_call

surf_vmap_call(xs: ArrayList, t: Array) -> ArrayList

Vectorized version of the call method over a list of surfaces.

Parameters:

Name	Type	Description	Default
xs	ArrayList	Spatial coordinates, (n_surf,)(..., 3)	required
t	Array	Time, ()	required

Returns:

Type	Description
ArrayList	Flow field values at the specified coordinates, (n_surf, )(..., 3)

Source code in src/flapjax/aero/flowfields.py

def surf_vmap_call(self, xs: ArrayList, t: Array) -> ArrayList:
    """
    Vectorized version of the __call__ method over a list of surfaces.
    :param xs: Spatial coordinates, ``(n_surf,)(..., 3)``
    :param t: Time, ()
    :return: Flow field values at the specified coordinates, ``(n_surf, )(..., 3)``
    """
    return ArrayList([self.vmap_call(x, t) for x in xs])

frequency_flowfields

FrequencyFlowField

FrequencyFlowField(
    sigma: float | Array,
    length_scale: float | Array,
    u_inf: float | Array,
)

Base class for turbulence spectra used in frequency-domain gust analysis.

Parameters:

Name	Type	Description	Default
sigma	float | Array	Turbulence intensity (RMS gust velocity), m/s.	required
length_scale	float | Array	Turbulence scale length, m.	required
u_inf	float | Array	Freestream velocity magnitude, m/s.	required

Source code in src/flapjax/aero/frequency_flowfields.py

def __init__(
    self,
    sigma: float | Array,
    length_scale: float | Array,
    u_inf: float | Array,
):
    r"""
    :param sigma: Turbulence intensity (RMS gust velocity), m/s.
    :param length_scale: Turbulence scale length, m.
    :param u_inf: Freestream velocity magnitude, m/s.
    """
    self.sigma: Array = jnp.array(sigma)
    self.length_scale: Array = jnp.array(length_scale)
    self.u_inf: Array = jnp.array(u_inf)

psd

psd(omega: Array) -> Array

Power spectral density of the vertical gust velocity.

Parameters:

Name	Type	Description	Default
omega	Array	Sampling requencies in rad/s, (n_freq, ).	required

Returns:

Type	Description
Array	PSD values, (n_freq, ).

Source code in src/flapjax/aero/frequency_flowfields.py

def psd(self, omega: Array) -> Array:
    r"""
    Power spectral density of the vertical gust velocity.

    :param omega: Sampling requencies in rad/s, ``(n_freq, )``.
    :return: PSD values, ``(n_freq, )``.
    """
    raise NotImplementedError("psd must be implemented in subclasses.")

VonKarmanFlowField

VonKarmanFlowField(
    sigma: float | Array,
    length_scale: float | Array,
    u_inf: float | Array,
)

Bases: FrequencyFlowField

Von Kármán continuous turbulence spectrum.

Parameters:

Name	Type	Description	Default
sigma	float | Array	Turbulence intensity (RMS gust velocity), m/s.	required
length_scale	float | Array	Turbulence scale length, m.	required
u_inf	float | Array	Freestream velocity magnitude, m/s.	required

Source code in src/flapjax/aero/frequency_flowfields.py

def __init__(
    self,
    sigma: float | Array,
    length_scale: float | Array,
    u_inf: float | Array,
):
    r"""
    :param sigma: Turbulence intensity (RMS gust velocity), m/s.
    :param length_scale: Turbulence scale length, m.
    :param u_inf: Freestream velocity magnitude, m/s.
    """
    self.sigma: Array = jnp.array(sigma)
    self.length_scale: Array = jnp.array(length_scale)
    self.u_inf: Array = jnp.array(u_inf)

DrydenFlowField

DrydenFlowField(
    sigma: float | Array,
    length_scale: float | Array,
    u_inf: float | Array,
)

Bases: FrequencyFlowField

Dryden continuous turbulence spectrum.

Parameters:

Name	Type	Description	Default
sigma	float | Array	Turbulence intensity (RMS gust velocity), m/s.	required
length_scale	float | Array	Turbulence scale length, m.	required
u_inf	float | Array	Freestream velocity magnitude, m/s.	required

Source code in src/flapjax/aero/frequency_flowfields.py

def __init__(
    self,
    sigma: float | Array,
    length_scale: float | Array,
    u_inf: float | Array,
):
    r"""
    :param sigma: Turbulence intensity (RMS gust velocity), m/s.
    :param length_scale: Turbulence scale length, m.
    :param u_inf: Freestream velocity magnitude, m/s.
    """
    self.sigma: Array = jnp.array(sigma)
    self.length_scale: Array = jnp.array(length_scale)
    self.u_inf: Array = jnp.array(u_inf)

gradients

data_structures

AeroGradsToCompute dataclass

AeroGradsToCompute(
    x0_aero: bool = True,
    flowfield: bool = False,
    cs_ang_t: bool = False,
    cs_vel_t: bool = False,
)

Class which contains flags to determine which gradients are to be computed for the aerodynamic problem during the adjoint solve. Defaults to computing only the aerodynamic grid gradients.

Parameters:

Name	Type	Description	Default
x0_aero	bool	Aerodynamic grid coordinates.	True
flowfield	bool	Flow field parameters.	False
cs_ang_t	bool	Control surface deflection angle time history.	False
cs_vel_t	bool	Control surface velocity time history.	False

AeroFullStates

AeroFullStates(
    gamma_b: ArrayList,
    gamma_w: ArrayList,
    gamma_b_dot: ArrayList,
    zeta_w: ArrayList,
)

Aerodynamic states used for the adjoint solve.

Parameters:

Name	Type	Description	Default
gamma_b	ArrayList	Bound panel circulation strengths. (n_surf, )(m, n)	required
gamma_w	ArrayList	Wake panel circulation strengths. (n_surf, )(m_star, n)	required
gamma_b_dot	ArrayList	Bound grid circulation time derivatives. (n_surf, )(m, n)	required
zeta_w	ArrayList	Wake grid coordinates. (n_surf, )(m_star + 1, n + 1, 3)	required

Source code in src/flapjax/aero/gradients/data_structures.py

def __init__(
    self,
    gamma_b: ArrayList,
    gamma_w: ArrayList,
    gamma_b_dot: ArrayList,
    zeta_w: ArrayList,
) -> None:
    r"""
    :param gamma_b: Bound panel circulation strengths. ``(n_surf, )(m, n)``
    :param gamma_w: Wake panel circulation strengths. ``(n_surf, )(m_star, n)``
    :param gamma_b_dot: Bound grid circulation time derivatives. ``(n_surf, )(m, n)``
    :param zeta_w: Wake grid coordinates. ``(n_surf, )(m_star + 1, n + 1, 3)``
    """
    self.gamma_b: ArrayList = gamma_b
    self.gamma_w: ArrayList = gamma_w
    self.gamma_b_dot: ArrayList = gamma_b_dot
    self.zeta_w: ArrayList = zeta_w

n_states property

n_states: int

Obtain the total number of states contained within the data structure, being the size of the vector obtained from self.ravel().

Returns:

Type	Description
int	Size of vector.

shapes

shapes() -> OrderedDict[
    str, tuple[int, ...] | ArrayListShape | None
]

Obtain the shapes of all arrays within the data structure.

Returns:

Type	Description
OrderedDict[str, tuple[int, ...] | ArrayListShape | None]	Dictionary of {name: shape} pairs of all arrays or ArrayLists within the data structure.

Source code in src/flapjax/aero/gradients/data_structures.py

def shapes(self) -> OrderedDict[str, tuple[int, ...] | ArrayListShape | None]:
    r"""
    Obtain the shapes of all arrays within the data structure.
    :return: Dictionary of {name: shape} pairs of all arrays or ArrayLists within the data structure.
    """
    return OrderedDict(
        gamma_b=self.gamma_b.shape,
        gamma_w=self.gamma_w.shape,
        gamma_b_dot=self.gamma_b_dot.shape,
        zeta_w=self.zeta_w.shape,
    )

from_vector staticmethod

from_vector(
    vect: Array,
    shapes: OrderedDict[
        str, tuple[int, ...] | ArrayListShape | None
    ],
) -> AeroFullStates

Construct an AeroFullStates object from a vector of data and a corresponding dictionary of shapes, being the inverse of self.ravel().

Parameters:

Name	Type	Description	Default
vect	Array	Aerodynamic state vector.	required
shapes	OrderedDict[str, tuple[int, ...] | ArrayListShape | None]	Dictionary of {name: shape} pairs of all arrays or ArrayLists within the data structure.	required

Returns:

Type	Description
AeroFullStates	AeroFullStates object.

Source code in src/flapjax/aero/gradients/data_structures.py

@staticmethod
def from_vector(
    vect: Array,
    shapes: OrderedDict[str, tuple[int, ...] | ArrayListShape | None],
) -> AeroFullStates:
    r"""
    Construct an AeroFullStates object from a vector of data and a corresponding dictionary of shapes, being the inverse
    of ``self.ravel()``.
    :param vect: Aerodynamic state vector.
    :param shapes: Dictionary of {name: shape} pairs of all arrays or ArrayLists within the data structure.
    :return: AeroFullStates object.
    """
    return AeroFullStates(**vect_to_arrs(vect, shapes))

ravel

ravel() -> Array

Ravel the data structure to a vector, being the inverse of cls.from_vector().

Returns:

Type	Description
Array	Data vector containing all states.

Source code in src/flapjax/aero/gradients/data_structures.py

def ravel(self) -> Array:
    r"""
    Ravel the data structure to a vector, being the inverse of ``cls.from_vector()``.
    :return: Data vector containing all states.
    """

    return jnp.concatenate(
        [
            self.gamma_b.ravel(),
            self.gamma_w.ravel(),
            self.gamma_b_dot.ravel(),
            self.zeta_w.ravel(),
        ]
    )

AeroDesignVariables

AeroDesignVariables(
    zeta_b0: ArrayList | None,
    flowfield: dict[str, Array] | None,
    cs_ang_t: dict[str, Array] | None,
    cs_vel_t: dict[str, Array] | None,
    f_shape: tuple[int, ...],
)

Bases: DesignVariables

Class to hold all differentiable aerodynamic design variables.

Parameters:

Name	Type	Description	Default
zeta_b0	ArrayList | None	Bound aerodynamic grid local reference coordinates. (n_surf, )(m+1, n+1, 3)	required
flowfield	dict[str, Array] | None	Dictionary of flowfield variables.	required
cs_ang_t	dict[str, Array] | None	Control surface angle time histories, {name: (n_tstep, )}.	required
cs_vel_t	dict[str, Array] | None	Control velocity time histories, {name: (n_tstep, )}.	required
f_shape	tuple[int, ...]	Shape of objective. As this class can hold both the primal design variables and the gradient of the objective with respect to design variables, the latter case results in arrays which have a shape which depends on the objective shape. In this case, all data has (f_shape, dv.shape) dimensionality.	required

Source code in src/flapjax/aero/gradients/data_structures.py

def __init__(
    self,
    zeta_b0: ArrayList | None,
    flowfield: dict[str, Array] | None,
    cs_ang_t: dict[str, Array] | None,
    cs_vel_t: dict[str, Array] | None,
    f_shape: tuple[int, ...],
):
    r"""
    :param zeta_b0: Bound aerodynamic grid local reference coordinates. ``(n_surf, )(m+1, n+1, 3)``
    :param flowfield: Dictionary of flowfield variables.
    :param cs_ang_t: Control surface angle time histories, {name: ``(n_tstep, )``}.
    :param cs_vel_t: Control velocity time histories, {name: ``(n_tstep, )``}.
    :param f_shape: Shape of objective. As this class can hold both the primal design variables and the gradient of
    the objective with respect to design variables, the latter case results in arrays which have a shape which
    depends on the objective shape. In this case, all data has (*f_shape, *dv.shape) dimensionality.
    """
    super().__init__()
    self.zeta_b0: ArrayList | None = zeta_b0
    self.flowfield: dict[str, Array] | None = flowfield
    self.cs_ang_t: dict[str, Array] | None = cs_ang_t
    self.cs_vel_t: dict[str, Array] | None = cs_vel_t

    self.f_shape: tuple[int, ...] = f_shape
    self.f_size: int = prod(f_shape)

    self.shapes: dict[
        str,
        tuple[int, ...]
        | ArrayListShape
        | dict[str, tuple[int, ...] | ArrayListShape]
        | None,
    ] = self.get_shapes()
    self.mapping, self.n_x = self.make_index_mapping()

get_cs_n

get_cs_n(
    i_ts: int | Array, dv_full: AeroDesignVariables
) -> tuple[dict[str, Array], dict[str, Array]]

Obtain the angles and velocities for all control surfaces at a given timestep.

Parameters:

Name	Type	Description	Default
i_ts	int | Array	Timestep index.	required
dv_full	AeroDesignVariables	Full design variables. This is used to substitute a non-differentiable value when the control inputs are chosen to be omitted from the design variables.	required

Returns:

Type	Description
tuple[dict[str, Array], dict[str, Array]]	Dictionary of {name: value} pairs for control angles and velocities.

Source code in src/flapjax/aero/gradients/data_structures.py

def get_cs_n(
    self,
    i_ts: int | Array,
    dv_full: AeroDesignVariables,
) -> tuple[dict[str, Array], dict[str, Array]]:
    r"""
    Obtain the angles and velocities for all control surfaces at a given timestep.
    :param i_ts: Timestep index.
    :param dv_full: Full design variables. This is used to substitute a non-differentiable value when the control
    inputs are chosen to be omitted from the design variables.
    :return: Dictionary of {name: value} pairs for control angles and velocities.
    """

    # get control surface deflections from design variables
    assert dv_full.cs_ang_t is not None and dv_full.cs_vel_t is not None
    cs_ang_n = {
        k: v[i_ts, ...]
        for k, v in (
            self.cs_ang_t if self.cs_ang_t is not None else dv_full.cs_ang_t
        ).items()
    }

    cs_vel_n = {
        k: v[i_ts, ...]
        for k, v in (
            self.cs_vel_t if self.cs_vel_t is not None else dv_full.cs_vel_t
        ).items()
    }

    return cs_ang_n, cs_vel_n

premultiply_adj

premultiply_adj(adj: Array) -> AeroDesignVariables

Premultiply all design gradients by the adjoint vector.

Parameters:

Name	Type	Description	Default
adj	Array	Adjoint vector.	required

Returns:

Type	Description
AeroDesignVariables	Adjoint-Jacobian product.

Source code in src/flapjax/aero/gradients/data_structures.py

def premultiply_adj(self, adj: Array) -> AeroDesignVariables:
    r"""
    Premultiply all design gradients by the adjoint vector.
    :param adj: Adjoint vector.
    :return: Adjoint-Jacobian product.
    """
    return AeroDesignVariables(
        zeta_b0=ArrayList(
            [
                jnp.einsum("ij,j...->i...", adj, self.zeta_b0[i])
                for i in range(len(self.zeta_b0))
            ]
        )
        if self.zeta_b0 is not None
        else None,
        flowfield={
            k: jnp.einsum("ij,j...->i...", adj, v)
            for k, v in self.flowfield.items()
        }
        if self.flowfield is not None
        else None,
        cs_ang_t={
            k: jnp.einsum("ij,j...->i...", adj, v) for k, v in self.cs_ang_t.items()
        }
        if self.cs_ang_t is not None
        else None,
        cs_vel_t={
            k: jnp.einsum("ij,j...->i...", adj, v) for k, v in self.cs_vel_t.items()
        }
        if self.cs_vel_t is not None
        else None,
        f_shape=(adj.shape[1],),
    )

to_dict

to_dict() -> dict[
    str, Array | ArrayList | dict[str, Array] | None
]

Extract the design variables as a dictionary.

Returns:

Type	Description
dict[str, Array | ArrayList | dict[str, Array] | None]	Dictionary of design variable name and value pairs.

Source code in src/flapjax/aero/gradients/data_structures.py

def to_dict(self) -> dict[str, Array | ArrayList | dict[str, Array] | None]:
    r"""
    Extract the design variables as a dictionary.
    :return: Dictionary of design variable name and value pairs.
    """
    return {
        "zeta_b0": self.zeta_b0,
        "flowfield": self.flowfield,
        "cs_ang_t": self.cs_ang_t,
        "cs_vel_t": self.cs_vel_t,
    }

plot

plot(
    case: AeroCase,
    rmat_nodal: ArrayList | None,
    directory: PathLike | str,
) -> Sequence[Path]

Plot the aerodynamic grid gradient for cases with a scalar objective.

Parameters:

Name	Type	Description	Default
case	AeroCase	Dynamic aerodynamic case object. This should only contain 1 time step, as the gradients for the grid are constant across all time steps.	required
rmat_nodal	ArrayList | None	Rotation matrices from beam. As the aerodynamic grid gradients are given in the local frame (as this is where zeta_b0 exists), this transformation is used to transform the gradients into the global frame for visualisation.	required
directory	PathLike | str	Directory in which to save the plots.	required

Returns:

Type	Description
Sequence[Path]	Sequence of paths of data created.

Source code in src/flapjax/aero/gradients/data_structures.py

def plot(
    self,
    case: AeroCase,
    rmat_nodal: ArrayList | None,
    directory: os.PathLike | str,
) -> Sequence[Path]:
    r"""
    Plot the aerodynamic grid gradient for cases with a scalar objective.
    :param case: Dynamic aerodynamic case object. This should only contain 1 time step, as the gradients for the
    grid are constant across all time steps.
    :param rmat_nodal: Rotation matrices from beam. As the aerodynamic grid gradients are given in the local frame
    (as this is where ``zeta_b0`` exists), this transformation is used to transform the gradients into the global frame
    for visualisation.
    :param directory: Directory in which to save the plots.
    :return: Sequence of paths of data created.
    """
    if self.f_size != 1:
        raise ValueError("Can only plot gradients for scalar objective functions.")

    if case.n_tstep != 1:
        raise ValueError("Can only plot gradients for singe timestep cases.")

    if self.zeta_b0 is None:
        warn(
            "Aerodynamic grid gradient not computed. Skipping grid gradient plotting."
        )
        return []

    directory_path = Path(directory)
    directory_path.mkdir(parents=True, exist_ok=True)
    paths = []

    for i_surf in range(case.n_surf):
        bound_filename = Path(directory).joinpath(
            case.surf_b_names[i_surf] + "_gradient"
        )

        if rmat_nodal is not None:
            d_x0_aero: Array = jnp.einsum(
                "ijk,...lik->lij", rmat_nodal[i_surf], self.zeta_b0[i_surf]
            )
        else:
            d_x0_aero = self.zeta_b0[i_surf]

        paths.append(
            plot_grid_to_vtk(
                case.zeta_b[i_surf],
                bound_filename,
                None,
                node_vector_data={
                    "zeta_b0": d_x0_aero,
                },
                cell_scalar_data={},
            )
        )

    return paths

linear

AeroLinearResult

AeroLinearResult(
    reference: AeroCase,
    u_t: AeroInputUnflattened,
    x_t: AeroStateUnflattened,
    y_t: AeroOutputUnflattened,
    u_t_tot: AeroInputUnflattened,
    x_t_tot: AeroStateUnflattened,
    y_t_tot: AeroOutputUnflattened,
    n_tstep: int,
    n_surf: int,
    t: Array,
    surf_b_names: list[str],
    surf_w_names: list[str],
)

Source code in src/flapjax/aero/linear/data_structures.py

def __init__(
    self,
    reference: AeroCase,
    u_t: AeroInputUnflattened,
    x_t: AeroStateUnflattened,
    y_t: AeroOutputUnflattened,
    u_t_tot: AeroInputUnflattened,
    x_t_tot: AeroStateUnflattened,
    y_t_tot: AeroOutputUnflattened,
    n_tstep: int,
    n_surf: int,
    t: Array,
    surf_b_names: list[str],
    surf_w_names: list[str],
) -> None:
    # system results, if simulated
    self.u_t: AeroInputUnflattened = u_t
    self.x_t: AeroStateUnflattened = x_t
    self.y_t: AeroOutputUnflattened = y_t
    self.u_t_tot: AeroInputUnflattened = u_t_tot
    self.x_t_tot: AeroStateUnflattened = x_t_tot
    self.y_t_tot: AeroOutputUnflattened = y_t_tot
    self.n_tstep: int = n_tstep
    self.n_surf: int = n_surf
    self.t: Array = t
    self.surf_b_names: list[str] = surf_b_names
    self.surf_w_names: list[str] = surf_w_names
    self.reference: AeroCase = reference

plot

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

Plot the aerodynamic grid at specified time steps.

Parameters:

Name	Type	Description	Default
directory	str | PathLike	Directory to save the plots to	required
index	slice | Sequence[int] | int | Array | None	Index or slice of time steps to plot. If None, plot all time steps.	None
plot_wake	bool	If True, plot the wake grid	True

Source code in src/flapjax/aero/linear/data_structures.py

def plot(
    self,
    directory: str | os.PathLike,
    index: slice | Sequence[int] | int | Array | None = None,
    plot_wake: bool = True,
) -> None:
    r"""
    Plot the aerodynamic grid at specified time steps.
    :param directory: Directory to save the plots to
    :param index: Index or slice of time steps to plot. If None, plot all time steps.
    :param plot_wake: If True, plot the wake grid
    """
    if isinstance(index, slice):
        index_ = jnp.arange(self.n_tstep)[index]
    elif isinstance(index, Sequence):
        index_ = jnp.array(index)
    elif isinstance(index, Array):
        index_ = index
    elif isinstance(index, int):
        index_ = (index,)
    elif index is None:
        index_ = jnp.arange(self.n_tstep)
    else:
        raise TypeError("index must be a slices, sequence of ints, or Array")

    directory_path = Path(directory).resolve()
    directory_path.mkdir(parents=True, exist_ok=True)

    paths: list[Sequence[Path]] = []
    for i_ts in index_:
        snapshot = self[i_ts]
        paths.append(snapshot.plot(directory, plot_wake=plot_wake))

    for i_surf in range(2 * self.n_surf):
        try:
            surf_paths = [paths[i][i_surf] for i in range(len(index_))]
            name = (self.surf_b_names + self.surf_w_names)[i_surf] + "_ts"
            write_pvd(directory, name, surf_paths, list(self.t[index_]))
        except IndexError:
            pass

LinearUVLM

LinearUVLM(
    case: UVLM,
    reference: AeroCase,
    wake_type: LinearWakeType = "frozen",
    bound_upwash: bool = True,
    wake_upwash: bool = True,
    unsteady_force: bool = True,
    *,
    skip_checks: bool = False,
    skip_linearisation: bool = False,
)

Bases: LinearModel[AeroCase, AeroInputUnflattened, AeroStateUnflattened, AeroOutputUnflattened, AeroLinearResult]

Class to represent a linearised UVLM aerodynamic system about a reference state.

Initialise linear UVLM system about a reference state.

Parameters:

Name	Type	Description	Default
case	UVLM	UVLM case object to linearise.	required
reference	AeroCase	StaticAero representing the reference state for linearisation.	required
wake_type	LinearWakeType	Instance of LinearWakeType enum to specify wake treatment.	'frozen'
bound_upwash	bool	If true, include bound surface upwash velocities as inputs.	True
wake_upwash	bool	If true, include wake surface upwash velocities as inputs.	True
unsteady_force	bool	If true, include unsteady force.	True

Source code in src/flapjax/aero/linear/linear_uvlm.py

def __init__(
    self,
    case: UVLM,
    reference: AeroCase,
    wake_type: LinearWakeType = "frozen",
    bound_upwash: bool = True,
    wake_upwash: bool = True,
    unsteady_force: bool = True,
    *,
    skip_checks: bool = False,
    skip_linearisation: bool = False,
):
    r"""
    Initialise linear UVLM system about a reference state.
    :param case: UVLM case object to linearise.
    :param reference: StaticAero representing the reference state for linearisation.
    :param wake_type: Instance of LinearWakeType enum to specify wake treatment.
    :param bound_upwash: If true, include bound surface upwash velocities as inputs.
    :param wake_upwash: If true, include wake surface upwash velocities as inputs.
    :param unsteady_force: If true, include unsteady force.
    """
    # options
    self.prescribed_wake, self.free_wake = {
        "frozen": (False, False),
        "prescribed": (True, False),
        "free": (True, True),
    }[wake_type]
    self.unsteady_force: bool = unsteady_force
    self.bound_upwash: bool = bound_upwash
    self.wake_upwash: bool = wake_upwash

    # time info
    super().__init__(reference=reference, dt=case.dt)

    # check that the reference state is steady
    # whilst linearisation can be performed about unsteady states, the current implementation omits some terms
    # required for this, however, cannot see a practical use case for such a model. Warn the user if the reference
    # state appears unsteady.
    if (
        not skip_checks
        and max([jnp.abs(zbd).max() for zbd in reference.zeta_b_dot]) > 1e-6
    ):
        warn(
            "Reference bound surface velocities are non-zero. Ensure that the reference state is steady for linearisation."
        )

    if (
        not skip_checks
        and max([jnp.abs(gbd).max() for gbd in reference.gamma_b_dot]) > 1e-6
    ):
        warn(
            "Reference bound circulation time derivative is non-zero. Ensure that the reference state is steady for linearisation."
        )

    # kernels
    self.kernels_b: Sequence[KernelFunction] = reference.n_surf * [
        biot_savart_cutoff
    ]
    self.kernels_w: Sequence[KernelFunction] = reference.n_surf * [
        biot_savart_cutoff
    ]

    # wake propagation deltas
    self.case: UVLM = case

    # linear system
    if not skip_linearisation:
        self.sys = self.linearise()

unpack_state_vector

unpack_state_vector(x: Array) -> S

Unpack a state vector into its components.

Parameters:

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

Returns:

Type	Description
S	StateUnflattened object.

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

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

unpack_output_vector

unpack_output_vector(y: Array) -> O

Unpack an output vector into its components.

Parameters:

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

Returns:

Type	Description
O	OutputUnflattened object.

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

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

pack_input_vector

pack_input_vector(u_input: InputUnflattened) -> Array

Pack an input unflattened object into a vector.

Parameters:

Name	Type	Description	Default
u_input	InputUnflattened	InputUnflattened object.	required

Returns:

Type	Description
Array	Input vector, (n_inputs, ).

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

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

pack_state_vector

pack_state_vector(x_state: StateUnflattened) -> Array

Pack a state unflattened object into a vector.

Parameters:

Name	Type	Description	Default
x_state	StateUnflattened	StateUnflattened object.	required

Returns:

Type	Description
Array	State vector, (n_states, ).

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

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

pack_output_vector

pack_output_vector(y_output: OutputUnflattened) -> Array

Pack an output unflattened object into a vector.

Parameters:

Name	Type	Description	Default
y_output	OutputUnflattened	OutputUnflattened object.	required

Returns:

Type	Description
Array	Output vector, (n_outputs, ).

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

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

get_total_input

get_total_input(u: InputUnflattened) -> InputUnflattened

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

Parameters:

Name	Type	Description	Default
u	InputUnflattened	InputUnflattened perturbation object.	required

Returns:

Type	Description
InputUnflattened	InputUnflattened total object.

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

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

get_total_state

get_total_state(x: StateUnflattened) -> StateUnflattened

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

Parameters:

Name	Type	Description	Default
x	StateUnflattened	StateUnflattened perturbation object.	required

Returns:

Type	Description
StateUnflattened	StateUnflattened total object.

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

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

get_total_output

get_total_output(y: OutputUnflattened) -> OutputUnflattened

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

Parameters:

Name	Type	Description	Default
y	OutputUnflattened	OutputUnflattened perturbation object.	required

Returns:

Type	Description
OutputUnflattened	OutputUnflattened total object.

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

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

get_total_input_t

get_total_input_t(
    u_t: InputUnflattened,
) -> InputUnflattened

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

Parameters:

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

Returns:

Type	Description
InputUnflattened	InputUnflattened total time history object.

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

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

get_total_state_t

get_total_state_t(
    x_t: StateUnflattened,
) -> StateUnflattened

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

Parameters:

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

Returns:

Type	Description
StateUnflattened	StateUnflattened total time history object.

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

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

get_total_output_t

get_total_output_t(
    y_t: OutputUnflattened,
) -> OutputUnflattened

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

Parameters:

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

Returns:

Type	Description
OutputUnflattened	OutputUnflattened total time history object.

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

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

get_zero_input

get_zero_input() -> InputUnflattened

Get a zero input unflattened object.

Returns:

Type	Description
InputUnflattened	InputUnflattened object with zero arrays.

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

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

get_zero_state

get_zero_state() -> StateUnflattened

Get a zero state unflattened object.

Returns:

Type	Description
StateUnflattened	StateUnflattened object with zero arrays.

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

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

get_zero_output

get_zero_output() -> OutputUnflattened

Get a zero output unflattened object.

Returns:

Type	Description
OutputUnflattened	OutputUnflattened object with zero arrays.

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

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

step_vec

step_vec(x_vec: Array, u_vec: Array) -> tuple[Array, Array]

Combined state/output step in vector form, operating on total (reference + perturbation) quantities and returning perturbations relative to the reference.

Parameters:

Name	Type	Description	Default
x_vec	Array	State vector, (n_states, )	required
u_vec	Array	Input vector, (n_inputs, )	required

Returns:

Type	Description
tuple[Array, Array]	Tuple of (state perturbation vector, output perturbation vector).

Source code in src/flapjax/aero/linear/linear_uvlm.py

def step_vec(self, x_vec: Array, u_vec: Array) -> tuple[Array, Array]:
    r"""
    Combined state/output step in vector form, operating on total (reference + perturbation) quantities and
    returning perturbations relative to the reference.
    :param x_vec: State vector, ``(n_states, )``
    :param u_vec: Input vector, ``(n_inputs, )``
    :return: Tuple of (state perturbation vector, output perturbation vector).
    """

    u_np1 = self._unpack_input_vector(u_vec)
    x_n = self.unpack_state_vector(x_vec)

    assert isinstance(u_np1, AeroInputUnflattened) and isinstance(
        x_n, AeroStateUnflattened
    ), (
        "Unpacked input and state must be of type AeroInputUnflattened and AeroStateUnflattened."
    )

    x_np1, y_n = self.step(u_np1=u_np1, x_n=x_n)

    return self.pack_state_vector(x_np1), self.pack_output_vector(y_n)

step

step(
    u_np1: AeroInputUnflattened, x_n: AeroStateUnflattened
) -> tuple[AeroStateUnflattened, AeroOutputUnflattened]

From total inputs u at timestep n+1 and states x at timestep n, compute the states at timestep n+1 and outputs at timestep n.

Source code in src/flapjax/aero/linear/linear_uvlm.py

def step(
    self,
    u_np1: AeroInputUnflattened,
    x_n: AeroStateUnflattened,
) -> tuple[AeroStateUnflattened, AeroOutputUnflattened]:
    r"""
    From total inputs `u` at timestep n+1 and states `x` at timestep n, compute the states at timestep n+1 and
    outputs at timestep n.
    """
    ref = self.reference

    zeta_b_np1: ArrayList = u_np1.zeta_b
    zeta_dot_b_np1: ArrayList = u_np1.zeta_b_dot
    gamma_b_n: ArrayList = x_n.gamma_b
    gamma_w_n: ArrayList = x_n.gamma_w

    if self.unsteady_force:
        assert x_n.gamma_b_nm1 is not None, "gamma_b_nm1 is None"
        gamma_b_dot_n: ArrayList = (gamma_b_n - x_n.gamma_b_nm1) / self.dt
    else:
        gamma_b_dot_n = ref.gamma_b_dot

    if self.prescribed_wake:
        assert x_n.zeta_b is not None, "zeta_b is None"
        assert x_n.zeta_w is not None, "zeta_w is None"
        zeta_b_n: ArrayList = x_n.zeta_b
        zeta_w_n: ArrayList = x_n.zeta_w
    else:
        zeta_b_n = ref.zeta_b
        zeta_w_n = ref.zeta_w

    q_n = AeroFullStates(
        gamma_b=gamma_b_n,
        gamma_w=gamma_w_n,
        gamma_b_dot=ref.gamma_b_dot,
        zeta_w=zeta_w_n,
    )
    (
        _,
        _,
        gamma_b_np1,
        gamma_w_np1,
        _,
        _,
        zeta_w_np1,
        *_,
    ) = self.case.base_solve_from_grid(
        q_nm1=q_n,
        t_n=ref.t,
        zeta_b_n=zeta_b_np1,
        zeta_b_nm1=zeta_b_n,
        zeta_b_dot_n=zeta_dot_b_np1,
        static=False,
        horseshoe=False,
        linearise_variable_wake=True,
        nu_b=u_np1.nu_b,
        nu_w=u_np1.nu_w,
    )

    # the forcing needs to be computed seperately to find its dependence on the current states
    rho = ref.flowfield.rho

    def v_out_func(x_target: Array) -> Array:
        return ref.flowfield.vmap_call(x=x_target, t=ref.t) + compute_v_ind(
            cs=x_target,
            zetas=ArrayList([*zeta_b_np1, *zeta_w_n]),
            gammas=ArrayList([*gamma_b_n, *gamma_w_n]),
            kernels=[*self.kernels_b, *self.kernels_w],
            batch_size=self.case.batch_size,
            mirror_normal=self.case.mirror_normal,
            mirror_point=self.case.mirror_point,
        )

    f_steady_n = compute_steady_forcing(
        zeta_b=zeta_b_np1,
        zeta_dot_b=zeta_dot_b_np1,
        gamma_b=gamma_b_n,
        gamma_w=gamma_w_n,
        rho=rho,
        v_func=v_out_func,
        v_inputs=u_np1.nu_b if self.bound_upwash else None,
        mirror_point=self.case.mirror_point,
        mirror_normal=self.case.mirror_normal,
        mirror_edge_low=self.case.mirror_edge_low,
        mirror_edge_high=self.case.mirror_edge_high,
    )

    normals = compute_nc(zetas=zeta_b_np1)
    if self.unsteady_force:
        f_unsteady_n = ArrayList(
            [
                split_to_vertex(
                    rho * gamma_b_dot_n[i][..., None] * normals[i], (0, 1)
                )
                for i in range(ref.n_surf)
            ]
        )
    else:
        f_unsteady_n = None

    x_np1 = AeroStateUnflattened(
        gamma_b=gamma_b_np1,
        gamma_w=gamma_w_np1,
        gamma_b_nm1=gamma_b_n if self.unsteady_force else None,
        zeta_w=zeta_w_np1 if self.prescribed_wake else None,
        zeta_b=zeta_b_np1 if self.prescribed_wake else None,
    )
    y_n = AeroOutputUnflattened(f_steady=f_steady_n, f_unsteady=f_unsteady_n)

    return x_np1, y_n

compute_jacobians

compute_jacobians(
    input_projection: LinearInputProjection | None = None,
    output_projection: LinearOutputProjection | None = None,
    residual_names: Sequence[str] | None = None,
) -> dict[
    str,
    tuple[
        Callable[..., Any], dict[str, Any], Sequence[str]
    ],
]

Build the Jacobians for the linear system.

Parameters:

Name	Type	Description	Default
input_projection	LinearInputProjection | None	If set, projects the inputs onto a different space.	None
output_projection	LinearOutputProjection | None	If set, projects the forcing outputs onto a different space.	None
residual_names	Sequence[str] | None	If set, restrict the returned residuals to this subset to avoid redundant computation.	None

Returns:

Type	Description
dict[str, tuple[Callable[..., Any], dict[str, Any], Sequence[str]]]	Mapping of residual name to Jacobian(s).

Source code in src/flapjax/aero/linear/linear_uvlm.py

def compute_jacobians(
    self,
    input_projection: LinearInputProjection | None = None,
    output_projection: LinearOutputProjection | None = None,
    residual_names: Sequence[str] | None = None,
) -> dict[str, tuple[Callable[..., Any], dict[str, Any], Sequence[str]]]:
    r"""
    Build the Jacobians for the linear system.
    :param input_projection: If set, projects the inputs onto a different space.
    :param output_projection: If set, projects the forcing outputs onto a different space.
    :param residual_names: If set, restrict the returned residuals to this subset to avoid redundant computation.
    :return: Mapping of residual name to Jacobian(s).
    """
    ref = self.reference

    # bound circulation
    gamma_b_args: dict[str, Any] = {
        "gamma_b_n_vec": ref.gamma_b.ravel(),
        "gamma_w_n_vec": ref.gamma_w.ravel(),
        "zeta_b_np1_vec": ref.zeta_b.ravel(),
        "zeta_b_dot_np1_vec": ref.zeta_b_dot.ravel(),
    }
    gamma_b_diff = [
        "gamma_b_n_vec",
        "gamma_w_n_vec",
        "zeta_b_np1_vec",
        "zeta_b_dot_np1_vec",
    ]
    if self.prescribed_wake:
        gamma_b_args["zeta_w_n_vec"] = ref.zeta_w.ravel()
        gamma_b_args["zeta_b_n_vec"] = ref.zeta_b.ravel()
        gamma_b_diff.extend(["zeta_w_n_vec", "zeta_b_n_vec"])
    if self.bound_upwash:
        gamma_b_args["nu_b_np1_vec"] = jnp.zeros(ref.zeta_b.size)
        gamma_b_diff.append("nu_b_np1_vec")
    if self.wake_upwash:
        gamma_b_args["nu_w_np1_vec"] = jnp.zeros(ref.zeta_w.size)
        gamma_b_diff.append("nu_w_np1_vec")

    # wake propagation (gamma_w and optionally zeta_w)
    wake_args: dict[str, Any] = {
        "gamma_b_n_vec": ref.gamma_b.ravel(),
        "gamma_w_n_vec": ref.gamma_w.ravel(),
    }
    gamma_w_diff = ["gamma_b_n_vec", "gamma_w_n_vec"]
    zeta_w_diff: list[str] = []
    if self.prescribed_wake:
        wake_args["zeta_w_n_vec"] = ref.zeta_w.ravel()
        wake_args["zeta_b_np1_vec"] = ref.zeta_b.ravel()
        zeta_w_diff.extend(["zeta_w_n_vec", "zeta_b_np1_vec"])
    if self.wake_upwash:
        wake_args["nu_w_np1_vec"] = jnp.zeros(ref.zeta_w.size)
        zeta_w_diff.append("nu_w_np1_vec")
    if self.free_wake:
        zeta_w_diff.extend(["gamma_b_n_vec", "gamma_w_n_vec"])

    # steady forcing
    f_steady_args: dict[str, Any] = {
        "gamma_b_n_vec": ref.gamma_b.ravel(),
        "gamma_w_n_vec": ref.gamma_w.ravel(),
        "zeta_b_np1_vec": ref.zeta_b.ravel(),
        "zeta_b_dot_np1_vec": ref.zeta_b_dot.ravel(),
    }
    f_steady_diff = [
        "gamma_b_n_vec",
        "gamma_w_n_vec",
        "zeta_b_np1_vec",
        "zeta_b_dot_np1_vec",
    ]
    if self.prescribed_wake:
        f_steady_args["zeta_w_n_vec"] = ref.zeta_w.ravel()
        f_steady_diff.append("zeta_w_n_vec")
    if self.bound_upwash:
        f_steady_args["nu_b_np1_vec"] = jnp.zeros(ref.zeta_b.size)
        f_steady_diff.append("nu_b_np1_vec")

    residuals: 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 **kw: self.wake_prop_step(**kw)[1],
            wake_args,
            gamma_w_diff,
        ),
    }
    if self.prescribed_wake:
        residuals["zeta_w"] = (
            lambda **kw: self.wake_prop_step(**kw)[0],
            wake_args,
            zeta_w_diff,
        )
    if self.unsteady_force:
        residuals["gamma_b_nm1"] = (
            lambda **kw: kw["gamma_b_n_vec"],
            {"gamma_b_n_vec": ref.gamma_b.ravel()},
            ["gamma_b_n_vec"],
        )
    if self.prescribed_wake:
        # zeta_b state at n+1 == zeta_b input at n+1
        residuals["zeta_b"] = (
            lambda **kw: kw["zeta_b_np1_vec"],
            {"zeta_b_np1_vec": ref.zeta_b.ravel()},
            ["zeta_b_np1_vec"],
        )
    residuals["f_steady"] = (self.f_steady_step, f_steady_args, f_steady_diff)
    if self.unsteady_force:
        f_unsteady_args: dict[str, Any] = {
            "gamma_b_n_vec": ref.gamma_b.ravel(),
            "gamma_b_nm1_vec": ref.gamma_b.ravel(),
            "zeta_b_np1_vec": ref.zeta_b.ravel(),
        }
        residuals["f_unsteady"] = (
            self.f_unsteady_step,
            f_unsteady_args,
            ["gamma_b_n_vec", "gamma_b_nm1_vec", "zeta_b_np1_vec"],
        )

    # apply I/O projections if provided
    if input_projection is not None:
        residuals = self._apply_input_projection(residuals, input_projection)
    if output_projection is not None:
        residuals = self._apply_output_projection(residuals, output_projection)

    if residual_names is not None:
        residuals = {k: v for k, v in residuals.items() if k in residual_names}

    return residuals

create_jacobians

create_jacobians(
    mode: ADMode | dict[str, ADMode] = "reverse",
    batch_size: int | None = None,
    input_projection: LinearInputProjection | None = None,
    output_projection: LinearOutputProjection | None = None,
    residual_names: Sequence[str] | None = None,
) -> dict[str, dict[str, Array]]

Compute the per-residual Jacobians using :func:jacrev_custom.

Source code in src/flapjax/aero/linear/linear_uvlm.py

def create_jacobians(
    self,
    mode: ADMode | dict[str, ADMode] = "reverse",
    batch_size: int | None = None,
    input_projection: LinearInputProjection | None = None,
    output_projection: LinearOutputProjection | None = None,
    residual_names: Sequence[str] | None = None,
) -> dict[str, dict[str, Array]]:
    r"""
    Compute the per-residual Jacobians using :func:`jacrev_custom`.
    """
    residuals = self.compute_jacobians(
        input_projection=input_projection,
        output_projection=output_projection,
        residual_names=residual_names,
    )

    jacobians: dict[str, dict[str, Array]] = {}
    for res_name, (res_func, args, diff_arg_names) in residuals.items():
        res_jac_options: dict[str, Callable[..., Array] | None] = {
            arg: None for arg in diff_arg_names
        }

        res_mode: ADMode = (
            mode.get(res_name, "reverse") if isinstance(mode, dict) else mode
        )

        jacs, _, _ = jacrev_custom(
            func=res_func,
            jac_options=res_jac_options,
            n_profile_loops=None,
            func_name=res_name,
            map_batch_size=batch_size,
            mode=res_mode,
        )(**args)
        jacobians[res_name] = jacs

    return jacobians

linearise

linearise(
    batch_size: int | None = None,
    *,
    input_projection: LinearInputProjection | None = None,
    output_projection: LinearOutputProjection | None = None,
) -> LinearSystem

Build the linear state-space system.

Parameters:

Name	Type	Description	Default
batch_size	int | None	If not None, batch the Jacobian passes to reduce memory.	None
input_projection	LinearInputProjection | None	Optional projection from lower-dim inputs (e.g. beam DOFs) to (zeta_b, zeta_b_dot).	None
output_projection	LinearOutputProjection | None	Optional projection from (f_steady, f_unsteady) to a smaller output (e.g. beam force); when set, C/D rows correspond to the projected output.	None

Returns:

Type	Description
LinearSystem	LinearSystem object.

Source code in src/flapjax/aero/linear/linear_uvlm.py

def linearise(
    self,
    batch_size: int | None = None,
    *,
    input_projection: LinearInputProjection | None = None,
    output_projection: LinearOutputProjection | None = None,
) -> LinearSystem:
    r"""
    Build the linear state-space system.
    :param batch_size: If not None, batch the Jacobian passes to reduce memory.
    :param input_projection: Optional projection from lower-dim inputs (e.g. beam DOFs) to
    ``(zeta_b, zeta_b_dot)``.
    :param output_projection: Optional projection from ``(f_steady, f_unsteady)`` to a
    smaller output (e.g. beam force); when set, C/D rows correspond to the projected output.
    :return: LinearSystem object.
    """
    jacobians = self.create_jacobians(
        mode="reverse",
        batch_size=batch_size,
        input_projection=input_projection,
        output_projection=output_projection,
    )

    # (row, column, size)
    ref = self.reference
    state_specs: list[tuple[str, str, int]] = [
        ("gamma_b", "gamma_b_n_vec", ref.gamma_b.size),
        ("gamma_w", "gamma_w_n_vec", ref.gamma_w.size),
    ]
    if self.unsteady_force:
        state_specs.append(("gamma_b_nm1", "gamma_b_nm1_vec", ref.gamma_b.size))
    if self.prescribed_wake:
        state_specs.append(("zeta_w", "zeta_w_n_vec", ref.zeta_w.size))
        state_specs.append(("zeta_b", "zeta_b_n_vec", ref.zeta_b.size))
    state_names = [s[0] for s in state_specs]
    state_arg_names = [s[1] for s in state_specs]
    state_sizes = [s[2] for s in state_specs]

    if input_projection is None:
        input_specs: list[tuple[str, int]] = [
            ("zeta_b_np1_vec", ref.zeta_b.size),
            ("zeta_b_dot_np1_vec", ref.zeta_b.size),
        ]
    else:
        input_specs = list(
            zip(input_projection.arg_names, input_projection.arg_sizes)
        )
    if self.bound_upwash:
        input_specs.append(("nu_b_np1_vec", ref.zeta_b.size))
    if self.wake_upwash:
        input_specs.append(("nu_w_np1_vec", ref.zeta_w.size))
    input_arg_names = [s[0] for s in input_specs]
    input_sizes = [s[1] for s in input_specs]

    if output_projection is None:
        n_fs = sum(3 * (m + 1) * (n + 1) for (m, n) in ref.gamma_b.shape)
        output_specs: list[tuple[str, int]] = [("f_steady", n_fs)]
        if self.unsteady_force:
            output_specs.append(("f_unsteady", n_fs))
    else:
        output_specs = [(output_projection.name, output_projection.size)]
    output_names = [s[0] for s in output_specs]
    output_sizes = [s[1] for s in output_specs]

    a = construct_named_block_jacobian(
        entries=tuple(jacobians[k] for k in state_names),
        keys=state_arg_names,
        widths=state_sizes,
        heights=state_sizes,
    )
    b = construct_named_block_jacobian(
        entries=tuple(jacobians[k] for k in state_names),
        keys=input_arg_names,
        widths=input_sizes,
        heights=state_sizes,
    )
    c = construct_named_block_jacobian(
        entries=tuple(jacobians[k] for k in output_names),
        keys=state_arg_names,
        widths=state_sizes,
        heights=output_sizes,
    )
    d = construct_named_block_jacobian(
        entries=tuple(jacobians[k] for k in output_names),
        keys=input_arg_names,
        widths=input_sizes,
        heights=output_sizes,
    )

    return LinearSystem(a=a, b=b, c=c, d=d, dt=self.dt)

run

run(
    u: AeroInputUnflattened,
    x0: AeroStateUnflattened | None = None,
    flowfield: FlowField | None = None,
) -> AeroLinearResult

Run the linear system.

Parameters:

Name	Type	Description	Default
u	AeroInputUnflattened	Total input over time (reference + pertubation).	required
x0	AeroStateUnflattened | None	Initial state perturbations, defaults to zero state.	None
flowfield	FlowField | None	FlowField object to provide flow velocities for bound and wake upwash, defaults to no flow.	None

Source code in src/flapjax/aero/linear/linear_uvlm.py

def run(
    self,
    u: AeroInputUnflattened,
    x0: AeroStateUnflattened | None = None,
    flowfield: FlowField | None = None,
) -> AeroLinearResult:
    r"""
    Run the linear system.
    :param u: Total input over time (reference + pertubation).
    :param x0: Initial state perturbations, defaults to zero state.
    :param flowfield: FlowField object to provide flow velocities for bound and wake upwash, defaults to no flow.
    """
    if self.prescribed_wake and self.sys.removed_u_np1:
        warn(
            "Wake perturbations coordinates at the trailing edge are zero when removing u_np1 from the system."
        )

    if x0 is None:
        x0_vec = None
    else:
        x0_vec = self._pack_state_vector_t(x0)

    n_tstep: int = u.zeta_b[0].shape[
        0
    ]  # number of time steps from first surface, first entry
    t = self.reference.t + jnp.arange(0, n_tstep) * self.dt  # time vector

    u_tot = u

    if self.bound_upwash and flowfield is None and u_tot.nu_b is None:
        warn(
            "No flowfield or bound upwash perturbations provided. Assuming zero bound upwash perturbations."
        )
        u_tot.nu_b = ArrayList(
            [jnp.zeros((n_tstep, *zb.shape)) for zb in self.reference.zeta_b]
        )

    if self.wake_upwash and flowfield is None and u_tot.nu_w is None:
        warn(
            "No flowfield or wake upwash perturbations provided. Assuming zero wake upwash perturbations."
        )
        u_tot.nu_w = ArrayList(
            [jnp.zeros((n_tstep, *zw.shape)) for zw in self.reference.zeta_w]
        )

    # add flowfield contributions to input upwash if provided
    if flowfield is not None:
        if self.bound_upwash:
            nu_b_flow = ArrayList([])
            for i_surf in range(self.reference.n_surf):
                nu_b_flow.append(
                    vmap(flowfield.vmap_call, in_axes=(None, 0), out_axes=0)(
                        self.reference.zeta_b[i_surf],
                        t,  # type: ignore
                    )
                    - flowfield.vmap_call(self.reference.zeta_b[i_surf], t[0])[
                        None, ...
                    ]
                )
            if u_tot.nu_b is None:
                u_tot.nu_b = nu_b_flow
            else:
                u_tot.nu_b += nu_b_flow
        if self.wake_upwash:
            nu_w_flow = ArrayList([])
            for i_surf in range(self.reference.n_surf):
                nu_w_flow.append(
                    vmap(flowfield.vmap_call, in_axes=(None, 0), out_axes=0)(
                        self.reference.zeta_w[i_surf],
                        t,  # type: ignore
                    )
                    - flowfield.vmap_call(self.reference.zeta_w[i_surf], t[0])[
                        None, ...
                    ]
                )
            if u_tot.nu_w is None:
                u_tot.nu_w = nu_w_flow
            else:
                u_tot.nu_w += nu_w_flow
    u_vec = self._pack_input_vector_t(u_tot)

    # run linear system
    x_t, y_t = self.sys.run(u_vec, x0_vec)

    x_t_obj = self._unpack_state_vector_t(x_t)
    y_t_obj = self._unpack_output_vector_t(y_t)

    assert isinstance(x_t_obj, AeroStateUnflattened) and isinstance(
        y_t_obj, AeroOutputUnflattened
    ), (
        "Unpacked state and output must be of type AeroStateUnflattened and AeroOutputUnflattened."
    )

    x_t_tot_obj = self.get_total_state_t(x_t_obj)
    y_t_tot_obj = self.get_total_output_t(y_t_obj)
    u_t_tot_obj = self.get_total_input_t(u_tot)

    assert (
        isinstance(u_t_tot_obj, AeroInputUnflattened)
        and isinstance(x_t_tot_obj, AeroStateUnflattened)
        and isinstance(y_t_tot_obj, AeroOutputUnflattened)
    ), (
        "Unpacked total state and output must be of type AeroStateUnflattened and AeroOutputUnflattened."
    )

    assert isinstance(self.reference, AeroCase), (
        "Reference state must be of type AeroCase."
    )

    # save results to object
    return AeroLinearResult(
        reference=self.reference,
        u_t=u,
        x_t=x_t_obj,
        y_t=y_t_obj,
        u_t_tot=u_t_tot_obj,
        x_t_tot=x_t_tot_obj,
        y_t_tot=y_t_tot_obj,
        n_tstep=n_tstep,
        t=t,
        n_surf=self.reference.n_surf,
        surf_b_names=self.case.surf_b_names,
        surf_w_names=self.case.surf_w_names,
    )

reference_snapshot

reference_snapshot() -> AeroCase

Get the reference (initial) initial_snapshot of the aerodynamic case. This will set the timestep as -1.

Returns:

Type	Description
AeroCase	StaticAero at reference state

Source code in src/flapjax/aero/linear/linear_uvlm.py

def reference_snapshot(self) -> AeroCase:
    r"""
    Get the reference (initial) initial_snapshot of the aerodynamic case. This will set the timestep as -1.
    :return: StaticAero at reference state
    """
    return AeroCase(
        zeta_b=self.reference.zeta_b,
        zeta_b_dot=self.reference.zeta_b_dot,
        zeta_w=self.reference.zeta_w,
        gamma_b=self.reference.gamma_b,
        gamma_b_dot=self.reference.gamma_b_dot,
        gamma_w=self.reference.gamma_w,
        f_steady=self.reference.f_steady,
        f_unsteady=self.reference.f_unsteady,
        cs_ang=self.reference.cs_ang,
        cs_vel=self.reference.cs_vel,
        surf_b_names=self.case.surf_b_names,
        surf_w_names=self.case.surf_w_names,
        i_ts=-1,
        t=jnp.array(0.0),
        c=self.reference.c,
        n=self.reference.nc,
        alpha=self.reference.alpha,
        cl=self.reference.cl,
        cd=self.reference.cd,
        cm=self.reference.cm,
        kernels=self.reference.kernels,
        mirror_normal=self.reference.mirror_normal,
        mirror_point=self.reference.mirror_point,
        mirror_edge_low=self.reference.mirror_edge_low,
        mirror_edge_high=self.reference.mirror_edge_high,
        flowfield=self.reference.flowfield,
        dof_mapping=self.reference.dof_mapping,
        free_wake=self.reference.free_wake,
        gamma_dot_relaxation=self.reference.gamma_dot_relaxation,
        static_horseshoe=self.reference.static_horseshoe,
        batch_size=self.case.batch_size,
    )

plot_reference

plot_reference(
    directory: PathLike, plot_wake: bool = True
) -> Sequence[Path]

Plot the reference (initial) initial_snapshot of the aerodynamic case. This will set the timestep as -1.

Parameters:

Name	Type	Description	Default
directory	PathLike	File path to save the plots to	required
plot_wake	bool	If True, plot the wake grid	True

Source code in src/flapjax/aero/linear/linear_uvlm.py

def plot_reference(
    self, directory: os.PathLike, plot_wake: bool = True
) -> Sequence[Path]:
    r"""
    Plot the reference (initial) initial_snapshot of the aerodynamic case. This will set the timestep as -1.
    :param directory: File path to save the plots to
    :param plot_wake: If True, plot the wake grid
    """
    return self.reference_snapshot().plot(
        Path(directory).resolve(), index=None, plot_wake=plot_wake
    )

data_structures

AeroLinearResult

AeroLinearResult(
    reference: AeroCase,
    u_t: AeroInputUnflattened,
    x_t: AeroStateUnflattened,
    y_t: AeroOutputUnflattened,
    u_t_tot: AeroInputUnflattened,
    x_t_tot: AeroStateUnflattened,
    y_t_tot: AeroOutputUnflattened,
    n_tstep: int,
    n_surf: int,
    t: Array,
    surf_b_names: list[str],
    surf_w_names: list[str],
)

Source code in src/flapjax/aero/linear/data_structures.py

def __init__(
    self,
    reference: AeroCase,
    u_t: AeroInputUnflattened,
    x_t: AeroStateUnflattened,
    y_t: AeroOutputUnflattened,
    u_t_tot: AeroInputUnflattened,
    x_t_tot: AeroStateUnflattened,
    y_t_tot: AeroOutputUnflattened,
    n_tstep: int,
    n_surf: int,
    t: Array,
    surf_b_names: list[str],
    surf_w_names: list[str],
) -> None:
    # system results, if simulated
    self.u_t: AeroInputUnflattened = u_t
    self.x_t: AeroStateUnflattened = x_t
    self.y_t: AeroOutputUnflattened = y_t
    self.u_t_tot: AeroInputUnflattened = u_t_tot
    self.x_t_tot: AeroStateUnflattened = x_t_tot
    self.y_t_tot: AeroOutputUnflattened = y_t_tot
    self.n_tstep: int = n_tstep
    self.n_surf: int = n_surf
    self.t: Array = t
    self.surf_b_names: list[str] = surf_b_names
    self.surf_w_names: list[str] = surf_w_names
    self.reference: AeroCase = reference

plot

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

Plot the aerodynamic grid at specified time steps.

Parameters:

Name	Type	Description	Default
directory	str | PathLike	Directory to save the plots to	required
index	slice | Sequence[int] | int | Array | None	Index or slice of time steps to plot. If None, plot all time steps.	None
plot_wake	bool	If True, plot the wake grid	True

Source code in src/flapjax/aero/linear/data_structures.py

def plot(
    self,
    directory: str | os.PathLike,
    index: slice | Sequence[int] | int | Array | None = None,
    plot_wake: bool = True,
) -> None:
    r"""
    Plot the aerodynamic grid at specified time steps.
    :param directory: Directory to save the plots to
    :param index: Index or slice of time steps to plot. If None, plot all time steps.
    :param plot_wake: If True, plot the wake grid
    """
    if isinstance(index, slice):
        index_ = jnp.arange(self.n_tstep)[index]
    elif isinstance(index, Sequence):
        index_ = jnp.array(index)
    elif isinstance(index, Array):
        index_ = index
    elif isinstance(index, int):
        index_ = (index,)
    elif index is None:
        index_ = jnp.arange(self.n_tstep)
    else:
        raise TypeError("index must be a slices, sequence of ints, or Array")

    directory_path = Path(directory).resolve()
    directory_path.mkdir(parents=True, exist_ok=True)

    paths: list[Sequence[Path]] = []
    for i_ts in index_:
        snapshot = self[i_ts]
        paths.append(snapshot.plot(directory, plot_wake=plot_wake))

    for i_surf in range(2 * self.n_surf):
        try:
            surf_paths = [paths[i][i_surf] for i in range(len(index_))]
            name = (self.surf_b_names + self.surf_w_names)[i_surf] + "_ts"
            write_pvd(directory, name, surf_paths, list(self.t[index_]))
        except IndexError:
            pass

linear_uvlm

LinearInputProjection dataclass

LinearInputProjection(
    arg_names: tuple[str, ...],
    arg_sizes: tuple[int, ...],
    arg_refs: tuple[Array, ...],
    to_zeta: Callable[..., tuple[Array, Array]],
)

Projection of grid coordinates and velocities from a lower-dimensional input space (e.g. beam DOFs) to (zeta_b_np1_vec, zeta_b_dot_np1_vec).

Parameters:

Name	Type	Description	Default
arg_names	tuple[str, ...]	Names of the projected input dofs (used as column labels).	required
arg_sizes	tuple[int, ...]	Sizes of each projected input dof block.	required
arg_refs	tuple[Array, ...]	Reference values for the projected inputs.	required
to_zeta	Callable[..., tuple[Array, Array]]	Callable (**{arg_name: vec}) -> (zeta_b_vec, zeta_b_dot_vec) returning their total values.	required

LinearOutputProjection dataclass

LinearOutputProjection(
    name: str,
    size: int,
    to_output: Callable[[Array, Array | None], Array],
)

Projection of forces from the full aerodynamic grid to a single output space.

Parameters:

Name	Type	Description	Default
name	str	Name of the projected output block (used as row label).	required
size	int	Size of the projected output.	required
to_output	Callable[[Array, Array | None], Array]	Callable (f_steady_vec, f_unsteady_vec | None) -> projected_vec.	required

LinearUVLM

LinearUVLM(
    case: UVLM,
    reference: AeroCase,
    wake_type: LinearWakeType = "frozen",
    bound_upwash: bool = True,
    wake_upwash: bool = True,
    unsteady_force: bool = True,
    *,
    skip_checks: bool = False,
    skip_linearisation: bool = False,
)

Bases: LinearModel[AeroCase, AeroInputUnflattened, AeroStateUnflattened, AeroOutputUnflattened, AeroLinearResult]

Class to represent a linearised UVLM aerodynamic system about a reference state.

Initialise linear UVLM system about a reference state.

Parameters:

Name	Type	Description	Default
case	UVLM	UVLM case object to linearise.	required
reference	AeroCase	StaticAero representing the reference state for linearisation.	required
wake_type	LinearWakeType	Instance of LinearWakeType enum to specify wake treatment.	'frozen'
bound_upwash	bool	If true, include bound surface upwash velocities as inputs.	True
wake_upwash	bool	If true, include wake surface upwash velocities as inputs.	True
unsteady_force	bool	If true, include unsteady force.	True

Source code in src/flapjax/aero/linear/linear_uvlm.py

def __init__(
    self,
    case: UVLM,
    reference: AeroCase,
    wake_type: LinearWakeType = "frozen",
    bound_upwash: bool = True,
    wake_upwash: bool = True,
    unsteady_force: bool = True,
    *,
    skip_checks: bool = False,
    skip_linearisation: bool = False,
):
    r"""
    Initialise linear UVLM system about a reference state.
    :param case: UVLM case object to linearise.
    :param reference: StaticAero representing the reference state for linearisation.
    :param wake_type: Instance of LinearWakeType enum to specify wake treatment.
    :param bound_upwash: If true, include bound surface upwash velocities as inputs.
    :param wake_upwash: If true, include wake surface upwash velocities as inputs.
    :param unsteady_force: If true, include unsteady force.
    """
    # options
    self.prescribed_wake, self.free_wake = {
        "frozen": (False, False),
        "prescribed": (True, False),
        "free": (True, True),
    }[wake_type]
    self.unsteady_force: bool = unsteady_force
    self.bound_upwash: bool = bound_upwash
    self.wake_upwash: bool = wake_upwash

    # time info
    super().__init__(reference=reference, dt=case.dt)

    # check that the reference state is steady
    # whilst linearisation can be performed about unsteady states, the current implementation omits some terms
    # required for this, however, cannot see a practical use case for such a model. Warn the user if the reference
    # state appears unsteady.
    if (
        not skip_checks
        and max([jnp.abs(zbd).max() for zbd in reference.zeta_b_dot]) > 1e-6
    ):
        warn(
            "Reference bound surface velocities are non-zero. Ensure that the reference state is steady for linearisation."
        )

    if (
        not skip_checks
        and max([jnp.abs(gbd).max() for gbd in reference.gamma_b_dot]) > 1e-6
    ):
        warn(
            "Reference bound circulation time derivative is non-zero. Ensure that the reference state is steady for linearisation."
        )

    # kernels
    self.kernels_b: Sequence[KernelFunction] = reference.n_surf * [
        biot_savart_cutoff
    ]
    self.kernels_w: Sequence[KernelFunction] = reference.n_surf * [
        biot_savart_cutoff
    ]

    # wake propagation deltas
    self.case: UVLM = case

    # linear system
    if not skip_linearisation:
        self.sys = self.linearise()

step_vec

step_vec(x_vec: Array, u_vec: Array) -> tuple[Array, Array]

Combined state/output step in vector form, operating on total (reference + perturbation) quantities and returning perturbations relative to the reference.

Parameters:

Name	Type	Description	Default
x_vec	Array	State vector, (n_states, )	required
u_vec	Array	Input vector, (n_inputs, )	required

Returns:

Type	Description
tuple[Array, Array]	Tuple of (state perturbation vector, output perturbation vector).

Source code in src/flapjax/aero/linear/linear_uvlm.py

def step_vec(self, x_vec: Array, u_vec: Array) -> tuple[Array, Array]:
    r"""
    Combined state/output step in vector form, operating on total (reference + perturbation) quantities and
    returning perturbations relative to the reference.
    :param x_vec: State vector, ``(n_states, )``
    :param u_vec: Input vector, ``(n_inputs, )``
    :return: Tuple of (state perturbation vector, output perturbation vector).
    """

    u_np1 = self._unpack_input_vector(u_vec)
    x_n = self.unpack_state_vector(x_vec)

    assert isinstance(u_np1, AeroInputUnflattened) and isinstance(
        x_n, AeroStateUnflattened
    ), (
        "Unpacked input and state must be of type AeroInputUnflattened and AeroStateUnflattened."
    )

    x_np1, y_n = self.step(u_np1=u_np1, x_n=x_n)

    return self.pack_state_vector(x_np1), self.pack_output_vector(y_n)

step

step(
    u_np1: AeroInputUnflattened, x_n: AeroStateUnflattened
) -> tuple[AeroStateUnflattened, AeroOutputUnflattened]

From total inputs u at timestep n+1 and states x at timestep n, compute the states at timestep n+1 and outputs at timestep n.

Source code in src/flapjax/aero/linear/linear_uvlm.py

def step(
    self,
    u_np1: AeroInputUnflattened,
    x_n: AeroStateUnflattened,
) -> tuple[AeroStateUnflattened, AeroOutputUnflattened]:
    r"""
    From total inputs `u` at timestep n+1 and states `x` at timestep n, compute the states at timestep n+1 and
    outputs at timestep n.
    """
    ref = self.reference

    zeta_b_np1: ArrayList = u_np1.zeta_b
    zeta_dot_b_np1: ArrayList = u_np1.zeta_b_dot
    gamma_b_n: ArrayList = x_n.gamma_b
    gamma_w_n: ArrayList = x_n.gamma_w

    if self.unsteady_force:
        assert x_n.gamma_b_nm1 is not None, "gamma_b_nm1 is None"
        gamma_b_dot_n: ArrayList = (gamma_b_n - x_n.gamma_b_nm1) / self.dt
    else:
        gamma_b_dot_n = ref.gamma_b_dot

    if self.prescribed_wake:
        assert x_n.zeta_b is not None, "zeta_b is None"
        assert x_n.zeta_w is not None, "zeta_w is None"
        zeta_b_n: ArrayList = x_n.zeta_b
        zeta_w_n: ArrayList = x_n.zeta_w
    else:
        zeta_b_n = ref.zeta_b
        zeta_w_n = ref.zeta_w

    q_n = AeroFullStates(
        gamma_b=gamma_b_n,
        gamma_w=gamma_w_n,
        gamma_b_dot=ref.gamma_b_dot,
        zeta_w=zeta_w_n,
    )
    (
        _,
        _,
        gamma_b_np1,
        gamma_w_np1,
        _,
        _,
        zeta_w_np1,
        *_,
    ) = self.case.base_solve_from_grid(
        q_nm1=q_n,
        t_n=ref.t,
        zeta_b_n=zeta_b_np1,
        zeta_b_nm1=zeta_b_n,
        zeta_b_dot_n=zeta_dot_b_np1,
        static=False,
        horseshoe=False,
        linearise_variable_wake=True,
        nu_b=u_np1.nu_b,
        nu_w=u_np1.nu_w,
    )

    # the forcing needs to be computed seperately to find its dependence on the current states
    rho = ref.flowfield.rho

    def v_out_func(x_target: Array) -> Array:
        return ref.flowfield.vmap_call(x=x_target, t=ref.t) + compute_v_ind(
            cs=x_target,
            zetas=ArrayList([*zeta_b_np1, *zeta_w_n]),
            gammas=ArrayList([*gamma_b_n, *gamma_w_n]),
            kernels=[*self.kernels_b, *self.kernels_w],
            batch_size=self.case.batch_size,
            mirror_normal=self.case.mirror_normal,
            mirror_point=self.case.mirror_point,
        )

    f_steady_n = compute_steady_forcing(
        zeta_b=zeta_b_np1,
        zeta_dot_b=zeta_dot_b_np1,
        gamma_b=gamma_b_n,
        gamma_w=gamma_w_n,
        rho=rho,
        v_func=v_out_func,
        v_inputs=u_np1.nu_b if self.bound_upwash else None,
        mirror_point=self.case.mirror_point,
        mirror_normal=self.case.mirror_normal,
        mirror_edge_low=self.case.mirror_edge_low,
        mirror_edge_high=self.case.mirror_edge_high,
    )

    normals = compute_nc(zetas=zeta_b_np1)
    if self.unsteady_force:
        f_unsteady_n = ArrayList(
            [
                split_to_vertex(
                    rho * gamma_b_dot_n[i][..., None] * normals[i], (0, 1)
                )
                for i in range(ref.n_surf)
            ]
        )
    else:
        f_unsteady_n = None

    x_np1 = AeroStateUnflattened(
        gamma_b=gamma_b_np1,
        gamma_w=gamma_w_np1,
        gamma_b_nm1=gamma_b_n if self.unsteady_force else None,
        zeta_w=zeta_w_np1 if self.prescribed_wake else None,
        zeta_b=zeta_b_np1 if self.prescribed_wake else None,
    )
    y_n = AeroOutputUnflattened(f_steady=f_steady_n, f_unsteady=f_unsteady_n)

    return x_np1, y_n

compute_jacobians

compute_jacobians(
    input_projection: LinearInputProjection | None = None,
    output_projection: LinearOutputProjection | None = None,
    residual_names: Sequence[str] | None = None,
) -> dict[
    str,
    tuple[
        Callable[..., Any], dict[str, Any], Sequence[str]
    ],
]

Build the Jacobians for the linear system.

Parameters:

Name	Type	Description	Default
input_projection	LinearInputProjection | None	If set, projects the inputs onto a different space.	None
output_projection	LinearOutputProjection | None	If set, projects the forcing outputs onto a different space.	None
residual_names	Sequence[str] | None	If set, restrict the returned residuals to this subset to avoid redundant computation.	None

Returns:

Type	Description
dict[str, tuple[Callable[..., Any], dict[str, Any], Sequence[str]]]	Mapping of residual name to Jacobian(s).

Source code in src/flapjax/aero/linear/linear_uvlm.py

def compute_jacobians(
    self,
    input_projection: LinearInputProjection | None = None,
    output_projection: LinearOutputProjection | None = None,
    residual_names: Sequence[str] | None = None,
) -> dict[str, tuple[Callable[..., Any], dict[str, Any], Sequence[str]]]:
    r"""
    Build the Jacobians for the linear system.
    :param input_projection: If set, projects the inputs onto a different space.
    :param output_projection: If set, projects the forcing outputs onto a different space.
    :param residual_names: If set, restrict the returned residuals to this subset to avoid redundant computation.
    :return: Mapping of residual name to Jacobian(s).
    """
    ref = self.reference

    # bound circulation
    gamma_b_args: dict[str, Any] = {
        "gamma_b_n_vec": ref.gamma_b.ravel(),
        "gamma_w_n_vec": ref.gamma_w.ravel(),
        "zeta_b_np1_vec": ref.zeta_b.ravel(),
        "zeta_b_dot_np1_vec": ref.zeta_b_dot.ravel(),
    }
    gamma_b_diff = [
        "gamma_b_n_vec",
        "gamma_w_n_vec",
        "zeta_b_np1_vec",
        "zeta_b_dot_np1_vec",
    ]
    if self.prescribed_wake:
        gamma_b_args["zeta_w_n_vec"] = ref.zeta_w.ravel()
        gamma_b_args["zeta_b_n_vec"] = ref.zeta_b.ravel()
        gamma_b_diff.extend(["zeta_w_n_vec", "zeta_b_n_vec"])
    if self.bound_upwash:
        gamma_b_args["nu_b_np1_vec"] = jnp.zeros(ref.zeta_b.size)
        gamma_b_diff.append("nu_b_np1_vec")
    if self.wake_upwash:
        gamma_b_args["nu_w_np1_vec"] = jnp.zeros(ref.zeta_w.size)
        gamma_b_diff.append("nu_w_np1_vec")

    # wake propagation (gamma_w and optionally zeta_w)
    wake_args: dict[str, Any] = {
        "gamma_b_n_vec": ref.gamma_b.ravel(),
        "gamma_w_n_vec": ref.gamma_w.ravel(),
    }
    gamma_w_diff = ["gamma_b_n_vec", "gamma_w_n_vec"]
    zeta_w_diff: list[str] = []
    if self.prescribed_wake:
        wake_args["zeta_w_n_vec"] = ref.zeta_w.ravel()
        wake_args["zeta_b_np1_vec"] = ref.zeta_b.ravel()
        zeta_w_diff.extend(["zeta_w_n_vec", "zeta_b_np1_vec"])
    if self.wake_upwash:
        wake_args["nu_w_np1_vec"] = jnp.zeros(ref.zeta_w.size)
        zeta_w_diff.append("nu_w_np1_vec")
    if self.free_wake:
        zeta_w_diff.extend(["gamma_b_n_vec", "gamma_w_n_vec"])

    # steady forcing
    f_steady_args: dict[str, Any] = {
        "gamma_b_n_vec": ref.gamma_b.ravel(),
        "gamma_w_n_vec": ref.gamma_w.ravel(),
        "zeta_b_np1_vec": ref.zeta_b.ravel(),
        "zeta_b_dot_np1_vec": ref.zeta_b_dot.ravel(),
    }
    f_steady_diff = [
        "gamma_b_n_vec",
        "gamma_w_n_vec",
        "zeta_b_np1_vec",
        "zeta_b_dot_np1_vec",
    ]
    if self.prescribed_wake:
        f_steady_args["zeta_w_n_vec"] = ref.zeta_w.ravel()
        f_steady_diff.append("zeta_w_n_vec")
    if self.bound_upwash:
        f_steady_args["nu_b_np1_vec"] = jnp.zeros(ref.zeta_b.size)
        f_steady_diff.append("nu_b_np1_vec")

    residuals: 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 **kw: self.wake_prop_step(**kw)[1],
            wake_args,
            gamma_w_diff,
        ),
    }
    if self.prescribed_wake:
        residuals["zeta_w"] = (
            lambda **kw: self.wake_prop_step(**kw)[0],
            wake_args,
            zeta_w_diff,
        )
    if self.unsteady_force:
        residuals["gamma_b_nm1"] = (
            lambda **kw: kw["gamma_b_n_vec"],
            {"gamma_b_n_vec": ref.gamma_b.ravel()},
            ["gamma_b_n_vec"],
        )
    if self.prescribed_wake:
        # zeta_b state at n+1 == zeta_b input at n+1
        residuals["zeta_b"] = (
            lambda **kw: kw["zeta_b_np1_vec"],
            {"zeta_b_np1_vec": ref.zeta_b.ravel()},
            ["zeta_b_np1_vec"],
        )
    residuals["f_steady"] = (self.f_steady_step, f_steady_args, f_steady_diff)
    if self.unsteady_force:
        f_unsteady_args: dict[str, Any] = {
            "gamma_b_n_vec": ref.gamma_b.ravel(),
            "gamma_b_nm1_vec": ref.gamma_b.ravel(),
            "zeta_b_np1_vec": ref.zeta_b.ravel(),
        }
        residuals["f_unsteady"] = (
            self.f_unsteady_step,
            f_unsteady_args,
            ["gamma_b_n_vec", "gamma_b_nm1_vec", "zeta_b_np1_vec"],
        )

    # apply I/O projections if provided
    if input_projection is not None:
        residuals = self._apply_input_projection(residuals, input_projection)
    if output_projection is not None:
        residuals = self._apply_output_projection(residuals, output_projection)

    if residual_names is not None:
        residuals = {k: v for k, v in residuals.items() if k in residual_names}

    return residuals

create_jacobians

create_jacobians(
    mode: ADMode | dict[str, ADMode] = "reverse",
    batch_size: int | None = None,
    input_projection: LinearInputProjection | None = None,
    output_projection: LinearOutputProjection | None = None,
    residual_names: Sequence[str] | None = None,
) -> dict[str, dict[str, Array]]

Compute the per-residual Jacobians using :func:jacrev_custom.

Source code in src/flapjax/aero/linear/linear_uvlm.py

def create_jacobians(
    self,
    mode: ADMode | dict[str, ADMode] = "reverse",
    batch_size: int | None = None,
    input_projection: LinearInputProjection | None = None,
    output_projection: LinearOutputProjection | None = None,
    residual_names: Sequence[str] | None = None,
) -> dict[str, dict[str, Array]]:
    r"""
    Compute the per-residual Jacobians using :func:`jacrev_custom`.
    """
    residuals = self.compute_jacobians(
        input_projection=input_projection,
        output_projection=output_projection,
        residual_names=residual_names,
    )

    jacobians: dict[str, dict[str, Array]] = {}
    for res_name, (res_func, args, diff_arg_names) in residuals.items():
        res_jac_options: dict[str, Callable[..., Array] | None] = {
            arg: None for arg in diff_arg_names
        }

        res_mode: ADMode = (
            mode.get(res_name, "reverse") if isinstance(mode, dict) else mode
        )

        jacs, _, _ = jacrev_custom(
            func=res_func,
            jac_options=res_jac_options,
            n_profile_loops=None,
            func_name=res_name,
            map_batch_size=batch_size,
            mode=res_mode,
        )(**args)
        jacobians[res_name] = jacs

    return jacobians

linearise

linearise(
    batch_size: int | None = None,
    *,
    input_projection: LinearInputProjection | None = None,
    output_projection: LinearOutputProjection | None = None,
) -> LinearSystem

Build the linear state-space system.

Parameters:

Name	Type	Description	Default
batch_size	int | None	If not None, batch the Jacobian passes to reduce memory.	None
input_projection	LinearInputProjection | None	Optional projection from lower-dim inputs (e.g. beam DOFs) to (zeta_b, zeta_b_dot).	None
output_projection	LinearOutputProjection | None	Optional projection from (f_steady, f_unsteady) to a smaller output (e.g. beam force); when set, C/D rows correspond to the projected output.	None

Returns:

Type	Description
LinearSystem	LinearSystem object.

Source code in src/flapjax/aero/linear/linear_uvlm.py

def linearise(
    self,
    batch_size: int | None = None,
    *,
    input_projection: LinearInputProjection | None = None,
    output_projection: LinearOutputProjection | None = None,
) -> LinearSystem:
    r"""
    Build the linear state-space system.
    :param batch_size: If not None, batch the Jacobian passes to reduce memory.
    :param input_projection: Optional projection from lower-dim inputs (e.g. beam DOFs) to
    ``(zeta_b, zeta_b_dot)``.
    :param output_projection: Optional projection from ``(f_steady, f_unsteady)`` to a
    smaller output (e.g. beam force); when set, C/D rows correspond to the projected output.
    :return: LinearSystem object.
    """
    jacobians = self.create_jacobians(
        mode="reverse",
        batch_size=batch_size,
        input_projection=input_projection,
        output_projection=output_projection,
    )

    # (row, column, size)
    ref = self.reference
    state_specs: list[tuple[str, str, int]] = [
        ("gamma_b", "gamma_b_n_vec", ref.gamma_b.size),
        ("gamma_w", "gamma_w_n_vec", ref.gamma_w.size),
    ]
    if self.unsteady_force:
        state_specs.append(("gamma_b_nm1", "gamma_b_nm1_vec", ref.gamma_b.size))
    if self.prescribed_wake:
        state_specs.append(("zeta_w", "zeta_w_n_vec", ref.zeta_w.size))
        state_specs.append(("zeta_b", "zeta_b_n_vec", ref.zeta_b.size))
    state_names = [s[0] for s in state_specs]
    state_arg_names = [s[1] for s in state_specs]
    state_sizes = [s[2] for s in state_specs]

    if input_projection is None:
        input_specs: list[tuple[str, int]] = [
            ("zeta_b_np1_vec", ref.zeta_b.size),
            ("zeta_b_dot_np1_vec", ref.zeta_b.size),
        ]
    else:
        input_specs = list(
            zip(input_projection.arg_names, input_projection.arg_sizes)
        )
    if self.bound_upwash:
        input_specs.append(("nu_b_np1_vec", ref.zeta_b.size))
    if self.wake_upwash:
        input_specs.append(("nu_w_np1_vec", ref.zeta_w.size))
    input_arg_names = [s[0] for s in input_specs]
    input_sizes = [s[1] for s in input_specs]

    if output_projection is None:
        n_fs = sum(3 * (m + 1) * (n + 1) for (m, n) in ref.gamma_b.shape)
        output_specs: list[tuple[str, int]] = [("f_steady", n_fs)]
        if self.unsteady_force:
            output_specs.append(("f_unsteady", n_fs))
    else:
        output_specs = [(output_projection.name, output_projection.size)]
    output_names = [s[0] for s in output_specs]
    output_sizes = [s[1] for s in output_specs]

    a = construct_named_block_jacobian(
        entries=tuple(jacobians[k] for k in state_names),
        keys=state_arg_names,
        widths=state_sizes,
        heights=state_sizes,
    )
    b = construct_named_block_jacobian(
        entries=tuple(jacobians[k] for k in state_names),
        keys=input_arg_names,
        widths=input_sizes,
        heights=state_sizes,
    )
    c = construct_named_block_jacobian(
        entries=tuple(jacobians[k] for k in output_names),
        keys=state_arg_names,
        widths=state_sizes,
        heights=output_sizes,
    )
    d = construct_named_block_jacobian(
        entries=tuple(jacobians[k] for k in output_names),
        keys=input_arg_names,
        widths=input_sizes,
        heights=output_sizes,
    )

    return LinearSystem(a=a, b=b, c=c, d=d, dt=self.dt)

run

run(
    u: AeroInputUnflattened,
    x0: AeroStateUnflattened | None = None,
    flowfield: FlowField | None = None,
) -> AeroLinearResult

Run the linear system.

Parameters:

Name	Type	Description	Default
u	AeroInputUnflattened	Total input over time (reference + pertubation).	required
x0	AeroStateUnflattened | None	Initial state perturbations, defaults to zero state.	None
flowfield	FlowField | None	FlowField object to provide flow velocities for bound and wake upwash, defaults to no flow.	None

Source code in src/flapjax/aero/linear/linear_uvlm.py

def run(
    self,
    u: AeroInputUnflattened,
    x0: AeroStateUnflattened | None = None,
    flowfield: FlowField | None = None,
) -> AeroLinearResult:
    r"""
    Run the linear system.
    :param u: Total input over time (reference + pertubation).
    :param x0: Initial state perturbations, defaults to zero state.
    :param flowfield: FlowField object to provide flow velocities for bound and wake upwash, defaults to no flow.
    """
    if self.prescribed_wake and self.sys.removed_u_np1:
        warn(
            "Wake perturbations coordinates at the trailing edge are zero when removing u_np1 from the system."
        )

    if x0 is None:
        x0_vec = None
    else:
        x0_vec = self._pack_state_vector_t(x0)

    n_tstep: int = u.zeta_b[0].shape[
        0
    ]  # number of time steps from first surface, first entry
    t = self.reference.t + jnp.arange(0, n_tstep) * self.dt  # time vector

    u_tot = u

    if self.bound_upwash and flowfield is None and u_tot.nu_b is None:
        warn(
            "No flowfield or bound upwash perturbations provided. Assuming zero bound upwash perturbations."
        )
        u_tot.nu_b = ArrayList(
            [jnp.zeros((n_tstep, *zb.shape)) for zb in self.reference.zeta_b]
        )

    if self.wake_upwash and flowfield is None and u_tot.nu_w is None:
        warn(
            "No flowfield or wake upwash perturbations provided. Assuming zero wake upwash perturbations."
        )
        u_tot.nu_w = ArrayList(
            [jnp.zeros((n_tstep, *zw.shape)) for zw in self.reference.zeta_w]
        )

    # add flowfield contributions to input upwash if provided
    if flowfield is not None:
        if self.bound_upwash:
            nu_b_flow = ArrayList([])
            for i_surf in range(self.reference.n_surf):
                nu_b_flow.append(
                    vmap(flowfield.vmap_call, in_axes=(None, 0), out_axes=0)(
                        self.reference.zeta_b[i_surf],
                        t,  # type: ignore
                    )
                    - flowfield.vmap_call(self.reference.zeta_b[i_surf], t[0])[
                        None, ...
                    ]
                )
            if u_tot.nu_b is None:
                u_tot.nu_b = nu_b_flow
            else:
                u_tot.nu_b += nu_b_flow
        if self.wake_upwash:
            nu_w_flow = ArrayList([])
            for i_surf in range(self.reference.n_surf):
                nu_w_flow.append(
                    vmap(flowfield.vmap_call, in_axes=(None, 0), out_axes=0)(
                        self.reference.zeta_w[i_surf],
                        t,  # type: ignore
                    )
                    - flowfield.vmap_call(self.reference.zeta_w[i_surf], t[0])[
                        None, ...
                    ]
                )
            if u_tot.nu_w is None:
                u_tot.nu_w = nu_w_flow
            else:
                u_tot.nu_w += nu_w_flow
    u_vec = self._pack_input_vector_t(u_tot)

    # run linear system
    x_t, y_t = self.sys.run(u_vec, x0_vec)

    x_t_obj = self._unpack_state_vector_t(x_t)
    y_t_obj = self._unpack_output_vector_t(y_t)

    assert isinstance(x_t_obj, AeroStateUnflattened) and isinstance(
        y_t_obj, AeroOutputUnflattened
    ), (
        "Unpacked state and output must be of type AeroStateUnflattened and AeroOutputUnflattened."
    )

    x_t_tot_obj = self.get_total_state_t(x_t_obj)
    y_t_tot_obj = self.get_total_output_t(y_t_obj)
    u_t_tot_obj = self.get_total_input_t(u_tot)

    assert (
        isinstance(u_t_tot_obj, AeroInputUnflattened)
        and isinstance(x_t_tot_obj, AeroStateUnflattened)
        and isinstance(y_t_tot_obj, AeroOutputUnflattened)
    ), (
        "Unpacked total state and output must be of type AeroStateUnflattened and AeroOutputUnflattened."
    )

    assert isinstance(self.reference, AeroCase), (
        "Reference state must be of type AeroCase."
    )

    # save results to object
    return AeroLinearResult(
        reference=self.reference,
        u_t=u,
        x_t=x_t_obj,
        y_t=y_t_obj,
        u_t_tot=u_t_tot_obj,
        x_t_tot=x_t_tot_obj,
        y_t_tot=y_t_tot_obj,
        n_tstep=n_tstep,
        t=t,
        n_surf=self.reference.n_surf,
        surf_b_names=self.case.surf_b_names,
        surf_w_names=self.case.surf_w_names,
    )

reference_snapshot

reference_snapshot() -> AeroCase

Get the reference (initial) initial_snapshot of the aerodynamic case. This will set the timestep as -1.

Returns:

Type	Description
AeroCase	StaticAero at reference state

Source code in src/flapjax/aero/linear/linear_uvlm.py

def reference_snapshot(self) -> AeroCase:
    r"""
    Get the reference (initial) initial_snapshot of the aerodynamic case. This will set the timestep as -1.
    :return: StaticAero at reference state
    """
    return AeroCase(
        zeta_b=self.reference.zeta_b,
        zeta_b_dot=self.reference.zeta_b_dot,
        zeta_w=self.reference.zeta_w,
        gamma_b=self.reference.gamma_b,
        gamma_b_dot=self.reference.gamma_b_dot,
        gamma_w=self.reference.gamma_w,
        f_steady=self.reference.f_steady,
        f_unsteady=self.reference.f_unsteady,
        cs_ang=self.reference.cs_ang,
        cs_vel=self.reference.cs_vel,
        surf_b_names=self.case.surf_b_names,
        surf_w_names=self.case.surf_w_names,
        i_ts=-1,
        t=jnp.array(0.0),
        c=self.reference.c,
        n=self.reference.nc,
        alpha=self.reference.alpha,
        cl=self.reference.cl,
        cd=self.reference.cd,
        cm=self.reference.cm,
        kernels=self.reference.kernels,
        mirror_normal=self.reference.mirror_normal,
        mirror_point=self.reference.mirror_point,
        mirror_edge_low=self.reference.mirror_edge_low,
        mirror_edge_high=self.reference.mirror_edge_high,
        flowfield=self.reference.flowfield,
        dof_mapping=self.reference.dof_mapping,
        free_wake=self.reference.free_wake,
        gamma_dot_relaxation=self.reference.gamma_dot_relaxation,
        static_horseshoe=self.reference.static_horseshoe,
        batch_size=self.case.batch_size,
    )

plot_reference

plot_reference(
    directory: PathLike, plot_wake: bool = True
) -> Sequence[Path]

Plot the reference (initial) initial_snapshot of the aerodynamic case. This will set the timestep as -1.

Parameters:

Name	Type	Description	Default
directory	PathLike	File path to save the plots to	required
plot_wake	bool	If True, plot the wake grid	True

Source code in src/flapjax/aero/linear/linear_uvlm.py

def plot_reference(
    self, directory: os.PathLike, plot_wake: bool = True
) -> Sequence[Path]:
    r"""
    Plot the reference (initial) initial_snapshot of the aerodynamic case. This will set the timestep as -1.
    :param directory: File path to save the plots to
    :param plot_wake: If True, plot the wake grid
    """
    return self.reference_snapshot().plot(
        Path(directory).resolve(), index=None, plot_wake=plot_wake
    )

unpack_state_vector

unpack_state_vector(x: Array) -> S

Unpack a state vector into its components.

Parameters:

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

Returns:

Type	Description
S	StateUnflattened object.

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

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

unpack_output_vector

unpack_output_vector(y: Array) -> O

Unpack an output vector into its components.

Parameters:

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

Returns:

Type	Description
O	OutputUnflattened object.

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

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

pack_input_vector

pack_input_vector(u_input: InputUnflattened) -> Array

Pack an input unflattened object into a vector.

Parameters:

Name	Type	Description	Default
u_input	InputUnflattened	InputUnflattened object.	required

Returns:

Type	Description
Array	Input vector, (n_inputs, ).

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

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

pack_state_vector

pack_state_vector(x_state: StateUnflattened) -> Array

Pack a state unflattened object into a vector.

Parameters:

Name	Type	Description	Default
x_state	StateUnflattened	StateUnflattened object.	required

Returns:

Type	Description
Array	State vector, (n_states, ).

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

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

pack_output_vector

pack_output_vector(y_output: OutputUnflattened) -> Array

Pack an output unflattened object into a vector.

Parameters:

Name	Type	Description	Default
y_output	OutputUnflattened	OutputUnflattened object.	required

Returns:

Type	Description
Array	Output vector, (n_outputs, ).

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

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

get_total_input

get_total_input(u: InputUnflattened) -> InputUnflattened

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

Parameters:

Name	Type	Description	Default
u	InputUnflattened	InputUnflattened perturbation object.	required

Returns:

Type	Description
InputUnflattened	InputUnflattened total object.

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

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

get_total_state

get_total_state(x: StateUnflattened) -> StateUnflattened

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

Parameters:

Name	Type	Description	Default
x	StateUnflattened	StateUnflattened perturbation object.	required

Returns:

Type	Description
StateUnflattened	StateUnflattened total object.

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

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

get_total_output

get_total_output(y: OutputUnflattened) -> OutputUnflattened

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

Parameters:

Name	Type	Description	Default
y	OutputUnflattened	OutputUnflattened perturbation object.	required

Returns:

Type	Description
OutputUnflattened	OutputUnflattened total object.

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

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

get_total_input_t

get_total_input_t(
    u_t: InputUnflattened,
) -> InputUnflattened

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

Parameters:

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

Returns:

Type	Description
InputUnflattened	InputUnflattened total time history object.

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

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

get_total_state_t

get_total_state_t(
    x_t: StateUnflattened,
) -> StateUnflattened

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

Parameters:

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

Returns:

Type	Description
StateUnflattened	StateUnflattened total time history object.

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

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

get_total_output_t

get_total_output_t(
    y_t: OutputUnflattened,
) -> OutputUnflattened

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

Parameters:

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

Returns:

Type	Description
OutputUnflattened	OutputUnflattened total time history object.

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

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

get_zero_input

get_zero_input() -> InputUnflattened

Get a zero input unflattened object.

Returns:

Type	Description
InputUnflattened	InputUnflattened object with zero arrays.

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

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

get_zero_state

get_zero_state() -> StateUnflattened

Get a zero state unflattened object.

Returns:

Type	Description
StateUnflattened	StateUnflattened object with zero arrays.

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

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

get_zero_output

get_zero_output() -> OutputUnflattened

Get a zero output unflattened object.

Returns:

Type	Description
OutputUnflattened	OutputUnflattened object with zero arrays.

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

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

utils

DynamicAeroSolver

Bases: Protocol

Required methods for a dynamic aerodynamic solver that can be coupled into the structural solver.

make_rectangular_grid

make_rectangular_grid(
    m: int,
    n: int,
    chord: Array | float,
    ea: Array | float,
    camber_line: tuple[Array, Array] | None = None,
    twist: Array | float = 0.0,
) -> Array

Create a rectangular aerodynamic grid.

Parameters:

Name	Type	Description	Default
m	int	Number of panels in the chordwise direction.	required
n	int	Number of panels in the spanwise direction.	required
chord	Array | float	Surface chord length.	required
ea	Array | float	Elastic axis location as fraction of chord.	required
camber_line	tuple[Array, Array] | None	Optional mean camber line as a pair (x/c, z/c) of equal-length vectors, giving the camber-line height at a set of chordwise stations, with x/c running from 0 (leading edge) to 1 (trailing edge). If None (default), the section is a flat plate.	None
twist	Array | float	Built-in geometric twist angle in radians, uniform over the whole grid, applied by rotating the local chord/camber section about the local spanwise axis (HINGE_AXIS_DEFAULT). This is a purely aerodynamic incidence offset: unlike a beam's y_vector-defined twist (which only reorients the structural cross-section's stiffness/mass axes), it actually changes the panels' angle of attack, since a uniform twist produces no curvature for the structural solver to pick up on its own. Default 0 (no twist).	0.0

Returns:

Type	Description
Array	Local grid points for planar wing, (zeta_m, zeta_n, 3).

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

def make_rectangular_grid(
    m: int,
    n: int,
    chord: Array | float,
    ea: Array | float,
    camber_line: tuple[Array, Array] | None = None,
    twist: Array | float = 0.0,
) -> Array:
    r"""
    Create a rectangular aerodynamic grid.
    :param m: Number of panels in the chordwise direction.
    :param n: Number of panels in the spanwise direction.
    :param chord: Surface chord length.
    :param ea: Elastic axis location as fraction of chord.
    :param camber_line: Optional mean camber line as a pair ``(x/c, z/c)`` of equal-length vectors, giving
    the camber-line height at a set of chordwise stations, with ``x/c`` running from 0 (leading edge)
    to 1 (trailing edge). If None (default), the section is a flat plate.
    :param twist: Built-in geometric twist angle in radians, uniform over the whole grid, applied by rotating the
    local chord/camber section about the local spanwise axis (``HINGE_AXIS_DEFAULT``). This is a purely
    aerodynamic incidence offset: unlike a beam's ``y_vector``-defined twist (which only reorients the
    structural cross-section's stiffness/mass axes), it actually changes the panels' angle of attack, since a
    uniform twist produces no curvature for the structural solver to pick up on its own. Default 0 (no twist).
    :return: Local grid points for planar wing, ``(zeta_m, zeta_n, 3)``.
    """

    x_over_c = jnp.linspace(0.0, 1.0, m + 1)
    grid = jnp.zeros((m + 1, n + 1, 3))
    grid = grid.at[..., 0].set((x_over_c * chord - ea * chord)[:, None])
    if camber_line is not None:
        camber_x, camber_z = camber_line
        z_over_c = jnp.interp(x_over_c, camber_x, camber_z)
        grid = grid.at[..., 2].set((z_over_c * chord)[:, None])
    rmat = exp_so3(HINGE_AXIS_DEFAULT * twist)
    grid = jnp.einsum("ij,mnj->mni", rmat, grid)
    return grid

add_control_surface

add_control_surface(
    grid: Array,
    angle: Array,
    m_slice: Array | Sequence[int] | slice,
    n_slice: Array | Sequence[int] | slice,
    hinge_axis: Array = HINGE_AXIS_DEFAULT,
) -> Array

Add a control surface to a panel grid.

Parameters:

Name	Type	Description	Default
grid	Array	Grid without deflection of this surface, (zeta_m, zeta_n, 3).	required
angle	Array	Angle in radians through which the control surface will be deflected.	required
m_slice	Array | Sequence[int] | slice	Slice of chordwise strips to include in the control surface.	required
n_slice	Array | Sequence[int] | slice	Slice of spanwise strips to include in the control surface.	required
hinge_axis	Array	Axis of the hinge surface in the local frame, (3, ).	HINGE_AXIS_DEFAULT

Returns:

Type	Description
Array	Deflected aerodynamic grid, (zeta_m, zeta_n, 3).

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

def add_control_surface(
    grid: Array,
    angle: Array,
    m_slice: Array | Sequence[int] | slice,
    n_slice: Array | Sequence[int] | slice,
    hinge_axis: Array = HINGE_AXIS_DEFAULT,
) -> Array:
    r"""
    Add a control surface to a panel grid.
    :param grid: Grid without deflection of this surface, ``(zeta_m, zeta_n, 3)``.
    :param angle: Angle in radians through which the control surface will be deflected.
    :param m_slice: Slice of chordwise strips to include in the control surface.
    :param n_slice: Slice of spanwise strips to include in the control surface.
    :param hinge_axis: Axis of the hinge surface in the local frame, ``(3, )``.
    :return: Deflected aerodynamic grid, ``(zeta_m, zeta_n, 3)``.
    """

    m_slice_arr: Array = index_to_arr(index=m_slice, n_entries=grid.shape[0])
    n_slice_arr: Array = index_to_arr(index=n_slice, n_entries=grid.shape[1])

    # grid for deflected surfaces
    grid_out = grid

    def inner_func(n_idx: Array) -> Array:
        hinge_point = grid[m_slice_arr[0], n_idx, :]  # (3, )

        crv = hinge_axis * angle  # cartesian rotation vector for surface, (3, ).
        rmat = exp_so3(crv)  # rotation matrix for rotating surface

        # transform coordinates to rotate control surface
        return (
            jnp.einsum(
                "ij,hj->hi",
                rmat,
                (grid[m_slice_arr, n_idx, :] - hinge_point[None, :]),
            )
            + hinge_point[None, :]
        )

    # update grid
    return grid_out.at[jnp.ix_(m_slice_arr, n_slice_arr, jnp.arange(3))].set(
        vmap(inner_func, in_axes=0, out_axes=1)(n_slice_arr)
    )

compute_surf_c

compute_surf_c(zeta: Array) -> Array

Compute the colocation points for a given grid of points on a single surface.

Parameters:

Name	Type	Description	Default
zeta	Array	Grid of points, (..., zeta_m, zeta_n, 3).	required

Returns:

Type	Description
Array	Colocation points, (..., m, n, 3).

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

def compute_surf_c(zeta: Array) -> Array:
    r"""
    Compute the colocation points for a given grid of points on a single surface.
    :param zeta: Grid of points, ``(..., zeta_m, zeta_n, 3)``.
    :return: Colocation points, ``(..., m, n, 3)``.
    """
    return neighbour_average(zeta, axes=(-3, -2))

compute_surf_nc

compute_surf_nc(zeta: Array) -> Array

Compute the surface normal vectors for a given grid of points on a single surface. These have length equal to the area of their corresponding panel.

Parameters:

Name	Type	Description	Default
zeta	Array	Grid of points, (..., zeta_m, zeta_n, 3).	required

Returns:

Type	Description
Array	Normal vectors, (..., m, n, 3).

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

def compute_surf_nc(zeta: Array) -> Array:
    r"""
    Compute the surface normal vectors for a given grid of points on a single surface. These have length equal to the
    area of their corresponding panel.
    :param zeta: Grid of points, ``(..., zeta_m, zeta_n, 3)``.
    :return: Normal vectors, ``(..., m, n, 3)``.
    """
    diag1 = zeta[..., 1:, 1:, :] - zeta[..., :-1, :-1, :]
    diag2 = zeta[..., 1:, :-1, :] - zeta[..., :-1, 1:, :]
    return jnp.cross(diag1, diag2)

compute_c

compute_c(zetas: ArrayList) -> ArrayList

Compute the colocation points for a list of surface grids.

Parameters:

Name	Type	Description	Default
zetas	ArrayList	Grids of points, (n_surf, )(zeta_m, zeta_n, 3).	required

Returns:

Type	Description
ArrayList	Colocation points (n_surf ,)(m, n, 3).

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

def compute_c(zetas: ArrayList) -> ArrayList:
    r"""
    Compute the colocation points for a list of surface grids.
    :param zetas: Grids of points, ``(n_surf, )(zeta_m, zeta_n, 3)``.
    :return: Colocation points ``(n_surf ,)(m, n, 3)``.
    """
    return ArrayList([compute_surf_c(zeta) for zeta in zetas])

compute_nc

compute_nc(zetas: ArrayList) -> ArrayList

Compute the surface normal vectors for a list of surface grids.

Parameters:

Name	Type	Description	Default
zetas	ArrayList	Grids of points, (n_surf, )(zeta_m, zeta_n, 3).	required

Returns:

Type	Description
ArrayList	Normal vectors (n_surf, )(m, n, 3).

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

def compute_nc(zetas: ArrayList) -> ArrayList:
    r"""
    Compute the surface normal vectors for a list of surface grids.
    :param zetas: Grids of points, ``(n_surf, )(zeta_m, zeta_n, 3)``.
    :return: Normal vectors ``(n_surf, )(m, n, 3)``.
    """
    return ArrayList([compute_surf_nc(zeta) for zeta in zetas])

compute_mirror_edges

compute_mirror_edges(
    zeta_b_ref: ArrayList,
    mirror_point: Array | None,
    mirror_normal: Array | None,
) -> tuple[ArrayList, ArrayList]

Determine which surfaces have an edge coinciding with a mirror plane.

Parameters:

Name	Type	Description	Default
zeta_b_ref	ArrayList	Reference bound grid coordinates, (n_surf, )(zeta_m, zeta_n, 3).	required
mirror_point	Array | None	Point on mirror plane, (3, ). If None (together with mirror_normal), no surface is considered to have a mirrored edge.	required
mirror_normal	Array | None	Normal vector to mirror across, (3, ).	required

Returns:

Type	Description
tuple[ArrayList, ArrayList]	Per-surface booleans for the low and high spanwise edges, each (n_surf, )().

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

def compute_mirror_edges(
    zeta_b_ref: ArrayList, mirror_point: Array | None, mirror_normal: Array | None
) -> tuple[ArrayList, ArrayList]:
    r"""
    Determine which surfaces have an edge coinciding with a mirror plane.
    :param zeta_b_ref: Reference bound grid coordinates, ``(n_surf, )(zeta_m, zeta_n, 3)``.
    :param mirror_point: Point on mirror plane, ``(3, )``. If None (together with ``mirror_normal``), no surface
    is considered to have a mirrored edge.
    :param mirror_normal: Normal vector to mirror across, ``(3, )``.
    :return: Per-surface booleans for the low and high spanwise edges, each ``(n_surf, )()``.
    """
    if mirror_point is None or mirror_normal is None:
        not_mirrored = jnp.array(False)
        return (
            ArrayList([not_mirrored for _ in zeta_b_ref]),
            ArrayList([not_mirrored for _ in zeta_b_ref]),
        )
    return (
        ArrayList(
            [
                _on_mirror_plane(zeta[:, 0, :], mirror_point, mirror_normal)
                for zeta in zeta_b_ref
            ]
        ),
        ArrayList(
            [
                _on_mirror_plane(zeta[:, -1, :], mirror_point, mirror_normal)
                for zeta in zeta_b_ref
            ]
        ),
    )

compute_steady_forcing

compute_steady_forcing(
    zeta_b: ArrayList,
    zeta_dot_b: ArrayList | None,
    gamma_b: ArrayList,
    gamma_w: ArrayList,
    rho: Array,
    v_func: Callable[[Array], Array],
    v_inputs: ArrayList | None,
    mirror_point: Array | None = None,
    mirror_normal: Array | None = None,
    mirror_edge_low: ArrayList | None = None,
    mirror_edge_high: ArrayList | None = None,
) -> ArrayList

Calculate steady aerodynamic forcing for all surfaces at specified time step.

Parameters:

Name	Type	Description	Default
zeta_b	ArrayList	Bound grid coordinates, (n_surf, )(zeta_m, zeta_n, 3).	required
zeta_dot_b	ArrayList | None	Bound grid velocities, (n_surf, )(zeta_m, zeta_n, 3).	required
gamma_b	ArrayList	Bound grid circulation, (n_surf, )(m, n)	required
gamma_w	ArrayList	Wake grid circulation, (n_surf, )(m, n)	required
rho	Array	Flow field density.	required
v_func	Callable[[Array], Array]	Total velocity as a function of coordinate.	required
v_inputs	ArrayList | None	Additive inputs for total velocity on bound grid vertex, used for the linear solver for custom perturbations.	required
mirror_point	Array | None	Point on mirror plane, (3, ).	None
mirror_normal	Array | None	Normal vector to mirror across, (3, ).	None
mirror_edge_low	ArrayList | None	Per-surface booleans marking whether that surface's n=0 edge lies on the mirror plane, (n_surf, )().	None
mirror_edge_high	ArrayList | None	As mirror_edge_low, for the n=-1 edge.	None

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

def compute_steady_forcing(
    zeta_b: ArrayList,
    zeta_dot_b: ArrayList | None,
    gamma_b: ArrayList,
    gamma_w: ArrayList,
    rho: Array,
    v_func: Callable[[Array], Array],
    v_inputs: ArrayList | None,
    mirror_point: Array | None = None,
    mirror_normal: Array | None = None,
    mirror_edge_low: ArrayList | None = None,
    mirror_edge_high: ArrayList | None = None,
) -> ArrayList:
    r"""
    Calculate steady aerodynamic forcing for all surfaces at specified time step.
    :param zeta_b: Bound grid coordinates, ``(n_surf, )(zeta_m, zeta_n, 3)``.
    :param zeta_dot_b: Bound grid velocities, ``(n_surf, )(zeta_m, zeta_n, 3)``.
    :param gamma_b: Bound grid circulation, ``(n_surf, )(m, n)``
    :param gamma_w: Wake grid circulation, ``(n_surf, )(m, n)``
    :param rho: Flow field density.
    :param v_func: Total velocity as a function of coordinate.
    :param v_inputs: Additive inputs for total velocity on bound grid vertex, used for the linear solver for custom
    perturbations.
    :param mirror_point: Point on mirror plane, ``(3, )``.
    :param mirror_normal: Normal vector to mirror across, ``(3, )``.
    :param mirror_edge_low: Per-surface booleans marking whether that surface's ``n=0`` edge lies on the mirror
    plane, ``(n_surf, )()``.
    :param mirror_edge_high: As ``mirror_edge_low``, for the ``n=-1`` edge.
    """

    has_mirror = mirror_point is not None and mirror_normal is not None
    if has_mirror and (mirror_edge_low is None or mirror_edge_high is None):
        raise ValueError(
            "mirror_edge_low/mirror_edge_high must be provided (e.g. via compute_mirror_edges) "
            "whenever mirror_point/mirror_normal are given."
        )
    f_steady = ArrayList([])

    if zeta_dot_b is None:
        zeta_dot_bs_: list[Array | None] = [None] * len(zeta_b)
    else:
        zeta_dot_bs_ = zeta_dot_b

    if v_inputs is None:
        v_inputs_ = [None] * len(zeta_b)
    else:
        v_inputs_ = v_inputs

    mirror_edge_low_ = (
        mirror_edge_low if mirror_edge_low is not None else [None] * len(zeta_b)
    )
    mirror_edge_high_ = (
        mirror_edge_high if mirror_edge_high is not None else [None] * len(zeta_b)
    )

    for (
        zeta_b_surf,
        zeta_dot_b_surf,
        gamma_b_surf,
        gamma_w_surf,
        v_input_surf,
        on_plane_low,
        on_plane_high,
    ) in zip(
        zeta_b,
        zeta_dot_bs_,
        gamma_b,
        gamma_w,
        v_inputs_,
        mirror_edge_low_,
        mirror_edge_high_,
    ):
        # compute midpoints
        mp_chordwise = neighbour_average(zeta_b_surf, axes=0)  # (gamma_m, gamma_n+1, 3)
        mp_spanwise = neighbour_average(zeta_b_surf, axes=1)  # (gamma_m+1, gamma_n, 3)

        assert zeta_dot_b_surf is not None

        mp_dot_chordwise = neighbour_average(
            zeta_dot_b_surf, axes=0
        )  # (gamma_m, gamma_n+1, 3)
        mp_dot_spanwise = neighbour_average(
            zeta_dot_b_surf, axes=1
        )  # (gamma_m+1, gamma_n, 3)

        # relative flow velocities at midpoints
        v_rel_chordwise = (
            v_func(mp_chordwise) - mp_dot_chordwise
        )  # (gamma_m, gamma_n+1, 3)
        v_rel_spanwise = (
            v_func(mp_spanwise) - mp_dot_spanwise
        )  # (gamma_m+1, gamma_n, 3)

        # add any input_ velocities
        if v_input_surf is not None:
            v_rel_chordwise += neighbour_average(v_input_surf, axes=0)
            v_rel_spanwise += neighbour_average(v_input_surf, axes=1)

        # equivalent strengths of filaments
        gamma_chordwise = jnp.zeros(
            v_rel_chordwise.shape[:-1]
        )  # (gamma_m, gamma_n+1, 3)
        gamma_chordwise = gamma_chordwise.at[:, :-1].set(gamma_b_surf)
        gamma_chordwise = gamma_chordwise.at[:, 1:].add(-gamma_b_surf)
        gamma_spanwise = jnp.zeros(v_rel_spanwise.shape[:-1])  # (gamma_m+1, gamma_n, 3)
        gamma_spanwise = gamma_spanwise.at[:-1, :].set(-gamma_b_surf)
        gamma_spanwise = gamma_spanwise.at[1:, :].add(gamma_b_surf)

        # add first wake gamma
        if gamma_w_surf.shape[0] > 0:
            gamma_spanwise = gamma_spanwise.at[-1, :].add(-gamma_w_surf[0, :])

        if has_mirror:
            assert mirror_point is not None and mirror_normal is not None
            # a spanwise edge lying exactly on the mirror plane carries no shed trailing vortex
            gamma_chordwise = gamma_chordwise.at[:, 0].set(
                jnp.where(on_plane_low, 0.0, gamma_chordwise[:, 0])
            )
            gamma_chordwise = gamma_chordwise.at[:, -1].set(
                jnp.where(on_plane_high, 0.0, gamma_chordwise[:, -1])
            )

        # filament vectors (from zeta_b_fil, which may differ from the midpoint geometry)
        r_chordwise = (
            zeta_b_surf[1:, :, :] - zeta_b_surf[:-1, :, :]
        )  # (gamma_m, gamma_n+1, 3)
        r_spanwise = (
            zeta_b_surf[:, 1:, :] - zeta_b_surf[:, :-1, :]
        )  # (gamma_m+1, gamma_n, 3)

        # forces from each set of filaments
        f_chordwise = rho * jnp.einsum(
            "ij,ijk->ijk",
            gamma_chordwise,
            jnp.cross(v_rel_chordwise, r_chordwise),
        )  # (gamma_m, gamma_n+1, 3)
        f_spanwise = rho * jnp.einsum(
            "ij,ijk->ijk", gamma_spanwise, jnp.cross(v_rel_spanwise, r_spanwise)
        )  # (gamma_m+1, gamma_n, 3)

        f_surf = split_to_vertex(f_chordwise, 0) + split_to_vertex(
            f_spanwise, 1
        )  # (gamma_m+1, gamma_n+1, 3)

        if has_mirror:
            assert mirror_point is not None and mirror_normal is not None
            # a spanwise edge on the mirror plane is shared with the image surface, and should receive the force from
            # the image's bound vortex
            f_ghost_low = _mirror_ghost_force(
                zeta_b_surf=zeta_b_surf,
                mp_dot_spanwise=mp_dot_spanwise,
                gamma_spanwise_edge=gamma_spanwise[:, 0],
                rho=rho,
                v_func=v_func,
                mirror_point=mirror_point,
                mirror_normal=mirror_normal,
                low=True,
            )
            f_surf = f_surf.at[:, 0, :].add(jnp.where(on_plane_low, f_ghost_low, 0.0))

            f_ghost_high = _mirror_ghost_force(
                zeta_b_surf=zeta_b_surf,
                mp_dot_spanwise=mp_dot_spanwise,
                gamma_spanwise_edge=gamma_spanwise[:, -1],
                rho=rho,
                v_func=v_func,
                mirror_point=mirror_point,
                mirror_normal=mirror_normal,
                low=False,
            )
            f_surf = f_surf.at[:, -1, :].add(
                jnp.where(on_plane_high, f_ghost_high, 0.0)
            )

        f_steady.append(f_surf)
    return f_steady

strip_alpha

strip_alpha(
    zeta_b: ArrayList,
    f_steady: ArrayList,
    v_func: Callable[[Array], Array],
    rho: Array,
    beta: Array,
) -> ArrayList

Compute the per-strip effective angle of attack from the UVLM sectional forcing. For each spanwise strip, the strip total force is projected onto the local lift direction to obtain the strip lift coefficient. The angle of attack is found by using a Prandtl-Glauert-corrected lift slope of 2 * pi / beta, since f_steady already reflects any compressibility correction applied to the circulation.

Parameters:

Name	Type	Description	Default
zeta_b	ArrayList	Bound grid coordinates, (n_surf, )(zeta_m, zeta_n, 3).	required
f_steady	ArrayList	Steady forcing, (n_surf, )(zeta_m, zeta_n, 3).	required
v_func	Callable[[Array], Array]	Reference velocity as a function of position, (..., 3) -> (..., 3).	required
rho	Array	Flow density.	required
beta	Array	Prandtl-Glauert compressibility factor, ().	required

Returns:

Type	Description
ArrayList	Per-surface strip angle of attack, (n_surf, )(n_strip,).

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

def strip_alpha(
    zeta_b: ArrayList,
    f_steady: ArrayList,
    v_func: Callable[[Array], Array],
    rho: Array,
    beta: Array,
) -> ArrayList:
    r"""
    Compute the per-strip effective angle of attack from the UVLM sectional forcing. For each spanwise strip, the strip
    total force is projected onto the local lift direction to obtain the strip lift coefficient. The angle of attack is
    found by using a Prandtl-Glauert-corrected lift slope of ``2 * pi / beta``, since
    ``f_steady`` already reflects any compressibility correction applied to the circulation.

    :param zeta_b: Bound grid coordinates, ``(n_surf, )(zeta_m, zeta_n, 3)``.
    :param f_steady: Steady forcing, ``(n_surf, )(zeta_m, zeta_n, 3)``.
    :param v_func: Reference velocity as a function of position, ``(..., 3) -> (..., 3)``.
    :param rho: Flow density.
    :param beta: Prandtl-Glauert compressibility factor, ``()``.
    :return: Per-surface strip angle of attack, ``(n_surf, )(n_strip,)``.
    """
    alphas = ArrayList([])
    for zeta_surf, f_surf in zip(zeta_b, f_steady):
        # find mid-panel chord length
        zeta_le = zeta_surf[0, :, :]
        zeta_te = zeta_surf[-1, :, :]
        le_j = 0.5 * (zeta_le[:-1, :] + zeta_le[1:, :])
        te_j = 0.5 * (zeta_te[:-1, :] + zeta_te[1:, :])

        chord_vec = te_j - le_j  # (n, 3)
        c_len = jnp.linalg.norm(chord_vec, axis=-1)  # (n)

        span_vec = 0.5 * (
            (zeta_le[1:, :] + zeta_te[1:, :]) - (zeta_le[:-1, :] + zeta_te[:-1, :])
        )  # (n, 3)
        b_len = jnp.linalg.norm(span_vec, axis=-1)  # (n)
        e_s = span_vec / b_len[:, None]  # unit vector in span direction

        strip_mid = 0.5 * (le_j + te_j)  # centre of strip
        v_ref = v_func(strip_mid)  # evaluate velocity at mid-strip
        v_mag2 = jnp.sum(v_ref * v_ref, axis=-1)

        e_l = jnp.cross(v_ref, e_s)  # vector in lift direction
        e_l /= jnp.linalg.norm(e_l, axis=-1, keepdims=True)  # make unit vector

        f_strip = neighbour_average(f_surf, axes=1).sum(axis=0)
        q = 0.5 * rho * v_mag2

        # correct from 2 pi / beta lift slope to custom input
        cl_uvlm = jnp.sum(f_strip * e_l, axis=-1) / (q * c_len * b_len)
        alphas.append(cl_uvlm * beta / (2.0 * jnp.pi))
    return alphas

apply_polar_correction

apply_polar_correction(
    zeta_b: ArrayList,
    f_steady: ArrayList,
    v_func: Callable[[Array], Array],
    rho: Array,
    beta: Array,
    polar_data: Sequence[Any | None],
    polar_function: Sequence[PolarFunction | None],
    alpha: ArrayList | None = None,
) -> tuple[
    ArrayList, ArrayList, ArrayList, ArrayList, ArrayList
]

Replace UVLM strip forcing with a sectional force built from tabulated airfoil polars. This method uses the UVLM to compute the strip-wise angles of attack, before obtaining the new forces from the passed polar databases.

For each surface where a polar database is provided, and for each spanwise strip, we firstly compute the strip forcing in the global frame. This forcing can then be converted into the local strip coordinate frame. We compute the velocity at the mid-strip point, from which we can calculate the lift coefficient. By assuming a Prandtl-Glauert-corrected 2*pi/beta potential flow lift slope, we can correct the forcing from polar_function's output for lift, drag and moment. To correct the output, we distribute the forcing back onto the grid in a way which satisfies the moment and force balance.

Parameters:

Name	Type	Description	Default
zeta_b	ArrayList	Bound grid coordinates, (n_surf, )(zeta_m, zeta_n, 3).	required
f_steady	ArrayList	Steady vertex forcing, (n_surf, )(zeta_m, zeta_n, 3).	required
v_func	Callable[[Array], Array]	Reference velocity as a function of position, (..., 3) -> (..., 3).	required
rho	Array	Flow density.	required
beta	Array	Prandtl-Glauert compressibility factor, ().	required
polar_data	Sequence[Any | None]	Per-surface polar database, length n_surf. Each entry is either None (no correction for that surface) or an arbitrary data structure.	required
polar_function	Sequence[PolarFunction | None]	Per-surface function mapping (alpha, polar_data) -> (cl, cd, cm) about the quarter-chord, length n_surf, where it expects database of the same type as the corresponding entry in polar_data. Set to None for surfaces with no polar correction.	required
alpha	ArrayList | None	Optional precomputed per-strip angle of attack, (n_surf, )(n_strip,). If None, computed internally.	None

Returns:

Type	Description
tuple[ArrayList, ArrayList, ArrayList, ArrayList, ArrayList]	Corrected vertex forcing, (n_surf, )(zeta_m, zeta_n, 3); per-strip lift scale factors, (n_surf, )(n, ), which can be used to scale the circulation strengths if requested; and the per-strip lift, drag and moment coefficients sampled from the polars, (n_surf, )(n, ) each.

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

def apply_polar_correction(
    zeta_b: ArrayList,
    f_steady: ArrayList,
    v_func: Callable[[Array], Array],
    rho: Array,
    beta: Array,
    polar_data: Sequence[Any | None],
    polar_function: Sequence[PolarFunction | None],
    alpha: ArrayList | None = None,
) -> tuple[ArrayList, ArrayList, ArrayList, ArrayList, ArrayList]:
    r"""
    Replace UVLM strip forcing with a sectional force built from tabulated airfoil polars. This method uses the UVLM to
    compute the strip-wise angles of attack, before obtaining the new forces from the passed polar databases.

    For each surface where a polar database is provided, and for each spanwise strip, we firstly compute the strip
    forcing in the global frame. This forcing can then be converted into the local strip coordinate frame. We
    compute the velocity at the mid-strip point, from which we can calculate the lift coefficient. By assuming a
    Prandtl-Glauert-corrected ``2*pi/beta`` potential flow lift slope, we can correct the forcing from
    ``polar_function``'s output for lift, drag and moment. To correct the output, we distribute the forcing back onto
    the grid in a way which satisfies the moment and force balance.
    :param zeta_b: Bound grid coordinates, ``(n_surf, )(zeta_m, zeta_n, 3)``.
    :param f_steady: Steady vertex forcing, ``(n_surf, )(zeta_m, zeta_n, 3)``.
    :param v_func: Reference velocity as a function of position, ``(..., 3) -> (..., 3)``.
    :param rho: Flow density.
    :param beta: Prandtl-Glauert compressibility factor, ``()``.
    :param polar_data: Per-surface polar database, length ``n_surf``. Each entry is either ``None`` (no
        correction for that surface) or an arbitrary data structure.
    :param polar_function: Per-surface function mapping ``(alpha, polar_data) -> (cl, cd, cm)`` about the
        quarter-chord, length ``n_surf``,  where it expects database of the same type as the corresponding entry in
        ``polar_data``. Set to ``None`` for surfaces with no polar correction.
    :param alpha: Optional precomputed per-strip angle of attack, ``(n_surf, )(n_strip,)``. If
        ``None``, computed internally.
    :return: Corrected vertex forcing, ``(n_surf, )(zeta_m, zeta_n, 3)``; per-strip lift scale factors,
        ``(n_surf, )(n, )``, which can be used to scale the circulation strengths if requested; and the
        per-strip lift, drag and moment coefficients sampled from the polars, ``(n_surf, )(n, )`` each.
    """
    if alpha is None:
        # compute the angles of attack for each strip if not passed
        alpha = strip_alpha(
            zeta_b=zeta_b, f_steady=f_steady, v_func=v_func, rho=rho, beta=beta
        )

    f_out = ArrayList([])
    lift_scale_out = ArrayList([])
    cl_out = ArrayList([])
    cd_out = ArrayList([])
    cm_out = ArrayList([])
    for i_surf, (zeta_surf, f_surf, database, func, alpha_surf) in enumerate(
        zip(zeta_b, f_steady, polar_data, polar_function, alpha)
    ):
        n = zeta_surf.shape[1] - 1

        if database is None:
            # no correction to apply
            f_out.append(f_surf)
            lift_scale_out.append(jnp.ones(n))
            cl_out.append(2.0 * jnp.pi * alpha_surf / beta)
            cd_out.append(jnp.zeros(n))
            cm_out.append(jnp.zeros(n))
            continue

        zeta_m = zeta_surf.shape[0]
        m = zeta_m - 1

        # find leading and trailing edges of centre of strip
        zeta_le = zeta_surf[0, :, :]
        zeta_te = zeta_surf[-1, :, :]
        le_j = 0.5 * (zeta_le[:-1, :] + zeta_le[1:, :])
        te_j = 0.5 * (zeta_te[:-1, :] + zeta_te[1:, :])

        chord_vec = te_j - le_j
        c_len = jnp.linalg.norm(chord_vec, axis=-1)  # chord of strips, (n, ).
        e_c = chord_vec / c_len[:, None]  # unit vector in chordwise direction, (n, 3)

        span_vec = 0.5 * (
            (zeta_le[1:, :] + zeta_te[1:, :]) - (zeta_le[:-1, :] + zeta_te[:-1, :])
        )
        b_len = jnp.linalg.norm(span_vec, axis=-1)  # span of strips, (n, ).
        e_s = span_vec / b_len[:, None]  # unit vector in spanwise direction, (n, 3)

        e_n = jnp.cross(e_c, e_s)
        e_n /= jnp.linalg.norm(
            e_n, axis=-1, keepdims=True
        )  # unit vector in normal direction, (n, 3)

        # compute dynamic pressure at centroid of strip
        strip_mid = 0.5 * (le_j + te_j)
        v_ref = v_func(strip_mid)
        v_mag2 = jnp.sum(v_ref * v_ref, axis=-1)
        v_mag = jnp.sqrt(v_mag2)
        q = 0.5 * rho * v_mag2

        e_l = jnp.cross(v_ref, e_s)
        e_l /= jnp.linalg.norm(
            e_l, axis=-1, keepdims=True
        )  # unit vector in flow direction, (n, 3)
        e_d = v_ref / v_mag[:, None]

        # sample the polar database via this surface's evaluation function
        assert func is not None
        cl_p, cd_p, cm_p = func(alpha_surf, database)

        # lift scale factor cl_polar / cl_uvlm with cl_uvlm = 2 pi alpha / beta
        cl_uvlm = 2.0 * jnp.pi * alpha_surf / beta
        lift_scale = jnp.where(jnp.abs(cl_uvlm) > EPSILON, cl_p / cl_uvlm, 1.0)
        lift_scale_out.append(lift_scale)
        cl_out.append(cl_p)
        cd_out.append(cd_p)
        cm_out.append(cm_p)

        qcb = q * c_len * b_len
        f_lump = qcb[:, None] * (cl_p[:, None] * e_l + cd_p[:, None] * e_d)
        f_couple = (qcb * cm_p)[:, None] * e_n  # rotate back into normal direction

        # chordwise triangular weights placing the lumped force at c/4
        chord_fracs = jnp.linspace(0.0, 1.0, zeta_m)
        w_chord = jnp.clip(1.0 - jnp.abs(chord_fracs - 0.25) * m, 0.0, 1.0)

        f_corrected = jnp.zeros_like(f_surf)

        # split force across the two bounding spanwise vertex columns
        lump_contrib = 0.5 * w_chord[:, None, None] * f_lump[None, :, :]
        f_corrected = f_corrected.at[:, :-1, :].add(lump_contrib)
        f_corrected = f_corrected.at[:, 1:, :].add(lump_contrib)

        # correct for moment with LE/TE opposing forces
        couple_contrib = 0.5 * f_couple
        f_corrected = f_corrected.at[0, :-1, :].add(couple_contrib)
        f_corrected = f_corrected.at[0, 1:, :].add(couple_contrib)
        f_corrected = f_corrected.at[-1, :-1, :].add(-couple_contrib)
        f_corrected = f_corrected.at[-1, 1:, :].add(-couple_contrib)

        f_out.append(f_corrected)

    return f_out, lift_scale_out, cl_out, cd_out, cm_out

propagate_surf_wake

propagate_surf_wake(
    gamma_b_nm1: Array,
    gamma_w_nm1: Array,
    zeta_b_n: Array,
    zeta_w_nm1: Array,
    delta_w: Array | None,
    v_func: Callable[[Array], Array],
    dt: Array,
    frozen_wake: bool,
    linearise_variable_wake: bool = False,
) -> tuple[Array | None, Array]

Convect the wake at some given velocity for a single surface from timestep n-1 to timestep n. This step includes convection from the trailing edge and culling the downstream data.

Parameters:

Name	Type	Description	Default
gamma_b_nm1	Array	Bound circulation at time step n-1, (m, n).	required
gamma_w_nm1	Array	Wake circulation at time step n-1, (m_star, n).	required
zeta_b_n	Array	Bound grid at time step n, (zeta_m, zeta_n, 3).	required
zeta_w_nm1	Array	Wake grid at time step n-1, (zeta_m_star, zeta_n, 3).	required
delta_w	Array | None	Desired wake discretisation, (zeta_m_star, 3), or None for uniform.	required
v_func	Callable[[Array], Array]	Function that computes the velocity as a function of coordinate, (3, ) -> (3, ).	required
dt	Array	Time step length.	required
frozen_wake	bool	If true, the grid stays constant with time. Used in the linearised case.	required
linearise_variable_wake	bool	If true, block gradients through the arc-length computation so that the re-discretisation is treated as a linear operator when differentiated.	False

Returns:

Type	Description
tuple[Array | None, Array]	New wake grid and circulation, (zeta_m_star, zeta_n, 3), (m_star, n).

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

def propagate_surf_wake(
    gamma_b_nm1: Array,
    gamma_w_nm1: Array,
    zeta_b_n: Array,
    zeta_w_nm1: Array,
    delta_w: Array | None,
    v_func: Callable[[Array], Array],
    dt: Array,
    frozen_wake: bool,
    linearise_variable_wake: bool = False,
) -> tuple[Array | None, Array]:
    r"""
    Convect the wake at some given velocity for a single surface from timestep n-1 to timestep n. This step includes
    convection from the trailing edge and culling the downstream data.
    :param gamma_b_nm1: Bound circulation at time step n-1, ``(m, n)``.
    :param gamma_w_nm1: Wake circulation at time step n-1, ``(m_star, n)``.
    :param zeta_b_n: Bound grid at time step n, ``(zeta_m, zeta_n, 3)``.
    :param zeta_w_nm1: Wake grid at time step n-1, ``(zeta_m_star, zeta_n, 3)``.
    :param delta_w: Desired wake discretisation, ``(zeta_m_star, 3)``, or None for uniform.
    :param v_func: Function that computes the velocity as a function of coordinate, ``(3, )`` -> ``(3, )``.
    :param dt: Time step length.
    :param frozen_wake: If true, the grid stays constant with time. Used in the linearised case.
    :param linearise_variable_wake: If true, block gradients through the arc-length computation so that
        the re-discretisation is treated as a linear operator when differentiated.
    :return: New wake grid and circulation, ``(zeta_m_star, zeta_n, 3)``, ``(m_star, n)``.
    """

    # trailing edge positions and circulations
    zeta_te = zeta_b_n[-1, ...]  # (zeta_n, 3)
    gamma_te = gamma_b_nm1[-1, ...]  # (gamma_n)

    # variable wake discretisation also depends on the final element
    if delta_w is not None:
        zeta_base = zeta_w_nm1  # (zeta_w_m, zeta_n, 3)
        gamma_base = gamma_w_nm1  # (gamma_w_m, gamma_n)
    else:
        zeta_base = zeta_w_nm1[:-1, ...]  # (zeta_w_m - 1, zeta_n, 3)
        gamma_base = gamma_w_nm1[:-1, ...]  # (gamma_w_m - 1, gamma_n)

    # values at t=n+1 before re-discretisation
    gamma_w_n = jnp.concatenate(
        (gamma_te[None, ...], gamma_base), axis=0
    )  # (gamma_w_m+1 | gamma_w_m, gamma_n)

    # if the wake is free, this should be embedded here
    v = v_func(zeta_base)  # (zeta_w_m | zeta_w_m-1, zeta_n, 3)

    # wake coordinates at t=n+1 before re-discretisation
    zeta_w_n = jnp.concatenate(
        (zeta_te[None, :, :], zeta_base + dt * v), axis=0
    )  # (zeta_w_m+1 | zeta_w_m, zeta_n, 3)

    if delta_w is not None:
        # streamline coordinates before re-discretisation
        s_zeta_w = jnp.concatenate(
            (
                jnp.zeros((1, zeta_te.shape[0])),  # (1, zeta_n)
                jnp.cumsum(
                    jnp.linalg.norm(
                        zeta_w_n[1:, ...] - zeta_w_n[:-1, ...], axis=-1
                    ),  # (zeta_w_m+1, zeta_n)
                    axis=0,
                ),  # (zeta_w_m, zeta_n)
            ),
            axis=0,
        )  # distance along each wake filament for each point (zeta_w_m + 1, zeta_n]

        if linearise_variable_wake:
            s_zeta_w = jax.lax.stop_gradient(s_zeta_w)

        # consider gamma to be at midpoints of zeta
        s_gamma_w = neighbour_average(s_zeta_w, axes=(0, 1))

        # vertex coordinates along desired discretised streamline, (m_star + 1)
        s_zeta_w_discretisation = jnp.concatenate((jnp.zeros(1), jnp.cumsum(delta_w)))

        # midpoint coordinates along desired discretised streamline, (m_star)
        s_gamma_w_discretisation = neighbour_average(s_zeta_w_discretisation, axes=(0,))

        # re-discretise coordinates onto desired grid
        zeta_w_n = vmap(
            vmap(jnp.interp, in_axes=(None, 0, 0), out_axes=1),
            in_axes=(None, None, 1),
            out_axes=2,
        )(
            s_zeta_w_discretisation, s_zeta_w.T, jnp.transpose(zeta_w_n, (1, 2, 0))
        )  # (zeta_w_m, zeta_n, 3)

        # re-discretise gamma onto desired grid
        gamma_w_n = vmap(jnp.interp, in_axes=(None, 0, 0), out_axes=1)(
            s_gamma_w_discretisation, s_gamma_w.T, gamma_w_n.T
        )  # (zeta_w_m, zeta_n, 3)

    # logic for edge case where there is no wake
    if gamma_w_nm1.size == 0:
        gamma_w_n = jnp.zeros_like(gamma_w_nm1)

    if frozen_wake:
        return None, gamma_w_n
    else:
        return zeta_w_n, gamma_w_n

propagate_wake

propagate_wake(
    gamma_b_nm1: ArrayList,
    gamma_w_nm1: ArrayList,
    zeta_b_n: ArrayList,
    zeta_w_nm1: ArrayList,
    delta_w: Sequence[Array | None],
    v_func: Callable[[Array], Array],
    dt: Array,
    frozen_wake: Literal[True],
    linearise_variable_wake: bool,
) -> tuple[None, ArrayList]

propagate_wake(
    gamma_b_nm1: ArrayList,
    gamma_w_nm1: ArrayList,
    zeta_b_n: ArrayList,
    zeta_w_nm1: ArrayList,
    delta_w: Sequence[Array | None],
    v_func: Callable[[Array], Array],
    dt: Array,
    frozen_wake: Literal[False],
    linearise_variable_wake: bool,
) -> tuple[ArrayList, ArrayList]

propagate_wake(
    gamma_b_nm1: ArrayList,
    gamma_w_nm1: ArrayList,
    zeta_b_n: ArrayList,
    zeta_w_nm1: ArrayList,
    delta_w: Sequence[Array | None],
    v_func: Callable[[Array], Array],
    dt: Array,
    frozen_wake: bool,
    linearise_variable_wake: bool = False,
) -> tuple[ArrayList | None, ArrayList]

Convect the wake for all surfaces.

Parameters:

Name	Type	Description	Default
gamma_b_nm1	ArrayList	Bound circulation at time step n-1, (n_surf, )(m, n).	required
gamma_w_nm1	ArrayList	Wake circulation at time step n-1, (n_surf, )(m_star, n).	required
zeta_b_n	ArrayList	Bound grid at time step n, (n_surf, )(zeta_m, zeta_n, 3).	required
zeta_w_nm1	ArrayList	Wake grid at time step n-1, (n_surf, )(zeta_m_star, zeta_n, 3).	required
delta_w	Sequence[Array | None]	Desired wake discretisation, (n_surf, )(zeta_m_star, 3) or None for uniform.	required
v_func	Callable[[Array], Array]	Function that computes the velocity, (3, ) -> (3, ).	required
dt	Array	Time step length.	required
frozen_wake	bool	If true, the grid stays constant with time, useful in the linearised case.	required
linearise_variable_wake	bool	If true, block gradients through the arc-length computation so that the re-discretisation is treated as a linear operator when differentiated with jax.jvp.	False

Returns:

Type	Description
tuple[ArrayList | None, ArrayList]	New wake grid and circulation, (n_surf, )(zeta_m_star, zeta_n, 3), (n_surf, )(m_star, n).

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

def propagate_wake(
    gamma_b_nm1: ArrayList,
    gamma_w_nm1: ArrayList,
    zeta_b_n: ArrayList,
    zeta_w_nm1: ArrayList,
    delta_w: Sequence[Array | None],
    v_func: Callable[[Array], Array],
    dt: Array,
    frozen_wake: bool,
    linearise_variable_wake: bool = False,
) -> tuple[ArrayList | None, ArrayList]:
    r"""
    Convect the wake for all surfaces.
    :param gamma_b_nm1: Bound circulation at time step n-1, ``(n_surf, )(m, n)``.
    :param gamma_w_nm1: Wake circulation at time step n-1, ``(n_surf, )(m_star, n)``.
    :param zeta_b_n: Bound grid at time step n, ``(n_surf, )(zeta_m, zeta_n, 3)``.
    :param zeta_w_nm1: Wake grid at time step n-1, ``(n_surf, )(zeta_m_star, zeta_n, 3)``.
    :param delta_w: Desired wake discretisation, ``(n_surf, )(zeta_m_star, 3)`` or None for uniform.
    :param v_func: Function that computes the velocity, ``(3, )`` -> ``(3, )``.
    :param dt: Time step length.
    :param frozen_wake: If true, the grid stays constant with time, useful in the linearised case.
    :param linearise_variable_wake: If true, block gradients through the arc-length computation so that
        the re-discretisation is treated as a linear operator when differentiated with jax.jvp.
    :return: New wake grid and circulation, ``(n_surf, )(zeta_m_star, zeta_n, 3)``, ``(n_surf, )(m_star, n)``.
    """

    n_surf = len(gamma_b_nm1)
    zeta_w_n: ArrayList | None = ArrayList([]) if not frozen_wake else None
    gamma_w_n = ArrayList([])

    for i_surf in range(n_surf):
        surf_zeta_w, surf_gamma_w = propagate_surf_wake(
            gamma_b_nm1=gamma_b_nm1[i_surf],
            gamma_w_nm1=gamma_w_nm1[i_surf],
            zeta_b_n=zeta_b_n[i_surf],
            zeta_w_nm1=zeta_w_nm1[i_surf],
            delta_w=delta_w[i_surf],
            v_func=v_func,
            dt=dt,
            frozen_wake=frozen_wake,
            linearise_variable_wake=linearise_variable_wake,
        )
        if zeta_w_n is not None:
            assert surf_zeta_w is not None
            zeta_w_n.append(surf_zeta_w)
        gamma_w_n.append(surf_gamma_w)
    return zeta_w_n, gamma_w_n

biot_savart

biot_savart(x: Array, y: Array) -> Array

Biot-Savart kernel without any smoothing or cutoff.

Parameters:

Name	Type	Description	Default
x	Array	Target point, (3, ).	required
y	Array	Filament endpoints, (2, 3).	required

Returns:

Type	Description
Array	Influence at target point, (3, ).

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

def biot_savart(x: Array, y: Array) -> Array:
    r"""
    Biot-Savart kernel without any smoothing or cutoff.
    :param x: Target point, ``(3, )``.
    :param y: Filament endpoints, ``(2, 3)``.
    :return: Influence at target point, ``(3, )``.
    """
    r0 = y[1, :] - y[0, :]
    r1 = x - y[0, :]
    r2 = x - y[1, :]
    r1_x_r2 = jnp.cross(r1, r2)
    diff_r = r1 / jnp.linalg.norm(r1) - r2 / jnp.linalg.norm(r2)
    return r1_x_r2 / (jnp.inner(r1_x_r2, r1_x_r2) * 4.0 * jnp.pi) * jnp.dot(r0, diff_r)

make_unit_epsilon

make_unit_epsilon(r: Array) -> Array

Differentiable function to obtain a unit vector that is defined for all r. As r -> 0, the output approaches zero instead of being undefined. Autodiff of this form matches a custom JVP to floating-point noise but transposes to a much cheaper VJP under reverse-mode.

Parameters:

Name	Type	Description	Default
r	Array	Vector to be normalised, (3, ).	required

Returns:

Type	Description
Array	Unit vector, (3, ).

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

def make_unit_epsilon(r: Array) -> Array:
    r"""
    Differentiable function to obtain a unit vector that is defined for all ``r``. As ``r`` -> 0, the output approaches
    zero instead of being undefined. Autodiff of this form matches a custom JVP to floating-point noise but transposes
    to a much cheaper VJP under reverse-mode.
    :param r: Vector to be normalised, ``(3, )``.
    :return: Unit vector, ``(3, )``.
    """
    return r / jnp.sqrt(jnp.sum(r**2) + EPSILON**2)

biot_savart_epsilon

biot_savart_epsilon(x: Array, y: Array) -> Array

Biot-Savart kernel with epsilon term added to remove singularity.

Parameters:

Name	Type	Description	Default
x	Array	Target point, (3, ).	required
y	Array	Filament endpoints, (2, 3).	required

Returns:

Type	Description
Array	Influence at target point, (3, ).

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

def biot_savart_epsilon(x: Array, y: Array) -> Array:
    r"""
    Biot-Savart kernel with epsilon term added to remove singularity.

    :param x: Target point, ``(3, )``.
    :param y: Filament endpoints, ``(2, 3)``.
    :return: Influence at target point, ``(3, )``.
    """
    r0 = y[1, :] - y[0, :]
    r1 = x - y[0, :]
    r2 = x - y[1, :]
    inv_n1 = jax.lax.rsqrt(jnp.sum(r1 * r1) + EPSILON * EPSILON)
    inv_n2 = jax.lax.rsqrt(jnp.sum(r2 * r2) + EPSILON * EPSILON)
    diff_r = r1 * inv_n1 - r2 * inv_n2
    r1_x_r2 = jnp.cross(r1, r2)
    r0_sq = jnp.sum(r0 * r0)
    denom = jnp.sum(r1_x_r2 * r1_x_r2) + EPSILON * r0_sq * r0_sq
    return r1_x_r2 * (jnp.dot(r0, diff_r) / (4.0 * jnp.pi * denom))

biot_savart_cutoff

biot_savart_cutoff(x: Array, y: Array) -> Array

Biot-Savart kernel with truncation radius to remove singularity.

Parameters:

Name	Type	Description	Default
x	Array	Target point, (3, ).	required
y	Array	Filament endpoints, (2, 3).	required

Returns:

Type	Description
Array	Influence at target point, (3, ).

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

def biot_savart_cutoff(x: Array, y: Array) -> Array:
    r"""
    Biot-Savart kernel with truncation radius to remove singularity.
    :param x: Target point, ``(3, )``.
    :param y: Filament endpoints, ``(2, 3)``.
    :return: Influence at target point, ``(3, )``.
    """
    r0 = y[1, :] - y[0, :]
    r1 = x - y[0, :]
    r2 = x - y[1, :]

    sm = jnp.inner(r0, r1) / jnp.inner(r0, y[1, :] - y[0, :])
    m = y[0, :] + sm * (y[1, :] - y[0, :])
    r = jnp.linalg.norm(x - m)  # radial distance

    def _kernel_value() -> Array:
        # Compute the standard Biot-Savart kernel, called only if r > R_CUTOFF
        r1_x_r2 = jnp.cross(r1, r2)
        r1_x_r2_unit2 = r1_x_r2 / (jnp.inner(r1_x_r2, r1_x_r2))
        diff_r = make_unit_epsilon(r1) - make_unit_epsilon(r2)
        return r1_x_r2_unit2 / (4.0 * jnp.pi) * jnp.dot(r0, diff_r)

    return cond((r > R_CUTOFF), _kernel_value, lambda: jnp.zeros(3))

mirror_grid

mirror_grid(
    zeta: Array, mirror_point: Array, mirror_normal: Array
) -> Array

Mirror a grid of points across a plane defined by a point and a normal vector.

Parameters:

Name	Type	Description	Default
zeta	Array	Grid of points, (zeta_m, zeta_n, 3).	required
mirror_point	Array	Point in mirror plane, (3, ).	required
mirror_normal	Array	Normal vector of mirror plane, (3, ). Should be normalised.	required

Returns:

Type	Description
Array	Mirrored grid of points, (zeta_m, zeta_n, 3).

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

def mirror_grid(zeta: Array, mirror_point: Array, mirror_normal: Array) -> Array:
    """
    Mirror a grid of points across a plane defined by a point and a normal vector.
    :param zeta: Grid of points, ``(zeta_m, zeta_n, 3)``.
    :param mirror_point: Point in mirror plane, ``(3, )``.
    :param mirror_normal: Normal vector of mirror plane, ``(3, )``. Should be normalised.
    :return: Mirrored grid of points, ``(zeta_m, zeta_n, 3)``.
    """
    diff = zeta - mirror_point[None, None, :]  # (zeta_m, zeta_n, 3)
    diff_n = jnp.einsum("ijk,k->ij", diff, mirror_normal)  # (zeta_m, zeta_n)
    return (
        zeta - 2.0 * diff_n[:, :, None] * mirror_normal[None, None, :]
    )  # (zeta_m, zeta_n, 3)

prandtl_glauert_transform

prandtl_glauert_transform(
    zeta: Array, u_inf_dir: Array, beta: Array
) -> Array

Apply the Prandtl-Glauert transform, used to map compressible-flow geometry onto its incompressible-equivalent for solving the ordinary (incompressible) UVLM equations. Components parallel to the freestream direction are unchanged; components perpendicular to the freestream are scaled by beta = sqrt(1 - M^2).

Parameters:

Name	Type	Description	Default
zeta	Array	Points or velocities to transform, (..., 3).	required
u_inf_dir	Array	Unit freestream direction, (3, ).	required
beta	Array	Prandtl-Glauert compressibility factor, ().	required

Returns:

Type	Description
Array	Transformed points or velocities, (..., 3).

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

def prandtl_glauert_transform(zeta: Array, u_inf_dir: Array, beta: Array) -> Array:
    r"""
    Apply the Prandtl-Glauert transform, used to map compressible-flow geometry onto its
    incompressible-equivalent for solving the ordinary (incompressible) UVLM equations. Components parallel to
    the freestream direction are unchanged; components perpendicular to the freestream are scaled by
    ``beta = sqrt(1 - M^2)``.
    :param zeta: Points or velocities to transform, ``(..., 3)``.
    :param u_inf_dir: Unit freestream direction, ``(3, )``.
    :param beta: Prandtl-Glauert compressibility factor, ``()``.
    :return: Transformed points or velocities, ``(..., 3)``.
    """
    r_par = jnp.einsum("...k,k->...", zeta, u_inf_dir)[..., None] * u_inf_dir
    return zeta + (beta - 1.0) * (zeta - r_par)

project_forcing_to_beam

project_forcing_to_beam(
    f_total: ArrayList,
    rmat: Array,
    dof_mapping: ArrayList,
    x0_aero: ArrayList,
    mirror_edge_low: ArrayList | None = None,
    mirror_edge_high: ArrayList | None = None,
) -> Array

Project aerodynamic forcing at specified time step onto the beam grid. Returned forces are in the global frame.

Parameters:

Name	Type	Description	Default
f_total	ArrayList	Total force on aerodynamic grid, (n_surf, )(m+1, n+1, 3)	required
rmat	Array	Rotation matrix for each node relative to reference, (n_nodes, 3, 3).	required
x0_aero	ArrayList	Reference coordinates for aerodynamic grid, (n_surf, )(zeta_m, zeta_n, 3).	required
dof_mapping	ArrayList	Mapping between aero and beam discretisations.	required
mirror_edge_low	ArrayList | None	Per-surface booleans marking whether that surface's n=0 edge lies on a mirror plane, (n_surf, )(). Where True, that vertex column's force is halved before being projected onto the beam.	None
mirror_edge_high	ArrayList | None	As mirror_edge_low, for the n=-1 edge.	None

Returns:

Type	Description
Array	Steady and unsteady forcing projected onto the beam grid, (n_nodes, 6)

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

def project_forcing_to_beam(
    f_total: ArrayList,
    rmat: Array,
    dof_mapping: ArrayList,
    x0_aero: ArrayList,
    mirror_edge_low: ArrayList | None = None,
    mirror_edge_high: ArrayList | None = None,
) -> Array:
    r"""
    Project aerodynamic forcing at specified time step onto the beam grid. Returned forces are in the global frame.
    :param f_total: Total force on aerodynamic grid, ``(n_surf, )(m+1, n+1, 3)``
    :param rmat: Rotation matrix for each node relative to reference, ``(n_nodes, 3, 3)``.
    :param x0_aero: Reference coordinates for aerodynamic grid, ``(n_surf, )(zeta_m, zeta_n, 3)``.
    :param dof_mapping: Mapping between aero and beam discretisations.
    :param mirror_edge_low: Per-surface booleans marking whether that surface's ``n=0`` edge lies on a mirror
    plane, ``(n_surf, )()``. Where True, that vertex column's force is halved before being projected onto the beam.
    :param mirror_edge_high: As ``mirror_edge_low``, for the ``n=-1`` edge.
    :return: Steady and unsteady forcing projected onto the beam grid, ``(n_nodes, 6)``
    """

    n_nodes = rmat.shape[0]
    result = jnp.zeros((n_nodes, 6))

    mirror_edge_low_ = (
        mirror_edge_low if mirror_edge_low is not None else [None] * len(f_total)
    )
    mirror_edge_high_ = (
        mirror_edge_high if mirror_edge_high is not None else [None] * len(f_total)
    )

    for i_surf, (f_surf, on_plane_low, on_plane_high) in enumerate(
        zip(f_total, mirror_edge_low_, mirror_edge_high_)
    ):
        if on_plane_low is not None:
            f_surf = f_surf.at[:, 0, :].multiply(jnp.where(on_plane_low, 0.5, 1.0))
        if on_plane_high is not None:
            f_surf = f_surf.at[:, -1, :].multiply(jnp.where(on_plane_high, 0.5, 1.0))

        # rotate relative distances to get moment arms
        this_rmat = rmat[dof_mapping[i_surf], ...]  # (zeta_n, 3, 3)
        r_x0 = jnp.einsum(
            "ijk,lik->lij", this_rmat, x0_aero[i_surf]
        )  # relative distance (zeta_n, zeta_m, 3)

        result = result.at[dof_mapping[i_surf], :3].add(
            f_surf.sum(axis=0)
        )  # forcing is sum along strip (zeta_n, 3)
        result = result.at[dof_mapping[i_surf], 3:].add(
            jnp.cross(r_x0, f_surf).sum(axis=0)
        )  # moment is cross(r, f) summed along strip (zeta_n, 3)
    return result

cs_ang_to_cs_vel

cs_ang_to_cs_vel(
    cs_ang_t: dict[str, Array], dt: float | Array
) -> dict[str, Array]

Approximate control surfaces velocities from the time series of their angles using finite differences.

Parameters:

Name	Type	Description	Default
cs_ang_t	dict[str, Array]	Time history of control surface angle, {name, (n_tstep, )}.	required
dt	float | Array	Time step length.	required

Returns:

Type	Description
dict[str, Array]	Control surface velocity, {name, (n_tstep, )}.

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

def cs_ang_to_cs_vel(cs_ang_t: dict[str, Array], dt: float | Array) -> dict[str, Array]:
    r"""
    Approximate control surfaces velocities from the time series of their angles using finite differences.
    :param cs_ang_t: Time history of control surface angle, ``{name, (n_tstep, )}``.
    :param dt: Time step length.
    :return: Control surface velocity, ``{name, (n_tstep, )}``.
    """
    cs_vel_t = {}
    for k, v in cs_ang_t.items():
        n_tstep = v.shape[0]
        cs_vel_t[k] = vmap(
            lambda i_ts: finite_difference(
                i_=i_ts,
                data=v,  # noqa: B023
                delta=jnp.array(dt),
                axis=0,
            ),
            in_axes=0,
            out_axes=0,
        )(jnp.arange(n_tstep))
    return cs_vel_t

cs_vel_to_cs_ang

cs_vel_to_cs_ang(
    cs_vel_t: dict[str, Array], dt: float | Array
) -> dict[str, Array]

Approximate control surfaces angles from the time series of their velocities using finite differences.

Parameters:

Name	Type	Description	Default
cs_vel_t	dict[str, Array]	Time history of control surface velocity, {name, (n_tstep, )}.	required
dt	float | Array	Time step length.	required

Returns:

Type	Description
dict[str, Array]	Control surface angle, {name, (n_tstep, )}.

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

def cs_vel_to_cs_ang(cs_vel_t: dict[str, Array], dt: float | Array) -> dict[str, Array]:
    r"""
    Approximate control surfaces angles from the time series of their velocities using finite differences.
    :param cs_vel_t: Time history of control surface velocity, ``{name, (n_tstep, )}``.
    :param dt: Time step length.
    :return: Control surface angle, ``{name, (n_tstep, )}``.
    """
    cs_ang_t = {}
    for k, v in cs_vel_t.items():
        cs_ang_t[k] = jnp.cumsum(v, axis=0) * dt
    return cs_ang_t

uvlm

UVLM

UVLM(
    grid_shapes: Sequence[
        GridDiscretisation | tuple[int, int, int]
    ],
    dof_mapping: ArrayList | Sequence[Array] | Array,
    variable_wake_disc: bool = False,
    mirror_point: Array | None = None,
    mirror_normal: Array | None = None,
    kernel: KernelFunction | None = None,
    grid_func: AeroGridFunction | None = None,
    free_wake: bool = False,
    gamma_dot_relaxation: float | Array = 0.7,
    include_unsteady_force: bool = True,
    batch_size: int | None = 64,
    polar_data: Sequence[Any | None] | None = None,
    polar_function: Sequence[PolarFunction | None]
    | None = None,
    polar_circulation_scale: float = 0.0,
)

Class to define an unsteady vortex lattice method aerodynamic case with arbitrary number of aerodynamic surfaces.

Initialise UVLM class with all non-design parameters.

Parameters:

Name	Type	Description	Default
grid_shapes	Sequence[GridDiscretisation | tuple[int, int, int]]	Discretisations for the number of chordwise, spanwise and wake-wise panels for each surface. May be passed as a sequence of either the GridDiscretisation class or a tuple of integers ordered as (m, n, m_star).	required
dof_mapping	ArrayList | Sequence[Array] | Array	Mapping from aerodynamic grid points to structure grid points for each surface.	required
variable_wake_disc	bool	If True, allow for variable wake discretisations.	False
mirror_point	Array | None	Optional point in mirror plane, (3, ). If provided, this will apply mirroring of the aerodynamic geometry and flow about the plane defined by this point and the mirror normal.	None
mirror_normal	Array | None	Optional normal vector for mirror plane, (3, ).	None
kernel	KernelFunction | None	Input for custom kernel function to use for induced velocity calculations.	None
grid_func	AeroGridFunction | None	Input functions used for defining surfaces with control surfaces. This function should take the reference local grid coordinates zeta_b0, as well as control inputs as keyword arguments, and return the deflected local grid coordinates as an ArrayList of equal dimensionality to zeta_b0.	None
free_wake	bool	If True, include the velocity induced from the aerodynamic elements for wake propagation.	False
gamma_dot_relaxation	float | Array	Filtering parameter used for obtaining the time derivative of the circulation strengths.	0.7
include_unsteady_force	bool	If True, include forces due to apparent mass for simulation.	True
batch_size	int | None	Batch size for vectorising AIC computations. Larger values may result in faster computations, at the expense of increased memory usage. Setting to None is equivelant to a vmap.	64
polar_data	Sequence[Any | None] | None	Optional per-surface polar database used to correct the UVLM sectional forcing, which can be an arbitrary data type.	None
polar_function	Sequence[PolarFunction | None] | None	Per-surface function mapping (alpha, database) -> (cl, cd, cm) about the quarter-chord. Entries of None in the sequence indicate no polar correction for that surface, or set the input value to None to apply no polar correction for any surface.	None
polar_circulation_scale	float	Factor in [0, 1] controlling how much of the per-strip lift correction factor cl_polar / cl_uvlm is applied to the bound circulation before it is stored and convected into the wake. A valud of 0 does not correct the circulation, whereas 1 fully rescales it so the shed vortex strength matches the polar-corrected lift.	0.0

Source code in src/flapjax/aero/uvlm.py

def __init__(
    self,
    grid_shapes: Sequence[GridDiscretisation | tuple[int, int, int]],
    dof_mapping: ArrayList | Sequence[Array] | Array,
    variable_wake_disc: bool = False,
    mirror_point: Array | None = None,
    mirror_normal: Array | None = None,
    kernel: KernelFunction | None = None,
    grid_func: AeroGridFunction | None = None,
    free_wake: bool = False,
    gamma_dot_relaxation: float | Array = 0.7,
    include_unsteady_force: bool = True,
    batch_size: int | None = 64,
    polar_data: Sequence[Any | None] | None = None,
    polar_function: Sequence[PolarFunction | None] | None = None,
    polar_circulation_scale: float = 0.0,
) -> None:
    r"""
    Initialise UVLM class with all non-design parameters.
    :param grid_shapes: Discretisations for the number of chordwise, spanwise and wake-wise panels for each surface.
    May be passed as a sequence of either the GridDiscretisation class or a tuple of integers ordered as ``(m, n, m_star)``.
    :param dof_mapping: Mapping from aerodynamic grid points to structure grid points for each surface.
    :param variable_wake_disc: If True, allow for variable wake discretisations.
    :param mirror_point: Optional point in mirror plane, ``(3, )``. If provided, this will apply mirroring of the aerodynamic
    geometry and flow about the plane defined by this point and the mirror normal.
    :param mirror_normal: Optional normal vector for mirror plane, ``(3, )``.
    :param kernel: Input for custom kernel function to use for induced velocity calculations.
    :param grid_func: Input functions used for defining surfaces with control surfaces. This function should take
    the reference local grid coordinates ``zeta_b0``, as well as control inputs as keyword arguments, and return the
    deflected local grid coordinates as an ArrayList of equal dimensionality to ``zeta_b0``.
    :param free_wake: If True, include the velocity induced from the aerodynamic elements for wake propagation.
    :param gamma_dot_relaxation: Filtering parameter used for obtaining the time derivative of the circulation
    strengths.
    :param include_unsteady_force: If True, include forces due to apparent mass for simulation.
    :param batch_size: Batch size for vectorising AIC computations. Larger values may result in faster computations,
    at the expense of increased memory usage. Setting to None is equivelant to a vmap.
    :param polar_data: Optional per-surface polar database used to correct the UVLM sectional forcing, which can
    be an arbitrary data type.
    :param polar_function: Per-surface function mapping ``(alpha, database) -> (cl, cd, cm)`` about the
    quarter-chord. Entries of ``None`` in the sequence indicate no polar correction for that surface, or set the
    input value to ``None`` to apply no polar correction for any surface.
    :param polar_circulation_scale: Factor in ``[0, 1]`` controlling how much of the per-strip lift
    correction factor ``cl_polar / cl_uvlm`` is applied to the bound circulation before it is stored and convected
    into the wake. A valud of 0 does not correct the circulation, whereas 1 fully rescales it so the shed vortex
    strength matches the polar-corrected lift.
    """

    # case for single inputs
    if isinstance(dof_mapping, Array):
        dof_mapping_arrlist: ArrayList = ArrayList([dof_mapping])
    elif isinstance(dof_mapping, Sequence):
        dof_mapping_arrlist = ArrayList(dof_mapping)
    elif isinstance(dof_mapping, ArrayList):
        dof_mapping_arrlist = dof_mapping
    else:
        raise TypeError("Invalid dof mapping type")
    self.dof_mapping: ArrayList = dof_mapping_arrlist

    # number of aerodynamic surfaces
    self.n_surf: int = len(grid_shapes)

    # set grid discretisations parameters for number of panels
    grid_disc = []

    for grid in grid_shapes:
        if isinstance(grid, Sequence):
            if len(grid) != 3:
                raise ValueError(
                    "Grid shape tuple must have exactly three elements (m, n, m_star)"
                )
            grid_disc.append(GridDiscretisation(*grid))
        elif isinstance(grid, GridDiscretisation):
            grid_disc.append(grid)
        else:
            raise TypeError(
                "Grid shape must be either a Sequence of three integers or a GridDiscretisation instance"
            )
    self.grid_disc: tuple[GridDiscretisation] = tuple(grid_disc)

    # count of number of panels
    self.n_bound_panels: tuple[int, ...] = tuple(
        [gd.m * gd.n for gd in self.grid_disc]
    )
    self.n_wake_panels: tuple[int, ...] = tuple(
        [gd.m_star * gd.n for gd in self.grid_disc]
    )
    self.n_panels_tot: int = sum(self.n_bound_panels) + sum(self.n_wake_panels)

    # placeholder for aerodynamic local grid coordinates, and global coordinates for wing and wake
    self.hg_ref = None
    self.zeta_b0 = None
    self.zeta_b_ref = None
    self.zeta_w_ref = None

    # placeholder for which surfaces have spanwise edges that lie on the mirror plane
    # needed for force corrections
    self.mirror_edge_low: ArrayList | None = None
    self.mirror_edge_high: ArrayList | None = None

    self.gamma_b_slice, self.gamma_w_slice = self._make_gamma_slices()

    # store DOF mapping
    if len(self.dof_mapping) != self.n_surf:
        raise ValueError(
            f"Expected {self.n_surf} DOF mapping arrays, got {len(self.dof_mapping)}"
        )
    for i_surf, map_ in enumerate(self.dof_mapping):
        check_arr_dtype(map_, int, "dof_mapping")
        check_arr_shape(map_, (self.grid_disc[i_surf].n + 1,), "grid_disc")

    # this must be optional as it is set as a design variable later
    self.flowfield = None

    # time step length
    self._dt: Array | None = None

    # wake discretisation parameters
    self.variable_wake_disc: bool = variable_wake_disc
    self.delta_w = None

    # kernel definitions per surface (separate for wing and wake)
    self.kernels_b: Sequence[KernelFunction] = self.n_surf * [
        kernel if kernel is not None else biot_savart_epsilon
    ]
    self.kernels_w: Sequence[KernelFunction] = self.n_surf * [
        kernel if kernel is not None else biot_savart_epsilon
    ]

    # settings for solvers
    self.free_wake: bool = free_wake
    self.gamma_dot_relaxation: float | Array = gamma_dot_relaxation
    self.include_unsteady_force: bool = include_unsteady_force
    self.batch_size: int | None = batch_size

    # optional per-surface polar database and evaluation function for sectional forcing correction
    polar_data_: list[Any | None] = (
        list(polar_data) if polar_data is not None else [None] * self.n_surf
    )
    if len(polar_data_) != self.n_surf:
        raise ValueError(
            f"Expected {self.n_surf} polar_data entries, got {len(polar_data_)}"
        )

    polar_function_: list[PolarFunction | None] = (
        list(polar_function) if polar_function is not None else [None] * self.n_surf
    )
    if len(polar_function_) != self.n_surf:
        raise ValueError(
            f"Expected {self.n_surf} polar_function entries, got {len(polar_function_)}"
        )

    for i_surf, (database, func) in enumerate(zip(polar_data_, polar_function_)):
        if (database is None) != (func is None):
            raise ValueError(
                f"Surface {i_surf}: polar_data and polar_function must either both be None or both be set"
            )

    self.polar_data: tuple[Any | None, ...] = tuple(polar_data_)
    self.polar_function: tuple[PolarFunction | None, ...] = tuple(polar_function_)

    if not 0.0 <= polar_circulation_scale <= 1.0:
        raise ValueError(
            f"polar_circulation_scale must be in [0, 1], got {polar_circulation_scale}."
        )
    self.polar_circulation_scale: float = float(polar_circulation_scale)

    # mirror definitions
    if (mirror_point is None) != (mirror_normal is None):
        raise ValueError(
            "Both mirror_point and mirror_normal must be provided to apply mirroring, or both must be None to "
            "apply no mirroring."
        )

    if mirror_point is None or mirror_normal is None:
        self.mirror_point: Array | None = None
        self.mirror_normal: Array | None = None
    else:
        self.mirror_point = mirror_point
        self.mirror_normal = mirror_normal / jnp.linalg.norm(
            mirror_normal
        )  # normalise

    # surface names used for plotting
    self.surf_b_names: list[str] = [f"surf_{i}_bound" for i in range(self.n_surf)]
    self.surf_w_names: list[str] = [f"surf_{i}_wake" for i in range(self.n_surf)]

    # control surface variables
    self.grid_func: AeroGridFunction = (
        grid_func if grid_func is not None else _identity_grid_func
    )

    self.cs_ang0: dict[str, Array] = {}
    self.cs_vel0: dict[str, Array] = {}

dt property writable

dt: Array

Get the time step length.

Returns:

Type	Description
Array	Time step length.

linearise

linearise(
    reference: AeroCase,
    wake_type: LinearWakeType,
    bound_upwash: bool = True,
    wake_upwash: bool = True,
    unsteady_force: bool = True,
) -> LinearUVLM

Create linearised aerodynamic model.

Parameters:

Name	Type	Description	Default
reference	AeroCase	Reference StaticAero around which to linearise.	required
wake_type	LinearWakeType	Type of wake model to use in linearisation, with options given from the LinearWakeType class (frozen, prescribed, or free). Value of None defaults to prescribed.	required
bound_upwash	bool	If true, linearise for flowfield perturbations at the bound vortex vertex.	True
wake_upwash	bool	If true, linearise for flowfield perturbations at the wake vortex vertex.	True
unsteady_force	bool	If true, include unsteady force terms in linearisation.	True

Returns:

Type	Description
LinearUVLM	LinearUVLM model, linearised at specified time step.

Source code in src/flapjax/aero/uvlm.py

def linearise(
    self,
    reference: AeroCase,
    wake_type: LinearWakeType,
    bound_upwash: bool = True,
    wake_upwash: bool = True,
    unsteady_force: bool = True,
) -> LinearUVLM:
    r"""
    Create linearised aerodynamic model.
    :param reference: Reference StaticAero around which to linearise.
    :param wake_type: Type of wake model to use in linearisation, with options given from the LinearWakeType class
     (frozen, prescribed, or free). Value of None defaults to prescribed.
    :param bound_upwash: If true, linearise for flowfield perturbations at the bound vortex vertex.
    :param wake_upwash: If true, linearise for flowfield perturbations at the wake vortex vertex.
    :param unsteady_force: If true, include unsteady force terms in linearisation.
    :return: LinearUVLM model, linearised at specified time step.
    """

    # local import used to prevent circular import issues
    from flapjax.aero.linear.linear_uvlm import LinearUVLM

    return LinearUVLM(
        self,
        reference=reference,
        wake_type=wake_type,
        bound_upwash=bound_upwash,
        wake_upwash=wake_upwash,
        unsteady_force=unsteady_force,
    )

set_design_variables

set_design_variables(
    dt: float | Array,
    flowfield: FlowField,
    zeta_b0: ArrayList | Sequence[Array] | Array,
    hg0: Array,
    delta_w: Sequence[Array | None] | Array | None = None,
    reference_cs_angles: dict[str, Array] | None = None,
) -> None

Set aerodynamic design variables for solution.

Parameters:

Name	Type	Description	Default
dt	float | Array	Time step length	required
flowfield	FlowField	FlowField object defining the background flow in space and time	required
delta_w	Sequence[Array | None] | Array | None	Vector to define segment lengths of a variable wake discretisation per surface. If None, this will use a uniform discretisation, as in the canonical UVLM.	None
zeta_b0	ArrayList | Sequence[Array] | Array	Aerodynamic local grid coordinates, (n_surf, )(zeta_m, zeta_n, 3).	required
hg0	Array	Beam reference global grid coordinates, (n_nodes, 4, 4). dictionary input for deflections, and returns an ArrayList of the local deflected grid.	required
reference_cs_angles	dict[str, Array] | None	Dictionary of {name: angle} for each control surface at the reference. If None, defaults to no control surfaces.	None

Source code in src/flapjax/aero/uvlm.py

def set_design_variables(
    self,
    dt: float | Array,
    flowfield: FlowField,
    zeta_b0: ArrayList | Sequence[Array] | Array,
    hg0: Array,
    delta_w: Sequence[Array | None] | Array | None = None,
    reference_cs_angles: dict[str, Array] | None = None,
) -> None:
    r"""
    Set aerodynamic design variables for solution.
    :param dt: Time step length
    :param flowfield: FlowField object defining the background flow in space and time
    :param delta_w: Vector to define segment lengths of a variable wake discretisation per surface. If None, this
    will use a uniform discretisation, as in the canonical UVLM.
    :param zeta_b0: Aerodynamic local grid coordinates, ``(n_surf, )(zeta_m, zeta_n, 3)``.
    :param hg0: Beam reference global grid coordinates, ``(n_nodes, 4, 4)``.
    dictionary input for deflections, and returns an ArrayList of the local deflected grid.
    :param reference_cs_angles: Dictionary of {name: angle} for each control surface at the reference. If None,
    defaults to no control surfaces.
    """

    if isinstance(delta_w, Array):
        delta_w_seq: Sequence[Array | None] = [
            delta_w if gd.m_star > 0 else None for gd in self.grid_disc
        ]
    elif delta_w is None:
        delta_w_seq = self.n_surf * [None]
    elif isinstance(delta_w, Sequence):
        if len(delta_w) != self.n_surf:
            raise ValueError(
                "Number of delta_w entries must match number of surfaces if passed as a Sequence"
            )
        delta_w_seq = delta_w
    else:
        raise TypeError("Invalid delta_w type")

    if isinstance(zeta_b0, Array):
        x0_aero_arraylist = ArrayList([zeta_b0])
    elif isinstance(zeta_b0, Sequence):
        x0_aero_arraylist = ArrayList(zeta_b0)
    elif isinstance(zeta_b0, ArrayList):
        x0_aero_arraylist = zeta_b0
    else:
        raise TypeError("Invalid zeta_b0 type")

    # set aerodynamic local coordinates
    if len(x0_aero_arraylist) != self.n_surf:
        raise ValueError(
            f"Expected {self.n_surf} aerodynamic grid coordinate arrays, got {len(zeta_b0)}"
        )

    for i_surf in range(self.n_surf):
        check_arr_shape(
            x0_aero_arraylist[i_surf],
            (self.grid_disc[i_surf].m + 1, self.grid_disc[i_surf].n + 1, 3),
            "zeta_b0",
        )
    self.zeta_b0 = x0_aero_arraylist

    if reference_cs_angles is not None:
        self.cs_ang0 = reference_cs_angles
        self.cs_vel0 = {
            k: jnp.zeros_like(v) for k, v in reference_cs_angles.items()
        }

    # set global grid coordinates for bound and wake
    check_arr_shape(hg0, (None, 4, 4), "hg0")
    self.hg_ref = hg0
    self.zeta_b_ref = self.hg_to_zeta_b(hg_n=hg0, cs_ang_n=self.cs_ang0)

    # surface spanwise edges that sit on the mirror plane, derived from reference configuration
    self.mirror_edge_low, self.mirror_edge_high = compute_mirror_edges(
        self.zeta_b_ref, self.mirror_point, self.mirror_normal
    )

    # set flowfield
    self.flowfield = flowfield

    # set timestep
    if isinstance(dt, float):
        self._dt = jnp.array(dt)
    elif isinstance(dt, Array):
        check_arr_shape(dt, (), "dt")
        self._dt = dt
    else:
        raise TypeError("dt must be either a float or an Array scalar")

    # set wake displacement
    self.delta_w = []
    for i_surf, dw_ in enumerate(delta_w_seq):
        if dw_ is None:
            self.delta_w.append(None)
        else:
            check_arr_shape(dw_, (self.grid_disc[i_surf].m_star,), "delta_w")
            self.delta_w.append(dw_)
    self.zeta_w_ref = self.initialise_wake()

case_from_dv

case_from_dv(dv: AeroDesignVariables) -> UVLM

Create a new UVLM instance as a function of design variables, allowing it to have defined gradients w.r.t. design variables.

Parameters:

Name	Type	Description	Default
dv	AeroDesignVariables	Design variables.	required

Returns:

Type	Description
UVLM	UVLM object with the same functionality as self.

Source code in src/flapjax/aero/uvlm.py

def case_from_dv(self, dv: AeroDesignVariables) -> UVLM:
    r"""
    Create a new UVLM instance as a function of design variables, allowing it to have defined gradients w.r.t.
    design variables.
    :param dv: Design variables.
    :return: UVLM object with the same functionality as ``self``.
    """
    inner_case = pytree_clone(self)
    flowfield = (
        inner_case.flowfield.from_design_variables(dv.flowfield)
        if dv.flowfield is not None
        else self.flowfield
    )
    cs_angles = (
        {k: jnp.atleast_1d(v)[0] for k, v in dv.cs_ang_t.items()}
        if dv.cs_ang_t is not None
        else self.cs_ang0
    )
    inner_case.set_design_variables(
        dt=self.dt,
        flowfield=flowfield,
        delta_w=self.delta_w,
        zeta_b0=dv_or(dv.zeta_b0, self.zeta_b0),
        hg0=self.hg_ref,
        reference_cs_angles=cs_angles,
    )

    return inner_case

get_design_variables

get_design_variables(
    cs_ang_t: dict[str, Array],
    cs_vel_t: dict[str, Array],
    grads_to_compute: AeroGradsToCompute | None,
) -> AeroDesignVariables

Extract design variables from the aerodynamic case. As the control input time histories are defined when initialising the simulation, they are not included in self and so are passed by argument.

Parameters:

Name	Type	Description	Default
cs_ang_t	dict[str, Array]	Time history of control surface angles, {keys: (n_tstep, )}.	required
cs_vel_t	dict[str, Array]	Time history of control surface velocities, {keys: (n_tstep, )}.	required
grads_to_compute	AeroGradsToCompute | None	Data structure which describes which design variables should be obtained. If None, all variables are obtained.	required

Returns:

Type	Description
AeroDesignVariables	Aerodynamic design variables.

Source code in src/flapjax/aero/uvlm.py

def get_design_variables(
    self,
    cs_ang_t: dict[str, Array],
    cs_vel_t: dict[str, Array],
    grads_to_compute: AeroGradsToCompute | None,
) -> AeroDesignVariables:
    r"""
    Extract design variables from the aerodynamic case. As the control input time histories are defined when
    initialising the simulation, they are not included in ``self`` and so are passed by argument.
    :param cs_ang_t: Time history of control surface angles, ``{keys: (n_tstep, )}``.
    :param cs_vel_t: Time history of control surface velocities, ``{keys: (n_tstep, )}``.
    :param grads_to_compute: Data structure which describes which design variables should be obtained. If None, all
    variables are obtained.
    :return: Aerodynamic design variables.
    """
    if isinstance(grads_to_compute, AeroGradsToCompute):
        return AeroDesignVariables(
            zeta_b0=self.zeta_b0 if grads_to_compute.x0_aero else None,
            flowfield=self.flowfield.to_design_variables()
            if grads_to_compute.flowfield
            else None,
            cs_ang_t=cs_ang_t if grads_to_compute.cs_ang_t else None,
            cs_vel_t=cs_vel_t if grads_to_compute.cs_vel_t else None,
            f_shape=(),
        )
    else:  # grads_to_compute is None
        return AeroDesignVariables(
            zeta_b0=self.zeta_b0,
            flowfield=self.flowfield.to_design_variables(),
            cs_ang_t=cs_ang_t,
            cs_vel_t=cs_vel_t,
            f_shape=(),
        )

hg_to_zeta_b

hg_to_zeta_b(
    hg_n: Array, cs_ang_n: dict[str, Array]
) -> ArrayList

Convert beam global grid coordinates to aerodynamic global grid coordinates.

Parameters:

Name	Type	Description	Default
hg_n	Array	Beam global grid coordinates at time step n, (n_nodes, 4, 4).	required
cs_ang_n	dict[str, Array]	Control surface angles as {surface_name: angle} pairs at time step n.	required

Returns:

Type	Description
ArrayList	Full aerodynamic global grid coordinates for each surface, (n_surf, )(zeta_m, zeta_n, 3).

Source code in src/flapjax/aero/uvlm.py

def hg_to_zeta_b(self, hg_n: Array, cs_ang_n: dict[str, Array]) -> ArrayList:
    r"""
    Convert beam global grid coordinates to aerodynamic global grid coordinates.
    :param hg_n: Beam global grid coordinates at time step n, ``(n_nodes, 4, 4)``.
    :param cs_ang_n: Control surface angles as {surface_name: angle} pairs at time step n.
    :return: Full aerodynamic global grid coordinates for each surface, ``(n_surf, )(zeta_m, zeta_n, 3)``.
    """

    zeta_b0_cs = self.grid_func(
        self.zeta_b0, **cs_ang_n
    )  # get local aerodynamic grid for control surface deflections.

    zetas = ArrayList([])
    for i_surf in range(self.n_surf):
        this_hg = jnp.take(
            hg_n, self.dof_mapping[i_surf], axis=0
        )  # (n_nodes, 4, 4)

        zetas.append(
            vmap(vmap(se3_vect_product, (None, 0), 0), (0, 1), 1)(
                this_hg, zeta_b0_cs[i_surf]
            )
        )
    return zetas

hg_dot_to_zeta_b_dot

hg_dot_to_zeta_b_dot(
    hg_n: Array,
    hg_dot_n: Array,
    cs_ang_n: dict[str, Array],
    cs_vel_n: dict[str, Array],
) -> ArrayList

Convert beam global grid velocities to aerodynamic global grid velocities.

Parameters:

Name	Type	Description	Default
hg_n	Array	Beam global grid coordinates, (n_nodes, 4, 4).	required
hg_dot_n	Array	Beam global grid velocities, (n_nodes, 4, 4).	required
cs_ang_n	dict[str, Array]	Control surface angles as {surface_name: angle} pairs.	required
cs_vel_n	dict[str, Array]	Control surface velocities as {surface_name: velocities} pairs.	required

Returns:

Type	Description
ArrayList	Full aerodynamic global grid velocities for each surface, (n_surf, )(zeta_m, zeta_n, 3).

Source code in src/flapjax/aero/uvlm.py

def hg_dot_to_zeta_b_dot(
    self,
    hg_n: Array,
    hg_dot_n: Array,
    cs_ang_n: dict[str, Array],
    cs_vel_n: dict[str, Array],
) -> ArrayList:
    r"""
    Convert beam global grid velocities to aerodynamic global grid velocities.
    :param hg_n: Beam global grid coordinates, ``(n_nodes, 4, 4)``.
    :param hg_dot_n: Beam global grid velocities, ``(n_nodes, 4, 4)``.
    :param cs_ang_n: Control surface angles as {surface_name: angle} pairs.
    :param cs_vel_n: Control surface velocities as {surface_name: velocities} pairs.
    :return: Full aerodynamic global grid velocities for each surface, ``(n_surf, )(zeta_m, zeta_n, 3)``.
    """
    zeta_b0_cs = self.grid_func(
        self.zeta_b0, **cs_ang_n
    )  # deflected local aerodynamic grid

    # as this is where the control velocities are used, they are checked here
    for key in set(cs_ang_n.keys()) | set(cs_vel_n.keys()):
        if key not in cs_vel_n or key not in cs_ang_n:
            raise ValueError(
                f"Missing pair of control angles and velocities for control surface key '{key}'"
            )

        if cs_vel_n[key].shape != cs_vel_n[key].shape:
            raise ValueError(
                f"Mismatched shape for control surface angles {cs_vel_n[key].shape} and velocities {cs_vel_n[key].shape}"
            )

    # use jvp to find the velocity of the aerodynamic grid due to control velocity.
    _, zeta_b0_dot_cs = jax.jvp(
        lambda angs: self.grid_func(
            self.zeta_b0, **angs
        ),  # local grid velocities due to control surface
        primals=(cs_ang_n,),
        tangents=(cs_vel_n,),
    )

    zeta_dots = ArrayList([])
    for i_surf in range(self.n_surf):
        this_hg = jnp.take(
            hg_n, self.dof_mapping[i_surf], axis=0
        )  # (n_nodes, 4, 4)
        this_hg_dot = jnp.take(
            hg_dot_n, self.dof_mapping[i_surf], axis=0
        )  # (n_nodes, 4, 4)
        this_rmat = this_hg[:, :3, :3]  # (n_span, 3, 3)
        zeta_dots.append(
            vmap(vmap(se3_vect_product, (None, 0), 0), (0, 1), 1)(
                this_hg_dot, zeta_b0_cs[i_surf]
            )
            + jnp.einsum(
                "njk,mnk->mnj",
                this_rmat,
                zeta_b0_dot_cs[i_surf],
            )
        )
    return zeta_dots

initialise_wake

initialise_wake(
    zeta_b: ArrayList | None = None,
) -> ArrayList

Generate initial wake grid coordinates, based on the bound grid coordinates and the freestream conditions.

Parameters:

Name	Type	Description	Default
zeta_b	ArrayList | None	Initial wake grid coordinates, (n_surf, )(zeta_m, zeta_n, 3). If None, this will use the initialised bound grid coordinates based on hg0.	None

Returns:

Type	Description
ArrayList	Initial wake grid coordinates, (zeta_m_star, zeta_n, 3)

Source code in src/flapjax/aero/uvlm.py

def initialise_wake(self, zeta_b: ArrayList | None = None) -> ArrayList:
    r"""
    Generate initial wake grid coordinates, based on the bound grid coordinates and the freestream conditions.
    :param zeta_b: Initial wake grid coordinates, ``(n_surf, )(zeta_m, zeta_n, 3)``. If None, this will use the
    initialised bound grid coordinates based on hg0.
    :return: Initial wake grid coordinates, ``(zeta_m_star, zeta_n, 3)``
    """
    zeta_b: ArrayList = zeta_b if zeta_b is not None else self.zeta_b_ref

    zeta0_w = ArrayList([])
    for i_surf, this_delta_w in enumerate(self.delta_w):
        # get bound grid coordinates
        zeta_te = zeta_b[i_surf][-1, :, :]  # (n+1, 3)

        # set wake grid coordinates as trailing edge + displacement
        if this_delta_w is None:
            this_delta_w = (
                jnp.ones(self.grid_disc[i_surf].m_star)
                * self.dt
                * self.flowfield.u_inf_mag
            )
        grid_s = jnp.concatenate((jnp.zeros(1), jnp.cumsum(this_delta_w)))

        zeta0_w.append(
            zeta_te[None, :, :]
            + jnp.outer(grid_s, self.flowfield.u_inf_dir)[:, None, :]
        )
    return zeta0_w

compute_gamma_dot staticmethod

compute_gamma_dot(
    gamma_b_n: ArrayList,
    gamma_b_nm1: ArrayList,
    gamma_b_dot_nm1: ArrayList,
    dt: Array,
    gamma_dot_relaxation: float | Array,
) -> ArrayList

Calculate time derivative of bound circulation strengths at specified time step using finite difference.

Parameters:

Name	Type	Description	Default
gamma_b_n	ArrayList	Bound circulation strengths at timestep n, (n_surf, )(m, n).	required
gamma_b_nm1	ArrayList	Bound circulation strengths at timestep n-1, (n_surf, )(m, n).	required
gamma_b_dot_nm1	ArrayList	Filtered bound circulation strengths time derivative at timestep n-1, (n_surf, )(m, n).	required
dt	Array	Time step length.	required
gamma_dot_relaxation	float | Array	Relaxation factor which filters the time derivative.	required

Source code in src/flapjax/aero/uvlm.py

@staticmethod
def compute_gamma_dot(
    gamma_b_n: ArrayList,
    gamma_b_nm1: ArrayList,
    gamma_b_dot_nm1: ArrayList,
    dt: Array,
    gamma_dot_relaxation: float | Array,
) -> ArrayList:
    r"""
    Calculate time derivative of bound circulation strengths at specified time step using finite difference.
    :param gamma_b_n: Bound circulation strengths at timestep n, ``(n_surf, )(m, n)``.
    :param gamma_b_nm1: Bound circulation strengths at timestep n-1, ``(n_surf, )(m, n)``.
    :param gamma_b_dot_nm1: Filtered bound circulation strengths time derivative at timestep n-1, ``(n_surf, )(m, n)``.
    :param dt: Time step length.
    :param gamma_dot_relaxation: Relaxation factor which filters the time derivative.
    """

    # first obtain the current unfiltered, and previous filtered values for gamma_dot
    gamma_b_dot_curr = (gamma_b_n - gamma_b_nm1) / dt

    # blend with relaxation parameter
    return gamma_b_dot_curr * gamma_dot_relaxation + gamma_b_dot_nm1 * (
        1.0 - gamma_dot_relaxation
    )

set_gamma_w

set_gamma_w(
    gamma_vec: Array, case: AeroCase, i_ts: int
) -> None

Set wake circulation strengths from total circulation strengths at specified time step. Can be passed either a full vector of strengths, or a sequence of strengths per surface.

Parameters:

Name	Type	Description	Default
case	AeroCase	AeroCase object.	required
gamma_vec	Array	Total circulation strengths vector, (gamma_w_tot, ).	required
i_ts	int	Timestep index.	required

Source code in src/flapjax/aero/uvlm.py

@singledispatchmethod
def set_gamma_w(self, gamma_vec: Array, case: AeroCase, i_ts: int) -> None:
    r"""
    Set wake circulation strengths from total circulation strengths at specified time step. Can be passed either a
    full vector of strengths, or a sequence of strengths per surface.
    :param case: AeroCase object.
    :param gamma_vec: Total circulation strengths vector, ``(gamma_w_tot, )``.
    :param i_ts: Timestep index.
    """
    for i_surf in range(self.n_surf):
        case.gamma_w[i_surf] = (
            case.gamma_w[i_surf]
            .at[i_ts, ...]
            .set(
                gamma_vec[self.gamma_w_slice[i_surf]].reshape(
                    self.grid_disc[i_surf].m_star, self.grid_disc[i_surf].n
                )
            )
        )

base_solve

base_solve(
    q_nm1: AeroFullStates | None,
    t_n: Array,
    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,
) -> tuple[
    ArrayList,
    ArrayList,
    ArrayList,
    ArrayList,
    ArrayList | None,
    ArrayList,
    ArrayList,
    ArrayList | None,
    ArrayList,
    ArrayList | None,
    ArrayList,
    ArrayList,
    ArrayList,
    ArrayList,
]

Solve the UVLM equations for a single time step from beam coordinate inputs.

Parameters:

Name	Type	Description	Default
q_nm1	AeroFullStates | None	Minimal aerodynamic states from timestep n-1.	required
t_n	Array	Time at timestep n.	required
hg_n	Array | None	Beam global grid coordinates at time step n, (n_nodes, 4, 4).	required
hg_nm1	Array | None	Beam global grid coordinates at time step n - 1, (n_nodes, 4, 4). Required for consistent free wake modelling.	required
hg_dot_n	Array | None	Beam global grid velocities, (n_nodes, 4, 4).	required
static	bool	If True, perform a static solve.	required
horseshoe	bool	If True, replace the wake with a horseshoe wake in static solve which extends a fixed distance.	required
cs_ang_n	dict[str, Array]	Control surface angle at timestep n, {name, ()}.	required
cs_ang_nm1	dict[str, Array] | None	Control surface angle at timestep n - 1, {name, ()}.	required
cs_vel_n	dict[str, Array] | None	Control surface velocity at timestep n, {name, ()}.	required

Returns:

Type	Description
tuple[ArrayList, ArrayList, ArrayList, ArrayList, ArrayList | None, ArrayList, ArrayList, ArrayList | None, ArrayList, ArrayList | None, ArrayList, ArrayList, ArrayList, ArrayList]	Collocation points, bound normals, bound circulation, wake circulation, bound circulation time derivative, bound grid, wake grid, bound grid time derivative, steady forcing, unsteady forcing, per-strip effective angle of attack, and per-strip lift, drag and moment coefficients.

Source code in src/flapjax/aero/uvlm.py

def base_solve(
    self,
    q_nm1: AeroFullStates | None,
    t_n: Array,
    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,
) -> tuple[
    ArrayList,
    ArrayList,
    ArrayList,
    ArrayList,
    ArrayList | None,
    ArrayList,
    ArrayList,
    ArrayList | None,
    ArrayList,
    ArrayList | None,
    ArrayList,
    ArrayList,
    ArrayList,
    ArrayList,
]:
    r"""
    Solve the UVLM equations for a single time step from beam coordinate inputs.
    :param q_nm1: Minimal aerodynamic states from timestep n-1.
    :param t_n: Time at timestep n.
    :param hg_n: Beam global grid coordinates at time step n, ``(n_nodes, 4, 4)``.
    :param hg_nm1: Beam global grid coordinates at time step n - 1, ``(n_nodes, 4, 4)``. Required for consistent free
    wake modelling.
    :param hg_dot_n: Beam global grid velocities, ``(n_nodes, 4, 4)``.
    :param static: If True, perform a static solve.
    :param horseshoe: If True, replace the wake with a horseshoe wake in static solve which extends a fixed
    distance.
    :param cs_ang_n: Control surface angle at timestep n, {name, ()}.
    :param cs_ang_nm1: Control surface angle at timestep n - 1, {name, ()}.
    :param cs_vel_n: Control surface velocity at timestep n, {name, ()}.
    :return: Collocation points, bound normals, bound circulation, wake circulation, bound circulation time
    derivative, bound grid, wake grid, bound grid time derivative, steady forcing, unsteady forcing, per-strip
    effective angle of attack, and per-strip lift, drag and moment coefficients.
    """

    zeta_b_n = self.hg_to_zeta_b(
        hg_n=hg_n if hg_n is not None else self.hg_ref, cs_ang_n=cs_ang_n
    )

    if hg_dot_n is None:
        zeta_b_dot_n: ArrayList | None = None
        zeta_b_nm1: ArrayList | None = None
    else:
        assert (
            hg_n is not None
            and cs_vel_n is not None
            and hg_nm1 is not None
            and cs_ang_nm1 is not None
        )
        zeta_b_dot_n = self.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
        )
        zeta_b_nm1 = self.hg_to_zeta_b(hg_n=hg_nm1, cs_ang_n=cs_ang_nm1)

    return self.base_solve_from_grid(
        q_nm1=q_nm1,
        t_n=t_n,
        zeta_b_n=zeta_b_n,
        zeta_b_nm1=zeta_b_nm1,
        zeta_b_dot_n=zeta_b_dot_n,
        static=static,
        horseshoe=horseshoe,
    )

base_solve_from_grid

base_solve_from_grid(
    q_nm1: AeroFullStates | None,
    t_n: Array,
    zeta_b_n: ArrayList,
    zeta_b_nm1: ArrayList | None,
    zeta_b_dot_n: ArrayList | None,
    static: bool,
    horseshoe: bool,
    *,
    linearise_variable_wake: bool = False,
    nu_b: ArrayList | None = None,
    nu_w: ArrayList | None = None,
) -> tuple[
    ArrayList,
    ArrayList,
    ArrayList,
    ArrayList,
    ArrayList | None,
    ArrayList,
    ArrayList,
    ArrayList | None,
    ArrayList,
    ArrayList | None,
    ArrayList,
    ArrayList,
    ArrayList,
    ArrayList,
]

Solve the UVLM equations for a single time step from aerodynamic grid inputs.

Parameters:

Name	Type	Description	Default
q_nm1	AeroFullStates | None	Aerodynamic states carried from timestep n-1. Required for dynamic (static=False) solves.	required
t_n	Array	Time at timestep n.	required
zeta_b_n	ArrayList	Bound aerodynamic grid at timestep n, (n_surf, )(zeta_m, zeta_n, 3).	required
zeta_b_nm1	ArrayList | None	Bound aerodynamic grid at timestep n-1, used to seed the wake-convection velocity for the free-wake case. Required for dynamic solves.	required
zeta_b_dot_n	ArrayList | None	Bound grid velocity at timestep n, or None for a static solve.	required
static	bool	If True, perform a static solve (initialise wake, skip wake propagation and unsteady forcing).	required
horseshoe	bool	If True, replace the wake with a horseshoe wake in the static solve.	required
linearise_variable_wake	bool	If True, block gradients through the arc-length discretisation in wake propagation so it acts as a linear operator when differentiated. Used by the linear system; default False.	False
nu_b	ArrayList | None	Optional additive bound upwash velocity at bound vertices, (n_surf, )(zeta_m, zeta_n, 3). Used by the linear system; ignored (equivalent to zero) if not supplied.	None
nu_w	ArrayList | None	Optional additive wake-convection velocity at wake vertices, (n_surf, )(zeta_m_star, zeta_n, 3). Used by the linear system; ignored if not supplied.	None

Returns:

Type	Description
tuple[ArrayList, ArrayList, ArrayList, ArrayList, ArrayList | None, ArrayList, ArrayList, ArrayList | None, ArrayList, ArrayList | None, ArrayList, ArrayList, ArrayList, ArrayList]	Collocation points, bound normals, bound circulation, wake circulation, bound circulation time derivative, bound grid, wake grid, bound grid velocity, steady forcing, unsteady forcing, per-strip effective angle of attack, and per-strip lift, drag and moment coefficients sampled from the polars.

Source code in src/flapjax/aero/uvlm.py

def base_solve_from_grid(
    self,
    q_nm1: AeroFullStates | None,
    t_n: Array,
    zeta_b_n: ArrayList,
    zeta_b_nm1: ArrayList | None,
    zeta_b_dot_n: ArrayList | None,
    static: bool,
    horseshoe: bool,
    *,
    linearise_variable_wake: bool = False,
    nu_b: ArrayList | None = None,
    nu_w: ArrayList | None = None,
) -> tuple[
    ArrayList,
    ArrayList,
    ArrayList,
    ArrayList,
    ArrayList | None,
    ArrayList,
    ArrayList,
    ArrayList | None,
    ArrayList,
    ArrayList | None,
    ArrayList,
    ArrayList,
    ArrayList,
    ArrayList,
]:
    r"""
    Solve the UVLM equations for a single time step from aerodynamic grid inputs.
    :param q_nm1: Aerodynamic states carried from timestep n-1. Required for dynamic (``static=False``) solves.
    :param t_n: Time at timestep n.
    :param zeta_b_n: Bound aerodynamic grid at timestep n, ``(n_surf, )(zeta_m, zeta_n, 3)``.
    :param zeta_b_nm1: Bound aerodynamic grid at timestep n-1, used to seed the wake-convection velocity for the
    free-wake case. Required for dynamic solves.
    :param zeta_b_dot_n: Bound grid velocity at timestep n, or ``None`` for a static solve.
    :param static: If True, perform a static solve (initialise wake, skip wake propagation and unsteady forcing).
    :param horseshoe: If True, replace the wake with a horseshoe wake in the static solve.
    :param linearise_variable_wake: If True, block gradients through the arc-length discretisation in wake
    propagation so it acts as a linear operator when differentiated. Used by the linear system; default False.
    :param nu_b: Optional additive bound upwash velocity at bound vertices, ``(n_surf, )(zeta_m, zeta_n, 3)``. Used by
    the linear system; ignored (equivalent to zero) if not supplied.
    :param nu_w: Optional additive wake-convection velocity at wake vertices, ``(n_surf, )(zeta_m_star, zeta_n, 3)``.
    Used by the linear system; ignored if not supplied.
    :return: Collocation points, bound normals, bound circulation, wake circulation, bound circulation time
    derivative, bound grid, wake grid, bound grid velocity, steady forcing, unsteady forcing, per-strip effective
    angle of attack, and per-strip lift, drag and moment coefficients sampled from the polars.
    """
    if isinstance(
        self.gamma_dot_relaxation,
        (int, float),  # prevents evaluating if value is traced
    ) and not (0.0 < self.gamma_dot_relaxation <= 1.0):
        warn("Gamma_dot relaxation factor not in (0, 1]")

    if not static and horseshoe:
        warn(
            "Horseshoe wake not compatible with non-static solve. Overriding horseshoe to False."
        )
        horseshoe = False

    if not static and q_nm1 is None:
        raise ValueError("q_nm1 needs to be specified for dynamic solve")

    c_n = compute_c(zetas=zeta_b_n)
    nc_n = compute_nc(zetas=zeta_b_n)

    # Prandtl-Glauert compressibility transform: components parallel to the freestream are unchanged, components
    # perpendicular to the freestream are scaled into new coordinates denoted with bar
    # physical circulation is recovered afterwardss as gamma = gamma_bar / beta**2
    x_hat = self.flowfield.u_inf_dir
    beta = self.flowfield.beta

    def _pg(z: Array) -> Array:
        return prandtl_glauert_transform(z, x_hat, beta)

    zeta_b_bar_n = ArrayList([_pg(z) for z in zeta_b_n])
    c_n_bar = compute_c(zetas=zeta_b_bar_n)
    nc_n_bar = compute_nc(zetas=zeta_b_bar_n)
    mirror_point_bar = (
        _pg(self.mirror_point) if self.mirror_point is not None else None
    )

    if zeta_b_dot_n is None:
        c_dot_n_bar: ArrayList | None = None
    else:
        c_dot_n = ArrayList(
            [neighbour_average(zeta_dot, axes=(0, 1)) for zeta_dot in zeta_b_dot_n]
        )
        c_dot_n_bar = ArrayList([_pg(cd) for cd in c_dot_n])

    if static:
        if horseshoe:
            zeta_w_n = ArrayList(
                [
                    self._make_surf_horseshoe_wake(
                        zeta_b=zeta_b_n[i_surf],
                        i_surf=i_surf,
                        horseshoe_length=HORSESHOE_LENGTH,
                    )
                    for i_surf in range(self.n_surf)
                ]
            )
        else:
            zeta_w_n = self.initialise_wake(zeta_b_n)

        gamma_w_n = None  # allocate later from gamma_b
        gamma_w_bar_n: ArrayList | None = None
        zeta_w_bar_n: ArrayList | None = ArrayList([_pg(z) for z in zeta_w_n])
    else:
        assert q_nm1 is not None and zeta_b_nm1 is not None

        zeta_full = ArrayList([*zeta_b_nm1, *q_nm1.zeta_w])
        gamma_full = ArrayList([*q_nm1.gamma_b, *q_nm1.gamma_w])

        def v_wake_prop(x_: Array) -> Array:
            v = self.flowfield.vmap_call(x=x_, t=t_n)
            if self.free_wake:
                v += compute_v_ind(
                    cs=x_,
                    zetas=zeta_full,
                    gammas=gamma_full,
                    kernels=[*self.kernels_b, *self.kernels_w],
                    batch_size=self.batch_size,
                    mirror_normal=self.mirror_normal,
                    mirror_point=self.mirror_point,
                )
            return v

        zeta_w_n, gamma_w_n = propagate_wake(
            gamma_b_nm1=q_nm1.gamma_b,
            gamma_w_nm1=q_nm1.gamma_w,
            zeta_b_n=zeta_b_n,
            zeta_w_nm1=q_nm1.zeta_w,
            delta_w=self.delta_w,
            v_func=v_wake_prop,
            dt=self.dt,
            frozen_wake=False,
            linearise_variable_wake=linearise_variable_wake,
        )

        # for the linearised case, add wake upwash from input
        if nu_w is not None:
            zeta_w_n = ArrayList(
                [zw + nub * self.dt for zw, nub in zip(zeta_w_n, nu_w)]
            )

        zeta_w_bar_n = ArrayList([_pg(z) for z in zeta_w_n])
        gamma_w_bar_n = ArrayList([beta**2 * gw for gw in gamma_w_n])

    aic_solve = compute_aic_solve(
        cs=c_n_bar,
        ns=nc_n_bar,
        zetas_b=zeta_b_bar_n,
        zetas_w=zeta_w_bar_n if static else None,
        kernels_b=self.kernels_b,
        kernels_w=self.kernels_w if static else None,
        batch_size=self.batch_size,
        mirror_normal=self.mirror_normal,
        mirror_point=mirror_point_bar,
    )

    # sampled at the physical collocation points, not transformed coordinates
    v_bc_n = self.flowfield.surf_vmap_call(xs=c_n, t=t_n)  # (n_surf, )(m, n, 3)

    if not static:
        assert c_dot_n_bar is not None
        v_bc_n -= c_dot_n_bar

        if zeta_w_bar_n is None or gamma_w_bar_n is None:
            raise ValueError("zeta_w_nm1 and gamma_w_nm1 are None")

        v_bc_n += compute_v_ind(
            cs=c_n_bar,
            zetas=zeta_w_bar_n,
            gammas=gamma_w_bar_n,
            kernels=self.kernels_w,
            batch_size=self.batch_size,
            mirror_normal=self.mirror_normal,
            mirror_point=mirror_point_bar,
        )

    # for linearised case, add bound grid upwash
    if nu_b is not None:
        v_bc_n += compute_c(ArrayList([_pg(nb) for nb in nu_b]))

    v_bc_n = ArrayList.einsum("ijk,ijk->ij", v_bc_n, nc_n_bar)  # (c_tot, )

    gamma_b_bar_vec_n = jnp.linalg.solve(aic_solve, -v_bc_n.ravel())
    gamma_b_n = ArrayList(
        [g / beta**2 for g in self._vec_to_gamma_b_list(gamma_b_bar_vec_n)]
    )

    def _static_wake_from_gamma_b(gamma_b: ArrayList) -> ArrayList:
        return ArrayList(
            [
                jnp.broadcast_to(
                    gb[[-1], ...],
                    shape=(
                        1
                        if (horseshoe and gd.m_star != 0)
                        else gd.m_star,  # m_star of 0 will override horseshoe
                        gd.n,
                    ),
                )
                for gb, gd in zip(gamma_b, self.grid_disc)
            ]
        )

    if static:
        gamma_w_n = _static_wake_from_gamma_b(gamma_b_n)

    assert gamma_w_n is not None
    assert zeta_w_n is not None

    zeta_b_dot_for_forces = (
        zeta_b_dot_n
        if zeta_b_dot_n is not None
        else ArrayList([jnp.zeros_like(zb) for zb in zeta_b_n])
    )

    def v_total_func(x_: Array) -> Array:
        assert gamma_w_n is not None
        return self.flowfield.vmap_call(x=x_, t=t_n) + compute_v_ind(
            cs=x_,
            zetas=ArrayList([*zeta_b_n, *zeta_w_n]),
            gammas=ArrayList([*gamma_b_n, *gamma_w_n]),
            kernels=[*self.kernels_b, *self.kernels_w],
            batch_size=self.batch_size,
            mirror_normal=self.mirror_normal,
            mirror_point=self.mirror_point,
        )

    f_steady = compute_steady_forcing(
        zeta_b=zeta_b_n,
        zeta_dot_b=zeta_b_dot_for_forces,
        gamma_b=gamma_b_n,
        gamma_w=gamma_w_n,
        rho=self.flowfield.rho,
        v_func=v_total_func,
        v_inputs=nu_b,
        mirror_point=self.mirror_point,
        mirror_normal=self.mirror_normal,
        mirror_edge_low=self.mirror_edge_low,
        mirror_edge_high=self.mirror_edge_high,
    )

    def v_freestream_func(x: Array) -> Array:
        return self.flowfield.vmap_call(x=x, t=t_n)

    alpha_n = strip_alpha(
        zeta_b=zeta_b_n,
        f_steady=f_steady,
        v_func=v_freestream_func,
        rho=self.flowfield.rho,
        beta=beta,
    )

    # update forcing with any polar corrections; surfaces with a ``None`` database report the Prandtl-Glauert
    # corrected flat-plate (2 pi alpha / beta) values implied by the UVLM itself
    f_steady, lift_scale, cl_n, cd_n, cm_n = apply_polar_correction(
        zeta_b=zeta_b_n,
        f_steady=f_steady,
        v_func=v_freestream_func,
        rho=self.flowfield.rho,
        beta=beta,
        polar_data=self.polar_data,
        polar_function=self.polar_function,
        alpha=alpha_n,
    )

    if self.polar_circulation_scale != 0.0:
        # blend bound (and, in the static case, wake) circulation towards the polar-corrected lift
        # scale=0 leaves gamma unchanged; scale=1 fully rescales it to lift_scale
        s = self.polar_circulation_scale
        gamma_b_n = ArrayList(
            [
                gb * (1.0 + s * (ls[None, :] - 1.0))
                for gb, ls in zip(gamma_b_n, lift_scale)
            ]
        )
        if static:
            gamma_w_n = _static_wake_from_gamma_b(gamma_b_n)

    if static:
        gamma_b_dot_n: ArrayList | None = None
        f_unsteady: ArrayList | None = None
    else:
        if q_nm1 is None:
            raise ValueError("q_nm1 needs to be specified for dynamic solve")

        gamma_b_dot_n = self.compute_gamma_dot(
            gamma_b_n=gamma_b_n,
            gamma_b_nm1=q_nm1.gamma_b,
            gamma_b_dot_nm1=q_nm1.gamma_b_dot,
            dt=self.dt,
            gamma_dot_relaxation=self.gamma_dot_relaxation,
        )
        f_unsteady = ArrayList(
            [
                split_to_vertex(
                    self.flowfield.rho
                    * gamma_b_dot_n[i_surf][..., None]
                    * nc_n[i_surf],
                    (0, 1),
                )
                for i_surf in range(self.n_surf)
            ]
        )

    return (
        c_n,
        nc_n,
        gamma_b_n,
        gamma_w_n,
        gamma_b_dot_n,
        zeta_b_n,
        zeta_w_n,
        zeta_b_dot_n,
        f_steady,
        f_unsteady,
        alpha_n,
        cl_n,
        cd_n,
        cm_n,
    )

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

Solve the UVLM equations for a single time step. Can be used for both static and dynamic solves. The solution is updated in-place in the case object.

Parameters:

Name	Type	Description	Default
case	AeroCase	Solution object.	required
i_ts	int	Timestep index to solve for.	required
hg_n	Array | None	Beam global grid coordinates at time step n, (zeta_n, 4, 4).	required
hg_nm1	Array | None	Beam global grid coordinates at time step n-1, (zeta_n, 4, 4).	required
hg_dot_n	Array | None	Beam global grid velocities at time step n, (zeta_n, 4, 4).	required
static	bool	If true, perform a static solve.	required
horseshoe	bool	If true, replace the wake with a static_horseshoe wake in static solve which extends a fixed distance.	required
cs_ang_n	dict[str, Array]	Control surface angle at timestep n, {name, ()}.	required
cs_ang_nm1	dict[str, Array] | None	Control surface angle at timestep n - 1, {name, ()}.	required
cs_vel_n	dict[str, Array] | None	Control surface velocity at timestep n, {name, ()}.	required

Returns:

Type	Description
AeroCase	Solution object with data for current time step added.

Source code in src/flapjax/aero/uvlm.py

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"""
    Solve the UVLM equations for a single time step. Can be used for both static and dynamic solves. The solution
    is updated in-place in the case object.
    :param case: Solution object.
    :param i_ts: Timestep index to solve for.
    :param hg_n: Beam global grid coordinates at time step n, ``(zeta_n, 4, 4)``.
    :param hg_nm1: Beam global grid coordinates at time step n-1, ``(zeta_n, 4, 4)``.
    :param hg_dot_n: Beam global grid velocities at time step n, ``(zeta_n, 4, 4)``.
    :param static: If true, perform a static solve.
    :param horseshoe: If true, replace the wake with a static_horseshoe wake in static solve which extends a fixed
    distance.
    :param cs_ang_n: Control surface angle at timestep n, {name, ()}.
    :param cs_ang_nm1: Control surface angle at timestep n - 1, {name, ()}.
    :param cs_vel_n: Control surface velocity at timestep n, {name, ()}.
    :return: Solution object with data for current time step added.
    """

    assert case.gamma_b_dot is not None and case.zeta_w is not None

    q_nm1 = AeroFullStates(
        gamma_b=case.gamma_b.index_all(i_ts - 1, ...),
        gamma_w=case.gamma_w.index_all(i_ts - 1, ...),
        gamma_b_dot=case.gamma_b_dot.index_all(i_ts - 1, ...),
        zeta_w=case.zeta_w.index_all(i_ts - 1, ...),
    )

    if not static:
        case.t = case.t.at[i_ts].set(
            jax.lax.select(i_ts, case.t[i_ts - 1] + self.dt, 0.0)
        )

    (
        c_n,
        nc_n,
        gamma_b_n,
        gamma_w_n,
        gamma_b_dot_n,
        zeta_b_n,
        zeta_w_n,
        zeta_b_dot_n,
        f_steady,
        f_unsteady,
        alpha_n,
        cl_n,
        cd_n,
        cm_n,
    ) = self.base_solve(
        q_nm1=q_nm1,
        t_n=case.t[i_ts, ...],
        hg_n=hg_n,
        hg_nm1=hg_nm1,
        hg_dot_n=hg_dot_n,
        static=static,
        horseshoe=horseshoe,
        cs_ang_n=cs_ang_n,
        cs_ang_nm1=cs_ang_nm1,
        cs_vel_n=cs_vel_n,
    )

    case.set_arraylist_at_ts("c", values=c_n, i_ts=i_ts)
    case.set_arraylist_at_ts("nc", values=nc_n, i_ts=i_ts)
    case.set_arraylist_at_ts("gamma_b", values=gamma_b_n, i_ts=i_ts)
    case.set_arraylist_at_ts("gamma_w", values=gamma_w_n, i_ts=i_ts)
    case.set_arraylist_at_ts("zeta_b", values=zeta_b_n, i_ts=i_ts)
    case.set_arraylist_at_ts("f_steady", values=f_steady, i_ts=i_ts)
    case.set_arraylist_at_ts("alpha", values=alpha_n, i_ts=i_ts)
    case.set_arraylist_at_ts("cl", values=cl_n, i_ts=i_ts)
    case.set_arraylist_at_ts("cd", values=cd_n, i_ts=i_ts)
    case.set_arraylist_at_ts("cm", values=cm_n, i_ts=i_ts)

    if not static:
        if gamma_b_dot_n is None:
            raise ValueError("gamma_b_dot_n is None")
        if zeta_b_dot_n is None:
            raise ValueError("zeta_b_dot_n is None")
        if f_unsteady is None:
            raise ValueError("f_unsteady is None")
        case.set_arraylist_at_ts("gamma_b_dot", values=gamma_b_dot_n, i_ts=i_ts)
        case.set_arraylist_at_ts("zeta_b_dot", values=zeta_b_dot_n, i_ts=i_ts)
        case.set_arraylist_at_ts("f_unsteady", values=f_unsteady, i_ts=i_ts)

    # set wake grid coordinates. If using static_horseshoe, it will still create a regular wake for plotting
    if horseshoe:
        case.set_arraylist_at_ts(
            "zeta_w", values=self.initialise_wake(zeta_w_n), i_ts=i_ts
        )
    else:
        case.set_arraylist_at_ts("zeta_w", values=zeta_w_n, i_ts=i_ts)

    return case

initialise_case_object

initialise_case_object(
    n_tstep: int,
    static_horseshoe: bool,
    free_wake: bool,
    gamma_dot_relaxation: float | Array,
    cs_ang_t: dict[str, Array],
    cs_vel_t: dict[str, Array],
) -> AeroCase

Initialise an AeroCase object to store the solution of the aerodynamic case for a given number of time steps. All solution data is initialised to zero.

Parameters:

Name	Type	Description	Default
n_tstep	int	Number of time steps to solve for in dynamic solution.	required
static_horseshoe	bool	Whether a horseshoe formulation was used for the static case.	required
free_wake	bool	Whether to use a free wake formulation.	required
gamma_dot_relaxation	float | Array	Relaxation factor for damping gamma_dot.	required
cs_ang_t	dict[str, Array]	Control surface angle time history, {name: (n_tstep, )}.	required
cs_vel_t	dict[str, Array]	Control surface velocity time history, {name: (n_tstep, )}.	required

Returns:

Type	Description
AeroCase	AeroCase object.

Source code in src/flapjax/aero/uvlm.py

def initialise_case_object(
    self,
    n_tstep: int,
    static_horseshoe: bool,
    free_wake: bool,
    gamma_dot_relaxation: float | Array,
    cs_ang_t: dict[str, Array],
    cs_vel_t: dict[str, Array],
) -> AeroCase:
    r"""
    Initialise an AeroCase object to store the solution of the aerodynamic case for a given number of time
    steps. All solution data is initialised to zero.
    :param n_tstep: Number of time steps to solve for in dynamic solution.
    :param static_horseshoe: Whether a horseshoe formulation was used for the static case.
    :param free_wake: Whether to use a free wake formulation.
    :param gamma_dot_relaxation: Relaxation factor for damping gamma_dot.
    :param cs_ang_t: Control surface angle time history, ``{name: (n_tstep, )}``.
    :param cs_vel_t: Control surface velocity time history, ``{name: (n_tstep, )}``.
    :return: AeroCase object.
    """
    # zero initialise
    return AeroCase(
        zeta_b=ArrayList(
            [jnp.zeros((n_tstep, gd.m + 1, gd.n + 1, 3)) for gd in self.grid_disc]
        ),
        zeta_b_dot=ArrayList(
            [jnp.zeros((n_tstep, gd.m + 1, gd.n + 1, 3)) for gd in self.grid_disc]
        ),
        zeta_w=ArrayList(
            [
                jnp.zeros((n_tstep, gd.m_star + 1, gd.n + 1, 3))
                for gd in self.grid_disc
            ]
        ),
        gamma_b=ArrayList(
            [jnp.zeros((n_tstep, gd.m, gd.n)) for gd in self.grid_disc]
        ),
        gamma_b_dot=ArrayList(
            [jnp.zeros((n_tstep, gd.m, gd.n)) for gd in self.grid_disc]
        ),
        gamma_w=ArrayList(
            [jnp.zeros((n_tstep, gd.m_star, gd.n)) for gd in self.grid_disc]
        ),
        f_steady=ArrayList(
            [jnp.zeros((n_tstep, gd.m + 1, gd.n + 1, 3)) for gd in self.grid_disc]
        ),
        f_unsteady=ArrayList(
            [jnp.zeros((n_tstep, gd.m + 1, gd.n + 1, 3)) for gd in self.grid_disc]
        ),
        alpha=ArrayList([jnp.zeros((n_tstep, gd.n)) for gd in self.grid_disc]),
        cl=ArrayList([jnp.zeros((n_tstep, gd.n)) for gd in self.grid_disc]),
        cd=ArrayList([jnp.zeros((n_tstep, gd.n)) for gd in self.grid_disc]),
        cm=ArrayList([jnp.zeros((n_tstep, gd.n)) for gd in self.grid_disc]),
        c=ArrayList([jnp.zeros((n_tstep, gd.m, gd.n, 3)) for gd in self.grid_disc]),
        n=ArrayList([jnp.zeros((n_tstep, gd.m, gd.n, 3)) for gd in self.grid_disc]),
        kernels=[*self.kernels_b, *self.kernels_w],
        mirror_point=self.mirror_point,
        mirror_normal=self.mirror_normal,
        mirror_edge_low=self.mirror_edge_low,
        mirror_edge_high=self.mirror_edge_high,
        flowfield=self.flowfield,
        surf_b_names=self.surf_b_names,
        surf_w_names=self.surf_w_names,
        i_ts=jnp.arange(n_tstep),
        t=jnp.zeros(n_tstep),
        dof_mapping=self.dof_mapping,
        static_horseshoe=static_horseshoe,
        free_wake=free_wake,
        gamma_dot_relaxation=gamma_dot_relaxation,
        cs_ang=cs_ang_t,
        cs_vel=cs_vel_t,
        batch_size=self.batch_size,
    )

static_solve

static_solve(
    hg: Array | None = None,
    t: Array | float = 0.0,
    horseshoe: bool = False,
    cs_ang: dict[str, Array] | None = None,
) -> AeroCase

Solve the VLM for given static beam coordinates. TODO: add free wake

Parameters:

Name	Type	Description	Default
hg	Array | None	Beam coordinates, (zeta_n, 4, 4).	None
t	Array | float	Time at which to solve static solution, used for background flowfield evaluation.	0.0
horseshoe	bool	If true, replace the wake with a horseshoe wake which extends a fixed distance.	False
cs_ang	dict[str, Array] | None	Control surface angles, {name: ()}.	None

Returns:

Type	Description
AeroCase	AeroCase solution object.

Source code in src/flapjax/aero/uvlm.py

def static_solve(
    self,
    hg: Array | None = None,
    t: Array | float = 0.0,
    horseshoe: bool = False,
    cs_ang: dict[str, Array] | None = None,
) -> AeroCase:
    r"""
    Solve the VLM for given static beam coordinates.
    TODO: add free wake
    :param hg: Beam coordinates, ``(zeta_n, 4, 4)``.
    :param t: Time at which to solve static solution, used for background flowfield evaluation.
    :param horseshoe: If true, replace the wake with a horseshoe wake which extends a fixed distance.
    :param cs_ang: Control surface angles, ``{name: ()}``.
    :return: AeroCase solution object.
    """

    case = self.initialise_case_object(
        1,
        static_horseshoe=horseshoe,
        gamma_dot_relaxation=0.7,
        free_wake=False,
        cs_ang_t=cs_ang if cs_ang is not None else self.cs_ang0,
        cs_vel_t={
            k: jnp.zeros_like(v)
            for k, v in (cs_ang if cs_ang is not None else self.cs_ang0).items()
        },
    )
    case.t = case.t.at[0].set(t)

    out_case = self.case_solve(
        case=case,
        i_ts=0,
        hg_n=hg,
        hg_nm1=None,
        hg_dot_n=None,
        static=True,
        horseshoe=horseshoe,
        cs_ang_n=cs_ang if cs_ang is not None else self.cs_ang0,
        cs_ang_nm1=None,
        cs_vel_n=None,
    )[0]

    if horseshoe:
        # if using a horseshoe wake, return the normal wake to prevent continuity errors
        out_case.zeta_w = self.initialise_wake(zeta_b=out_case.zeta_b)

    return out_case

prescribed_dynamic_solve

prescribed_dynamic_solve(
    init_case: AeroCase,
    hg_t: Array,
    hg_dot_t: Array,
    cs_ang_t: dict[str, Array] | None = None,
    cs_vel_t: dict[str, Array] | None = None,
) -> AeroCase

Solve the UVLM for prescribed grid motions.

Parameters:

Name	Type	Description	Default
init_case	AeroCase	StaticAero object containing initial conditions for the solution at time step 0.	required
hg_t	Array	Beam coordinates over time, (n_tstep, n_nodes, 4, 4).	required
hg_dot_t	Array	Beam coordinate time derivative over time, (n_tstep, n_nodes, 4, 4).	required
cs_ang_t	dict[str, Array] | None	Control surface angle time history, {name, (n_tstep, )}.	None
cs_vel_t	dict[str, Array] | None	Control surface velocity time history, {name, (n_tstep, )}.	None

Returns:

Type	Description
AeroCase	AeroCase solution object.

Source code in src/flapjax/aero/uvlm.py

def prescribed_dynamic_solve(
    self,
    init_case: AeroCase,
    hg_t: Array,
    hg_dot_t: Array,
    cs_ang_t: dict[str, Array] | None = None,
    cs_vel_t: dict[str, Array] | None = None,
) -> AeroCase:
    r"""
    Solve the UVLM for prescribed grid motions.
    :param init_case: StaticAero object containing initial conditions for the solution at time step 0.
    :param hg_t: Beam coordinates over time, ``(n_tstep, n_nodes, 4, 4)``.
    :param hg_dot_t: Beam coordinate time derivative over time, ``(n_tstep, n_nodes, 4, 4)``.
    :param cs_ang_t: Control surface angle time history, {name, ``(n_tstep, )``}.
    :param cs_vel_t: Control surface velocity time history, {name, ``(n_tstep, )``}.
    :return: AeroCase solution object.
    """
    check_arr_shape(hg_t, (None, None, 4, 4), "hg_n")
    check_if_all_se3_g(hg_t, True)

    if hg_t.shape != hg_dot_t.shape:
        raise ValueError(
            f"hg_dot_n must have the same shape as hg_n, got {hg_dot_t.shape} vs {hg_t.shape}"
        )

    check_if_all_se3_a(hg_dot_t, True)

    n_tstep = hg_t.shape[0]

    case = init_case.to_dynamic(i_ts=0, n_tstep=n_tstep)

    def _step_func(i_ts_: int, case_: AeroCase) -> AeroCase:
        cs_angle_nm1, cs_angle_n = (
            (
                {k: v[i_ts__] for k, v in cs_ang_t.items()}
                if cs_ang_t is not None
                else {}
            )
            for i_ts__ in (i_ts_ - 1, i_ts_)
        )
        cs_velocity_n = (
            {k: v[i_ts_] for k, v in cs_vel_t.items()}
            if cs_vel_t is not None
            else {}
        )

        case_ = self.case_solve(
            case=case_,
            i_ts=i_ts_,
            hg_n=hg_t[i_ts_, ...],
            hg_nm1=hg_t[i_ts_ - 1, ...],
            hg_dot_n=hg_dot_t[i_ts_, ...],
            static=False,
            horseshoe=False,
            cs_ang_n=cs_angle_n,
            cs_ang_nm1=cs_angle_nm1,
            cs_vel_n=cs_velocity_n,
        )
        jax_print(
            "UVLM timestep {i_ts_}",
            i_ts_=i_ts_,
            verbose_level="normal",
        )
        return case_

    case = fori_loop(
        1,
        n_tstep,
        _step_func,
        init_val=case,
    )
    return case

reference_configuration

reference_configuration() -> AeroCase

Get the reference (initial) snapshot of the aerodynamic case. This will set the timestep as -1.

Returns:

Type	Description
AeroCase	StaticAero object at initial time step.

Source code in src/flapjax/aero/uvlm.py

def reference_configuration(self) -> AeroCase:
    r"""
    Get the reference (initial) snapshot of the aerodynamic case. This will set the timestep as -1.
    :return: StaticAero object at initial time step.
    """
    return AeroCase(
        zeta_b=self.zeta_b_ref,
        zeta_b_dot=ArrayList(
            [jnp.zeros((gd.m + 1, gd.n + 1, 3)) for gd in self.grid_disc]
        ),
        zeta_w=self.zeta_w_ref,
        c=compute_c(self.zeta_b_ref),
        n=compute_nc(self.zeta_b_ref),
        gamma_b=ArrayList([jnp.zeros((gd.m, gd.n)) for gd in self.grid_disc]),
        gamma_b_dot=ArrayList([jnp.zeros((gd.m, gd.n)) for gd in self.grid_disc]),
        gamma_w=ArrayList([jnp.zeros((gd.m_star, gd.n)) for gd in self.grid_disc]),
        f_steady=ArrayList(
            [jnp.zeros((gd.m + 1, gd.n + 1, 3)) for gd in self.grid_disc]
        ),
        f_unsteady=ArrayList(
            [jnp.zeros((gd.m + 1, gd.n + 1, 3)) for gd in self.grid_disc]
        ),
        alpha=ArrayList([jnp.zeros((gd.n,)) for gd in self.grid_disc]),
        cl=ArrayList([jnp.zeros((gd.n,)) for gd in self.grid_disc]),
        cd=ArrayList([jnp.zeros((gd.n,)) for gd in self.grid_disc]),
        cm=ArrayList([jnp.zeros((gd.n,)) for gd in self.grid_disc]),
        surf_b_names=self.surf_b_names,
        surf_w_names=self.surf_w_names,
        i_ts=-1,
        t=jnp.array(0.0),
        dof_mapping=self.dof_mapping,
        flowfield=self.flowfield,
        mirror_point=self.mirror_point,
        mirror_normal=self.mirror_normal,
        mirror_edge_low=self.mirror_edge_low,
        mirror_edge_high=self.mirror_edge_high,
        kernels=[*self.kernels_b, *self.kernels_w],
        static_horseshoe=False,
        gamma_dot_relaxation=0.0,
        free_wake=False,
        cs_ang=self.cs_ang0,
        cs_vel=self.cs_vel0,
        batch_size=self.batch_size,
    )

plot_reference

plot_reference(
    directory: PathLike | str, plot_wake: bool = True
) -> Sequence[Path]

Plot the reference (initial) snapshot of the aerodynamic case. This will set the timestep as -1.

Parameters:

Name	Type	Description	Default
directory	PathLike | str	Path to write files to.	required
plot_wake	bool	If True, plot the wake grid.	True

Source code in src/flapjax/aero/uvlm.py

def plot_reference(
    self, directory: os.PathLike | str, plot_wake: bool = True
) -> Sequence[Path]:
    r"""
    Plot the reference (initial) snapshot of the aerodynamic case. This will set the timestep as -1.
    :param directory: Path to write files to.
    :param plot_wake: If True, plot the wake grid.
    """
    return self.reference_configuration().plot(
        Path(directory).resolve(), plot_wake=plot_wake
    )

gamma_b_res_func

gamma_b_res_func(
    i_ts: int | Array,
    t_n: Array,
    varphi_n: Array,
    v_n: Array,
    gamma_b_n: Array,
    gamma_w_n: Array,
    zeta_w_n: Array,
    dv: AeroelasticDesignVariables,
    dv_full: AeroelasticDesignVariables,
    struct_obj: BeamStructure,
) -> Array

Bound circulation residual used for adjoint computations.

:math:\mathbf{r}_{\Gamma_b} = \left(\boldsymbol{\mathcal{A}}_{b, n} \cdot \mathbf{n}_n\right)^{-1} \left[\left( \boldsymbol{\mathcal{A}}_{w, n} \boldsymbol{\Gamma}_{w, n} +\mathbf{v}_{bc, n} - \dot{\boldsymbol{\zeta}}_c\right) \cdot \mathbf{n}_n\right] + \boldsymbol{\Gamma}_{b, n}

Parameters:

Name	Type	Description	Default
i_ts	int | Array	Time step index.	required
t_n	Array	Time at step n.	required
varphi_n	Array	varphi vector at timestep n.	required
v_n	Array	Beam velocity vector at timestep n.	required
gamma_b_n	Array	Bound circulation vector at timestep n.	required
gamma_w_n	Array	Wake circulation vector at timestep n.	required
zeta_w_n	Array	Wake grid vector at timestep n.	required
dv	AeroelasticDesignVariables	Aeroelastic design variables.	required
dv_full	AeroelasticDesignVariables	Aeroelastic design variables without omissions for the variables where gradients aren't requested.	required
struct_obj	BeamStructure	Beam structure.	required

Returns:

Type	Description
Array	Bound circulation residual.

Source code in src/flapjax/aero/uvlm.py

def gamma_b_res_func(
    self,
    i_ts: int | Array,
    t_n: Array,
    varphi_n: Array,
    v_n: Array,
    gamma_b_n: Array,
    gamma_w_n: Array,
    zeta_w_n: Array,
    dv: AeroelasticDesignVariables,
    dv_full: AeroelasticDesignVariables,
    struct_obj: BeamStructure,
) -> Array:
    r"""
    Bound circulation residual used for adjoint computations.

    :math:`\mathbf{r}_{\Gamma_b} = \left(\boldsymbol{\mathcal{A}}_{b, n} \cdot \mathbf{n}_n\right)^{-1} \left[\left(
    \boldsymbol{\mathcal{A}}_{w, n} \boldsymbol{\Gamma}_{w, n} +\mathbf{v}_{bc, n} - \dot{\boldsymbol{\zeta}}_c\right)
    \cdot \mathbf{n}_n\right] + \boldsymbol{\Gamma}_{b, n}`

    :param i_ts: Time step index.
    :param t_n: Time at step n.
    :param varphi_n: varphi vector at timestep n.
    :param v_n: Beam velocity vector at timestep n.
    :param gamma_b_n: Bound circulation vector at timestep n.
    :param gamma_w_n: Wake circulation vector at timestep n.
    :param zeta_w_n: Wake grid vector at timestep n.
    :param dv: Aeroelastic design variables.
    :param dv_full: Aeroelastic design variables without omissions for the variables where gradients aren't
    requested.
    :param struct_obj: Beam structure.
    :return: Bound circulation residual.
    """

    varphi_n = varphi_n.reshape(-1, 6)
    v_n = v_n.reshape(-1, 6)
    gamma_b_n = ArrayList.from_vector(
        vect=gamma_b_n,
        shapes=ArrayListShape([(gd.m, gd.n) for gd in self.grid_disc]),
    )
    gamma_w_n = ArrayList.from_vector(
        vect=gamma_w_n,
        shapes=ArrayListShape([(gd.m_star, gd.n) for gd in self.grid_disc]),
    )
    zeta_w_n = ArrayList.from_vector(
        vect=zeta_w_n,
        shapes=ArrayListShape(
            [(gd.m_star + 1, gd.n + 1, 3) for gd in self.grid_disc]
        ),
    )

    inner_struct = struct_obj.case_from_dv(dv=dv.structure)
    hg_n = inner_struct.compute_hg_from_varphi(varphi=varphi_n)
    hg_dot_n = inner_struct.make_hg_dot(hg=hg_n, v=v_n)

    inner_case = self.case_from_dv(dv=dv.aero)

    # get control surface deflections from design variables
    cs_ang_n, cs_vel_n = dv.aero.get_cs_n(i_ts=i_ts, dv_full=dv_full.aero)

    zeta_b_n = inner_case.hg_to_zeta_b(hg_n=hg_n, cs_ang_n=cs_ang_n)
    c_n = compute_c(
        zetas=zeta_b_n
    )  # physical collocation points, used for freestream sampling only

    # Prandtl-Glauert compressibility transformforward solve.
    x_hat = inner_case.flowfield.u_inf_dir
    beta = inner_case.flowfield.beta

    def _pg(z: Array) -> Array:
        return prandtl_glauert_transform(z, x_hat, beta)

    zeta_b_bar_n = ArrayList([_pg(z) for z in zeta_b_n])
    c_n_bar = compute_c(zetas=zeta_b_bar_n)
    nc_n_bar = compute_nc(zetas=zeta_b_bar_n)
    mirror_point_bar = (
        _pg(inner_case.mirror_point)
        if inner_case.mirror_point is not None
        else None
    )

    zeta_b_dot_n = inner_case.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
    )

    c_dot_n = ArrayList(
        [neighbour_average(zeta_dot, axes=(0, 1)) for zeta_dot in zeta_b_dot_n]
    )
    c_dot_n_bar = ArrayList([_pg(cd) for cd in c_dot_n])

    zeta_w_bar_n = ArrayList([_pg(z) for z in zeta_w_n])
    gamma_w_bar_n = ArrayList([beta**2 * gw for gw in gamma_w_n])

    aic_solve = compute_aic_solve(
        cs=c_n_bar,
        ns=nc_n_bar,
        zetas_b=zeta_b_bar_n,
        zetas_w=None,
        kernels_b=inner_case.kernels_b,
        kernels_w=None,
        batch_size=self.batch_size,
        mirror_normal=inner_case.mirror_normal,
        mirror_point=mirror_point_bar,
    )

    v_bc_n = inner_case.flowfield.surf_vmap_call(
        xs=c_n, t=t_n
    )  # (n_surf, )(m, n, 3)

    # structural component
    v_bc_n -= c_dot_n_bar

    # find wake component
    v_bc_n += compute_v_ind(
        cs=c_n_bar,
        zetas=zeta_w_bar_n,
        gammas=gamma_w_bar_n,
        kernels=inner_case.kernels_w,
        batch_size=self.batch_size,
        mirror_normal=inner_case.mirror_normal,
        mirror_point=mirror_point_bar,
    )

    v_bc_n = ArrayList.einsum("ijk,ijk->ij", v_bc_n, nc_n_bar)  # (c_tot, )

    gamma_b_bar_vec_n = jnp.linalg.solve(aic_solve, -v_bc_n.ravel())
    gamma_b_nm1_update = ArrayList(
        [g / beta**2 for g in self._vec_to_gamma_b_list(gamma_b_bar_vec_n)]
    )

    return (gamma_b_nm1_update - gamma_b_n).ravel()

wake_prop_res_func

wake_prop_res_func(
    i_ts: int | Array,
    t_n: Array,
    varphi_nm1: Array,
    varphi_n: Array,
    gamma_b_nm1: Array,
    gamma_w_nm1: Array,
    gamma_w_n: Array,
    zeta_w_nm1: Array,
    zeta_w_n: Array,
    dv: AeroelasticDesignVariables,
    dv_full: AeroelasticDesignVariables,
    struct_obj: BeamStructure,
) -> tuple[Array, Array]

Wake propagation residual for both grid coordinates and circulation strengths.

Parameters:

Name	Type	Description	Default
i_ts	int | Array	Time step index.	required
t_n	Array	Time at timestep n.	required
varphi_nm1	Array	Beam minimal coordinates vector at timestep n-1.	required
varphi_n	Array	Beam maximal coordinates vector at timestep n.	required
gamma_b_nm1	Array	Bound circulation vector at timestep n-1.	required
gamma_w_nm1	Array	Wake circulation vector at timestep n-1.	required
gamma_w_n	Array	Wake circulation vector at timestep n.	required
zeta_w_nm1	Array	Wake grid vector at timestep n-1.	required
zeta_w_n	Array	Wake grid vector at timestep n.	required
dv	AeroelasticDesignVariables	Aeroelastic design variables.	required
dv_full	AeroelasticDesignVariables	Aeroelastic design variables without omissions.	required
struct_obj	BeamStructure	Beam structure.	required

Returns:

Type	Description
tuple[Array, Array]	Wake grid and circulation residuals.

Source code in src/flapjax/aero/uvlm.py

def wake_prop_res_func(
    self,
    i_ts: int | Array,
    t_n: Array,
    varphi_nm1: Array,
    varphi_n: Array,
    gamma_b_nm1: Array,
    gamma_w_nm1: Array,
    gamma_w_n: Array,
    zeta_w_nm1: Array,
    zeta_w_n: Array,
    dv: AeroelasticDesignVariables,
    dv_full: AeroelasticDesignVariables,
    struct_obj: BeamStructure,
) -> tuple[Array, Array]:
    r"""
    Wake propagation residual for both grid coordinates and circulation strengths.

    :param i_ts: Time step index.
    :param t_n: Time at timestep n.
    :param varphi_nm1: Beam minimal coordinates vector at timestep n-1.
    :param varphi_n: Beam maximal coordinates vector at timestep n.
    :param gamma_b_nm1: Bound circulation vector at timestep n-1.
    :param gamma_w_nm1: Wake circulation vector at timestep n-1.
    :param gamma_w_n: Wake circulation vector at timestep n.
    :param zeta_w_nm1: Wake grid vector at timestep n-1.
    :param zeta_w_n: Wake grid vector at timestep n.
    :param dv: Aeroelastic design variables.
    :param dv_full: Aeroelastic design variables without omissions.
    :param struct_obj: Beam structure.
    :return: Wake grid and circulation residuals.
    """

    varphi_nm1 = varphi_nm1.reshape(-1, 6)
    varphi_n = varphi_n.reshape(-1, 6)

    gamma_b_nm1 = ArrayList.from_vector(
        vect=gamma_b_nm1,
        shapes=ArrayListShape([(gd.m, gd.n) for gd in self.grid_disc]),
    )

    gamma_w_nm1 = ArrayList.from_vector(
        vect=gamma_w_nm1,
        shapes=ArrayListShape([(gd.m_star, gd.n) for gd in self.grid_disc]),
    )

    gamma_w_n = ArrayList.from_vector(
        vect=gamma_w_n,
        shapes=ArrayListShape([(gd.m_star, gd.n) for gd in self.grid_disc]),
    )
    zeta_w_nm1 = ArrayList.from_vector(
        vect=zeta_w_nm1,
        shapes=ArrayListShape(
            [(gd.m_star + 1, gd.n + 1, 3) for gd in self.grid_disc]
        ),
    )

    zeta_w_n = ArrayList.from_vector(
        vect=zeta_w_n,
        shapes=ArrayListShape(
            [(gd.m_star + 1, gd.n + 1, 3) for gd in self.grid_disc]
        ),
    )

    inner_struct = struct_obj.case_from_dv(dv=dv.structure)
    hg_nm1 = inner_struct.compute_hg_from_varphi(varphi=varphi_nm1)
    hg_n = inner_struct.compute_hg_from_varphi(varphi=varphi_n)

    inner_case = self.case_from_dv(dv=dv.aero)

    # get control surface deflections from design variables
    cs_ang_nm1, _ = dv.aero.get_cs_n(i_ts=i_ts - 1, dv_full=dv_full.aero)
    cs_ang_n, _ = dv.aero.get_cs_n(i_ts=i_ts, dv_full=dv_full.aero)
    zeta_b_nm1 = inner_case.hg_to_zeta_b(hg_n=hg_nm1, cs_ang_n=cs_ang_nm1)
    zeta_b_n = inner_case.hg_to_zeta_b(hg_n=hg_n, cs_ang_n=cs_ang_n)

    def v_wake_prop(x_: Array) -> Array:
        v = inner_case.flowfield.vmap_call(x=x_, t=t_n)
        if self.free_wake:
            v += compute_v_ind(
                cs=x_,
                zetas=ArrayList([*zeta_b_nm1, *zeta_w_nm1]),
                gammas=ArrayList([*gamma_b_nm1, *gamma_w_nm1]),
                kernels=[*inner_case.kernels_b, *inner_case.kernels_w],
                batch_size=self.batch_size,
                mirror_normal=inner_case.mirror_normal,
                mirror_point=inner_case.mirror_point,
            )
        return v

    zeta_w_nm1_update, gamma_w_nm1_update = propagate_wake(
        gamma_b_nm1=gamma_b_nm1,
        gamma_w_nm1=gamma_w_nm1,
        zeta_b_n=zeta_b_n,
        zeta_w_nm1=zeta_w_nm1,
        delta_w=inner_case.delta_w,
        v_func=v_wake_prop,
        dt=inner_case.dt,
        frozen_wake=False,
        linearise_variable_wake=False,
    )

    return (zeta_w_nm1_update - zeta_w_n).ravel(), (
        gamma_w_nm1_update - gamma_w_n
    ).ravel()

gamma_b_dot_res_func

gamma_b_dot_res_func(
    gamma_b_nm1: Array,
    gamma_b_n: Array,
    gamma_b_dot_nm1: Array,
    gamma_b_dot_n: Array,
    dv: AeroelasticDesignVariables,
) -> Array

Bound circulation time derivative residual.

:math:\frac{g}{h} \left[\mathbf{\Gamma}_{b, n} - \mathbf{\Gamma}_{b, n-1}\right] + (1-g) \dot{\mathbf{\Gamma}}_{b, n-1} - \dot{\mathbf{\Gamma}}_{b, n}

Parameters:

Name	Type	Description	Default
gamma_b_nm1	Array	Bound circulation vector at timestep n-1.	required
gamma_b_n	Array	Bound circulation vector at timestep n.	required
gamma_b_dot_nm1	Array	Bound circulation time derivative vector at timestep n-1.	required
gamma_b_dot_n	Array	Bound circulation time derivative vector at timestep n.	required
dv	AeroelasticDesignVariables	Design variables. Whilst this function does not depend upon it, including it simplified obtaining the residual design gradient.	required

Returns:

Type	Description
Array	Bound circulation time derivative residual.

Source code in src/flapjax/aero/uvlm.py

def gamma_b_dot_res_func(
    self,
    gamma_b_nm1: Array,
    gamma_b_n: Array,
    gamma_b_dot_nm1: Array,
    gamma_b_dot_n: Array,
    dv: AeroelasticDesignVariables,
) -> Array:
    r"""
    Bound circulation time derivative residual.

    :math:`\frac{g}{h} \left[\mathbf{\Gamma}_{b, n} - \mathbf{\Gamma}_{b, n-1}\right] + (1-g)
    \dot{\mathbf{\Gamma}}_{b, n-1} - \dot{\mathbf{\Gamma}}_{b, n}`

    :param gamma_b_nm1: Bound circulation vector at timestep n-1.
    :param gamma_b_n: Bound circulation vector at timestep n.
    :param gamma_b_dot_nm1: Bound circulation time derivative vector at timestep n-1.
    :param gamma_b_dot_n: Bound circulation time derivative vector at timestep n.
    :param dv: Design variables. Whilst this function does not depend upon it, including it simplified obtaining
    the residual design gradient.
    :return: Bound circulation time derivative residual.
    """
    del dv  # intentionally unused

    gamma_b_nm1 = ArrayList.from_vector(
        vect=gamma_b_nm1,
        shapes=ArrayListShape([(gd.m, gd.n) for gd in self.grid_disc]),
    )

    gamma_b_n = ArrayList.from_vector(
        vect=gamma_b_n,
        shapes=ArrayListShape([(gd.m, gd.n) for gd in self.grid_disc]),
    )

    gamma_b_dot_nm1 = ArrayList.from_vector(
        vect=gamma_b_dot_nm1,
        shapes=ArrayListShape([(gd.m, gd.n) for gd in self.grid_disc]),
    )

    gamma_b_dot_n = ArrayList.from_vector(
        vect=gamma_b_dot_n,
        shapes=ArrayListShape([(gd.m, gd.n) for gd in self.grid_disc]),
    )

    return (
        self.gamma_dot_relaxation / self.dt * (gamma_b_n - gamma_b_nm1)
        + (1.0 - self.gamma_dot_relaxation) * gamma_b_dot_nm1
        - gamma_b_dot_n
    ).ravel()

f_aero_res_func

f_aero_res_func(
    i_ts: int | Array,
    t_n: Array,
    varphi_n: Array,
    v_n: Array,
    gamma_b_n: Array,
    gamma_w_n: Array,
    gamma_b_dot_n: Array,
    zeta_w_n: Array,
    dv: AeroelasticDesignVariables,
    dv_full: AeroelasticDesignVariables,
    struct_obj: BeamStructure,
    f_aero_beam_n: Array,
    block_grid_gradients: bool,
    solve_dofs: tuple[int, ...],
) -> Array

Aerodynamic forcing residual, compared in the local frame for compatibility with the structure.

:math:\boldsymbol{\mathcal{F}}_{\text{aero}, n}(\mathbf{\Gamma}_{b, n}, \mathbf{\Gamma}_{w, n}, \dot{\mathbf{\Gamma}}_{b, n}, \boldsymbol{\zeta}_{b, n}, \dot{\boldsymbol{\zeta}}_{b, n}, \boldsymbol{\zeta}_{w, n}) - \mathbf{f}_{\text{aero}, n}

Parameters:

Name	Type	Description	Default
i_ts	int | Array	Time step index.	required
t_n	Array	Time at timestep n.	required
varphi_n	Array	Beam minimal coordinates vector at timestep n.	required
v_n	Array	Beam velocity vector at timestep n.	required
gamma_b_n	Array	Bound circulation vector at timestep n.	required
gamma_w_n	Array	Wake circulation vector at timestep n.	required
gamma_b_dot_n	Array	Bound circulation vector time derivative at timestep n.	required
zeta_w_n	Array	Wake grid vector at timestep n.	required
dv	AeroelasticDesignVariables	Aeroelastic design variables.	required
dv_full	AeroelasticDesignVariables	Aeroelastic design variables without omissions.	required
struct_obj	BeamStructure	Beam structure.	required
f_aero_beam_n	Array	Local aerodynamic forcing projected onto beam.	required
block_grid_gradients	bool	If true, blocks the gradient path for the dependency of the steady aerodynamic forcing on the bound grid.	required
solve_dofs	tuple[int, ...]	Index of forces to keep for solution.	required

Source code in src/flapjax/aero/uvlm.py

def f_aero_res_func(
    self,
    i_ts: int | Array,
    t_n: Array,
    varphi_n: Array,
    v_n: Array,
    gamma_b_n: Array,
    gamma_w_n: Array,
    gamma_b_dot_n: Array,
    zeta_w_n: Array,
    dv: AeroelasticDesignVariables,
    dv_full: AeroelasticDesignVariables,
    struct_obj: BeamStructure,
    f_aero_beam_n: Array,
    block_grid_gradients: bool,
    solve_dofs: tuple[int, ...],
) -> Array:
    r"""
    Aerodynamic forcing residual, compared in the local frame for compatibility with the structure.

    :math:`\boldsymbol{\mathcal{F}}_{\text{aero}, n}(\mathbf{\Gamma}_{b, n}, \mathbf{\Gamma}_{w, n},
    \dot{\mathbf{\Gamma}}_{b, n}, \boldsymbol{\zeta}_{b, n}, \dot{\boldsymbol{\zeta}}_{b, n},
    \boldsymbol{\zeta}_{w, n}) - \mathbf{f}_{\text{aero}, n}`

    :param i_ts: Time step index.
    :param t_n: Time at timestep n.
    :param varphi_n: Beam minimal coordinates vector at timestep n.
    :param v_n: Beam velocity vector at timestep n.
    :param gamma_b_n: Bound circulation vector at timestep n.
    :param gamma_w_n: Wake circulation vector at timestep n.
    :param gamma_b_dot_n: Bound circulation vector time derivative at timestep n.
    :param zeta_w_n: Wake grid vector at timestep n.
    :param dv: Aeroelastic design variables.
    :param dv_full: Aeroelastic design variables without omissions.
    :param struct_obj: Beam structure.
    :param f_aero_beam_n: Local aerodynamic forcing projected onto beam.
    :param block_grid_gradients: If true, blocks the gradient path for the dependency of the steady aerodynamic
    forcing on the bound grid.
    :param solve_dofs: Index of forces to keep for solution.
    """

    varphi_n = varphi_n.reshape(-1, 6)
    v_n = v_n.reshape(-1, 6)

    gamma_b_n = ArrayList.from_vector(
        vect=gamma_b_n,
        shapes=ArrayListShape([(gd.m, gd.n) for gd in self.grid_disc]),
    )

    gamma_w_n = ArrayList.from_vector(
        vect=gamma_w_n,
        shapes=ArrayListShape([(gd.m_star, gd.n) for gd in self.grid_disc]),
    )

    gamma_b_dot_n = ArrayList.from_vector(
        vect=gamma_b_dot_n,
        shapes=ArrayListShape([(gd.m, gd.n) for gd in self.grid_disc]),
    )

    zeta_w_n = ArrayList.from_vector(
        vect=zeta_w_n,
        shapes=ArrayListShape(
            [(gd.m_star + 1, gd.n + 1, 3) for gd in self.grid_disc]
        ),
    )

    f_aero_beam_n = f_aero_beam_n.reshape(-1, 6)

    inner_struct = struct_obj.case_from_dv(dv=dv.structure)
    hg_n = inner_struct.compute_hg_from_varphi(varphi=varphi_n)
    hg_dot_n = inner_struct.make_hg_dot(hg=hg_n, v=v_n)

    inner_case = self.case_from_dv(dv=dv.aero)

    # create grid
    cs_ang_n, cs_vel_n = dv.aero.get_cs_n(i_ts=i_ts, dv_full=dv_full.aero)

    zeta_b_n = inner_case.hg_to_zeta_b(hg_n=hg_n, cs_ang_n=cs_ang_n)
    if block_grid_gradients:
        zeta_b_n = jax.lax.stop_gradient(zeta_b_n)
        zeta_w_n = jax.lax.stop_gradient(zeta_w_n)

    zeta_b_dot_n = inner_case.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
    )
    nc_n = compute_nc(zetas=zeta_b_n)

    def v_total_func(x_: Array) -> Array:
        return inner_case.flowfield.vmap_call(x=x_, t=t_n) + compute_v_ind(
            cs=x_,
            zetas=ArrayList([*zeta_b_n, *zeta_w_n]),
            gammas=ArrayList([*gamma_b_n, *gamma_w_n]),
            kernels=[*inner_case.kernels_b, *inner_case.kernels_w],
            batch_size=self.batch_size,
            mirror_normal=inner_case.mirror_normal,
            mirror_point=inner_case.mirror_point,
        )

    f_steady = compute_steady_forcing(
        zeta_b=zeta_b_n,
        zeta_dot_b=zeta_b_dot_n,
        gamma_b=gamma_b_n,
        gamma_w=gamma_w_n,
        rho=inner_case.flowfield.rho,
        v_func=v_total_func,
        v_inputs=None,
        mirror_point=inner_case.mirror_point,
        mirror_normal=inner_case.mirror_normal,
        mirror_edge_low=inner_case.mirror_edge_low,
        mirror_edge_high=inner_case.mirror_edge_high,
    )

    if self.include_unsteady_force:
        f_unsteady: ArrayList = ArrayList(
            [
                split_to_vertex(
                    inner_case.flowfield.rho
                    * gamma_b_dot_n[i_surf][..., None]
                    * nc_n[i_surf],
                    (0, 1),
                )
                for i_surf in range(inner_case.n_surf)
            ]
        )
        f_tot = f_steady + f_unsteady
    else:
        f_tot = f_steady

    # project forcing to beam (global frame)
    f_tot_beam_global = project_forcing_to_beam(
        f_total=f_tot,
        rmat=hg_n[:, :3, :3],
        dof_mapping=inner_case.dof_mapping,
        x0_aero=inner_case.zeta_b0,
        mirror_edge_low=inner_case.mirror_edge_low,
        mirror_edge_high=inner_case.mirror_edge_high,
    )

    # transform to local frame to match f_aero_beam_n
    f_tot_beam_local = transform_nodal_vect(
        vect=f_tot_beam_global, rmat=jnp.swapaxes(hg_n[:, :3, :3], -2, -1)
    )

    return (f_tot_beam_local - f_aero_beam_n).ravel()[jnp.array(solve_dofs)]

timestep_residual

timestep_residual(
    i_ts: int | Array,
    varphi_nm1: Array,
    varphi_n: Array,
    v_n: Array,
    t_n: Array,
    q_n: AeroFullStates,
    q_nm1: AeroFullStates,
    dv: AeroelasticDesignVariables,
    dv_full: AeroelasticDesignVariables,
    f_aero_beam_n: Array,
    struct_obj: BeamStructure,
    approx_grads: bool,
) -> Array

Compute the residual vector to the UVLM equations. These are given as:

:math:\left(\boldsymbol{\mathcal{A}}_{b, n} \cdot \mathbf{n}_n\right)^{-1} \left[\left( \boldsymbol{\mathcal{A}}_{w, n} \boldsymbol{\Gamma}_{w, n} +\mathbf{v}_{bc, n} - \dot{\mathbf{c}}\right) \cdot \mathbf{n}_n\right] + \boldsymbol{\Gamma}_{b, n}

:math:\boldsymbol{\mathcal{W}}_{\Gamma}(\mathbf{\Gamma}_{b, {n-1}}, \mathbf{\Gamma}_{w, {n-1}}) - \mathbf{\Gamma}_{w, n}

:math:\frac{g}{h} \left[\mathbf{\Gamma}_{b, n} - \mathbf{\Gamma}_{b, n-1}\right] + (1-g) \dot{\mathbf{\Gamma}}_{b, n-1} - \dot{\mathbf{\Gamma}}_{b, n}

:math:\boldsymbol{\mathcal{W}}_{\zeta}(\boldsymbol{\zeta}_{b, n}, \boldsymbol{\zeta}_{w, n-1}) - \boldsymbol{\zeta}_{w, n}

:math:\boldsymbol{\mathcal{F}}_{\text{aero}, n}(\mathbf{\Gamma}_{b, n}, \mathbf{\Gamma}_{w, n}, \dot{\mathbf{\Gamma}}_{b, n}, \boldsymbol{\zeta}_{b, n}, \dot{\boldsymbol{\zeta}}_{b, n}, \boldsymbol{\zeta}_{w, n}) - \mathbf{f}_{\text{aero}, n}

Parameters:

Name	Type	Description	Default
i_ts	int | Array	Time step index.	required
varphi_nm1	Array	Beam minimal coordinates at timestep n-1, (n_nodes, 6).	required
varphi_n	Array	Beam maximal coordinates at timestep n, (n_nodes, 6).	required
v_n	Array	Beam velocity at timestep n, (n_nodes, 6).	required
t_n	Array	Time at step n.	required
q_n	AeroFullStates	Aero minimal states at timestep n.	required
q_nm1	AeroFullStates	Aero minimal states at timestep n-1.	required
dv	AeroelasticDesignVariables	Aeroelastic design variables.	required
dv_full	AeroelasticDesignVariables	Aeroelastic design variables without omissions.	required
f_aero_beam_n	Array	Aerodynamic forcing for the beam at timestep n, in the local frame.	required
struct_obj	BeamStructure	Beam structure.	required
approx_grads	bool	If true, eliminate grid gradients from the aerodynamic force residual.	required

Returns:

Type	Description
Array	Residual vector.

Source code in src/flapjax/aero/uvlm.py

def timestep_residual(
    self,
    i_ts: int | Array,
    varphi_nm1: Array,
    varphi_n: Array,
    v_n: Array,
    t_n: Array,
    q_n: AeroFullStates,
    q_nm1: AeroFullStates,
    dv: AeroelasticDesignVariables,
    dv_full: AeroelasticDesignVariables,
    f_aero_beam_n: Array,
    struct_obj: BeamStructure,
    approx_grads: bool,
) -> Array:
    r"""
    Compute the residual vector to the UVLM equations. These are given as:

    :math:`\left(\boldsymbol{\mathcal{A}}_{b, n} \cdot \mathbf{n}_n\right)^{-1} \left[\left(
    \boldsymbol{\mathcal{A}}_{w, n} \boldsymbol{\Gamma}_{w, n} +\mathbf{v}_{bc, n} - \dot{\mathbf{c}}\right)
    \cdot \mathbf{n}_n\right] + \boldsymbol{\Gamma}_{b, n}`

    :math:`\boldsymbol{\mathcal{W}}_{\Gamma}(\mathbf{\Gamma}_{b, {n-1}}, \mathbf{\Gamma}_{w, {n-1}})
    - \mathbf{\Gamma}_{w, n}`

    :math:`\frac{g}{h} \left[\mathbf{\Gamma}_{b, n} - \mathbf{\Gamma}_{b, n-1}\right] + (1-g)
    \dot{\mathbf{\Gamma}}_{b, n-1} - \dot{\mathbf{\Gamma}}_{b, n}`

    :math:`\boldsymbol{\mathcal{W}}_{\zeta}(\boldsymbol{\zeta}_{b, n}, \boldsymbol{\zeta}_{w, n-1})
    - \boldsymbol{\zeta}_{w, n}`

    :math:`\boldsymbol{\mathcal{F}}_{\text{aero}, n}(\mathbf{\Gamma}_{b, n}, \mathbf{\Gamma}_{w, n},
    \dot{\mathbf{\Gamma}}_{b, n}, \boldsymbol{\zeta}_{b, n}, \dot{\boldsymbol{\zeta}}_{b, n},
    \boldsymbol{\zeta}_{w, n}) - \mathbf{f}_{\text{aero}, n}`

    :param i_ts: Time step index.
    :param varphi_nm1: Beam minimal coordinates at timestep n-1, ``(n_nodes, 6)``.
    :param varphi_n: Beam maximal coordinates at timestep n, ``(n_nodes, 6)``.
    :param v_n: Beam velocity at timestep n, ``(n_nodes, 6)``.
    :param t_n: Time at step n.
    :param q_n: Aero minimal states at timestep n.
    :param q_nm1: Aero minimal states at timestep n-1.
    :param dv: Aeroelastic design variables.
    :param dv_full: Aeroelastic design variables without omissions.
    :param f_aero_beam_n: Aerodynamic forcing for the beam at timestep n, in the local frame.
    :param struct_obj: Beam structure.
    :param approx_grads: If true, eliminate grid gradients from the aerodynamic force residual.
    :return: Residual vector.
    """

    zeta_w_res, gamma_w_res = self.wake_prop_res_func(
        i_ts=i_ts,
        varphi_n=varphi_n.ravel(),
        varphi_nm1=varphi_nm1.ravel(),
        t_n=t_n,
        dv=dv,
        dv_full=dv_full,
        gamma_b_nm1=q_nm1.gamma_b.ravel(),
        gamma_w_nm1=q_nm1.gamma_w.ravel(),
        gamma_w_n=q_n.gamma_w.ravel(),
        zeta_w_nm1=q_nm1.zeta_w.ravel(),
        zeta_w_n=q_n.zeta_w.ravel(),
        struct_obj=struct_obj,
    )

    return jnp.concatenate(
        (
            self.gamma_b_res_func(
                i_ts=i_ts,
                varphi_n=varphi_n.ravel(),
                v_n=v_n.ravel(),
                t_n=t_n,
                dv=dv,
                dv_full=dv_full,
                gamma_b_n=q_n.gamma_b.ravel(),
                gamma_w_n=q_n.gamma_w.ravel(),
                zeta_w_n=q_n.zeta_w.ravel(),
                struct_obj=struct_obj,
            ),
            gamma_w_res,
            self.gamma_b_dot_res_func(
                gamma_b_nm1=q_nm1.gamma_b.ravel(),
                gamma_b_n=q_n.gamma_b.ravel(),
                gamma_b_dot_nm1=q_nm1.gamma_b_dot.ravel(),
                gamma_b_dot_n=q_n.gamma_b_dot.ravel(),
                dv=dv,
            ),
            zeta_w_res,
            self.f_aero_res_func(
                i_ts=i_ts,
                varphi_n=varphi_n.ravel(),
                v_n=v_n.ravel(),
                t_n=t_n,
                dv=dv,
                dv_full=dv_full,
                gamma_b_n=q_n.gamma_b.ravel(),
                gamma_w_n=q_n.gamma_w.ravel(),
                gamma_b_dot_n=q_n.gamma_b_dot.ravel(),
                zeta_w_n=q_n.zeta_w.ravel(),
                f_aero_beam_n=f_aero_beam_n.ravel(),
                struct_obj=struct_obj,
                block_grid_gradients=approx_grads,
                solve_dofs=tuple(range(struct_obj.n_dof)),
            ),
        )
    )

construct_approximate_jacobians

construct_approximate_jacobians(
    aero_sol: AeroCase,
    structure_sol: StructureCase,
    struct_obj: BeamStructure,
    dv: AeroelasticDesignVariables,
    dv_full: AeroelasticDesignVariables,
    solve_dofs: tuple[int, ...],
    jacobian_approximations: AeroJacobianApproximations,
) -> dict[str, dict[str, Callable[..., Any] | None]]

Compute approximations for the aerodynamic residual Jacobians specified in the jacobian_approximations data structure. The residuals covered are the bound circulation, wake propagation, bound circulation rate, and aerodynamic forcing.

Parameters:

Name	Type	Description	Default
aero_sol	AeroCase	Aerodynamic solution from which to extract states for the initial time step.	required
structure_sol	StructureCase	Structural solution from which to extract states for the initial time step.	required
struct_obj	BeamStructure	Beam structure used by the residual functions for the kinematic transformations.	required
dv	AeroelasticDesignVariables	Aeroelastic design variables.	required
dv_full	AeroelasticDesignVariables	Aeroelastic design variables without omissions for the variables where gradients aren't requested.	required
solve_dofs	tuple[int, ...]	Active structural degrees of freedom.	required
jacobian_approximations	AeroJacobianApproximations	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/aero/uvlm.py

def construct_approximate_jacobians(
    self,
    aero_sol: AeroCase,
    structure_sol: StructureCase,
    struct_obj: BeamStructure,
    dv: AeroelasticDesignVariables,
    dv_full: AeroelasticDesignVariables,
    solve_dofs: tuple[int, ...],
    jacobian_approximations: AeroJacobianApproximations,
) -> dict[str, dict[str, Callable[..., Any] | None]]:
    r"""
    Compute approximations for the aerodynamic residual Jacobians specified in the ``jacobian_approximations`` data
    structure. The residuals covered are the bound circulation, wake propagation, bound circulation rate, and
    aerodynamic forcing.
    :param aero_sol: Aerodynamic solution from which to extract states for the initial time step.
    :param structure_sol: Structural solution from which to extract states for the initial time step.
    :param struct_obj: Beam structure used by the residual functions for the kinematic transformations.
    :param dv: Aeroelastic design variables.
    :param dv_full: Aeroelastic design variables without omissions for the variables where gradients aren't
    requested.
    :param solve_dofs: Active structural degrees of freedom.
    :param jacobian_approximations: Data structure which defines which approximations to create.
    :return: Dictionary of approximations keyed by residual name.
    """
    q_nm1 = aero_sol.get_states(0)
    q_n = aero_sol.get_states(1)
    struct_nm1 = structure_sol.get_minimal_states(0)
    struct_n = structure_sol.get_minimal_states(1)

    varphi_nm1 = struct_nm1.varphi.ravel()
    varphi_n = struct_n.varphi.ravel()
    v_n = struct_n.v.ravel()
    gamma_b_nm1 = q_nm1.gamma_b.ravel()
    gamma_b_n = q_n.gamma_b.ravel()
    gamma_w_nm1 = q_nm1.gamma_w.ravel()
    gamma_w_n = q_n.gamma_w.ravel()
    gamma_b_dot_nm1 = q_nm1.gamma_b_dot.ravel()
    gamma_b_dot_n = q_n.gamma_b_dot.ravel()
    zeta_w_nm1 = q_nm1.zeta_w.ravel()
    zeta_w_n = q_n.zeta_w.ravel()

    if struct_n.f_ext_aero is None:
        raise ValueError("Missing aerodynamic forcing states")
    f_aero_beam_n = struct_n.f_ext_aero.ravel()

    t_n = aero_sol.t[1]

    gamma_b_args = {
        "i_ts": 1,
        "t_n": t_n,
        "varphi_n": varphi_n,
        "v_n": v_n,
        "gamma_b_n": gamma_b_n,
        "gamma_w_n": gamma_w_n,
        "zeta_w_n": zeta_w_n,
        "dv": dv,
        "dv_full": dv_full,
        "struct_obj": struct_obj,
    }

    wake_args = {
        "i_ts": 1,
        "t_n": t_n,
        "varphi_nm1": varphi_nm1,
        "varphi_n": varphi_n,
        "gamma_b_nm1": gamma_b_nm1,
        "gamma_w_nm1": gamma_w_nm1,
        "gamma_w_n": gamma_w_n,
        "zeta_w_nm1": zeta_w_nm1,
        "zeta_w_n": zeta_w_n,
        "dv": dv,
        "dv_full": dv_full,
        "struct_obj": struct_obj,
    }

    gamma_b_dot_args = {
        "gamma_b_nm1": gamma_b_nm1,
        "gamma_b_n": gamma_b_n,
        "gamma_b_dot_nm1": gamma_b_dot_nm1,
        "gamma_b_dot_n": gamma_b_dot_n,
        "dv": dv,
    }

    f_aero_args = {
        "i_ts": 1,
        "t_n": t_n,
        "varphi_n": varphi_n,
        "v_n": v_n,
        "gamma_b_n": gamma_b_n,
        "gamma_w_n": gamma_w_n,
        "gamma_b_dot_n": gamma_b_dot_n,
        "zeta_w_n": zeta_w_n,
        "dv": dv,
        "dv_full": dv_full,
        "struct_obj": struct_obj,
        "f_aero_beam_n": f_aero_beam_n,
        "block_grid_gradients": True,
        "solve_dofs": solve_dofs,
    }

    res_args: dict[
        str, tuple[Callable[..., Array], dict[str, Any], Sequence[str]]
    ] = {
        "gamma_b": (
            self.gamma_b_res_func,
            gamma_b_args,
            [f.name for f in fields(GammaBApprox)],
        ),
        "gamma_w": (
            lambda **kwargs: self.wake_prop_res_func(**kwargs)[1],
            wake_args,
            [f.name for f in fields(GammaWApprox)],
        ),
        "zeta_w": (
            lambda **kwargs: self.wake_prop_res_func(**kwargs)[0],
            wake_args,
            [f.name for f in fields(ZetaWApprox)],
        ),
        "gamma_b_dot": (
            self.gamma_b_dot_res_func,
            gamma_b_dot_args,
            [f.name for f in fields(GammaBDotApprox)],
        ),
        "f_aero": (
            self.f_aero_res_func,
            f_aero_args,
            [f.name for f in fields(FAeroApprox)],
        ),
    }

    return construct_approximation(
        res_args=res_args, jacobian_approximations=jacobian_approximations
    )

timestep_residual_jacobians

timestep_residual_jacobians(
    i_ts: int | Array,
    varphi_nm1: Array,
    varphi_n: Array,
    v_n: Array,
    t_n: Array,
    q_n: AeroFullStates,
    q_nm1: AeroFullStates,
    dv: AeroelasticDesignVariables,
    dv_full: AeroelasticDesignVariables,
    f_aero_beam_n: Array,
    struct_obj: BeamStructure,
    approx_grads: bool,
    solve_dofs: tuple[int, ...],
    n_profile_loops: int | None,
    jac_options: dict[str, dict[str, Any]],
    compute_wake_gradients: bool = True,
    mode: ADMode = "reverse",
    map_batch_size: int | None = None,
) -> tuple[
    Array,
    Array,
    AeroelasticDesignVariables,
    Array,
    Array,
    dict[str, dict[str, float]] | None,
    dict[str, dict[str, float]] | None,
]

Compute the Jacobians of the aerodynamic problem.

Parameters:

Name	Type	Description	Default
i_ts	int | Array	Time step index.	required
varphi_nm1	Array	Minimal structural coordinates at timestep n-1, (n_nodes, 6).	required
varphi_n	Array	Minimal structural coordinates at timestep n, (n_nodes, 6).	required
v_n	Array	Structural velocity at timestep n, (n_nodes, 6).	required
t_n	Array	Time at timestep n.	required
q_n	AeroFullStates	Aerodynamic minimal states at timestep n.	required
q_nm1	AeroFullStates	Aerodynamic minimal states at timestep n-1.	required
dv	AeroelasticDesignVariables	Aeroelastic design variables.	required
dv_full	AeroelasticDesignVariables	Aeroelastic design variables without omissions.	required
f_aero_beam_n	Array	Aerodynamic forcing in local frame of reference at timestep n, (n_nodes, 6).	required
struct_obj	BeamStructure	Structural object.	required
approx_grads	bool	If true, eliminate grid gradients from force computation.	required
solve_dofs	tuple[int, ...]	Degrees of freedom to solve for. This removes non-active forcing entries.	required
n_profile_loops	int | None	Optional number of loops for profiling function.	required
jac_options	dict[str, dict[str, Any]]	Dictionary of optional functions which can be used to substitute Jacobian evaluations from AD.	required
compute_wake_gradients	bool	If False, skip Jacobians involving the wake for zeta_w and gamma_w.	True
mode	ADMode	AD mode for obtaining gradients, either forward or reverse.	'reverse'
map_batch_size	int | None	Batch size for vectorising Jacobian construction.	None

Returns:

Type	Description
tuple[Array, Array, AeroelasticDesignVariables, Array, Array, dict[str, dict[str, float]] | None, dict[str, dict[str, float]] | None]	Gradients of aerodynamic residual with respect to previous states, current states, and design variables. Additionally, includes Jacobians of the aerodynamic residual with respect to the structural displacement and velocity, and profiling times for compilation and run time.

Source code in src/flapjax/aero/uvlm.py

def timestep_residual_jacobians(
    self,
    i_ts: int | Array,
    varphi_nm1: Array,
    varphi_n: Array,
    v_n: Array,
    t_n: Array,
    q_n: AeroFullStates,
    q_nm1: AeroFullStates,
    dv: AeroelasticDesignVariables,
    dv_full: AeroelasticDesignVariables,
    f_aero_beam_n: Array,
    struct_obj: BeamStructure,
    approx_grads: bool,
    solve_dofs: tuple[int, ...],
    n_profile_loops: int | None,
    jac_options: dict[str, dict[str, Any]],
    compute_wake_gradients: bool = True,
    mode: ADMode = "reverse",
    map_batch_size: int | None = None,
) -> tuple[
    Array,
    Array,
    AeroelasticDesignVariables,
    Array,
    Array,
    dict[str, dict[str, float]] | None,
    dict[str, dict[str, float]] | None,
]:
    r"""
    Compute the Jacobians of the aerodynamic problem.
    :param i_ts: Time step index.
    :param varphi_nm1: Minimal structural coordinates at timestep n-1, ``(n_nodes, 6)``.
    :param varphi_n: Minimal structural coordinates at timestep n, ``(n_nodes, 6)``.
    :param v_n: Structural velocity at timestep n, ``(n_nodes, 6)``.
    :param t_n: Time at timestep n.
    :param q_n: Aerodynamic minimal states at timestep n.
    :param q_nm1: Aerodynamic minimal states at timestep n-1.
    :param dv: Aeroelastic design variables.
    :param dv_full: Aeroelastic design variables without omissions.
    :param f_aero_beam_n: Aerodynamic forcing in local frame of reference at timestep n, ``(n_nodes, 6)``.
    :param struct_obj: Structural object.
    :param approx_grads: If true, eliminate grid gradients from force computation.
    :param solve_dofs: Degrees of freedom to solve for. This removes non-active forcing entries.
    :param n_profile_loops: Optional number of loops for profiling function.
    :param jac_options: Dictionary of optional functions which can be used to substitute Jacobian evaluations from
    AD.
    :param compute_wake_gradients: If False, skip Jacobians involving the wake for ``zeta_w`` and ``gamma_w``.
    :param mode: AD mode for obtaining gradients, either ``forward`` or ``reverse``.
    :param map_batch_size: Batch size for vectorising Jacobian construction.
    :return: Gradients of aerodynamic residual with respect to previous states, current states, and design
    variables. Additionally, includes Jacobians of the aerodynamic residual with respect to the structural
    displacement and velocity, and profiling times for compilation and run time.
    """

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

    varphi_nm1 = varphi_nm1.ravel()
    varphi_n = varphi_n.ravel()
    v_n = v_n.ravel()
    gamma_b_nm1 = q_nm1.gamma_b.ravel()
    gamma_b_n = q_n.gamma_b.ravel()
    gamma_w_nm1 = q_nm1.gamma_w.ravel()
    gamma_w_n = q_n.gamma_w.ravel()
    gamma_b_dot_nm1 = q_nm1.gamma_b_dot.ravel()
    gamma_b_dot_n = q_n.gamma_b_dot.ravel()
    zeta_w_nm1 = q_nm1.zeta_w.ravel()
    zeta_w_n = q_n.zeta_w.ravel()

    if not compute_wake_gradients:
        # optionally omit wake computations
        jac_options: dict[str, dict[str, Callable[..., Array] | None]] = {
            k: dict(v) for k, v in jac_options.items()
        }

        # remove entries that involve the wake
        for res_name, wake_keys in (
            ("gamma_b", ("zeta_w_n", "gamma_w_n")),
            ("gamma_w", ("zeta_w_n", "zeta_w_nm1")),
            ("f_aero", ("zeta_w_n", "gamma_w_n")),
        ):
            for key in wake_keys:
                jac_options[res_name].pop(key, None)

    # stop AD tape where required
    gamma_b_gamma_w_n_input = (
        jax.lax.stop_gradient(gamma_w_n)
        if not compute_wake_gradients
        else gamma_w_n
    )
    gamma_b_zeta_w_n_input = (
        jax.lax.stop_gradient(zeta_w_n) if not compute_wake_gradients else zeta_w_n
    )

    d_gamma_b, compile_time["gamma_b"], run_time["gamma_b"] = jacrev_custom(
        func=self.gamma_b_res_func,
        jac_options=jac_options["gamma_b"],
        n_profile_loops=n_profile_loops,
        func_name="gamma_b",
        mode=mode,
        map_batch_size=map_batch_size,
    )(
        i_ts=i_ts,
        t_n=t_n,
        varphi_n=varphi_n,
        v_n=v_n,
        gamma_b_n=gamma_b_n,
        gamma_w_n=gamma_b_gamma_w_n_input,
        zeta_w_n=gamma_b_zeta_w_n_input,
        dv=dv,
        dv_full=dv_full,
        struct_obj=struct_obj,
    )

    if compute_wake_gradients:
        d_gamma_w, compile_time["gamma_w"], run_time["gamma_w"] = jacrev_custom(
            func=lambda **kwargs: self.wake_prop_res_func(**kwargs)[1],
            jac_options=jac_options["gamma_w"],
            n_profile_loops=n_profile_loops,
            func_name="gamma_w",
            mode=mode,
            map_batch_size=map_batch_size,
        )(
            i_ts=i_ts,
            t_n=t_n,
            varphi_nm1=varphi_nm1,
            varphi_n=varphi_n,
            gamma_b_nm1=gamma_b_nm1,
            gamma_w_nm1=gamma_w_nm1,
            gamma_w_n=gamma_w_n,
            zeta_w_nm1=zeta_w_nm1,
            zeta_w_n=zeta_w_n,
            dv=dv,
            dv_full=dv_full,
            struct_obj=struct_obj,
        )

        d_zeta_w, compile_time["zeta_w"], run_time["zeta_w"] = jacrev_custom(
            func=lambda **kwargs: self.wake_prop_res_func(**kwargs)[0],
            jac_options=jac_options["zeta_w"],
            n_profile_loops=n_profile_loops,
            func_name="zeta_w",
            mode=mode,
            map_batch_size=map_batch_size,
        )(
            i_ts=i_ts,
            t_n=t_n,
            varphi_nm1=varphi_nm1,
            varphi_n=varphi_n,
            gamma_b_nm1=gamma_b_nm1,
            gamma_w_nm1=gamma_w_nm1,
            gamma_w_n=gamma_w_n,
            zeta_w_nm1=zeta_w_nm1,
            zeta_w_n=zeta_w_n,
            dv=dv,
            dv_full=dv_full,
            struct_obj=struct_obj,
        )
    else:
        # skip computations
        d_gamma_w = {}
        d_zeta_w = {}

    d_gamma_b_dot, compile_time["gamma_b_dot"], run_time["gamma_b_dot"] = (
        jacrev_custom(
            func=self.gamma_b_dot_res_func,
            jac_options=jac_options["gamma_b_dot"],
            n_profile_loops=n_profile_loops,
            func_name="gamma_b_dot",
            mode=mode,
            map_batch_size=map_batch_size,
        )(
            gamma_b_nm1=gamma_b_nm1,
            gamma_b_n=gamma_b_n,
            gamma_b_dot_nm1=gamma_b_dot_nm1,
            gamma_b_dot_n=gamma_b_dot_n,
            dv=dv,
        )
    )

    f_aero_gamma_w_n_input = (
        jax.lax.stop_gradient(gamma_w_n)
        if not compute_wake_gradients
        else gamma_w_n
    )
    f_aero_zeta_w_n_input = (
        jax.lax.stop_gradient(zeta_w_n) if not compute_wake_gradients else zeta_w_n
    )

    d_f_aero, compile_time["f_aero"], run_time["f_aero"] = jacrev_custom(
        func=self.f_aero_res_func,
        jac_options=jac_options["f_aero"],
        n_profile_loops=n_profile_loops,
        func_name="f_aero",
        static_argnames=("block_grid_gradients", "solve_dofs"),
        mode=mode,
        map_batch_size=map_batch_size,
    )(
        i_ts=i_ts,
        t_n=t_n,
        varphi_n=varphi_n,
        v_n=v_n,
        gamma_b_n=gamma_b_n,
        gamma_w_n=f_aero_gamma_w_n_input,
        gamma_b_dot_n=gamma_b_dot_n,
        zeta_w_n=f_aero_zeta_w_n_input,
        dv=dv,
        dv_full=dv_full,
        struct_obj=struct_obj,
        f_aero_beam_n=f_aero_beam_n.ravel(),
        block_grid_gradients=approx_grads,
        solve_dofs=solve_dofs,
    )
    # slice f_aero Jacobian to get forcing only on active degrees of freedom
    d_f_aero["f_aero_beam_n"] = d_f_aero["f_aero_beam_n"][:, jnp.array(solve_dofs)]

    # Jacobians block widths and heights, assembled in degree of freedom order
    n_solve_dof = len(solve_dofs)
    aero_entries_list: list[dict[str, Any]] = [d_gamma_b]
    aero_heights_list: list[int] = [gamma_b_n.size]
    aero_n_keys_list: list[str] = ["gamma_b_n"]
    aero_nm1_keys_list: list[str] = ["gamma_b_nm1"]
    if compute_wake_gradients:
        aero_entries_list.append(d_gamma_w)
        aero_heights_list.append(gamma_w_n.size)
        aero_n_keys_list.append("gamma_w_n")
        aero_nm1_keys_list.append("gamma_w_nm1")
    aero_entries_list.append(d_gamma_b_dot)
    aero_heights_list.append(gamma_b_n.size)
    aero_n_keys_list.append("gamma_b_dot_n")
    aero_nm1_keys_list.append("gamma_b_dot_nm1")

    if compute_wake_gradients:
        aero_entries_list.append(d_zeta_w)
        aero_heights_list.append(zeta_w_n.size)
        aero_n_keys_list.append("zeta_w_n")
        aero_nm1_keys_list.append("zeta_w_nm1")
    aero_entries_list.append(d_f_aero)
    aero_heights_list.append(n_solve_dof)
    aero_n_keys_list.append("f_aero_beam_n")
    aero_nm1_keys_list.append("f_aero_beam_nm1")

    aero_entries = tuple(aero_entries_list)
    aero_heights: tuple[int, ...] = tuple(aero_heights_list)
    aero_n_keys = tuple(aero_n_keys_list)
    aero_nm1_keys = tuple(aero_nm1_keys_list)
    aero_widths = aero_heights

    struct_sizes = (
        struct_obj.n_dof,
        struct_obj.n_dof,
        struct_obj.n_dof,
        struct_obj.n_dof,
    )

    d_aero_res_d_q_nm1 = construct_named_block_jacobian(
        entries=aero_entries,
        keys=aero_nm1_keys,
        widths=aero_widths,
        heights=aero_heights,
    )

    d_aero_res_d_q_n = construct_named_block_jacobian(
        entries=aero_entries,
        keys=aero_n_keys,
        widths=aero_widths,
        heights=aero_heights,
    )

    # residual of aero problem w.r.t. structural states
    d_struct_res_d_q_nm1 = construct_named_block_jacobian(
        entries=aero_entries,
        keys=("varphi_nm1", "v_nm1", "v_dot_nm1", "a_nm1"),
        widths=struct_sizes,
        heights=aero_heights,
    )

    d_struct_res_d_q_n = construct_named_block_jacobian(
        entries=aero_entries,
        keys=("varphi_n", "v_n", "v_dot_n", "a_n"),
        widths=struct_sizes,
        heights=aero_heights,
    )

    # handle design gradients: include a row per included residual, mirroring
    # the aero block layout above.
    dv_rows: list[Any] = [d_gamma_b["dv"]]
    if compute_wake_gradients:
        dv_rows.append(d_gamma_w["dv"])
    dv_rows.append(d_gamma_b_dot["dv"])
    if compute_wake_gradients:
        dv_rows.append(d_zeta_w["dv"])
    dv_rows.append(d_f_aero["dv"])

    from flapjax.coupled.data_structures import AeroelasticDesignVariables

    d_res_d_dv = AeroelasticDesignVariables.concatenate(*dv_rows)

    return (
        d_aero_res_d_q_nm1,
        d_aero_res_d_q_n,
        d_res_d_dv,
        d_struct_res_d_q_nm1,
        d_struct_res_d_q_n,
        compile_time if n_profile_loops is not None else None,
        run_time if n_profile_loops is not None else None,
    )