Utils

flapjax.utils

data_structures

ConvergenceStatus

ConvergenceStatus(
    convergence_settings: ConvergenceSettings,
)

Object to track convergence status of an iterative solver based on absolute and relative tolerances. Absolute convergence is measured for the delta vector, while relative convergence is measured as the ratio of the maximum element in the delta vector to the total vector.

Parameters:

Name	Type	Description	Default
convergence_settings	ConvergenceSettings	Convergence settings object containing tolerances and maximum iteration count for convergence failure.	required

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

def __init__(self, convergence_settings: ConvergenceSettings):
    r"""
    Object to track convergence status of an iterative solver based on absolute and relative tolerances. Absolute
    convergence is measured for the delta vector, while relative convergence is measured as the ratio of the maximum
    element in the delta vector to the total vector.
    :param convergence_settings: Convergence settings object containing tolerances and maximum iteration count for
    convergence failure.
    """

    # make sure settings allow for loop to be broken
    if (
        convergence_settings.rel_disp_tol is None
        and convergence_settings.abs_disp_tol is None
        and convergence_settings.rel_force_tol is None
        and convergence_settings.abs_force_tol is None
    ):
        if convergence_settings.max_n_iter is None:
            raise ValueError(
                "No convergence criteria provided, at least one tolerance or maximum iteration count "
                "must be specified."
            )
        warn(
            "No convergence tolerances provided, will iterate until maximum iteration counter."
        )

    # base parameters
    self.i_iter: Array = jnp.zeros((), dtype=int)
    self.convergence_settings: ConvergenceSettings = convergence_settings

    # store residual values
    self.rel_disp_val: Array = jnp.zeros(())
    self.abs_disp_val: Array = jnp.zeros(())
    self.rel_force_val: Array = jnp.zeros(())
    self.abs_force_val: Array = jnp.zeros(())

    # convergence status
    self.converged: Array = jnp.zeros((), dtype=bool)
    self.converged_abs_disp: Array = jnp.zeros((), dtype=bool)
    self.converged_rel_disp: Array = jnp.zeros((), dtype=bool)
    self.converged_rel_force: Array = jnp.zeros((), dtype=bool)
    self.converged_abs_force: Array = jnp.zeros((), dtype=bool)

    # flags for other convergence failure modes
    self.final_iter: Array = jnp.zeros((), dtype=bool)
    self.has_nan: Array = jnp.zeros((), dtype=bool)

update

update(
    delta_disp: Array | None,
    total_disp: Array | None,
    delta_force: Array | None,
    total_force: Array | None,
) -> None

Parameters:

Name	Type	Description	Default
delta_disp	Array | None	Difference in displacement vector between current and previous iteration.	required
total_disp	Array | None	Total displacement for time step, used for relative convergence calculation, typically the current solution vector.	required
delta_force	Array | None	Residual force vector, used for force convergence calculation.	required
total_force	Array | None	Total force vector for time step, used for relative convergence calculation.	required

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

