Algebra

flapjax.algebra

array_utils

ArrayList

ArrayList(arrs: Sequence[Array])

Class to hold a sequence of arrays, with overloaded arithmetic operations. This allows for more elegant handling of non-uniform arrays in various calculations.

Parameters:

Name	Type	Description	Default
arrs	Sequence[Array]	Sequence of arrays to hold.	required

Source code in src/flapjax/algebra/array_utils.py

def __init__(self, arrs: Sequence[Array]) -> None:
    self.data: list[Array] = list(arrs)

shape property

shape: ArrayListShape

Get the shapes of the arrays in the ArrayList.

Returns:

Type	Description
ArrayListShape	ArrayListShape containing the shapes of the arrays in the ArrayList.

to_list

to_list() -> list[Array]

Convert the ArrayList to a standard Python list of arrays.

Returns:

Type	Description
list[Array]	List of arrays.

Source code in src/flapjax/algebra/array_utils.py

def to_list(self) -> list[Array]:
    r"""
    Convert the ArrayList to a standard Python list of arrays.
    :return: List of arrays.
    """
    return list(self.data)

at

at(idx: int) -> Array

Get the array at the given index.

Parameters:

Name	Type	Description	Default
idx	int	Index of the array to get.	required

Returns:

Type	Description
Array	Array at the given index.

Source code in src/flapjax/algebra/array_utils.py

def at(self, idx: int) -> Array:
    r"""
    Get the array at the given index.
    :param idx: Index of the array to get.
    :return: Array at the given index.
    """
    return self.data[idx]

ravel

ravel() -> Array

Flatten the sequence of arrays into a single 1D array.

Returns:

Type	Description
Array	Flattened 1D array.

Source code in src/flapjax/algebra/array_utils.py

def ravel(self) -> Array:
    r"""
    Flatten the sequence of arrays into a single 1D array.
    :return: Flattened 1D array.
    """
    return flatten_to_1d(self.data)

from_vector classmethod

from_vector(
    vect: Array, shapes: ArrayListShape
) -> ArrayList

Unravel a 1D vector into a sequence of arrays with the given shapes.

Parameters:

Name	Type	Description	Default
vect	Array	Input 1D vector to unravel.	required
shapes	ArrayListShape	ArrayListShape containing the shapes of the arrays to unravel into.	required

Returns:

Type	Description
ArrayList	ArrayList containing the unravelled arrays.

Source code in src/flapjax/algebra/array_utils.py

@classmethod
def from_vector(cls, vect: Array, shapes: ArrayListShape) -> ArrayList:
    r"""
    Unravel a 1D vector into a sequence of arrays with the given shapes.
    :param vect: Input 1D vector to unravel.
    :param shapes: ArrayListShape containing the shapes of the arrays to unravel into.
    :return: ArrayList containing the unravelled arrays.
    """

    if vect.size != sum(shapes.sizes):
        raise ValueError(
            f"Input vector must have the same number of elements as shapes. Vector has "
            f"{vect.size} elements, but shape requires {sum(shapes.sizes)}."
        )

    arrs = []
    idx = 0
    for shape in shapes.shapes:
        size = math.prod(shape)
        arrs.append(vect[idx : idx + size].reshape(shape))
        idx += size
    return cls(arrs)

index_all

index_all(
    *idx: EllipsisType | int | slice | Array | None,
) -> ArrayList

Get the value of all arrays at the given index. This is equivalent to self[i][idx] for i in range(varphi).

Source code in src/flapjax/algebra/array_utils.py

def index_all(
    self,
    *idx: EllipsisType | int | slice | Array | None,
) -> ArrayList:
    r"""
    Get the value of all arrays at the given index. This is equivalent to `self[i][idx] for i in range(varphi)`.
    """
    return ArrayList([self.data[i][idx] for i in range(len(self.data))])

einsum staticmethod

einsum(subscript: str, *operands: ArrayList) -> ArrayList

Perform Einstein summation on sequences of arrays.

Parameters:

Name	Type	Description	Default
subscript	str	Subscript for Einstein summation. This does not include the indices for the sequence dimension.	required
operands	ArrayList	Sequences of arrays to perform Einstein summation on.	()

Returns:

Type	Description
ArrayList	Sequence of arrays resulting from Einstein summation.

Source code in src/flapjax/algebra/array_utils.py

@staticmethod
def einsum(subscript: str, *operands: ArrayList) -> ArrayList:
    r"""
    Perform Einstein summation on sequences of arrays.
    :param subscript: Subscript for Einstein summation. This does not include the indices for the sequence dimension.
    :param operands: Sequences of arrays to perform Einstein summation on.
    :return: Sequence of arrays resulting from Einstein summation.
    """
    n_arrays = len(operands[0].data)
    for op in operands:
        if len(op.data) != n_arrays:
            raise ValueError("All ArrayLists must have the same length.")

    return ArrayList(
        [
            jnp.einsum(subscript, *(op.data[i] for op in operands))
            for i in range(n_arrays)
        ]
    )

zeros_like staticmethod

zeros_like(arr: ArrayList) -> ArrayList

Create a new ArrayList with the same shape as the input_, but filled with zeros.

Parameters:

Name	Type	Description	Default
arr	ArrayList	Input ArrayList to create zeros like.	required

Returns:

Type	Description
ArrayList	New ArrayList filled with zeros.

Source code in src/flapjax/algebra/array_utils.py

@staticmethod
def zeros_like(arr: ArrayList) -> ArrayList:
    r"""
    Create a new ArrayList with the same shape as the input_, but filled with zeros.
    :param arr: Input ArrayList to create zeros like.
    :return: New ArrayList filled with zeros.
    """
    return ArrayList([jnp.zeros_like(a) for a in arr.data])

ArrayListShape

ArrayListShape(shapes: Sequence[tuple[int, ...]])

Class to hold the shapes of the arrays in an ArrayList. This is used for indexing and reshaping operations.

Source code in src/flapjax/algebra/array_utils.py

def __init__(self, shapes: Sequence[tuple[int, ...]]) -> None:
    self.shapes: Sequence[tuple[int, ...]] = shapes
    self.n_arrays: int = len(self.shapes)
    self.sizes: Sequence[int] = [prod(shape) for shape in self.shapes]

total_size

total_size() -> int

Get the total number of entries in the ArrayList.

Source code in src/flapjax/algebra/array_utils.py

def total_size(self) -> int:
    r"""
    Get the total number of entries in the ArrayList.
    """
    return sum(self.sizes)

check_arr_shape

check_arr_shape(
    arr: Array,
    expected_shape: tuple[int | None, ...],
    name: str | None,
) -> None

Asserts that the shape of the given array matches the expected shape.

Parameters:

Name	Type	Description	Default
arr	Array	Input array to check.	required
expected_shape	tuple[int | None, ...]	Expected shape of the array, as a tuple of integers, with None used for dimensions that can be of any size.	required
name	str | None	Name of the input array. This is used to provide more informative error messages.	required

Raises:

Type	Description
ValueError	If the shape of the array does not match the expected shape.

Source code in src/flapjax/algebra/array_utils.py

def check_arr_shape(
    arr: Array, expected_shape: tuple[int | None, ...], name: str | None
) -> None:
    """Asserts that the shape of the given array matches the expected shape.
    :param arr: Input array to check.
    :param expected_shape: Expected shape of the array, as a tuple of integers, with None used for dimensions that can
    be of any size.
    :param name: Name of the input array. This is used to provide more informative error messages.
    :raises ValueError: If the shape of the array does not match the expected shape.
    """
    actual_shape = arr.shape
    if len(actual_shape) == len(expected_shape):
        for i_dim in range(len(expected_shape)):
            if expected_shape[i_dim] is None:
                continue
            if actual_shape[i_dim] != expected_shape[i_dim]:
                break
        else:
            return

    message = f"Expected shape {expected_shape}, but got shape {actual_shape}."
    if name is not None:
        message += f"Issue with input '{name}'"
    raise ValueError(message)

check_arr_ndim

check_arr_ndim(
    arr: Array, expected_ndim: int, name: str | None
) -> None

Asserts that the number of dimensions of the given array matches the expected number.

Parameters:

Name	Type	Description	Default
arr	Array	Input array to check.	required
expected_ndim	int	Expected number of dimensions of the array.	required
name	str | None	Name of the input array. This is used to provide more informative error messages.	required

Raises:

Type	Description
ValueError	If the number of dimensions of the array does not match the expected value.

Source code in src/flapjax/algebra/array_utils.py

def check_arr_ndim(arr: Array, expected_ndim: int, name: str | None) -> None:
    """Asserts that the number of dimensions of the given array matches the expected number.
    :param arr: Input array to check.
    :param expected_ndim: Expected number of dimensions of the array.
    :param name: Name of the input array. This is used to provide more informative error messages.
    :raises ValueError: If the number of dimensions of the array does not match the expected value.
    """
    actual_ndim = arr.ndim
    if actual_ndim != expected_ndim:
        message = f"Expected {expected_ndim} dimensions, but got {actual_ndim}."
        if name is not None:
            message += f"Issue with input '{name}'"
        raise ValueError(message)

check_arr_dtype

check_arr_dtype(
    arr: Array, expected_dtype: type, name: str | None
) -> None

Asserts that the data type of the given array matches the expected type.

Parameters:

Name	Type	Description	Default
arr	Array	Input array to check.	required
expected_dtype	type	Expected underlying data type of the array.	required
name	str | None	Name of the input array. This is used to provide more informative error messages.	required

Raises:

Type	Description
ValueError	If the data type of the array does not match the expected type.

Source code in src/flapjax/algebra/array_utils.py

def check_arr_dtype(arr: Array, expected_dtype: type, name: str | None) -> None:
    """Asserts that the data type of the given array matches the expected type.
    :param arr: Input array to check.
    :param expected_dtype: Expected underlying data type of the array.
    :param name: Name of the input array. This is used to provide more informative error messages.
    :raises ValueError: If the data type of the array does not match the expected type.
    """

    if expected_dtype is int:
        jax_dtype = jnp.integer
    elif expected_dtype is float:
        jax_dtype = jnp.floating
    else:
        jax_dtype = expected_dtype

    actual_dtype = arr.dtype
    if not jnp.issubdtype(actual_dtype, jax_dtype):
        message = f"Expected {jax_dtype}, but got {actual_dtype}."
        if name is not None:
            message += f"Issue with input '{name}'"
        raise ValueError(message)

flatten_to_1d

flatten_to_1d(arrs: Sequence[Array]) -> Array

Convert a list of ND arrays into a single 1D vector by flattening and concatenating

Parameters:

Name	Type	Description	Default
arrs	Sequence[Array]	List of arrays to flatten and concatenate	required

Returns:

Type	Description
Array	Single 1D vector

Source code in src/flapjax/algebra/array_utils.py

def flatten_to_1d(arrs: Sequence[Array]) -> Array:
    r"""
    Convert a list of ND arrays into a single 1D vector by flattening and concatenating
    :param arrs: List of arrays to flatten and concatenate
    :return: Single 1D vector
    """
    return jnp.concatenate([arr.ravel() for arr in arrs])

block_axis

block_axis(
    arrs: Sequence[Sequence[Array]], axes: Sequence[int]
) -> Array

Form a block matrix along two given axes

Parameters:

Name	Type	Description	Default
arrs	Sequence[Sequence[Array]]	Double nested sequence of arrays ()()(n, m)	required
axes	Sequence[int]	Axes along which to concatenate the arrays	required

Returns:

Type	Description
Array	Block matrix, (n_total, m_total)

Source code in src/flapjax/algebra/array_utils.py

def block_axis(arrs: Sequence[Sequence[Array]], axes: Sequence[int]) -> Array:
    r"""
    Form a block matrix along two given axes
    :param arrs: Double nested sequence of arrays ``()()(n, m)``
    :param axes: Axes along which to concatenate the arrays
    :return: Block matrix, ``(n_total, m_total)``
    """
    # obtain the number of levels in the nested sequence
    if len(axes) != 2:
        raise ValueError("axes must be a sequence of two integers.")

    return jnp.concatenate(
        [jnp.concatenate(arrs1, axis=axes[1]) for arrs1 in arrs], axis=axes[0]
    )

neighbour_average

neighbour_average(
    arr: Array, axes: int | Sequence[int]
) -> Array

Find the pairwise average of the array along the specified axes.

Parameters:

Name	Type	Description	Default
arr	Array	Input array to average.	required
axes	int | Sequence[int]	Axis or axes along which to average.	required

Returns:

Type	Description
Array	Averaged array.

Source code in src/flapjax/algebra/array_utils.py

def neighbour_average(arr: Array, axes: int | Sequence[int]) -> Array:
    r"""
    Find the pairwise average of the array along the specified axes.
    :param arr: Input array to average.
    :param axes: Axis or axes along which to average.
    :return: Averaged array.
    """

    def _single_neighbour_average(arr_: Array, axis_: int) -> Array:
        r"""
        Average the values of the array along the specified axes, considering the neighbouring elements.
        :param arr_: Input array to average.
        :param axis_: Axis along which to average.
        :return: Averaged array.
        """
        index1: list[slice] = [slice(None, None)] * arr_.ndim
        index1[axis_] = slice(None, -1)
        index2: list[slice] = [slice(None, None)] * arr_.ndim
        index2[axis_] = slice(1, None)

        return 0.5 * (arr_[tuple(index1)] + arr_[tuple(index2)])

    if isinstance(axes, int):
        return _single_neighbour_average(arr, axes)
    elif isinstance(axes, Sequence):
        for ax in axes:
            arr = _single_neighbour_average(arr, ax)
        return arr
    else:
        raise TypeError("Axis must be an int or a sequence of ints.")

