kups.core.utils.math ¶

@jit
@vectorize(signature="(4)->(3)")
def cubic_roots(coefficients: Array) -> Array:
    """Find all roots of a cubic polynomial using the companion matrix method.

    Solves `ax³ + bx² + cx + d = 0` by computing eigenvalues of the companion
    matrix. Returns all three roots (real or complex).

    Args:
        coefficients: Array of shape `(..., 4)` containing `[a, b, c, d]` for each
            cubic polynomial.

    Returns:
        Array of shape `(..., 3)` containing the three roots. May be complex-valued.

    Example:
        ```python
        # Solve x^3 - 6x^2 + 11x - 6 = 0 (roots: 1, 2, 3)
        coeffs = jnp.array([1.0, -6.0, 11.0, -6.0])
        roots = cubic_roots(coeffs)
        ```

    Note:
        This method is numerically stable and handles multiple polynomials in
        parallel via vectorization.
    """
    a, b, c, d = coefficients
    C = jnp.array([[0, 0, -d / a], [1, 0, -c / a], [0, 1, -b / a]], dtype=jnp.float_)
    solutions = jnp.linalg.eigvals(C)
    return solutions

@jit
@vectorize(signature="(3,3)->(),(3,3)")
def det_and_inverse_3x3(A: jax.Array) -> tuple[jax.Array, jax.Array]:
    """Compute determinant and inverse of 3×3 matrices via the adjugate method.

    Efficiently computes both the determinant and inverse of 3×3 matrices using
    explicit formulas for the cofactor matrix. More efficient than general
    matrix inversion for small matrices.

    Args:
        A: Array of shape `(..., 3, 3)` containing 3×3 matrices.

    Returns:
        Tuple of (determinant, inverse):
            - determinant: Array of shape `(...)` containing scalar determinants.
            - inverse: Array of shape `(..., 3, 3)` containing inverted matrices.

    Example:
        ```python
        A = jnp.eye(3)
        det, A_inv = det_and_inverse_3x3(A)
        # det = 1.0, A_inv = identity matrix
        ```
    """
    assert A.shape == (3, 3), "Input must be a 3x3 matrix"

    # Determinant
    det = (
        A[0, 0] * (A[1, 1] * A[2, 2] - A[1, 2] * A[2, 1])
        - A[0, 1] * (A[1, 0] * A[2, 2] - A[1, 2] * A[2, 0])
        + A[0, 2] * (A[1, 0] * A[2, 1] - A[1, 1] * A[2, 0])
    )
    # Cofactor matrix
    cofactor = jnp.array(
        [
            [
                (A[1, 1] * A[2, 2] - A[1, 2] * A[2, 1]),
                -(A[1, 0] * A[2, 2] - A[1, 2] * A[2, 0]),
                (A[1, 0] * A[2, 1] - A[1, 1] * A[2, 0]),
            ],
            [
                -(A[0, 1] * A[2, 2] - A[0, 2] * A[2, 1]),
                (A[0, 0] * A[2, 2] - A[0, 2] * A[2, 0]),
                -(A[0, 0] * A[2, 1] - A[0, 1] * A[2, 0]),
            ],
            [
                (A[0, 1] * A[1, 2] - A[0, 2] * A[1, 1]),
                -(A[0, 0] * A[1, 2] - A[0, 2] * A[1, 0]),
                (A[0, 0] * A[1, 1] - A[0, 1] * A[1, 0]),
            ],
        ]
    )
    # Adjugate is transpose of cofactor
    adjugate = cofactor.T
    # Inverse
    return det, adjugate / det

@jit
def log_factorial_ratio(N: Array, M: Array) -> Array:
    """Compute log(N!/M!) efficiently using the log-gamma function.

    Uses the identity `log(n!) = lgamma(n+1)` to compute the ratio directly.

    Args:
        N: Integer array (numerator factorials).
        M: Integer array (denominator factorials).

    Returns:
        Result containing `log(N!/M!)` for each pair of elements.

    Example:
        ```python
        N = jnp.array([10, 5])
        M = jnp.array([5, 8])
        result = log_factorial_ratio(N, M)
        # Computes [log(10!/5!), log(5!/8!)]
        ```
    """
    return jax.lax.lgamma(N + 1.0) - jax.lax.lgamma(M + 1.0)