def update(
    self,
    delta_disp: Array | None,
    total_disp: Array | None,
    delta_force: Array | None,
    total_force: Array | None,
) -> None:
    r"""
    :param delta_disp: Difference in displacement vector between current and previous iteration.
    :param total_disp: Total displacement for time step, used for relative convergence calculation, typically the
    current solution vector.
    :param delta_force: Residual force vector, used for force convergence calculation.
    :param total_force: Total force vector for time step, used for relative convergence calculation.
    """
    # update iteration counter
    self.i_iter += 1

    # check absolute displacement convergence
    if delta_disp is not None:
        self.abs_disp_val = jnp.linalg.norm(delta_disp)

        # NaNs are checked for with displacement magnitude
        self.has_nan = jnp.isnan(delta_disp).any()

    if self.convergence_settings.abs_disp_tol is not None:
        self.converged_abs_disp = (
            self.abs_disp_val < self.convergence_settings.abs_disp_tol
        )

    # check relative displacement convergence:
    if self.convergence_settings.rel_disp_tol is not None:
        if total_disp is None:
            raise ValueError("total_disp cannot be None")
        max_total_elem = jnp.linalg.norm(total_disp)
        self.rel_disp_val = self.abs_disp_val / max_total_elem
        self.converged_rel_disp = (
            jnp.nan_to_num(self.rel_disp_val, True, jnp.inf)
            < self.convergence_settings.rel_disp_tol
        )

    # check absolute force convergence
    if self.convergence_settings.abs_force_tol is not None:
        if delta_force is None:
            raise ValueError("delta_force cannot be None")
        self.abs_force_val = jnp.linalg.norm(delta_force)
        self.converged_abs_force = (
            self.abs_force_val < self.convergence_settings.abs_force_tol
        )

    # check relative force convergence:
    if self.convergence_settings.rel_force_tol is not None:
        if total_force is None:
            raise ValueError("total_force cannot be None")
        max_total_elem = jnp.abs(total_force).max()
        self.rel_force_val = self.abs_force_val / max_total_elem
        self.converged_rel_force = (
            jnp.nan_to_num(self.rel_force_val, True, jnp.inf)
            < self.convergence_settings.rel_force_tol
        )

    # find convergence status of numerics (excluding failure modes such as max iterations or nans)
    self.converged = (
        self.converged_rel_disp
        | self.converged_abs_disp
        | self.converged_rel_force
        | self.converged_abs_force
    ) & (self.i_iter > 0)

    # check for failure modes
    if self.convergence_settings.max_n_iter is not None:
        self.final_iter = self.i_iter >= self.convergence_settings.max_n_iter

get_status

get_status() -> Array

Get overall convergence status.

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

def get_status(self) -> Array:
    """Get overall convergence status."""
    return self.converged | self.has_nan | self.final_iter

reset_status

reset_status() -> None

Reset convergence status for next load step, setting all convergence flags to False.

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

def reset_status(self) -> None:
    r"""
    Reset convergence status for next load step, setting all convergence flags to False.
    """
    false_ = jnp.zeros((), dtype=bool)
    zero_ = jnp.zeros(())
    self.converged = false_
    self.converged_abs_disp = false_
    self.converged_rel_disp = false_
    self.converged_rel_force = false_
    self.converged_abs_force = false_
    self.rel_disp_val = zero_
    self.abs_disp_val = zero_
    self.rel_force_val = zero_
    self.abs_force_val = zero_
    self.has_nan = false_
    self.final_iter = false_
    self.i_iter = jnp.zeros((), dtype=int)

print_struct_message

print_struct_message(
    i_ts: int | None, t: Array | None, i_load_step: int
) -> None

Print convergence message for structure based on status.

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

def print_struct_message(
    self, i_ts: int | None, t: Array | None, i_load_step: int
) -> None:
    """Print convergence message for structure based on status."""
    self._print_convergence_message("Struct", i_ts, t, i_load_step)

print_fsi_message

print_fsi_message(
    i_ts: int | None, t: Array | None
) -> None

Print convergence message for FSI based on status.

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

def print_fsi_message(self, i_ts: int | None, t: Array | None) -> None:
    """Print convergence message for FSI based on status."""
    self._print_convergence_message("FSI", i_ts, t, None)

linear

LinearModel

LinearModel(reference: R, dt: float | Array)

Bases: ABC

Base class to represent a linearised system.

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

def __init__(self, reference: R, dt: float | Array):
    # slices of individual surface components in full vector
    self.input_slices, self.n_inputs = self._make_input_slices(reference=reference)
    self.state_slices, self.n_states = self._make_state_slices(reference=reference)
    self.output_slices, self.n_outputs = self._make_output_slices(
        reference=reference
    )
    self.dt: Array = jnp.array(dt)

    self._reference: R = reference

    self._reference_inputs: dict[str, Array | ArrayList | None] = (
        self.extract_reference_inputs()
    )
    self._reference_states: dict[str, Array | ArrayList | None] = (
        self.extract_reference_states()
    )
    self._reference_outputs: dict[str, Array | ArrayList | None] = (
        self.extract_reference_outputs()
    )

    self._sys: LinearSystem | None = None

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)))

LinearSystem

LinearSystem(
    a: Array,
    b: Array,
    c: Array,
    d: Array,
    dt: float | Array,
    continuous_time: bool = False,
    removed_u_np1: bool = False,
)