split_to_vertex

split_to_vertex(
    arr: Array, axes: int | Sequence[int]
) -> Array

Split the array into its vertex components along the specified axes. This corresponds to the process of splitting the forcing generated by a panel into its four corners.

Parameters:

Name	Type	Description	Default
arr	Array	Input array to split, or equivalent ArrayList. (..., n, ..., m, ...)	required
axes	int | Sequence[int]	Axis or axes along which to split.	required

Returns:

Type	Description
Array	Array with vertex components. (..., n+1, ..., m+1, ...)

Source code in src/flapjax/algebra/array_utils.py

@singledispatch
def split_to_vertex(arr: Array, axes: int | Sequence[int]) -> Array:
    r"""
    Split the array into its vertex components along the specified axes. This corresponds to the process of splitting
    the forcing generated by a panel into its four corners.
    :param arr: Input array to split, or equivalent ArrayList. ``(..., n, ..., m, ...)``
    :param axes: Axis or axes along which to split.
    :return: Array with vertex components. ``(..., n+1, ..., m+1, ...)``
    """

    def _single_split_to_vertex(arr_: Array, axis_: int) -> Array:
        r"""
        Split the array into its vertex components along the specified axis.
        :param arr_: Input array to split.
        :param axis_: Axis along which to split.
        :return: Array with vertex components, with dimension increased by one along the specified axis.
        """

        shape = list(arr_.shape)
        shape[axis_] += 1

        new_arr_ = jnp.empty(shape, dtype=arr_.dtype)

        index1: list[slice] = [slice(None, None)] * len(shape)
        index1[axis_] = slice(None, -1)

        index2: list[slice] = [slice(None, None)] * len(shape)
        index2[axis_] = slice(1, None)

        new_arr_ = new_arr_.at[tuple(index1)].set(0.5 * arr_)
        new_arr_ = new_arr_.at[tuple(index2)].add(0.5 * arr_)
        return new_arr_

    for ax in axes if isinstance(axes, Sequence) else [axes]:
        arr = _single_split_to_vertex(arr, ax)
    return arr

vect_to_arrs

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

Reconstruct a dictionary with a combination of key-Array and key-ArrayList pairs. The shapes of the arrays and array lists are specified in the shapes argument, which is an ordered dictionary mapping

Parameters:

Name	Type	Description	Default
vect	Array	Data vector.	required
shapes	OrderedDict[str, tuple[int, ...] | ArrayListShape | None]	Shapes for unflattened data.	required

Returns:

Type	Description
OrderedDict[str, Array | ArrayList | None]	Ordered dictionary of unflattened data.

Source code in src/flapjax/algebra/array_utils.py

def vect_to_arrs(
    vect: Array, shapes: OrderedDict[str, tuple[int, ...] | ArrayListShape | None]
) -> OrderedDict[str, Array | ArrayList | None]:
    r"""
    Reconstruct a dictionary with a combination of key-Array and key-ArrayList pairs. The shapes of the arrays and array
    lists are specified in the shapes argument, which is an ordered dictionary mapping
    :param vect: Data vector.
    :param shapes: Shapes for unflattened data.
    :return: Ordered dictionary of unflattened data.
    """

    out_vals = OrderedDict()
    cnt: int = 0

    for name, shape in shapes.items():
        if isinstance(shape, tuple):
            sz = prod(shape)
            out_vals[name] = vect[cnt : cnt + sz].reshape(shape)
            cnt += sz
        elif isinstance(shape, ArrayListShape):
            sz = shape.total_size()
            out_vals[name] = ArrayList.from_vector(
                vect=vect[cnt : cnt + sz], shapes=shape
            )
            cnt += sz
        elif shape is None:
            out_vals[name] = None
        else:
            raise TypeError("Shape must be a tuple or an ArrayListShape.")
    return out_vals

construct_named_block_jacobian

construct_named_block_jacobian(
    entries: tuple[dict, ...],
    keys: Sequence[str],
    widths: Sequence[int],
    heights: Sequence[int],
) -> Array

Assemble a block Jacobian matrix from named partial derivatives.

Parameters:

Name	Type	Description	Default
entries	tuple[dict, ...]	One dict per residual, mapping variable names to Jacobians.	required
keys	Sequence[str]	Variable names that define the column blocks.	required
widths	Sequence[int]	Widths of the blocks.	required
heights	Sequence[int]	Heights of the blocks.	required

Returns:

Type	Description
Array	Dense 2-D Jacobian assembled from the blocks.

Source code in src/flapjax/algebra/array_utils.py

def construct_named_block_jacobian(
    entries: tuple[dict, ...],
    keys: Sequence[str],
    widths: Sequence[int],
    heights: Sequence[int],
) -> Array:
    r"""
    Assemble a block Jacobian matrix from named partial derivatives.
    :param entries: One dict per residual, mapping variable names to Jacobians.
    :param keys: Variable names that define the column blocks.
    :param widths: Widths of the blocks.
    :param heights: Heights of the blocks.
    :return: Dense 2-D Jacobian assembled from the blocks.
    """
    n_block_rows = len(entries)
    n_block_cols = len(keys)

    if len(entries) != len(heights):
        raise ValueError(
            f"Height dimensions do not match. Has {len(entries)} entries versus {len(heights)} heights."
        )
    if len(keys) != len(widths):
        raise ValueError(
            f"Width dimensions do not match. Has {len(keys)} keys versus {len(widths)} widths."
        )

    # Build grid of 2-D blocks, substituting absent entries for zero blocks.
    grid: list[list[Array]] = []
    for i in range(n_block_rows):
        block_height = heights[i]
        row = []
        for j in range(n_block_cols):
            block_width = widths[j]

            if keys[j] in entries[i]:
                row.append(entries[i][keys[j]])
            else:
                row.append(jnp.zeros((block_height, block_width)))
        grid.append(row)
    return jnp.block(grid)

base

matrix2

matrix2(mat: Array) -> Array

Computes the square of a matrix.

Parameters:

Name	Type	Description	Default
mat	Array	Matrix, (n, n).	required

Returns:

Type	Description
Array	Matrix squared, (n, n).

Source code in src/flapjax/algebra/base.py

def matrix2(mat: Array) -> Array:
    r"""
    Computes the square of a matrix.
    :param mat: Matrix, ``(n, n)``.
    :return: Matrix squared, ``(n, n)``.
    """
    return mat @ mat

clip_to_pi

clip_to_pi(val: float | Array)

Clips an angle value to be within [-pi, pi].

Parameters:

Name	Type	Description	Default
val	float | Array	Scalar to bound.	required

Returns:

Type	Description
	Bounded scalar within [-pi, pi].

Source code in src/flapjax/algebra/base.py

def clip_to_pi(val: float | Array):
    r"""
    Clips an angle value to be within `[-pi, pi]`.
    :param val: Scalar to bound.
    :return: Bounded scalar within `[-pi, pi]`.
    """
    return jnp.arctan2(jnp.sin(val), jnp.cos(val))

chi

chi(rmat: Array) -> Array

Converts a 3x3 rotation matrix to a 6x6 matrix used in spatial transformations.

Parameters:

Name	Type	Description	Default
rmat	Array	Rotation matrix, (a, b).	required

Returns:

Type	Description
Array	Block matrix with diagonal rotation matrices, (2a, 2b).

Source code in src/flapjax/algebra/base.py

def chi(rmat: Array) -> Array:
    r"""
    Converts a 3x3 rotation matrix to a 6x6 matrix used in spatial transformations.
    :param rmat: Rotation matrix, ``(a, b)``.
    :return: Block matrix with diagonal rotation matrices, ``(2a, 2b)``.
    """
    return jnp.block([[rmat, jnp.zeros_like(rmat)], [jnp.zeros_like(rmat), rmat]])

finite_difference

finite_difference(
    i_: int,
    data: Array,
    delta: Array,
    axis: int,
    order: int = 1,
) -> Array

Compute the finite difference of the data at a given time step. This assumes that data[:i_+1] is available.

Parameters:

Name	Type	Description	Default
i_	int	Index of derivative to obtain.	required
data	Array	Data to compute the finite difference on, (...).	required
delta	Array	Small perturbation value for finite difference, which divides the difference.	required
axis	int	Axis along which to compute the finite difference.	required
order	int	Order of the finite difference (1 or 2).	1

Returns:

Type	Description
Array	Finite difference of the data at the specified time step.

Source code in src/flapjax/algebra/base.py

def finite_difference(
    i_: int, data: Array, delta: Array, axis: int, order: int = 1
) -> Array:
    r"""
    Compute the finite difference of the data at a given time step. This assumes that ``data[:i_+1]`` is available.
    :param i_: Index of derivative to obtain.
    :param data: Data to compute the finite difference on, (...).
    :param delta: Small perturbation value for finite difference, which divides the difference.
    :param axis: Axis along which to compute the finite difference.
    :param order: Order of the finite difference (1 or 2).
    :return: Finite difference of the data at the specified time step.
    """

    if order not in (0, 1, 2):
        raise ValueError("Order must be 0, 1, or 2.")

    def _slice_order(shift_: int) -> tuple[slice | int, ...]:
        sl: list[slice | int] = [slice(None)] * data.ndim
        sl[axis] = i_ - shift_
        return tuple(sl)

    def _order0() -> Array:
        return jnp.zeros([n for i, n in enumerate(data.shape) if i != axis])

    def _order1() -> Array:
        return (data[_slice_order(0)] - data[_slice_order(1)]) / delta

    def _order2() -> Array:
        return (
            3.0 * data[_slice_order(0)]
            - 4.0 * data[_slice_order(1)]
            + data[_slice_order(2)]
        ) / (2.0 * delta)

    def _err() -> Array:
        return jnp.full([n for i, n in enumerate(data.shape) if i != axis], jnp.nan)

    # use lower int_order when not enough data is available
    # for the instance where only a single data point is available, gradient is set to zero
    order: Array = jnp.array((order, i_)).min()
    return cond(
        order == 0,
        _order0,
        lambda: cond(order == 1, _order1, lambda: cond(order == 2, _order2, _err)),
    )

exp_sum

exp_sum(
    a: Array, order: int = BASE_SUMMATION_ORDER
) -> Array

Computes the matrix exponential using truncated summation.

Parameters:

Name	Type	Description	Default
a	Array	Algebra matrix to exponentiate, (n, n).	required
order	int	Order of summation.	BASE_SUMMATION_ORDER

Returns:

Type	Description
Array	Exponential of matrix, (n, n)

Source code in src/flapjax/algebra/base.py

def exp_sum(a: Array, order: int = BASE_SUMMATION_ORDER) -> Array:
    r"""
    Computes the matrix exponential using truncated summation.
    :param a: Algebra matrix to exponentiate, ``(n, n)``.
    :param order: Order of summation.
    :return: Exponential of matrix, ``(n, n)``
    """

    if a.ndim != 2 or a.shape[0] != a.shape[1]:
        raise ValueError("Input must be a square matrix")

    result = jnp.eye(a.shape[0])
    for i in range(1, order + 1):
        result += jnp.linalg.matrix_power(a, i) / factorial(i)
    return result

log_sum

log_sum(
    g: Array, order: int = BASE_SUMMATION_ORDER
) -> Array

Computes the matrix logarithm using truncated summation.

Parameters:

Name	Type	Description	Default
g	Array	Group matrix to exponentiate, (n, n).	required
order	int	Order of summation.	BASE_SUMMATION_ORDER

Returns:

Type	Description
Array	Logarithm of matrix, (n, n)

Source code in src/flapjax/algebra/base.py

def log_sum(g: Array, order: int = BASE_SUMMATION_ORDER) -> Array:
    r"""
    Computes the matrix logarithm using truncated summation.
    :param g: Group matrix to exponentiate, ``(n, n)``.
    :param order: Order of summation.
    :return: Logarithm of matrix, ``(n, n)``
    """

    if g.ndim != 2 or g.shape[0] != g.shape[1]:
        raise ValueError("Input must be a square matrix")

    g_e = g - jnp.eye(g.shape[0])
    result = g_e

    for i in range(2, order + 1):
        result += (-1.0) ** (i + 1) * jnp.linalg.matrix_power(g_e, i) / i
    return result

t_sum

t_sum(a: Array, order: int = BASE_SUMMATION_ORDER) -> Array

Computes the tangent operator truncated summation. This is used to validate other implementations.

Parameters:

Name	Type	Description	Default
a	Array	Adjoint action matrix, (n, n)	required
order	int	Order of summation.	BASE_SUMMATION_ORDER

Returns:

Type	Description
Array	Tangent operator, (n, n)

Source code in src/flapjax/algebra/base.py

def t_sum(a: Array, order: int = BASE_SUMMATION_ORDER) -> Array:
    r"""
    Computes the tangent operator truncated summation. This is used to validate other implementations.
    :param a: Adjoint action matrix, ``(n, n)``
    :param order: Order of summation.
    :return: Tangent operator, ``(n, n)``
    """

    if a.ndim != 2 or a.shape[0] != a.shape[1]:
        raise ValueError("Input must be a square matrix")

    result = jnp.eye(a.shape[0])
    for i in range(1, order + 1):
        result += (-1.0) ** i * jnp.linalg.matrix_power(a, i) / factorial(i + 1)
    return result