@jit(static_argnames=("hermitian",))
def logm(A: Array, *, hermitian: bool = False) -> Array:
    r"""Compute the principal matrix logarithm via eigendecomposition.

    Inverse of `jax.scipy.linalg.expm`: returns `L` with `expm(L) == A` for
    diagonalizable `A`. Diagonalizing $A = V\,\mathrm{diag}(\lambda)\,V^{-1}$
    gives $\log A = V\,\mathrm{diag}(\log\lambda)\,V^{-1}$.

    Args:
        A: Array of shape `(..., n, n)`. Must be diagonalizable; the result is
            inaccurate for defective matrices (e.g. non-trivial Jordan blocks).
        hermitian: If True, treat `A` as Hermitian/symmetric and use `eigh`.
            This gives a real result for positive-definite `A` and is
            differentiable. If False (default), `eig` handles general matrices
            but the result is complex and is not differentiable with respect to
            `A`, since JAX does not implement derivatives of non-symmetric
            eigenvectors.

    Returns:
        Array of shape `(..., n, n)` holding the principal logarithm. Real when
        `hermitian` and `A` is positive-definite, complex otherwise.

    Example:
        ```python
        A = jax.scipy.linalg.expm(B)
        B_recovered = logm(A)  # ≈ B
        ```
    """
    if hermitian:
        w, V = jnp.linalg.eigh(A)
        return (V * jnp.log(w)[..., None, :]) @ jnp.conjugate(jnp.swapaxes(V, -1, -2))
    w, V = jnp.linalg.eig(A)
    return (V * jnp.log(w)[..., None, :]) @ jnp.linalg.inv(V)

def next_higher_power(value: Array, base: Array | float = 2.0) -> Array:
    """Compute the next higher power of a given base for each element.

    For each element in `value`, finds the smallest power of `base` that is
    greater than or equal to that element.

    Args:
        value: Array of shape `(...)` containing input values.
        base: Base for the power calculation (default is 2.0).

    Returns:
        Array of shape `(...)` containing the next higher powers.

    Example:
        ```python
        values = jnp.array([3, 5, 10])
        result = next_higher_power(values, base=2)
        # result = [4, 8, 16]
        ```
    """
    value = jnp.ceil(value).astype(int)
    log_base = jnp.log(base)
    exponents = jnp.ceil(jnp.log(value) / log_base)
    result = jnp.power(base, exponents)
    # Linear growth for base <= 1
    return jnp.where(base <= 1, value, result).astype(int)

@jit
def solve_affine_ode(A: Array, b: Array, x0: Array, dt: Array | float) -> Array:
    r"""Solve the affine ODE $\dot{x} = A\,x + b$ exactly over $[0, \Delta t]$.

    Computes

    $$x(\Delta t) = e^{A\,\Delta t}\,x_0 + \varphi_1(A\,\Delta t)\,\Delta t\,b,
    \qquad \varphi_1(M) = M^{-1}(e^{M} - I).$$

    Uses the augmented-matrix trick (Al-Mohy & Higham 2011) so that a single
    matrix exponential delivers both terms and the formula is well-defined
    even when ``A`` has zero eigenvalues:

    $$\exp\!\left(\Delta t \begin{bmatrix} A & b \\ 0 & 0 \end{bmatrix}\right)
      = \begin{bmatrix} e^{A\,\Delta t} & \varphi_1(A\,\Delta t)\,\Delta t\,b \\ 0 & 1 \end{bmatrix}.$$

    Args:
        A: Coefficient matrix, shape ``(..., n, n)``. Any shape (triangular
            or general) is accepted; ``expm`` handles all cases.
        b: Constant inhomogeneous term, shape ``(..., n)``.
        x0: Initial state, shape ``(..., n)``.
        dt: Scalar timestep (broadcastable to the batch shape).

    Returns:
        State at ``dt``, shape ``(..., n)``.
    """
    n = x0.shape[-1]
    # Augmented matrix M = [[A, b], [0, 0]], shape (..., n+1, n+1)
    A_padded = jnp.concatenate([A, b[..., None]], axis=-1)  # (..., n, n+1)
    zero_row = jnp.zeros((*A.shape[:-2], 1, n + 1), dtype=A.dtype)
    M = jnp.concatenate([A_padded, zero_row], axis=-2)  # (..., n+1, n+1)
    E = jax.scipy.linalg.expm(M * dt)
    return jnp.einsum("...ij,...j->...i", E[..., :n, :n], x0) + E[..., :n, n]