Linear system represented in state-space form, with tools for time-stepping

Initialise the LinearSystem with state-space linear operators.

Parameters:

Name	Type	Description	Default
a	Array	System matrix A	required
b	Array	Input matrix B	required
c	Array	Output matrix C	required
d	Array	Feedthrough matrix D	required
removed_u_np1	bool	If true, indicates that the system is in terms of inputs at time step n only.	False

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

def __init__(
    self,
    a: Array,
    b: Array,
    c: Array,
    d: Array,
    dt: float | Array,
    continuous_time: bool = False,
    removed_u_np1: bool = False,
) -> None:
    r"""
    Initialise the LinearSystem with state-space linear operators.
    :param a: System matrix A
    :param b: Input matrix B
    :param c: Output matrix C
    :param d: Feedthrough matrix D
    :param removed_u_np1: If true, indicates that the system is in terms of inputs at time step n only.
    """
    self.a: Array = a
    self.b: Array = b
    self.c: Array = c
    self.d: Array = d
    self.dt: float | Array = dt
    self.n_inputs: int = b.shape[1]
    self.n_states: int = a.shape[0]
    self.n_outputs: int = c.shape[0]
    self.continuous_time: bool = continuous_time
    self.removed_u_np1: bool = removed_u_np1

remove_u_np1

remove_u_np1() -> None

Remove the dependence on u at time varphi+1 from the linear system, modifying b and d accordingly. :math:D_{new} = C B + D and :math:B_{new} = A B

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

def remove_u_np1(self) -> None:
    r"""
    Remove the dependence on u at time varphi+1 from the linear system, modifying b and d accordingly.
    :math:`D_{new} = C B + D` and :math:`B_{new} = A B`
    """
    if self.removed_u_np1:
        warn("u_np1 has already been removed from the system. Skipping.")
    else:
        self.d = (self.c @ self.b) + self.d
        self.b = self.a @ self.b
        self.removed_u_np1 = True

continuous_to_discrete

continuous_to_discrete(
    method: Literal["zoh", "tustin"] = "tustin",
) -> LinearSystem

Convert the continuous-time linear system to a discrete-time linear system.

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

def continuous_to_discrete(
    self, method: Literal["zoh", "tustin"] = "tustin"
) -> LinearSystem:
    r"""
    Convert the continuous-time linear system to a discrete-time linear system.
    """
    jax_print(
        "Converting continuous-time system to discrete-time system.",
        verbose_level="normal",
    )
    if not self.continuous_time:
        warn("System is already discrete-time.")
        return self

    a_c = self.a
    b_c = self.b

    match method:
        case "zoh":
            # perform a ZOH conversion
            # create a matrix aug = [[A, B], [0, 0]], exp(aug * dt) = [[Ad, Bd], [0, I]]
            aug = jnp.block(
                [
                    [a_c, b_c],
                    [
                        jnp.zeros(
                            (self.b.shape[1], self.a.shape[0] + self.b.shape[1])
                        )
                    ],
                ]
            )
            assert aug.shape[0] == aug.shape[1], (
                "Augmented matrix must be square for matrix exponential."
            )

            exp_aug = jsp.linalg.expm(self.dt * aug)

            a_d = exp_aug[: self.a.shape[0], : self.a.shape[1]]
            b_d = exp_aug[: self.a.shape[0], self.a.shape[1] :]
        case "tustin":
            # Tustin bilinear transformation
            mat_inv = jnp.linalg.inv(jnp.eye(self.a.shape[0]) - a_c * 0.5 * self.dt)
            a_d = mat_inv @ (jnp.eye(self.a.shape[0]) + a_c * 0.5 * self.dt)
            b_d = mat_inv @ (b_c * self.dt)

    return LinearSystem(
        a=a_d, b=b_d, c=self.c, d=self.d, dt=self.dt, continuous_time=False
    )

run

run(
    u: Array, x0: Array | None = None
) -> tuple[Array, Array]

Run the linear system for a time history of input vector u.

Parameters:

Name	Type	Description	Default
u	Array	Time history of input vectors, shape (n_tstep, n_inputs)	required
x0	Array | None	Initial state vector, (n_states, ). If None, assumed to be zero.	None