t_inv_sum

t_inv_sum(
    a: Array, order: int = BASE_SUMMATION_ORDER
) -> Array

Computes the inverse tangent operator truncated summation.

Parameters:

Name	Type	Description	Default
a	Array	Adjoint action matrix, (n, n)	required
order	int	Order of summation.	BASE_SUMMATION_ORDER

Returns:

Type	Description
Array	Inverse tangent operator, (n, n)

Source code in src/flapjax/algebra/base.py

def t_inv_sum(a: Array, order: int = BASE_SUMMATION_ORDER) -> Array:
    r"""
    Computes the inverse tangent operator truncated summation.
    :param a: Adjoint action matrix, ``(n, n)``
    :param order: Order of summation.
    :return: Inverse tangent operator, ``(n, n)``
    """

    if a.ndim != 2 or a.shape[0] != a.shape[1]:
        raise ValueError("Input must be a square matrix")

    b = bernoulli(order)

    result = jnp.eye(a.shape[0])
    for i in range(1, order + 1):
        result += (-1.0) ** i * b[i] * jnp.linalg.matrix_power(a, i) / factorial(i)
    return result

jacrev_kwargs

jacrev_kwargs(
    func: Callable[..., Array],
    argnames: str | Sequence[str],
    allow_int: bool = True,
) -> Callable[..., dict[str, Any]]

Custom reverse Jacobian routine which allows for keyword arguments.

Parameters:

Name	Type	Description	Default
func	Callable[..., Array]	Function for which to obtain Jacobians. Must be callable with the keyword arguments later passed to the returned function.	required
argnames	str | Sequence[str]	Argument names of variables for which to obtain Jacobians. These must be a subset of the keyword argument names provided to the returned function.	required
allow_int	bool	As jax.jacrev.	True

Returns:

Type	Description
Callable[..., dict[str, Any]]	Function that accepts keyword arguments and returns a dictionary of argname: Jacobian pairs.

Source code in src/flapjax/algebra/base.py

def jacrev_kwargs(
    func: Callable[..., Array],
    argnames: str | Sequence[str],
    allow_int: bool = True,
) -> Callable[..., dict[str, Any]]:
    r"""
    Custom reverse Jacobian routine which allows for keyword arguments.
    :param func: Function for which to obtain Jacobians. Must be callable with the keyword arguments later passed to
    the returned function.
    :param argnames: Argument names of variables for which to obtain Jacobians. These must be a subset of the keyword
    argument names provided to the returned function.
    :param allow_int: As `jax.jacrev`.
    :return: Function that accepts keyword arguments and returns a dictionary of argname: Jacobian pairs.
    """
    return _jac_kwargs(func, argnames=argnames, mode="reverse", allow_int=allow_int)

jacfwd_kwargs

jacfwd_kwargs(
    func: Callable[..., Array],
    argnames: str | Sequence[str],
) -> Callable[..., dict[str, Any]]

Forward-mode counterpart of :func:jacrev_kwargs. Prefer this when the input dimension for the selected arguments is smaller than the output dimension of func, which happens for design-variable Jacobians and for dynamic-adjoint residual blocks whose input state size is smaller than the residual dimension.

Parameters:

Name	Type	Description	Default
func	Callable[..., Array]	Function for which to obtain Jacobians.	required
argnames	str | Sequence[str]	Argument names of variables for which to obtain Jacobians.	required

Returns:

Type	Description
Callable[..., dict[str, Any]]	Function that accepts keyword arguments and returns a dictionary of argname: Jacobian pairs.

Source code in src/flapjax/algebra/base.py

def jacfwd_kwargs(
    func: Callable[..., Array],
    argnames: str | Sequence[str],
) -> Callable[..., dict[str, Any]]:
    r"""
    Forward-mode counterpart of :func:`jacrev_kwargs`. Prefer this when the input dimension for the selected
    arguments is smaller than the output dimension of ``func``, which happens for design-variable Jacobians and for
    dynamic-adjoint residual blocks whose input state size is smaller than the residual dimension.
    :param func: Function for which to obtain Jacobians.
    :param argnames: Argument names of variables for which to obtain Jacobians.
    :return: Function that accepts keyword arguments and returns a dictionary of argname: Jacobian pairs.
    """
    return _jac_kwargs(func, argnames=argnames, mode="forward", allow_int=False)

jacrev_custom

jacrev_custom(
    func: Callable[..., Array],
    jac_options: dict[
        str, Callable[[dict[str, Any]], Array] | None
    ],
    n_profile_loops: int | None,
    func_name: str,
    static_argnames: Sequence[str] = (),
    mode: ADMode = "reverse",
    map_batch_size: int | None = None,
) -> Callable[
    ...,
    tuple[
        dict[str, Any],
        dict[str, float] | None,
        dict[str, float] | None,
    ],
]

Obtain the Jacobians of the function func with respect to a chosen set of arguments.

Parameters:

Name	Type	Description	Default
func	Callable[..., Array]	Function for which to obtain the Jacobians.	required
jac_options	dict[str, Callable[[dict[str, Any]], Array] | None]	Dictionary with variable names as keys, with values being a tuple of the argument number in func, and an optional function which can be used to compute the given Jacobian. If this is None, the Jacobian is computed by applying AD with no approximations.	required
n_profile_loops	int | None	Number of profile loops. If None, no profiling is done.	required
func_name	str	Name of the function to be called, used for console prints during profiling.	required
static_argnames	Sequence[str]	Argument names to treat as static when JIT-compiling the AD Jacobian. Needed so the profile loop hits a cached jaxpr instead of re-tracing on every call.	()
mode	ADMode	"reverse" (default) or "forward".	'reverse'
map_batch_size	int | None	When set, batch passes to obtain Jacobian rather than vmapping, reducing memory at the expense of computation time.	None

Returns:

Type	Description
Callable[..., tuple[dict[str, Any], dict[str, float] | None, dict[str, float] | None]]	Function to obtain Jacobians, as well as respective compile and run times for Jacobians if n_profile_loops is not None.

Source code in src/flapjax/algebra/base.py

def jacrev_custom(
    func: Callable[..., Array],
    jac_options: dict[str, Callable[[dict[str, Any]], Array] | None],
    n_profile_loops: int | None,
    func_name: str,
    static_argnames: Sequence[str] = (),
    mode: ADMode = "reverse",
    map_batch_size: int | None = None,
) -> Callable[
    ..., tuple[dict[str, Any], dict[str, float] | None, dict[str, float] | None]
]:
    r"""
    Obtain the Jacobians of the function `func` with respect to a chosen set of arguments.

    :param func: Function for which to obtain the Jacobians.
    :param jac_options: Dictionary with variable names as keys, with values being a tuple of the argument number in
    `func`, and an optional function which can be used to compute the given Jacobian. If this is None, the Jacobian is
    computed by applying AD with no approximations.
    :param n_profile_loops: Number of profile loops. If None, no profiling is done.
    :param func_name: Name of the function to be called, used for console prints during profiling.
    :param static_argnames: Argument names to treat as static when JIT-compiling the AD Jacobian. Needed so the
    profile loop hits a cached jaxpr instead of re-tracing on every call.
    :param mode: ``"reverse"`` (default) or ``"forward"``.
    :param map_batch_size: When set, batch passes to obtain Jacobian rather than vmapping, reducing memory at the
    expense of computation time.
    :return: Function to obtain Jacobians, as well as respective compile and run times for
    Jacobians if `n_profile_loops` is not None.
    """

    static_argnames_t = tuple(static_argnames)

    if mode == "reverse":
        if map_batch_size is not None:

            def _build_jac(argnames):
                return _jacrev_via_map_kwargs(
                    func, argnames=argnames, map_batch_size=map_batch_size
                )
        else:

            def _build_jac(argnames):
                return jacrev_kwargs(func, argnames=argnames, allow_int=True)
    elif mode == "forward":
        if map_batch_size is not None:

            def _build_jac(argnames):
                return _jacfwd_via_map_kwargs(
                    func, argnames=argnames, map_batch_size=map_batch_size
                )
        else:

            def _build_jac(argnames):
                return jacfwd_kwargs(func, argnames=argnames)
    else:
        raise ValueError('Invalid mode, use "forward" or "reverse"')

    def inner_func(
        **args: Any,
    ) -> tuple[dict[str, Any], dict[str, float] | None, dict[str, float] | None]:
        jacobians: dict[str, Any] = {}
        ad_args: list[
            str
        ] = []  # accumulate list of argument names that we will use for AD
        compile_time: dict[str, float | None] = {}
        run_time: dict[str, float | None] = {}

        # compute cases which have a passed approximation function for finding the Jacobian
        for arg, jac_func in jac_options.items():
            if jac_func is not None:
                jacobians[arg], compile_time[arg], run_time[arg] = conditional_profile(
                    func=jac_func,
                    n_loops=n_profile_loops,
                    func_name=func_name,
                    arg_name=arg,
                )(args)  # evaluate function
            else:
                ad_args.append(arg)  # indicate that we will use AD to compute

        # compute cases where we perform AD directly
        if ad_args:
            per_arg_jacobians: dict[str, Any] = {}
            if n_profile_loops is not None:
                for ad_arg in ad_args:
                    # profile for individual Jacobians
                    jac_func = jax.jit(
                        _build_jac(ad_arg),
                        static_argnames=static_argnames_t,
                    )
                    val, compile_time[ad_arg], run_time[ad_arg] = conditional_profile(
                        func=jac_func,
                        n_loops=n_profile_loops,
                        func_name=func_name,
                        arg_name=ad_arg,
                    )(**args)
                    per_arg_jacobians.update(val)

            if mode == "forward":
                all_jacobians = (
                    per_arg_jacobians
                    if n_profile_loops is not None
                    else _build_jac(ad_args)(**args)
                )
            else:
                jac_func = _build_jac(ad_args)
                if n_profile_loops is not None:
                    jac_func = jax.jit(jac_func, static_argnames=static_argnames_t)

                (
                    all_jacobians,
                    compile_time["all"],
                    run_time["all"],
                ) = conditional_profile(
                    func=jac_func,
                    n_loops=n_profile_loops,
                    func_name=func_name,
                    arg_name="all",
                )(**args)

            jacobians.update(all_jacobians)

        if n_profile_loops is not None:
            return jacobians, compile_time, run_time  # type: ignore
        else:
            return jacobians, None, None

    return inner_func

jacobian_approximation

jacobian_approximation(
    func: Callable[..., Any],
    args: Any,
    approx_type: Literal[
        "zero",
        "constant",
        "identity",
        "dense_linear",
        "lazy_linear",
    ]
    | None,
    jacobian_argname: str,
    hessian_argnames: str | Sequence[str],
) -> Callable[..., Any] | None

Compute approximations of Jacobians. The following options for approximation are available: None - No approximation is computed. zero - Assume that the Jacobian is zero. identity - Assume that the Jacobian is the identity matrix. Valid only for 2D square entries. constant - Assume that the Jacobian is constant across all time steps. Alongside these, options are also available to compute the Jacobians by a linear approximation by using a Hessian-vector product, where the Hessian is constant. This introduces additional options: dense_linear - The dense Hessian is explicitly computed, which increases memory cost but reduces computation cost. lazy_linear - Uses the JAX linearise routine, which avoids explicitly computing the dense Hessian, which reduces memory cost at the expense of increased computation cost.

Parameters:

Name	Type	Description	Default
func	Callable[..., Any]	Function for which Jacobians are to be computed.	required
args	Any	Arguments for func which define the state around which we want to make the approximation.	required
approx_type	Literal['zero', 'constant', 'identity', 'dense_linear', 'lazy_linear'] | None	Type of approximation to be used to compute Jacobians. Options are "constant", "dense_linear" or "lazy_linear". If None, no approximation is computed.	required
jacobian_argname	str	Argument which the function Jacobian is to be computed with respect to.	required
hessian_argnames	str | Sequence[str]	Argument names which the Jacobian is linearised with respect to for computing Hessian-vector products.	required

Returns:

Type	Description
Callable[..., Any] | None	Function which approximates the Jacobian which takes the same arguments as func, or None.

Source code in src/flapjax/algebra/base.py