@jit(static_argnames=("lower",))
def triangular_3x3_det_and_inverse(
    A: jax.Array, *, lower: bool = True
) -> tuple[jax.Array, jax.Array]:
    """Compute determinant and inverse of triangular 3×3 matrices.

    Exploits triangular structure with closed-form expressions: the determinant
    is the product of diagonal elements and the inverse follows from forward
    substitution. Upper-triangular inputs are handled by transposition.

    Args:
        A: Array of shape `(..., 3, 3)` containing triangular matrices.
        lower: Whether matrices are lower (True) or upper (False) triangular.

    Returns:
        Tuple of (determinant, inverse):
            - determinant: Array of shape `(...)` containing diagonal products.
            - inverse: Array of shape `(..., 3, 3)` containing inverted matrices.

    Example:
        ```python
        L = jnp.array([[1, 0, 0], [2, 3, 0], [4, 5, 6]])
        det, L_inv = triangular_3x3_det_and_inverse(L, lower=True)
        # det = 18 (product of diagonals: 1*3*6)
        ```
    """
    if not lower:
        det, inv = triangular_3x3_det_and_inverse(jnp.swapaxes(A, -1, -2), lower=True)
        return det, jnp.swapaxes(inv, -1, -2)

    a = A[..., 0, 0]
    b = A[..., 1, 0]
    c = A[..., 1, 1]
    d = A[..., 2, 0]
    e = A[..., 2, 1]
    f = A[..., 2, 2]

    det = a * c * f
    m00, m11, m22 = 1.0 / a, 1.0 / c, 1.0 / f
    m10 = -b / (a * c)
    m21 = -e / (c * f)
    m20 = (e * b - c * d) / (a * c * f)

    zero = jnp.zeros_like(a)
    inv = jnp.stack(
        [
            jnp.stack([m00, zero, zero], axis=-1),
            jnp.stack([m10, m11, zero], axis=-1),
            jnp.stack([m20, m21, m22], axis=-1),
        ],
        axis=-2,
    )
    return det, inv

@jit(static_argnames=("lower",))
def triangular_3x3_expm(A: Array, *, lower: bool = True) -> Array:
    r"""Matrix exponential of triangular 3×3 matrices.

    Closed form from the fact that ``B`` and ``A = expm(B)`` commute: the diagonal
    of ``A`` is ``exp`` of ``B``'s diagonal, and the off-diagonal couplings are
    divided differences of ``exp`` over the diagonal entries. Using divided
    differences keeps repeated diagonal entries (e.g. cubic / tetragonal cells)
    well-conditioned.

    Real and gradient-safe inverse of
    [triangular_3x3_logm][kups.core.utils.math.triangular_3x3_logm]; agrees with
    ``jax.scipy.linalg.expm`` on triangular matrices but avoids its Padé
    scaling-and-squaring iteration.

    Args:
        A: Array of shape `(..., 3, 3)` containing triangular matrices.
        lower: Whether matrices are lower (True) or upper (False) triangular.

    Returns:
        Array of shape `(..., 3, 3)` containing the triangular exponentials.
    """
    M = A if lower else jnp.swapaxes(A, -1, -2)
    d0, d1, d2 = M[..., 0, 0], M[..., 1, 1], M[..., 2, 2]
    a, b, c = M[..., 1, 0], M[..., 2, 0], M[..., 2, 1]
    zero = jnp.zeros_like(d0)
    out = jnp.stack(
        [
            jnp.stack([jnp.exp(d0), zero, zero], axis=-1),
            jnp.stack([a * _ddexp(d0, d1), jnp.exp(d1), zero], axis=-1),
            jnp.stack(
                [
                    b * _ddexp(d0, d2) + a * c * _ddexp2(d0, d1, d2),
                    c * _ddexp(d1, d2),
                    jnp.exp(d2),
                ],
                axis=-1,
            ),
        ],
        axis=-2,
    )
    return out if lower else jnp.swapaxes(out, -1, -2)

def triangular_3x3_from_tril(tril: Array) -> Array:
    """Assemble a lower-triangular 3×3 matrix from its 6 elements.

    Args:
        tril: Array of shape `(..., 6)` holding ``[m00, m10, m11, m20, m21, m22]``.

    Returns:
        Array of shape `(..., 3, 3)`::

            [[m00,   0,   0],
             [m10, m11,   0],
             [m20, m21, m22]]
    """
    zero = jnp.zeros_like(tril[..., :1])
    return jnp.stack(
        [
            jnp.concatenate([tril[..., 0:1], zero, zero], axis=-1),
            jnp.concatenate([tril[..., 1:3], zero], axis=-1),
            tril[..., 3:6],
        ],
        axis=-2,
    )