Returns:

Type	Description
tuple[Array, Array]	State history and output history, (n_tstep, n_states) and (n_tstep, n_outputs)

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

def run(self, u: Array, x0: Array | None = None) -> tuple[Array, Array]:
    r"""
    Run the linear system for a time history of input vector u.
    :param u: Time history of input vectors, shape ``(n_tstep, n_inputs)``
    :param x0: Initial state vector, ``(n_states, )``. If None, assumed to be zero.
    :return: State history and output history, ``(n_tstep, n_states)`` and ``(n_tstep, n_outputs)``
    """
    if x0 is not None:
        check_arr_shape(x0, (self.n_states,), "x0")
    check_arr_shape(u, (None, self.n_inputs), "u")
    n_tstep = u.shape[0]

    if self.continuous_time:
        self.continuous_to_discrete()

    def state_func(i_ts: int, x_: Array) -> Array:
        r"""
        State update function for time step i_ts, given as :math:`x_{varphi} = A x_{varphi-1} + B u_{varphi-1}` or :math:`x_n = A x_{varphi-1} + B u_n`.
        :param i_ts: Time step index to obtain new states for.
        :param x_: State history array being updated, ``(n_tstep, n_states)``
        :return: Updated state history array, ``(n_tstep, n_states)``
        """
        jax_print(
            "Linear system state step {i_ts}",
            i_ts=i_ts,
            verbose_level="normal",
        )
        this_u = u[i_ts - 1, ...] if self.removed_u_np1 else u[i_ts, ...]
        return x_.at[i_ts, ...].set(self.a @ x_[i_ts - 1, ...] + self.b @ this_u)

    x = jnp.zeros((n_tstep, self.n_states))
    if x0 is not None:
        x = x.at[0, ...].set(x0)
    x = jax.lax.fori_loop(1, n_tstep, state_func, x)

    def output_func(i_ts: int, y_: Array) -> Array:
        r"""
        Output computation function for time step i_ts, given as :math:`y_n = C x_n + D u_n`.
        :param i_ts: Time step index to obtain outputs for.
        :param y_: Output history array being updated, ``(n_tstep, n_outputs)``
        :return: Updated output history array, ``(n_tstep, n_outputs)``
        """
        jax_print("Linear system output step {i_ts}", i_ts=i_ts)
        return y_.at[i_ts, ...].set(self.c @ x[i_ts, ...] + self.d @ u[i_ts, ...])

    y = jnp.zeros((n_tstep, self.n_outputs))
    y = jax.lax.fori_loop(0, n_tstep, output_func, y)
    return x, y

conjugate_partner_mask

conjugate_partner_mask(
    freq_hz: Array, damping: Array, tiebreaker: Array
) -> Array

Return a boolean mask identifying the second member of each complex-conjugate mode pair.

Modes are first sorted by frequency, and falls back to a tiebreaker (typically the real part of the eigenvalue) as a secondary key. Each mode whose predecessor matches in frequency and |damping| is then flagged.

Parameters:

Name	Type	Description	Default
freq_hz	Array	Natural frequency of each mode (n_modes, ).	required
damping	Array	Damping ratio of each mode (n_modes, ).	required
tiebreaker	Array	Secondary sort key used to keep conjugate partners adjacent (n_modes, ).	required

Returns:

Type	Description
Array	Boolean mask, True where the mode is the redundant partner in a conjugate pair (n_modes, ).

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

def conjugate_partner_mask(
    freq_hz: Array,
    damping: Array,
    tiebreaker: Array,
) -> Array:
    r"""
    Return a boolean mask identifying the second member of each complex-conjugate mode pair.

    Modes are first sorted by frequency, and falls back to a tiebreaker (typically the real part of the eigenvalue) as a
    secondary key. Each mode whose predecessor matches in frequency and \|damping\|
    is then flagged.
    :param freq_hz: Natural frequency of each mode ``(n_modes, )``.
    :param damping: Damping ratio of each mode ``(n_modes, )``.
    :param tiebreaker: Secondary sort key used to keep conjugate partners adjacent ``(n_modes, )``.
    :return: Boolean mask, True where the mode is the redundant partner in a conjugate pair ``(n_modes, )``.
    """
    idx = jnp.lexsort((tiebreaker, freq_hz))
    freq_sorted = freq_hz[idx]
    damp_sorted = damping[idx]
    partner_sorted = jnp.concatenate(
        [
            jnp.array([False]),
            jnp.isclose(freq_sorted[1:], freq_sorted[:-1], rtol=1e-6, atol=0.0)
            & jnp.isclose(
                jnp.abs(damp_sorted[1:]),
                jnp.abs(damp_sorted[:-1]),
                rtol=1e-6,
                atol=1e-9,
            ),
        ]
    )
    # scatter back to original ordering
    return jnp.zeros_like(partner_sorted).at[idx].set(partner_sorted)