def jacobian_approximation(
    func: Callable[..., Any],
    args: Any,
    approx_type: Literal["zero", "constant", "identity", "dense_linear", "lazy_linear"]
    | None,
    jacobian_argname: str,
    hessian_argnames: str | Sequence[str],
) -> Callable[..., Any] | None:
    r"""
    Compute approximations of Jacobians. The following options for approximation are available:
    `None` - No approximation is computed.
    `zero` - Assume that the Jacobian is zero.
    `identity` - Assume that the Jacobian is the identity matrix. Valid only for 2D square entries.
    `constant` - Assume that the Jacobian is constant across all time steps.
     Alongside these, options are also available to compute the Jacobians by a linear approximation by using a
     Hessian-vector product, where the Hessian is constant. This introduces additional options:
     `dense_linear` - The dense Hessian is explicitly computed, which increases memory cost but reduces computation cost.
     `lazy_linear` - Uses the JAX linearise routine, which avoids explicitly computing the dense Hessian, which reduces
     memory cost at the expense of increased computation cost.
    :param func: Function for which Jacobians are to be computed.
    :param args: Arguments for `func` which define the state around which we want to make the approximation.
    :param approx_type: Type of approximation to be used to compute Jacobians. Options are "constant", "dense_linear" or
    "lazy_linear". If None, no approximation is computed.
    :param jacobian_argname: Argument which the function Jacobian is to be computed with respect to.
    :param hessian_argnames: Argument names which the Jacobian is linearised with respect to for computing
    Hessian-vector products.
    :return: Function which approximates the Jacobian which takes the same arguments as `func`, or None.
    """

    hessian_argnames_t = (
        (hessian_argnames,)
        if isinstance(hessian_argnames, str)
        else tuple(hessian_argnames)
    )

    # closure over non-array arguments
    _dyn_args, _static_args = _split_dyn_static(
        args, keep_dynamic=(jacobian_argname, *hessian_argnames_t)
    )
    _func = func if not _static_args else (lambda **_dk: func(**_dk, **_static_args))

    match approx_type:
        case None:
            # no approximation, and so we will compute the exact Jacobian using AD later
            return None
        case "zero":
            # assume that the Jacobian is zero, avoiding any Jacobian computation. In practice, this is only useful
            # when the `constant` path would be slow to linearise.
            jac_shape = jax.eval_shape(
                jacrev_kwargs(_func, argnames=jacobian_argname, allow_int=True),
                **_dyn_args,
            )[jacobian_argname]
            zero_jac = jax.tree.map(
                lambda s: jnp.zeros(s.shape, dtype=s.dtype), jac_shape
            )
            return lambda *_, **__: zero_jac
        case "identity":
            # assume that the Jacobian is the identity matrix
            jac_shape = jax.eval_shape(
                jacrev_kwargs(_func, argnames=jacobian_argname, allow_int=True),
                **_dyn_args,
            )[jacobian_argname]
            if not isinstance(jac_shape, jax.ShapeDtypeStruct):
                raise TypeError(
                    "identity approximation requires an Array-valued Jacobian, got a pytree"
                )
            shape = jac_shape.shape
            if len(shape) != 2 or shape[0] != shape[1]:
                raise ValueError(
                    f"identity approximation expected a 2-D square Jacobian for "
                    f"{jacobian_argname}, got shape {shape}"
                )
            identity_jac = jnp.eye(shape[0], dtype=jac_shape.dtype)
            return lambda *_, **__: identity_jac
        case "constant":
            # assume that the Jacobian stays constant for all time
            jac = jacrev_kwargs(_func, argnames=jacobian_argname, allow_int=True)(
                **_dyn_args
            )[jacobian_argname]
            return lambda *_, **__: jac
        case "lazy_linear":
            # Jacobian is computed using a Hessian-vector product where the Hessian is chosen to be constant
            # uses the JAX linearise function which means that the dense Hessian is never computed
            # reduced memory cost at the expense of increased computation cost
            def _jac_at(*y_vals: Any) -> Array:
                new_args = dict(_dyn_args)
                for name, val in zip(hessian_argnames_t, y_vals):
                    new_args[name] = val
                return jacrev_kwargs(_func, argnames=jacobian_argname, allow_int=True)(
                    **new_args
                )[jacobian_argname]

            y0_vals = tuple(_dyn_args[name] for name in hessian_argnames_t)
            j0, jac_jvp = jax.linearize(_jac_at, *y0_vals)

            def lazy_linear_approx(
                new_args: dict[str, Any], *_: Any, **__: Any
            ) -> Array:
                dy_vals = tuple(
                    new_args[name] - _dyn_args[name] for name in hessian_argnames_t
                )
                return j0 + jac_jvp(*dy_vals)

            return lazy_linear_approx
        case "dense_linear":
            # as `lazy_linear`, except that the dense Hessian is explicitly computed.
            def _jac_at_kwargs(**kw: Any) -> Array:
                new_args = dict(_dyn_args)
                new_args.update(kw)
                return jacrev_kwargs(_func, argnames=jacobian_argname, allow_int=True)(
                    **new_args
                )[jacobian_argname]

            y0_kwargs = {name: _dyn_args[name] for name in hessian_argnames_t}
            j0 = _jac_at_kwargs(**y0_kwargs)
            hessians = jacrev_kwargs(
                _jac_at_kwargs, argnames=hessian_argnames_t, allow_int=True
            )(**y0_kwargs)

            def dense_linear_approx(
                new_args: dict[str, Any], *_: Any, **__: Any
            ) -> Array:
                result = j0
                for name in hessian_argnames_t:
                    dy = new_args[name] - _dyn_args[name]
                    result += jnp.tensordot(hessians[name], dy, axes=dy.ndim)
                return result

            return dense_linear_approx
        case _:
            raise ValueError(f"Invalid approximation type: {approx_type}")

integration

gauss_lobatto

gauss_lobatto(
    f: Callable[[Array], Array],
    bounds: Array,
    f_bounds: Array,
    int_order: Literal[3, 4, 5],
) -> Array

Integrate using quadrature with Gauss-Lobatto points. Makes use of function values at the bounds. See https://en.wikipedia.org/wiki/Gaussian_quadrature.

Parameters:

Name	Type	Description	Default
f	Callable[[Array], Array]	Function to integrate (must support vector mapping), () -> (...).	required
bounds	Array	Scalar bounds of integration in function space, (2, ).	required
f_bounds	Array	values of function at the bounds, (2, ...).	required
int_order	Literal[3, 4, 5]	Order of integration, 3, 4, or 5.	required

Returns:

Type	Description
Array	Integrated value, (...).

Source code in src/flapjax/algebra/integration.py

def gauss_lobatto(
    f: Callable[[Array], Array],
    bounds: Array,
    f_bounds: Array,
    int_order: Literal[3, 4, 5],
) -> Array:
    r"""
    Integrate using quadrature with Gauss-Lobatto points. Makes use of function values at the bounds.
    See https://en.wikipedia.org/wiki/Gaussian_quadrature.
    :param f: Function to integrate (must support vector mapping), () -> (...).
    :param bounds: Scalar bounds of integration in function space, ``(2, )``.
    :param f_bounds: values of function at the bounds, ``(2, ...)``.
    :param int_order: Order of integration, 3, 4, or 5.
    :return: Integrated value, (...).
    """
    match int_order:
        case 3:
            x_i = jnp.array((0.0,))
            w_i = jnp.array((4.0 / 3.0,))
        case 4:
            x_i = jnp.array((-1.0 / jnp.sqrt(5.0), 1.0 / jnp.sqrt(5.0)))
            w_i = jnp.array((5.0 / 6.0, 5.0 / 6.0))
        case 5:
            x_i = jnp.array((-jnp.sqrt(3.0 / 7.0), 0.0, jnp.sqrt(3.0 / 7.0)))
            w_i = jnp.array((49.0 / 90.0, 32.0 / 45.0, 49.0 / 90.0))
        case _:
            raise ValueError("Order must be one of 3, 4, or 5.")

    range_ = bounds[1] - bounds[0]
    x_i_scaled = bounds[0] + 0.5 * (x_i + 1.0) * range_  # (n_i, )
    f_i = vmap(f, 0, 0)(x_i_scaled)  # (n_i, ...)

    return (
        range_
        / 2.0
        * (
            2.0 / (int_order * (int_order - 1)) * jnp.sum(f_bounds, axis=0)
            + jnp.einsum("i,i...->...", w_i, f_i)
        )
    )

gauss_legendre

gauss_legendre(
    f: Callable[[Array], Array],
    bounds: Array,
    int_order: Literal[1, 2, 3],
) -> Array

Integrate using quadrature with Gauss-Legendre points. See https://en.wikipedia.org/wiki/Gaussian_quadrature.

Parameters:

Name	Type	Description	Default
f	Callable[[Array], Array]	Function to integrate (must support vector mapping), () -> (...).	required
bounds	Array	Scalar bounds of integration in function space, (2, ).	required
int_order	Literal[1, 2, 3]	Order of integration.	required

Returns:

Type	Description
Array	Integrated value, (...).

Source code in src/flapjax/algebra/integration.py

def gauss_legendre(
    f: Callable[[Array], Array],
    bounds: Array,
    int_order: Literal[1, 2, 3],
) -> Array:
    r"""
    Integrate using quadrature with Gauss-Legendre points. See https://en.wikipedia.org/wiki/Gaussian_quadrature.
    :param f: Function to integrate (must support vector mapping), () -> (...).
    :param bounds: Scalar bounds of integration in function space, ``(2, )``.
    :param int_order: Order of integration.
    :return: Integrated value, (...).
    """

    match int_order:
        case 1:
            x_i = jnp.array((0.0,))
            w_i = jnp.array((2.0,))
        case 2:
            x_i = jnp.array((-1.0 / jnp.sqrt(3.0), 1.0 / jnp.sqrt(3.0)))
            w_i = jnp.array((1.0, 1.0))
        case 3:
            x_i = jnp.array((-jnp.sqrt(3.0 / 5.0), 0.0, jnp.sqrt(3.0 / 5.0)))
            w_i = jnp.array((5.0 / 9.0, 8.0 / 9.0, 5.0 / 9.0))
        case _:
            raise ValueError("Order must be one of 1, 2, or 3.")

    range_ = bounds[1] - bounds[0]
    x_i_scaled = bounds[0] + 0.5 * (x_i + 1.0) * range_  # (n_i, )
    f_i = vmap(f, 0, 0)(x_i_scaled)  # (n_i, ...)

    return range_ / 2.0 * jnp.einsum("i,i...->...", w_i, f_i)

se3

bracket_se3

bracket_se3(a_vec: Array, b_vec: Array) -> Array

Computes the Lie bracket of two se(3) elements, :math:\tilde{a}\tilde{b} - \tilde{b}\tilde{a}.

Parameters:

Name	Type	Description	Default
a_vec	Array	Lie algebra vector in se(3), (6, ).	required
b_vec	Array	Lie algebra vector in se(3), (6, ).	required

Returns:

Type	Description
Array	Lie bracket, (4, 4).

Source code in src/flapjax/algebra/se3.py

def bracket_se3(a_vec: Array, b_vec: Array) -> Array:
    r"""
    Computes the Lie bracket of two se(3) elements, :math:`\tilde{a}\tilde{b} - \tilde{b}\tilde{a}`.
    :param a_vec: Lie algebra vector in se(3), ``(6, )``.
    :param b_vec: Lie algebra vector in se(3), ``(6, )``.
    :return: Lie bracket, ``(4, 4)``.
    """
    mat1 = ha_to_ha_tilde(a_vec)
    mat2 = ha_to_ha_tilde(b_vec)

    return mat1 @ mat2 - mat2 @ mat1

bracket_neg_se3

bracket_neg_se3(a_vec: Array, b_vec: Array) -> Array

Computes the negative Lie bracket of two se(3) elements, :math:\tilde{a}\tilde{b} + \tilde{b}\tilde{a}.

Parameters:

Name	Type	Description	Default
a_vec	Array	Lie algebra vector in se(3), (6, ).	required
b_vec	Array	Lie algebra vector in se(3), (6, ).	required

Returns:

Type	Description
Array	Negative lie bracket, (4, 4).

Source code in src/flapjax/algebra/se3.py

def bracket_neg_se3(a_vec: Array, b_vec: Array) -> Array:
    r"""
    Computes the negative Lie bracket of two se(3) elements, :math:`\tilde{a}\tilde{b} + \tilde{b}\tilde{a}`.
    :param a_vec: Lie algebra vector in se(3), ``(6, )``.
    :param b_vec: Lie algebra vector in se(3), ``(6, )``.
    :return: Negative lie bracket, ``(4, 4)``.
    """
    mat1 = ha_to_ha_tilde(a_vec)
    mat2 = ha_to_ha_tilde(b_vec)

    return mat1 @ mat2 + mat2 @ mat1

t_u_omega_plus

t_u_omega_plus(ha: Array) -> Array

Computes the :math:\mathbf{T}_{U \omega+} matrix, used for computing the tangent application for SE(3). Formulation from Geometrically exact beam finite element formulated on the special Euclidean group SE(3), by Sonneville et al., 2013, Eq A.12

Parameters:

Name	Type	Description	Default
ha	Array	Vector in se(3), (6, ).	required

Returns:

Type	Description
Array	Operator, (3, 3).

Source code in src/flapjax/algebra/se3.py

def t_u_omega_plus(ha: Array) -> Array:
    r"""
    Computes the :math:`\mathbf{T}_{U \omega+}` matrix, used for computing the tangent application for SE(3). Formulation
    from Geometrically exact beam finite element formulated on the special Euclidean group SE(3), by Sonneville et al., 2013, Eq A.12
    :param ha: Vector in se(3), ``(6, )``.
    :return: Operator, ``(3, 3)``.
    """
    # (6,) -> (3, 3)
    a = ha[:3]
    b = ha[3:]
    b_norm2 = jnp.inner(b, b)

    def t_u_omega_plus_full() -> Array:
        r"""Full computation of :math:`\mathbf{T}_{U \omega+}` for non-small angles."""
        alpha_ = alpha(b)
        beta_ = beta(b)

        return (
            -0.5 * beta_ * vec_to_skew(a)
            + (1.0 - alpha_) / b_norm2 * bracket_neg_so3(a, b)
            + jnp.inner(b, a)
            / b_norm2
            * (
                (beta_ - alpha_) * vec_to_skew(b)
                + (0.5 * beta_ - 3.0 * (1.0 - alpha_) / b_norm2)
                * matrix2(vec_to_skew(b))
            )
        )

    def t_u_omega_plus_small_angle() -> Array:
        r"""Computation of :math:`\mathbf{T}_{U \omega+}` when the rotation angle is small."""
        return -0.5 * vec_to_skew(a)

    return cond(
        b_norm2 > SMALL_ANG_THRESH, t_u_omega_plus_full, t_u_omega_plus_small_angle
    )