@jit(static_argnames=("lower",))
def triangular_3x3_logm(A: Array, *, lower: bool = True) -> Array:
    r"""Principal matrix logarithm of triangular 3×3 matrices with positive diagonal.

    Closed form from the fact that ``L`` and ``B = logm(L)`` commute (``BL = LB``):
    the diagonal of ``B`` is ``log`` of ``L``'s diagonal, and the off-diagonal
    couplings are divided differences of ``log`` over the diagonal entries. Using
    divided differences keeps repeated diagonal entries (e.g. cubic / tetragonal
    cells) well-conditioned.

    Real and gradient-safe, unlike the eigendecomposition-based
    [logm][kups.core.utils.math.logm] which is complex and inaccurate for these
    (defective) matrices. Inverse of ``jax.scipy.linalg.expm`` on triangular
    matrices.

    Args:
        A: Array of shape `(..., 3, 3)` containing triangular matrices with
            strictly positive diagonal.
        lower: Whether matrices are lower (True) or upper (False) triangular.

    Returns:
        Array of shape `(..., 3, 3)` containing the triangular logarithms.
    """
    M = A if lower else jnp.swapaxes(A, -1, -2)
    l0, l1, l2 = M[..., 0, 0], M[..., 1, 1], M[..., 2, 2]
    a, b, c = M[..., 1, 0], M[..., 2, 0], M[..., 2, 1]
    zero = jnp.zeros_like(l0)
    out = jnp.stack(
        [
            jnp.stack([jnp.log(l0), zero, zero], axis=-1),
            jnp.stack([a * _ddlog(l0, l1), jnp.log(l1), zero], axis=-1),
            jnp.stack(
                [
                    b * _ddlog(l0, l2) + a * c * _ddlog2(l0, l1, l2),
                    c * _ddlog(l1, l2),
                    jnp.log(l2),
                ],
                axis=-1,
            ),
        ],
        axis=-2,
    )
    return out if lower else jnp.swapaxes(out, -1, -2)

@jit(static_argnames=("lower", "side"))
def triangular_3x3_matmul(
    L: Array, x: Array, *, lower: bool = True, side: MatmulSide | str = MatmulSide.RIGHT
) -> Array:
    """Optimized matrix-vector multiplication for triangular 3×3 matrices.

    Specialized implementation that exploits triangular structure to avoid
    computing with known-zero elements. On CPU, uses unrolled loops for better
    performance than einsum.

    Args:
        L: Array of shape `(..., 3, 3)` containing triangular matrices.
        x: Array of shape `(..., 3)` containing vectors to multiply.
        lower: Whether `L` is lower (True) or upper (False) triangular.
        side: Multiplication side:
            - `"right"`: Computes `xL`
            - `"left"`: Computes `Lx`

    Returns:
        Shape `(..., 3)` containing the result.

    Example:
        ```python
        L = jnp.array([[1, 0, 0], [2, 3, 0], [4, 5, 6]])
        x = jnp.array([1.0, 2.0, 3.0])
        result = triangular_3x3_matmul(L, x, lower=True, side="right")
        # Computes L @ x efficiently
        ```

    Note:
        Automatically selects between einsum (GPU) and unrolled loops (CPU) for
        optimal performance.
    """
    side = MatmulSide(side)
    cuda = jax.devices()[0].device_kind == "cuda"

    @vectorize(signature="(3,3),(3)->(3)")
    def inner(L: Array, x: Array) -> Array:
        if cuda:
            if side is MatmulSide.RIGHT:
                return jnp.einsum("ji,j->i", L, x)
            else:
                return jnp.einsum("ij,j->i", L, x)
        # The following unrolled implementation is faster on CPU for small matrices
        if side is MatmulSide.RIGHT:
            if lower:
                L_0, L_1, L_2 = L[:, 0], L[1:, 1], L[2:, 2]
                x_0, x_1, x_2 = x, x[1:], x[2:]
            else:
                L_0, L_1, L_2 = L[:1, 0], L[:2, 1], L[:3, 2]
                x_0, x_1, x_2 = x[:1], x[:2], x
        elif side is MatmulSide.LEFT:
            if lower:
                L_0, L_1, L_2 = L[0, :1], L[1, :2], L[2, :3]
                x_0, x_1, x_2 = x[:1], x[:2], x
            else:
                L_0, L_1, L_2 = L[0, :], L[1, 1:], L[2, 2:]
                x_0, x_1, x_2 = x, x[1:], x[2:]
        else:
            raise ValueError(f"Invalid side argument: {side}")
        return jnp.stack([L_0 @ x_0, L_1 @ x_1, L_2 @ x_2])

    return inner(L, x)