print_utils

get_verbosity

get_verbosity() -> VerbosityLevel

Return the current verbosity level.

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

def get_verbosity() -> VerbosityLevel:
    """
    Return the current verbosity level.
    """
    return VERBOSITY_LEVEL

verbosity

verbosity(
    level: VerbosityLevel,
) -> Generator[None, None, None]

Context manager to temporarily change the verbosity level.

Parameters:

Name	Type	Description	Default
level	VerbosityLevel	Custom verbosity to use in context.	required

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

@contextmanager
def verbosity(level: VerbosityLevel) -> Generator[None, None, None]:
    r"""
    Context manager to temporarily change the verbosity level.
    :param level: Custom verbosity to use in context.
    """
    old = VERBOSITY_LEVEL
    set_verbosity(level)
    try:
        yield
    finally:
        set_verbosity(old)

utils

make_pytree

make_pytree(cls: type[T]) -> type[T]

Register a class as a JAX pytree.

When applying to classes, define _static: ClassVar[tuple[str, ...]], listing field names whose values are hashable / non-traced. Dynamic children are auto-derived from vars(self) at flatten time. Instance attributes set outside __init__ become part of the pytree structure.

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

def make_pytree[T](cls: type[T]) -> type[T]:
    """
    Register a class as a JAX pytree.

    When applying to classes, define ``_static: ClassVar[tuple[str, ...]]``, listing
    field names whose values are hashable / non-traced. Dynamic children are
    auto-derived from ``vars(self)`` at flatten time. Instance attributes set
    outside ``__init__`` become part of the pytree structure.

    """
    static_attr = getattr(cls, "_static", None)

    if isinstance(static_attr, tuple):
        static_names: tuple[str, ...] = tuple(static_attr)  # type: ignore[arg-type]
        static_set = frozenset(static_names)

        def flatten_func(self: T) -> tuple[tuple[Any, ...], tuple[Any, ...]]:
            state = vars(self)
            dynamic_names = tuple(k for k in state if k not in static_set)
            children = tuple(state[k] for k in dynamic_names)
            static_vals = tuple(state[k] for k in static_names)
            return children, (dynamic_names, static_vals)

        def unflatten_func(aux_data: tuple[Any, ...], children: tuple[Any, ...]) -> T:
            dynamic_names, static_vals = aux_data
            obj = cls.__new__(cls)
            for name, val in zip(dynamic_names, children):
                setattr(obj, name, val)
            for name, val in zip(static_names, static_vals):
                setattr(obj, name, val)
            return obj

    else:
        raise TypeError(
            f"@make_pytree on {cls.__name__}: needs a `_static: "
            f"ClassVar[tuple[str, ...]]` class attribute"
        )

    tree_util.register_pytree_node(cls, flatten_func, unflatten_func)
    return cls

dv_or

dv_or(dv_val: V | None, fallback: V) -> V

Returndv_val if it is not None, otherwise return fallback. Used to collapse the dv.x if dv.x is not None else self.x calls when combining two classes of design variables.

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

def dv_or[V](dv_val: V | None, fallback: V) -> V:
    r"""
    Return``dv_val`` if it is not None, otherwise return ``fallback``. Used to collapse the
    ``dv.x if dv.x is not None else self.x`` calls when combining two classes of design variables.
    """
    return dv_val if dv_val is not None else fallback

pytree_clone

pytree_clone(obj: V) -> V

Cheap clone of a pytree-registered object

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