t_u_omega_minus

t_u_omega_minus(ha: Array) -> Array

Computes the :math:\mathbf{T}_{U \omega-} matrix, used for computing the inverse tangent application for se(3). Formulation from Geometrically exact beam finite element formulated on the special Euclidean group SE(3), by Sonneville et al., 2013, Eq A.14. This can be represented in terms of :math:\mathbf{T}_{U \omega+} and the inverse of the SO(3) tangent operator.

Parameters:

Name	Type	Description	Default
ha	Array	Vector in se(3), (6, )	required

Returns:

Type	Description
Array	Operator, (3, 3).

Source code in src/flapjax/algebra/se3.py

def t_u_omega_minus(ha: Array) -> Array:
    r"""
    Computes the :math:`\mathbf{T}_{U \omega-}` matrix, used for computing the inverse tangent application for se(3).
    Formulation from Geometrically exact beam finite element formulated on the special Euclidean group SE(3), by
    Sonneville et al., 2013, Eq A.14. This can be represented in terms of :math:`\mathbf{T}_{U \omega+}` and the
    inverse of the SO(3) tangent operator.
    :param ha: Vector in se(3), ``(6, )``
    :return: Operator, ``(3, 3)``.
    """

    t_inv_ = t_inv_so3(ha[3:])
    return -t_inv_ @ t_u_omega_plus(ha) @ t_inv_

t_se3

t_se3(ha: Array) -> Array

Computes the tangent operator for se(3). Formulation from Geometrically exact beam finite element formulated on the special Euclidean group SE(3), by Sonneville et al., 2013, Eq A.11.

Parameters:

Name	Type	Description	Default
ha	Array	se(3) vector, (6, ).	required

Returns:

Type	Description
Array	Tangent operator, (6, 6).

Source code in src/flapjax/algebra/se3.py

def t_se3(ha: Array) -> Array:
    r"""
    Computes the tangent operator for se(3). Formulation from Geometrically exact beam finite element formulated on the
    special Euclidean group SE(3), by Sonneville et al., 2013, Eq A.11.
    :param ha: se(3) vector, ``(6, )``.
    :return: Tangent operator, ``(6, 6)``.
    """

    def t_se3_full() -> Array:
        t_ = t_so3(ha[3:])
        return jnp.block([[t_, t_u_omega_plus(ha)], [jnp.zeros((3, 3)), t_]])

    def t_se3_small_angle() -> Array:
        return t_sum(ha_to_ha_hat(ha), 2)

    ang_mag2 = jnp.inner(ha[3:], ha[3:])
    return cond(ang_mag2 > SMALL_ANG_THRESH, t_se3_full, t_se3_small_angle)

t_inv_se3

t_inv_se3(ha: Array) -> Array

Computes the inverse tangent operator for se(3). Formulation from Geometrically exact beam finite element formulated on the special Euclidean group SE(3), by Sonneville et al., 2013, Eq A.13.

Parameters:

Name	Type	Description	Default
ha	Array	se(3) algebra vector, (6, ).	required

Returns:

Type	Description
Array	Inverse angent operator, (6, 6).

Source code in src/flapjax/algebra/se3.py

def t_inv_se3(ha: Array) -> Array:
    r"""
    Computes the inverse tangent operator for se(3). Formulation from Geometrically exact beam finite element formulated
    on the special Euclidean group SE(3), by Sonneville et al., 2013, Eq A.13.
    :param ha: se(3) algebra vector, ``(6, )``.
    :return: Inverse angent operator, ``(6, 6)``.
    """

    def t_inv_se3_full() -> Array:
        t_ = t_inv_so3(ha[3:])
        return jnp.block([[t_, t_u_omega_minus(ha)], [jnp.zeros((3, 3)), t_]])

    def t_inv_se3_small_angle() -> Array:
        return t_inv_sum(ha_to_ha_hat(ha), 2)

    ang_mag2 = jnp.inner(ha[3:], ha[3:])
    return cond(ang_mag2 > SMALL_ANG_THRESH, t_inv_se3_full, t_inv_se3_small_angle)

log_se3

log_se3(hg: Array) -> Array

Computes the logarithm map from SE(3) to se(3). Formulation from Geometrically exact beam finite element formulated on the special Euclidean group SE(3), by Sonneville et al., 2013, Eq A.15.

Parameters:

Name	Type	Description	Default
hg	Array	SE(3) group element, (4, 4).	required

Returns:

Type	Description
Array	se(3) algebra vector, (6, ).

Source code in src/flapjax/algebra/se3.py

def log_se3(hg: Array) -> Array:
    r"""
    Computes the logarithm map from SE(3) to se(3). Formulation from Geometrically exact beam finite element formulated
    on the special Euclidean group SE(3), by Sonneville et al., 2013, Eq A.15.
    :param hg: SE(3) group element, ``(4, 4)``.
    :return: se(3) algebra vector, ``(6, )``.
    """

    omega = log_so3(hg[:3, :3])

    def log_se3_full() -> Array:
        return jnp.concatenate((t_inv_so3(omega).T @ hg[:3, 3], omega))

    def log_se3_small_angle() -> Array:
        return ha_tilde_to_ha(log_sum(hg, 2))

    ang_mag2 = jnp.inner(omega, omega)
    return cond(ang_mag2 > SMALL_ANG_THRESH, log_se3_full, log_se3_small_angle)

exp_se3

exp_se3(ha: Array) -> Array

Computes the exponential map from se(3) to SE(3). Formulation from Geometrically exact beam finite element formulated on the special Euclidean group SE(3), by Sonneville et al., 2013, Eq A.10.

Parameters:

Name	Type	Description	Default
ha	Array	se(3) algebra vector, (6, ).	required

Returns:

Type	Description
Array	SE(3) group element, (4, 4).

Source code in src/flapjax/algebra/se3.py

def exp_se3(ha: Array) -> Array:
    r"""
    Computes the exponential map from se(3) to SE(3). Formulation from Geometrically exact beam finite element formulated
    on the special Euclidean group SE(3), by Sonneville et al., 2013, Eq A.10.
    :param ha: se(3) algebra vector, ``(6, )``.
    :return: SE(3) group element, ``(4, 4)``.
    """

    def exp_se3_full() -> Array:
        return jnp.block(
            [
                [exp_so3(ha[3:]), (t_so3(ha[3:]).T @ ha[:3])[:, None]],
                [jnp.zeros((1, 3)), jnp.ones((1, 1))],
            ]
        )

    def exp_se3_small_angle() -> Array:
        return exp_sum(ha_to_ha_tilde(ha), 2)

    ang_mag2 = jnp.inner(ha[3:], ha[3:])
    return cond(ang_mag2 > SMALL_ANG_THRESH, exp_se3_full, exp_se3_small_angle)

x_rmat_to_hg

x_rmat_to_hg(x: Array, rmat: Array) -> Array

Combines a translation vector and rotation matrix into an element of the SE(3) group.

Parameters:

Name	Type	Description	Default
x	Array	Translation vector, (3, ).	required
rmat	Array	Rotation matrix, (3, 3).	required

Returns:

Type	Description
Array	SE(3) group element, (4, 4).

Source code in src/flapjax/algebra/se3.py

def x_rmat_to_hg(x: Array, rmat: Array) -> Array:
    r"""
    Combines a translation vector and rotation matrix into an element of the SE(3) group.
    :param x: Translation vector, ``(3, )``.
    :param rmat: Rotation matrix, ``(3, 3)``.
    :return: SE(3) group element, ``(4, 4)``.
    """
    return jnp.block([[rmat, x[:, None]], [jnp.zeros((1, 3)), jnp.ones((1, 1))]])

x_rmat_to_ha

x_rmat_to_ha(x: Array, rmat: Array) -> Array

Combines a translation vector and rotation matrix into an element of the se(3) algebra.

Parameters:

Name	Type	Description	Default
x	Array	Translation vector, (3, ).	required
rmat	Array	Rotation matrix, (3, 3).	required

Returns:

Type	Description
Array	se(3) algebra vector, (6, ).

Source code in src/flapjax/algebra/se3.py

def x_rmat_to_ha(x: Array, rmat: Array) -> Array:
    r"""
    Combines a translation vector and rotation matrix into an element of the se(3) algebra.
    :param x: Translation vector, ``(3, )``.
    :param rmat: Rotation matrix, ``(3, 3)``.
    :return: se(3) algebra vector, ``(6, )``.
    """
    return log_se3(x_rmat_to_hg(x, rmat))

hg_to_x_rmat

hg_to_x_rmat(hg: Array) -> tuple[Array, Array]

Decomposes an SE(3) group element into a translation vector and rotation matrix.

Parameters:

Name	Type	Description	Default
hg	Array	SE(3) group element, (4, 4).	required

Returns:

Type	Description
tuple[Array, Array]	Translation vector, (3, ), and rotation matrix, (3, 3).

Source code in src/flapjax/algebra/se3.py

def hg_to_x_rmat(hg: Array) -> tuple[Array, Array]:
    r"""
    Decomposes an SE(3) group element into a translation vector and rotation matrix.
    :param hg: SE(3) group element, ``(4, 4)``.
    :return: Translation vector, ``(3, )``, and rotation matrix, ``(3, 3)``.
    """
    return hg[:3, 3], hg[:3, :3]

vect_product

vect_product(hg: Array, x: Array) -> Array

Computes the resulting vector of an SE(3) group element and a 3D translation vector.

Parameters:

Name	Type	Description	Default
hg	Array	SE(3) group element, (4, 4).	required
x	Array	Translation vector, (3, ).	required

Returns:

Type	Description
Array	Resulting translation vector, (3, ).

Source code in src/flapjax/algebra/se3.py

def vect_product(hg: Array, x: Array) -> Array:
    r"""
    Computes the resulting vector of an SE(3) group element and a 3D translation vector.
    :param hg: SE(3) group element, ``(4, 4)``.
    :param x: Translation vector, ``(3, )``.
    :return: Resulting translation vector, ``(3, )``.
    """
    return hg[:3, :3] @ x + hg[:3, 3]

hg_inv

hg_inv(hg: Array) -> Array

Computes the inverse of an SE(3) group element.

Parameters:

Name	Type	Description	Default
hg	Array	SE(3) group element, (4, 4).	required

Returns:

Type	Description
Array	Inverse SE(3) group element, (4, 4).

Source code in src/flapjax/algebra/se3.py

def hg_inv(hg: Array) -> Array:
    r"""
    Computes the inverse of an SE(3) group element.
    :param hg: SE(3) group element, ``(4, 4)``.
    :return: Inverse SE(3) group element, ``(4, 4)``.
    """
    x, rmat = hg_to_x_rmat(hg)
    return jnp.block(
        [[rmat.T, -(rmat.T @ x)[:, None]], [jnp.zeros((1, 3)), jnp.ones((1, 1))]]
    )

ha_to_ha_tilde

ha_to_ha_tilde(ha: Array) -> Array

Converts a se(3) vector into its matrix representation.

Parameters:

Name	Type	Description	Default
ha	Array	se(3) algebra vector, (6, ).	required

Returns:

Type	Description
Array	se(3) algebra element in matrix form, (4, 4).

Source code in src/flapjax/algebra/se3.py

def ha_to_ha_tilde(ha: Array) -> Array:
    r"""
    Converts a se(3) vector into its matrix representation.
    :param ha: se(3) algebra vector, ``(6, )``.
    :return: se(3) algebra element in matrix form, ``(4, 4)``.
    """
    return jnp.block([[vec_to_skew(ha[3:]), ha[:3, None]], [jnp.zeros((1, 4))]])

ha_tilde_to_ha

ha_tilde_to_ha(ha_tilde: Array) -> Array

Converts a se(3) element matrix into its vector representation.

Parameters:

Name	Type	Description	Default
ha_tilde	Array	se(3) algebra element in matrix form, (4, 4).	required

Returns:

Type	Description
Array	se(3) algebra vector, (6, ).

Source code in src/flapjax/algebra/se3.py

def ha_tilde_to_ha(ha_tilde: Array) -> Array:
    r"""
    Converts a se(3) element matrix into its vector representation.
    :param ha_tilde: se(3) algebra element in matrix form, ``(4, 4)``.
    :return: se(3) algebra vector, ``(6, )``.
    """
    ha_u = ha_tilde[:3, 3]
    ha_omega = skew_to_vec(ha_tilde[:3, :3])
    return jnp.concatenate((ha_u, ha_omega), axis=-1)

ha_to_ha_hat

ha_to_ha_hat(ha: Array) -> Array

Converts a se(3) vector into its hat matrix representation. Formulation from Geometrically exact beam finite element formulated on the special Euclidean group SE(3), by Sonneville et al., 2013, Eq 15.

Parameters:

Name	Type	Description	Default
ha	Array	se(3) algebra vector, (6, ).	required

Returns:

Type	Description
Array	Hat matrix representation, (6, 6).

Source code in src/flapjax/algebra/se3.py