def pytree_clone[V](obj: V) -> V:
    r"""
    Cheap clone of a pytree-registered object
    """
    leaves, treedef = tree_util.tree_flatten(obj)
    return tree_util.tree_unflatten(treedef, leaves)

index_to_arr

index_to_arr(
    index: int | Array | Sequence[int] | slice | None,
    n_entries: int,
) -> Array

Convert an input index to an Array index.

Parameters:

Name	Type	Description	Default
index	int | Array | Sequence[int] | slice | None	Index to apply to a sequence	required
n_entries	int	Number of entries in the full un-indexed sequence	required

Returns:

Type	Description
Array	Array index corresponding to the input index

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

def index_to_arr(
    index: int | Array | Sequence[int] | slice | None, n_entries: int
) -> Array:
    r"""
    Convert an input index to an Array index.
    :param index: Index to apply to a sequence
    :param n_entries: Number of entries in the full un-indexed sequence
    :return: Array index corresponding to the input index
    """
    if isinstance(index, slice):
        return jnp.arange(n_entries)[index]
    elif isinstance(index, Sequence):
        return jnp.array(index)
    elif isinstance(index, Array):
        return index
    elif isinstance(index, int):
        return jnp.array([index])
    elif index is None:
        return jnp.arange(n_entries)
    else:
        raise TypeError("index must be a slices, sequence of ints, or Array")

nested_list_to_tuple

nested_list_to_tuple(a: list[Any]) -> tuple[Any, ...]

Convert nested lists to nested tuples.

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

def nested_list_to_tuple(a: list[Any]) -> tuple[Any, ...]:
    r"""
    Convert nested lists to nested tuples.
    """

    def inner_func(a_):
        try:
            return tuple(nested_list_to_tuple(i) for i in a_)
        except TypeError:
            return a_

    return inner_func(a)

conditional_profile

conditional_profile(
    func: Callable[..., U],
    n_loops: int | None,
    func_name: str,
    arg_name: str,
) -> Callable[..., tuple[U, float | None, float | None]]

Function wrapper which optionally profiles it.

Parameters:

Name	Type	Description	Default
func	Callable[..., U]	Function to be profiled.	required
n_loops	int | None	Number of times to run the function, where the average time will be taken. If None, no profiling is done.	required
func_name	str	Name of the function to be profiled, used for console printing.	required
arg_name	str	Name of the argument to be profiled, used for console printing.	required

Returns:

Type	Description
Callable[..., tuple[U, float | None, float | None]]	New function which returns the same value as func, alongside the compile time and average run time of func, which are substituted for None when no profiling is done.

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

def conditional_profile[U](
    func: Callable[..., U],
    n_loops: int | None,
    func_name: str,
    arg_name: str,
) -> Callable[..., tuple[U, float | None, float | None]]:
    r"""
    Function wrapper which optionally profiles it.
    :param func: Function to be profiled.
    :param n_loops: Number of times to run the function, where the average time will be taken. If None, no profiling is
    done.
    :param func_name: Name of the function to be profiled, used for console printing.
    :param arg_name: Name of the argument to be profiled, used for console printing.
    :return: New function which returns the same value as `func`, alongside the compile time and average run time of
    `func`, which are substituted for None when no profiling is done.
    """

    if n_loops is None:
        return lambda *args, **kwargs: (func(*args, **kwargs), None, None)
    else:
        if not isinstance(n_loops, int) or n_loops <= 0:
            raise ValueError("n_loops must be a positive int")
        loops: int = n_loops

        def inner_func(*args_: Any, **a_: Any) -> tuple[U, float, float]:
            # warm-up: includes compile + first run
            start_time = time.time()
            out = jax.block_until_ready(func(*args_, **a_))
            compile_run_time = time.time() - start_time

            # timed run on the compiled function
            start_time = time.time()
            for _ in range(loops):
                jax.block_until_ready(func(*args_, **a_))

            run_time = (time.time() - start_time) / loops
            compile_time = compile_run_time - run_time

            jax_print(
                f"| Function: {func_name:<14} Argument: {arg_name:<16} Compile time: {compile_time:<8.4f} Run time: {run_time:<8.4f} |",
                verbose_level="normal",
            )

            return out, compile_time, run_time

        return inner_func