def ha_to_ha_hat(ha: Array) -> Array:
    r"""
    Converts a se(3) vector into its hat matrix representation. Formulation from Geometrically exact beam finite element
    formulated on the special Euclidean group SE(3), by Sonneville et al., 2013, Eq 15.
    :param ha: se(3) algebra vector, ``(6, )``.
    :return: Hat matrix representation, ``(6, 6)``.
    """
    return jnp.block(
        [
            [vec_to_skew(ha[3:]), vec_to_skew(ha[:3])],
            [jnp.zeros((3, 3)), vec_to_skew(ha[3:])],
        ]
    )

rmat_to_ha_hat

rmat_to_ha_hat(rmat: Array) -> Array

Converts a rotation matrix into its hat matrix representation in se(3) with zero translation. Formulation from "A geometric local frame approach for flexible multibody systems", by Sonneville, 2015, Eq 1.50, p. 15.

Parameters:

Name	Type	Description	Default
rmat	Array	Rotation matrix, (3, 3).	required

Returns:

Type	Description
Array	Hat matrix representation, (6, 6).

Source code in src/flapjax/algebra/se3.py

def rmat_to_ha_hat(rmat: Array) -> Array:
    r"""
    Converts a rotation matrix into its hat matrix representation in se(3) with zero translation. Formulation from
    "A geometric local frame approach for flexible multibody systems", by Sonneville, 2015, Eq 1.50, p. 15.
    :param rmat: Rotation matrix, ``(3, 3)``.
    :return: Hat matrix representation, ``(6, 6)``.
    """
    return chi(rmat)

ha_hat_to_ha

ha_hat_to_ha(ha_hat: Array) -> Array

Converts a se(3) hat matrix representation into se(3) vector. Formulation from Geometrically exact beam finite element formulated on the special Euclidean group SE(3), by Sonneville et al., 2013, Eq 15.

Parameters:

Name	Type	Description	Default
ha_hat	Array	Hat matrix representation, (6, 6).	required

Returns:

Type	Description
Array	se(3) vector, (6, ).

Source code in src/flapjax/algebra/se3.py

def ha_hat_to_ha(ha_hat: Array) -> Array:
    r"""
    Converts a se(3) hat matrix representation into se(3) vector. Formulation from Geometrically exact beam finite
    element formulated on the special Euclidean group SE(3), by Sonneville et al., 2013, Eq 15.
    :param ha_hat: Hat matrix representation, ``(6, 6)``.
    :return: se(3) vector, ``(6, )``.
    """
    ha_u = skew_to_vec(ha_hat[:3, 3:])
    ha_omega = skew_to_vec(ha_hat[:3, :3])
    return jnp.concatenate((ha_u, ha_omega), axis=-1)

ha_to_ha_check

ha_to_ha_check(ha: Array) -> Array

Converts a se(3) vector into its check matrix representation. Formulation from "Geometrically exact beam finite element formulated on the special Euclidean group SE(3)", by Sonneville, 2014, Eq 16.

Parameters:

Name	Type	Description	Default
ha	Array	se(3) algebra vector, (6, ).	required

Returns:

Type	Description
Array	Check matrix representation, (6, 6).

Source code in src/flapjax/algebra/se3.py

def ha_to_ha_check(ha: Array) -> Array:
    r"""
    Converts a se(3) vector into its check matrix representation. Formulation from "Geometrically exact beam finite element formulated
    on the special Euclidean group SE(3)", by Sonneville, 2014, Eq 16.
    :param ha: se(3) algebra vector, ``(6, )``.
    :return: Check matrix representation, ``(6, 6)``.
    """
    return -jnp.block(
        [
            [jnp.zeros((3, 3)), vec_to_skew(ha[:3])],
            [vec_to_skew(ha[:3]), vec_to_skew(ha[3:])],
        ]
    )

hg_to_ha_hat

hg_to_ha_hat(hg: Array) -> Array

Converts an SE(3) group element into its se(3) hat matrix representation.

Parameters:

Name	Type	Description	Default
hg	Array	SE(3) group element, (4, 4).	required

Returns:

Type	Description
Array	se(3) hat matrix representation, (6, 6).

Source code in src/flapjax/algebra/se3.py

def hg_to_ha_hat(hg: Array) -> Array:
    r"""
    Converts an SE(3) group element into its se(3) hat matrix representation.
    :param hg: SE(3) group element, ``(4, 4)``.
    :return: se(3) hat matrix representation, ``(6, 6)``.
    """
    rmat = hg[:3, :3]
    x = hg[:3, 3]

    return jnp.block(
        [
            [rmat, vec_to_skew(x) @ rmat],
            [jnp.zeros((3, 3)), rmat],
        ]
    )

hg_to_d

hg_to_d(hg1: Array, hg2: Array) -> Array

Obtains the relative configuration vector between two SE(3) group elements. Formulation from Geometrically exact beam finite element formulated on the special Euclidean group SE(3), by Sonneville et al., 2013, Eq 56.

Parameters:

Name	Type	Description	Default
hg1	Array	Base SE(3) group element at s=0, (4, 4).	required
hg2	Array	Tip SE(3) group element at s=L, (4, 4).	required

Returns:

Type	Description
Array	se(3) relative configuration vector, (6, ).

Source code in src/flapjax/algebra/se3.py

def hg_to_d(hg1: Array, hg2: Array) -> Array:
    r"""
    Obtains the relative configuration vector between two SE(3) group elements. Formulation from Geometrically exact
    beam finite element formulated on the special Euclidean group SE(3), by Sonneville et al., 2013, Eq 56.
    :param hg1: Base SE(3) group element at s=0, ``(4, 4)``.
    :param hg2: Tip SE(3) group element at s=L, ``(4, 4)``.
    :return: se(3) relative configuration vector, ``(6, )``.
    """
    return log_se3(hg_inv(hg1) @ hg2)

p

p(d: Array, ad_inv: Array) -> Array

Computes the :math:\mathbf{P}(\mathbf{d}) = \frac{d \mathbf{d}}{d \mathbf{h}_{AB}} matrix. Formulation from "A geometric local frame approach for flexible multibody systems", by Sonneville, 2015, Eq 6.141, p. 90.

Parameters:

Name	Type	Description	Default
d	Array	Relative se(3) configuration vector, (6, ).	required
ad_inv	Array	Adjoint action for rotation, (6, 6).	required

Returns:

Type	Description
Array	Matrix, (6, 12).

Source code in src/flapjax/algebra/se3.py

def p(
    d: Array,
    ad_inv: Array,
) -> Array:
    r"""
    Computes the :math:`\mathbf{P}(\mathbf{d}) = \frac{d \mathbf{d}}{d \mathbf{h}_{AB}}` matrix. Formulation
    from "A geometric local frame approach for flexible multibody systems", by Sonneville, 2015, Eq 6.141, p. 90.
    :param d: Relative se(3) configuration vector, ``(6, )``.
    :param ad_inv: Adjoint action for rotation, ``(6, 6)``.
    :return: Matrix, ``(6, 12)``.
    """

    t_inv_pos = t_inv_se3(d)
    t_inv_neg = t_inv_pos - ha_to_ha_hat(d)  # t_inv(-d)

    return jnp.concatenate((-t_inv_neg @ ad_inv, t_inv_pos @ ad_inv), axis=1)

t_star

t_star(s_l: Array, d: Array) -> Array

Matrix which described perturbations in the algebra element along an element with respect to the algebra elements at both ends of the element, :math:T^*(s, \mathbf{d}) = \frac{d \mathbf{h}(s)}{d \mathbf{h}_A} or :math:\frac{d \mathbf{h}(s)}{d \mathbf{h}_B}. Formulation from Geometrically exact beam finite element formulated on the special Euclidean group SE(3), by Sonneville et al., 2013, Eq 70.

Parameters:

Name	Type	Description	Default
s_l	Array	Relative position along the element :math:\frac{s}{l0} \in [0, 1], ().	required
d	Array	Relative se(3) configuration vector, (6, ).	required

Returns:

Type	Description
Array	:math:T^*(s, \mathbf{d}) matrix, (6, 6).

Source code in src/flapjax/algebra/se3.py

def t_star(s_l: Array, d: Array) -> Array:
    r"""
    Matrix which described perturbations in the algebra element along an element with respect to the algebra elements at
    both ends of the element, :math:`T^*(s, \mathbf{d}) = \frac{d \mathbf{h}(s)}{d \mathbf{h}_A}` or
    :math:`\frac{d \mathbf{h}(s)}{d \mathbf{h}_B}`. Formulation from Geometrically exact beam finite element
    formulated on the special Euclidean group SE(3), by Sonneville et al., 2013, Eq 70.
    :param s_l: Relative position along the element :math:`\frac{s}{l0} \in [0, 1]`, ().
    :param d: Relative se(3) configuration vector, ``(6, )``.
    :return: :math:`T^*(s, \mathbf{d})` matrix, ``(6, 6)``.
    """
    return s_l * t_se3(s_l * d) @ t_inv_se3(d)

q

q(s_l: Array, d: Array, ad_inv: Array) -> Array

Matrix which described perturbations in the algebra element along an element with respect to the algebra elements at both ends of the element, :math:Q(s, \mathbf{d}) = [\mathbf{I}_{6 \times 6} - T^*(s, \mathbf{d}) & T^*(s, \mathbf{d})]. Formulation from "A geometric local frame approach for flexible multibody systems", by Sonneville, 2015, Eq 6.145, p. 90.

Parameters:

Name	Type	Description	Default
s_l	Array	Relative position along the element :math:\frac{s}{l0} \in [0, 1], ().	required
d	Array	Relative se(3) configuration vector, (6, ).	required
ad_inv	Array	Adjoint action for base rotation, (6, 6).	required

Returns:

Type	Description
Array	:math:Q(s, \mathbf{d}) matrix, (6, 12).

Source code in src/flapjax/algebra/se3.py

def q(s_l: Array, d: Array, ad_inv: Array) -> Array:
    r"""
    Matrix which described perturbations in the algebra element along an element with respect to the algebra elements at
    both ends of the element, :math:`Q(s, \mathbf{d}) = [\mathbf{I}_{6 \times 6} - T^*(s, \mathbf{d}) &
    T^*(s, \mathbf{d})]`. Formulation from "A geometric local frame approach for flexible multibody systems",
    by Sonneville, 2015, Eq 6.145, p. 90.
    :param s_l: Relative position along the element :math:`\frac{s}{l0} \in [0, 1]`, ().
    :param d: Relative se(3) configuration vector, ``(6, )``.
    :param ad_inv: Adjoint action for base rotation, ``(6, 6)``.
    :return: :math:`Q(s, \mathbf{d})` matrix, ``(6, 12)``.
    """
    t_star_ = t_star(s_l, d)

    return jnp.concatenate(((jnp.eye(6) - t_star_) @ ad_inv, t_star_ @ ad_inv), axis=1)

q_dot

q_dot(
    s_l: Array, d: Array, d_dot: Array, ad_inv: Array
) -> Array

Time derivative of the matrix which described perturbations in the algebra element along an element with respect to the algebra elements at both ends of the element, :math:\dot{Q}(s, \mathbf{d}).

Parameters:

Name	Type	Description	Default
s_l	Array	Relative position along the element :math:\frac{s}{l0} \in [0, 1], ().	required
d	Array	Relative se(3) configuration vector, (6, ).	required
d_dot	Array	Time derivative of relative se(3) configuration vector, (6, ).	required
ad_inv	Array	Adjoint action for base rotation, (6, 6).	required

Returns:

Type	Description
Array	Time derivative of :math:Q(s, \mathbf{d}) matrix, (6, 12).

Source code in src/flapjax/algebra/se3.py

def q_dot(s_l: Array, d: Array, d_dot: Array, ad_inv: Array) -> Array:
    r"""
    Time derivative of the matrix which described perturbations in the algebra element along an element with respect to
    the algebra elements at both ends of the element, :math:`\dot{Q}(s, \mathbf{d})`.
    :param s_l: Relative position along the element :math:`\frac{s}{l0} \in [0, 1]`, ().
    :param d: Relative se(3) configuration vector, ``(6, )``.
    :param d_dot: Time derivative of relative se(3) configuration vector, ``(6, )``.
    :param ad_inv: Adjoint action for base rotation, ``(6, 6)``.
    :return: Time derivative of :math:`Q(s, \mathbf{d})` matrix, ``(6, 12)``.
    """

    _, tangents = jax.jvp(
        lambda d_: q(s_l, d_, ad_inv),
        primals=[d],
        tangents=[d_dot],
    )

    return tangents

so3

vec_to_skew

vec_to_skew(vec: Array) -> Array

Converts a 3D vector to a skew-symmetric matrix.

Parameters:

Name	Type	Description	Default
vec	Array	3D vector, (3, ).	required

Returns:

Type	Description
Array	Skew-symmetric matrix, (3, 3)

Source code in src/flapjax/algebra/so3.py

def vec_to_skew(vec: Array) -> Array:
    r"""
    Converts a 3D vector to a skew-symmetric matrix.
    :param vec: 3D vector, ``(3, )``.
    :return: Skew-symmetric matrix, ``(3, 3)``
    """
    return jnp.array(
        ((0.0, -vec[2], vec[1]), (vec[2], 0.0, -vec[0]), (-vec[1], vec[0], 0.0))
    )

skew_to_vec

skew_to_vec(mat: Array) -> Array

Converts a skew-symmetric matrix to a 3D vector. Note this refers to both skew symmetric entries for consistent gradients.

Parameters:

Name	Type	Description	Default
mat	Array	Skew-symmetric matrix, (3, 3)	required

Returns:

Type	Description
Array	3D vector, (3, ).

Source code in src/flapjax/algebra/so3.py

def skew_to_vec(mat: Array) -> Array:
    r"""
    Converts a skew-symmetric matrix to a 3D vector. Note this refers to both skew symmetric entries for consistent
    gradients.
    :param mat: Skew-symmetric matrix, ``(3, 3)``
    :return: 3D vector, ``(3, )``.
    """
    a1 = 0.5 * (mat[2, 1] - mat[1, 2])
    a2 = 0.5 * (mat[0, 2] - mat[2, 0])
    a3 = 0.5 * (mat[1, 0] - mat[0, 1])

    return jnp.array((a1, a2, a3))

alpha

alpha(b: Array) -> Array

Computes the alpha function for SO(3) operations. This includes a small angle approximation as b approaches zero. Formulation from Geometrically exact beam finite element formulated on the special Euclidean group SE(3), by Sonneville et al., 2013, Eq A.4

Parameters:

Name	Type	Description	Default
b	Array	Input vector, (3, ).	required

Returns:

Type	Description
Array	Alpha value, ().

Source code in src/flapjax/algebra/so3.py

def alpha(b: Array) -> Array:
    r"""
    Computes the alpha function for SO(3) operations. This includes a small angle approximation as b approaches zero.
    Formulation from Geometrically exact beam finite element formulated on the special Euclidean group SE(3), by
    Sonneville et al., 2013, Eq A.4
    :param b: Input vector, ``(3, )``.
    :return: Alpha value, ().
    """
    b_norm = jnp.linalg.norm(b)

    def alpha_full() -> Array:
        return jnp.sin(b_norm) / b_norm

    def alpha_small_angle() -> Array:
        return 1.0 - b_norm**2 / 6.0

    return cond(b_norm > SMALL_ANG_THRESH, alpha_full, alpha_small_angle)

beta

beta(b: Array) -> Array

Computes the beta function for SO(3) operations. This includes a small angle approximation as b approaches zero. Formulation from Geometrically exact beam finite element formulated on the special Euclidean group SE(3), by Sonneville et al., 2013, Eq A.4.

Parameters:

Name	Type	Description	Default
b	Array	Input vector, (3, ).	required

Returns:

Type	Description
Array	Beta value, ()

Source code in src/flapjax/algebra/so3.py

def beta(b: Array) -> Array:
    r"""
    Computes the beta function for SO(3) operations. This includes a small angle approximation as b approaches zero.
    Formulation from Geometrically exact beam finite element formulated on the special Euclidean group SE(3), by
    Sonneville et al., 2013, Eq A.4.
    :param b: Input vector, ``(3, )``.
    :return: Beta value, ()
    """
    b_norm2 = jnp.inner(b, b)
    b_norm = jnp.sqrt(b_norm2)

    def beta_full() -> Array:
        return 2.0 * (1.0 - jnp.cos(b_norm)) / b_norm2

    def beta_small_angle() -> Array:
        return 0.5 - b_norm2 / 24.0

    return cond(b_norm > SMALL_ANG_THRESH, beta_full, beta_small_angle)

bound_h_omega

bound_h_omega(h_omega: Array) -> Array

Bounds the angle of a rotation vector to be within [-pi, pi].

Parameters:

Name	Type	Description	Default
h_omega	Array	Cartesian rotation vector, (3, ).	required

Returns:

Type	Description
Array	Bounded Cartesian rotation vector, (3, ).

Source code in src/flapjax/algebra/so3.py

def bound_h_omega(h_omega: Array) -> Array:
    r"""
    Bounds the angle of a rotation vector to be within `[-pi, pi]`.
    :param h_omega: Cartesian rotation vector, ``(3, )``.
    :return: Bounded Cartesian rotation vector, ``(3, )``.
    """
    ang = jnp.linalg.norm(h_omega)

    def nonzero_ang() -> Array:
        n = h_omega / ang
        bounded_ang = clip_to_pi(ang)
        return bounded_ang * n

    def small_ang() -> Array:
        return h_omega

    return cond(ang > SMALL_ANG_THRESH, nonzero_ang, small_ang)

bracket_so3

bracket_so3(vec1: Array, vec2: Array) -> Array

Computes the Lie bracket of two so(3) elements represented as vectors.

Parameters:

Name	Type	Description	Default
vec1	Array	Vector 1, (3, ).	required
vec2	Array	Vector 2, (3, ).	required

Returns:

Type	Description
Array	Lie bracket, (3, 3).

Source code in src/flapjax/algebra/so3.py

def bracket_so3(vec1: Array, vec2: Array) -> Array:
    r"""
    Computes the Lie bracket of two so(3) elements represented as vectors.
    :param vec1: Vector 1, ``(3, )``.
    :param vec2: Vector 2, ``(3, )``.
    :return: Lie bracket, ``(3, 3)``.
    """
    mat1 = vec_to_skew(vec1)
    mat2 = vec_to_skew(vec2)

    return mat1 @ mat2 - mat2 @ mat1

bracket_neg_so3

bracket_neg_so3(vec1: Array, vec2: Array) -> Array

Computes the negative Lie bracket of two so(3) elements represented as vectors.

Parameters:

Name	Type	Description	Default
vec1	Array	Vector 1, (3, )	required
vec2	Array	Vector 2, (3, )	required

Returns:

Type	Description
Array	Negative Lie bracket, (3, 3)

Source code in src/flapjax/algebra/so3.py

def bracket_neg_so3(vec1: Array, vec2: Array) -> Array:
    r"""
    Computes the negative Lie bracket of two so(3) elements represented as vectors.
    :param vec1: Vector 1, ``(3, )``
    :param vec2: Vector 2, ``(3, )``
    :return: Negative Lie bracket, ``(3, 3)``
    """
    mat1 = vec_to_skew(vec1)
    mat2 = vec_to_skew(vec2)

    return mat1 @ mat2 + mat2 @ mat1

t_so3

t_so3(ha_omega: Array) -> Array

Computes the tangent operator for SO(3) given a rotation vector. Includes a small angle approximation as the rotation approaches zero. Formulation from Geometrically exact beam finite element formulated on the special Euclidean group SE(3), by Sonneville et al., 2013, Eq A.6

Parameters:

Name	Type	Description	Default
ha_omega	Array	Rotation vector, (3, ).	required

Returns:

Type	Description
Array	Tangent operator, (3, 3)

Source code in src/flapjax/algebra/so3.py

def t_so3(ha_omega: Array) -> Array:
    r"""
    Computes the tangent operator for SO(3) given a rotation vector. Includes a small angle approximation as the
    rotation approaches zero. Formulation from Geometrically exact beam finite element formulated on the special
    Euclidean group SE(3), by Sonneville et al., 2013, Eq A.6
    :param ha_omega: Rotation vector, ``(3, )``.
    :return: Tangent operator, ``(3, 3)``
    """

    def t_so3_full() -> Array:
        return (
            jnp.eye(3)
            - 0.5 * beta(ha_omega) * vec_to_skew(ha_omega)
            + (1.0 - alpha(ha_omega))
            / jnp.inner(ha_omega, ha_omega)
            * matrix2(vec_to_skew(ha_omega))
        )

    def t_so3_small_angle() -> Array:
        return t_sum(vec_to_skew(ha_omega), 2)

    ang_mag2 = jnp.inner(ha_omega, ha_omega)
    return cond(ang_mag2 > SMALL_ANG_THRESH, t_so3_full, t_so3_small_angle)

t_inv_so3

t_inv_so3(ha_omega: Array) -> Array

Computes the inverse tangent operator for SO(3) given a rotation vector. Includes a small angle approximation as the rotation approaches zero. Formulation from Geometrically exact beam finite element formulated on the special Euclidean group SE(3), by Sonneville et al., 2013, Eq A.7

Parameters:

Name	Type	Description	Default
ha_omega	Array	Rotation vector, (3, )	required

Returns:

Type	Description
Array	Inverse tangent operator, (3, 3)

Source code in src/flapjax/algebra/so3.py

def t_inv_so3(ha_omega: Array) -> Array:
    r"""
    Computes the inverse tangent operator for SO(3) given a rotation vector. Includes a small angle approximation as the
    rotation approaches zero. Formulation from Geometrically exact beam finite element formulated on the special
    Euclidean group SE(3), by Sonneville et al., 2013, Eq A.7
    :param ha_omega: Rotation vector, ``(3, )``
    :return: Inverse tangent operator, ``(3, 3)``
    """

    def t_inv_so3_full() -> Array:
        return (
            jnp.eye(3)
            + 0.5 * vec_to_skew(ha_omega)
            + (1.0 - alpha(ha_omega) / beta(ha_omega))
            / jnp.inner(ha_omega, ha_omega)
            * matrix2(vec_to_skew(ha_omega))
        )

    def t_inv_so3_small_angle() -> Array:
        return t_inv_sum(vec_to_skew(ha_omega), 2)

    ang_mag2 = jnp.linalg.norm(ha_omega)
    return cond(ang_mag2 > SMALL_ANG_THRESH, t_inv_so3_full, t_inv_so3_small_angle)

exp_so3

exp_so3(ha_omega: Array) -> Array

Computes the exponential map from so(3) to SO(3) given a rotation vector. Includes a small angle approximation as the angle approaches zero.

Parameters:

Name	Type	Description	Default
ha_omega	Array	Rotation vector, (3, )	required

Returns:

Type	Description
Array	Rotation matrix, (3, 3)

Source code in src/flapjax/algebra/so3.py

def exp_so3(ha_omega: Array) -> Array:
    r"""
    Computes the exponential map from so(3) to SO(3) given a rotation vector. Includes a small angle approximation as
    the angle approaches zero.
    :param ha_omega: Rotation vector, ``(3, )``
    :return: Rotation matrix, ``(3, 3)``
    """

    def exp_so3_full() -> Array:
        # has a singularity as ha_omega -> 0
        ang = jnp.linalg.norm(ha_omega)
        return (
            jnp.eye(3)
            + jnp.sin(ang) / ang * vec_to_skew(ha_omega)
            + (1.0 - jnp.cos(ang)) / ang**2 * matrix2(vec_to_skew(ha_omega))
        )

    def exp_so3_small_angle() -> Array:
        return exp_sum(vec_to_skew(ha_omega), 2)

    ang_mag2 = jnp.inner(ha_omega, ha_omega)
    return cond(ang_mag2 > SMALL_ANG_THRESH, exp_so3_full, exp_so3_small_angle)

log_so3

log_so3(rmat: Array) -> Array

Computes the logarithmic map from SO(3) to so(3) given a rotation matrix. Includes a small angle approximation as the angle approaches zero.

Parameters:

Name	Type	Description	Default
rmat	Array	Rotation matrix, (3, 3)	required

Returns:

Type	Description
Array	Rotation vector, (3, ).

Source code in src/flapjax/algebra/so3.py

def log_so3(rmat: Array) -> Array:
    r"""
    Computes the logarithmic map from SO(3) to so(3) given a rotation matrix. Includes a small angle approximation as
    the angle approaches zero.
    :param rmat: Rotation matrix, ``(3, 3)``
    :return: Rotation vector, ``(3, )``.
    """

    cos_theta = 0.5 * (jnp.trace(rmat) - 1.0)

    def log_so3_full() -> Array:
        # acos is only computed here, where theta is large enough that its gradient is finite
        theta = jnp.acos(cos_theta)
        return skew_to_vec(theta / (2.0 * jnp.sin(theta)) * (rmat - rmat.T))

    def log_so3_small_angle() -> Array:
        return skew_to_vec(log_sum(rmat, 2))

    return cond(
        cos_theta < jnp.cos(SMALL_ANG_THRESH), log_so3_full, log_so3_small_angle
    )

test_routines

check_if_so3_g

check_if_so3_g(
    rmat: Array, raise_if_false: bool = True
) -> bool

Check if rotation matrix is a valid SO3 group element

Parameters:

Name	Type	Description	Default
rmat	Array	Rotation matrix, (3, 3)	required
raise_if_false	bool	If the check fails, raise ValueError	True

Returns:

Type	Description
bool	Boolean indicating if matrix is SO3

Source code in src/flapjax/algebra/test_routines.py

def check_if_so3_g(
    rmat: Array,
    raise_if_false: bool = True,
) -> bool:
    r"""
    Check if rotation matrix is a valid SO3 group element
    :param rmat: Rotation matrix, ``(3, 3)``
    :param raise_if_false: If the check fails, raise ValueError
    :return: Boolean indicating if matrix is SO3
    """
    column_mags = jnp.linalg.norm(rmat, axis=0)
    row_mags = jnp.linalg.norm(rmat, axis=1)

    # if not correct shape
    if rmat.shape != (3, 3):
        if raise_if_false:
            raise ValueError("Matrix not SO3 as shape is not (3, 3)")
        return False

    # if not unit magnitude
    if not jnp.allclose(jnp.concatenate((column_mags, row_mags)), 1.0):
        if raise_if_false:
            raise ValueError(
                "Matrix is not SO3 as rows or columns are not of unit magnitude"
            )
        return False

    # if not orthogonal
    if not (
        jnp.allclose(rmat.T @ rmat, jnp.eye(3))
        and jnp.allclose(rmat @ rmat.T, jnp.eye(3))
    ):
        if raise_if_false:
            raise ValueError("Matrix is not SO3 as it is not orthogonal")
        return False
    return True

check_if_so3_a

check_if_so3_a(
    h_tilde: Array, raise_if_false: bool = True
) -> bool

Check if rotation matrix is a valid so3 algebra element

Parameters:

Name	Type	Description	Default
h_tilde	Array	Algebra matrix, (3, 3)	required
raise_if_false	bool	If the check fails, raise ValueError	True

Returns:

Type	Description
bool	Boolean indicating if matrix is so3

Source code in src/flapjax/algebra/test_routines.py

def check_if_so3_a(h_tilde: Array, raise_if_false: bool = True) -> bool:
    r"""
    Check if rotation matrix is a valid so3 algebra element
    :param h_tilde: Algebra matrix, ``(3, 3)``
    :param raise_if_false: If the check fails, raise ValueError
    :return: Boolean indicating if matrix is so3
    """
    # check shape
    if h_tilde.shape != (3, 3):
        if raise_if_false:
            raise ValueError("Matrix not so3 as shape is not (3, 3)")
        return False

    # check for nonzero diagonal elements
    if jnp.any(jnp.diagonal(h_tilde)) != 0.0:
        if raise_if_false:
            raise ValueError("Matrix not so3 as diagonal elements are not zero")
        return False

    # check for skew symmetry
    if jnp.any(h_tilde + h_tilde.T):
        if raise_if_false:
            raise ValueError("Matrix not so3 as it is not skew symmetric")
        return False
    return True

check_if_se3_g

check_if_se3_g(
    hg: Array, raise_if_false: bool = True
) -> bool

Check if matrix is a valid SE3 group element

Parameters:

Name	Type	Description	Default
hg	Array	SE(3) matrix, (4, 4)	required
raise_if_false	bool	If the check fails, raise ValueError	True

Returns:

Type	Description
bool	Boolean indicating if matrix is SE(3)

Source code in src/flapjax/algebra/test_routines.py

def check_if_se3_g(hg: Array, raise_if_false: bool = True) -> bool:
    r"""
    Check if matrix is a valid SE3 group element
    :param hg: SE(3) matrix, ``(4, 4)``
    :param raise_if_false: If the check fails, raise ValueError
    :return: Boolean indicating if matrix is SE(3)
    """
    # check shape
    if hg.shape != (4, 4):
        if raise_if_false:
            raise ValueError("Matrix not SE3 as shape is not (4, 4)")
        return False

    # check rotational component
    if not check_if_so3_g(hg[:3, :3], raise_if_false=raise_if_false):
        return False

    # check last row
    if not jnp.allclose(hg[3, :], jnp.array([0.0, 0.0, 0.0, 1.0])):
        if raise_if_false:
            raise ValueError("Matrix not SE3 as last row is not [0, 0, 0, 1]")
        return False
    return True

check_if_all_se3_g

check_if_all_se3_g(
    hgs: Array, raise_if_false: bool = True
) -> bool

Check if array of matrices are valid SE3 group elements

Parameters:

Name	Type	Description	Default
hgs	Array	SE(3) matrices, (..., 4, 4)	required
raise_if_false	bool	If the check fails, raise ValueError	True

Returns:

Type	Description
bool	Boolean indicating if all matrices are SE(3)

Source code in src/flapjax/algebra/test_routines.py

def check_if_all_se3_g(hgs: Array, raise_if_false: bool = True) -> bool:
    r"""
    Check if array of matrices are valid SE3 group elements
    :param hgs: SE(3) matrices, ``(..., 4, 4)``
    :param raise_if_false: If the check fails, raise ValueError
    :return: Boolean indicating if all matrices are SE(3)
    """
    # check shape
    if hgs.shape[-2:] != (4, 4):
        if raise_if_false:
            raise ValueError("Input not se3 as last two dimensions are not (4, 4)")
        return False

    def check_if_so3_g_jittable(rmat: Array) -> Array:
        column_mags = jnp.linalg.norm(rmat, axis=0)
        row_mags = jnp.linalg.norm(rmat, axis=1)

        # check if unit magnitude
        out = jnp.all(jnp.allclose(jnp.concatenate((column_mags, row_mags)), 1.0))

        # check if orthogonal
        out &= jnp.all(jnp.allclose(rmat.T @ rmat, jnp.eye(3), atol=1e-5, rtol=1e-3))
        out &= jnp.all(jnp.allclose(rmat @ rmat.T, jnp.eye(3), atol=1e-5, rtol=1e-3))
        return out

    hgs_flat = hgs.reshape(-1, 4, 4)

    results = jnp.all(
        vmap(check_if_so3_g_jittable, in_axes=0, out_axes=0)(hgs_flat[:, :3, :3])
    )
    results &= jnp.all(jnp.allclose(hgs_flat[:, 3, :3], 0.0))
    results &= jnp.all(jnp.allclose(hgs_flat[:, 3, 3], 1.0))

    if not results:
        if raise_if_false:
            raise ValueError("Not all matrices are se3 elements")
        return False
    return True

check_if_se3_a

check_if_se3_a(
    h_tilde: Array, raise_if_false: bool = True
) -> bool

Check if matrix is a valid se(3) group element

Parameters:

Name	Type	Description	Default
h_tilde	Array	se(3) matrix, (4, 4)	required
raise_if_false	bool	If the check fails, raise ValueError	True

Returns:

Type	Description
bool	Boolean indicating if matrix is se(3)

Source code in src/flapjax/algebra/test_routines.py

def check_if_se3_a(h_tilde: Array, raise_if_false: bool = True) -> bool:
    r"""
    Check if matrix is a valid se(3) group element
    :param h_tilde: se(3) matrix, ``(4, 4)``
    :param raise_if_false: If the check fails, raise ValueError
    :return: Boolean indicating if matrix is se(3)
    """
    # check shape
    if h_tilde.shape != (4, 4):
        if raise_if_false:
            raise ValueError("Matrix not se3 as shape is not (4, 4)")
        return False

    # check so3 component
    if not check_if_so3_a(h_tilde[:3, :3], raise_if_false=raise_if_false):
        return False

    # check last row and column
    if jnp.any(h_tilde[3, :]):
        if raise_if_false:
            raise ValueError("Matrix not se3 as last row is not zero")
        return False
    return True

check_if_all_se3_a

check_if_all_se3_a(
    h_tildes: Array, raise_if_false: bool = True
) -> bool

Check if array of matrices are valid se(3) algebra elements

Parameters:

Name	Type	Description	Default
h_tildes	Array	se(3) matrices, (..., 4, 4)	required
raise_if_false	bool	If the check fails, raise ValueError	True

Returns:

Type	Description
bool	Boolean indicating if all matrices are se(3)

Source code in src/flapjax/algebra/test_routines.py

def check_if_all_se3_a(h_tildes: Array, raise_if_false: bool = True) -> bool:
    r"""
    Check if array of matrices are valid se(3) algebra elements
    :param h_tildes: se(3) matrices, ``(..., 4, 4)``
    :param raise_if_false: If the check fails, raise ValueError
    :return: Boolean indicating if all matrices are se(3)
    """
    h_tildes_flat = h_tildes.reshape(-1, 4, 4)

    # check bottom row
    results = jnp.all(jnp.allclose(h_tildes_flat[:, 3, :], 0.0))

    # check diagonal
    results &= jnp.all(jnp.allclose(h_tildes_flat[:, (0, 1, 2), (0, 1, 2)], 0.0))

    # check skew symmetry of so3 part
    results &= jnp.all(
        jnp.allclose(
            h_tildes_flat[:, :3, :3],
            -jnp.transpose(h_tildes_flat[:, :3, :3], (0, 2, 1)),
        )
    )

    if not results:
        if raise_if_false:
            raise ValueError("Not all matrices are se3 algebra elements")
        return False
    return True

k_t_expected

k_t_expected(
    coeffs: Array | Sequence[float], length: Array | float
) -> Array

Compute expected two-node beam undeformed element stiffness matrix given coefficients and length

Parameters:

Name	Type	Description	Default
coeffs	Array | Sequence[float]	Stiffness coefficients which make up the diagonal of the local stiffness matrix, (6, ).	required
length	Array | float	Beam length, ()	required

Returns:

Type	Description
Array	Beam tangent stiffness matrix, (12, 12)

Source code in src/flapjax/algebra/test_routines.py

def k_t_expected(coeffs: Array | Sequence[float], length: Array | float) -> Array:
    r"""
    Compute expected two-node beam undeformed element stiffness matrix given coefficients and length
    :param coeffs: Stiffness coefficients which make up the diagonal of the local stiffness matrix, ``(6, )``.
    :param length: Beam length, ()
    :return: Beam tangent stiffness matrix, ``(12, 12)``
    """
    if (isinstance(coeffs, Array) and coeffs.shape != (6,)) or (
        isinstance(coeffs, Sequence) and len(coeffs) != 6
    ):
        raise ValueError("Coefficients array must be of shape (6, )")

    if isinstance(length, Array) and not jnp.isscalar(length):
        raise ValueError("Length l0 must be a scalar value")

    eax, *_, gjx, eiy, eiz = coeffs

    k_upper_left = jnp.array(
        [
            [eax / length, 0.0, 0.0, 0.0, 0.0, 0.0],
            [0.0, 12.0 * eiz / length**3, 0.0, 0.0, 0.0, 6.0 * eiz / length**2],
            [0.0, 0.0, 12.0 * eiy / length**3, 0.0, -6.0 * eiy / length**2, 0.0],
            [0.0, 0.0, 0.0, gjx / length, 0.0, 0.0],
            [0.0, 0.0, -6.0 * eiy / length**2, 0.0, 4.0 * eiy / length, 0.0],
            [0.0, 6.0 * eiz / length**2, 0.0, 0.0, 0.0, 4.0 * eiz / length],
        ]
    )

    k_upper_right = jnp.array(
        [
            [-eax / length, 0.0, 0.0, 0.0, 0.0, 0.0],
            [0.0, -12.0 * eiz / length**3, 0.0, 0.0, 0.0, 6.0 * eiz / length**2],
            [0.0, 0.0, -12.0 * eiy / length**3, 0.0, -6.0 * eiy / length**2, 0.0],
            [0.0, 0.0, 0.0, -gjx / length, 0.0, 0.0],
            [0.0, 0.0, 6.0 * eiy / length**2, 0.0, 2.0 * eiy / length, 0.0],
            [0.0, -6.0 * eiz / length**2, 0.0, 0.0, 0.0, 2.0 * eiz / length],
        ]
    )

    k_lower_left = k_upper_right.T

    k_lower_right = k_upper_left
    k_lower_right = k_lower_right.at[1:3, 4:6].mul(-1.0)
    k_lower_right = k_lower_right.at[4:6, 1:3].mul(-1.0)

    return jnp.block([[k_upper_left, k_upper_right], [k_lower_left, k_lower_right]])

const_curvature_beam

const_curvature_beam(
    kappa: float | Array,
    s: float | Array,
    direction: Literal["y", "z"],
) -> Array

For a beam with constant curvature, with base node at the origin and curvature in the positive z direction (i.e., existing in the x_target-y plane with z=0), obtain the coordinates along the beam length for

Parameters:

Name	Type	Description	Default
kappa	float | Array	Curvature of the element, ()	required
s	float | Array	Position along the beam length, ()	required
direction	Literal['y', 'z']	Direction of moment applied, either y or z	required

Returns:

Type	Description
Array	Coordinate of point along the beam, (3, ).

Source code in src/flapjax/algebra/test_routines.py

def const_curvature_beam(
    kappa: float | Array, s: float | Array, direction: Literal["y", "z"]
) -> Array:
    r"""
    For a beam with constant curvature, with base node at the origin and curvature in the positive z direction
    (i.e., existing in the x_target-y plane with z=0), obtain the coordinates along the beam length for
    :param kappa: Curvature of the element, ``()``
    :param s: Position along the beam length, ``()``
    :param direction: Direction of moment applied, either ``y`` or ``z``
    :return: Coordinate of point along the beam, ``(3, )``.
    """

    x = jnp.sin(s * kappa) / kappa
    v_deflection = (1.0 - jnp.cos(s * kappa)) / kappa

    match direction:
        case "y":
            return jnp.array([x, 0.0, -v_deflection])
        case "z":
            return jnp.array([x, v_deflection, 0.0])
        case _:
            raise ValueError("Direction must be 'y' or 'z'")

get_curvature

get_curvature(d: Array) -> Array

Obtain curvature from relative configuration vector

Parameters:

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

Returns:

Type	Description
Array	Curvature of neutral axis, ()

Source code in src/flapjax/algebra/test_routines.py

def get_curvature(d: Array) -> Array:
    r"""
    Obtain curvature from relative configuration vector
    :param d: Relative configuration vector, ``(6, )``
    :return: Curvature of neutral axis, ()
    """
    return jnp.linalg.norm(jnp.cross(d[3:], d[:3])) / jnp.linalg.norm(d[:3])

get_torsion

get_torsion(d: Array) -> Array

Obtain torsion from relative configuration vector

Parameters:

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

Returns:

Type	Description
Array	Torsion of neutral axis, ()

Source code in src/flapjax/algebra/test_routines.py

def get_torsion(d: Array) -> Array:
    r"""
    Obtain torsion from relative configuration vector
    :param d: Relative configuration vector, ``(6, )``
    :return: Torsion of neutral axis, ()
    """
    return -jnp.inner(d[3:], d[:3]) / jnp.linalg.norm(d[:3])