bijx.cg¶

Crouch-Grossmann integration methods for Lie group ordinary differential equations.

This module implements geometric integration schemes for ODEs on matrix Lie groups, using the Crouch-Grossmann family of methods. These ensure solutions remain on the manifold throughout integration (up to numerical accuracy).

Key concepts:

Lie group ODEs: Differential equations \(\dot{g} = A(t,g) g\) where \(g(t) \in G\)
Crouch-Grossmann schemes: Runge-Kutta-type methods using matrix exponentials
Butcher tableaux: Coefficient arrays defining integration schemes

Mathematical background: For ODEs on matrix Lie groups of the form \(\dot{g} = f(t,g)\) where \(f\) takes values in the Lie algebra, Crouch-Grossmann methods approximate: \(g_{n+1} = \exp(\sum_i b_i k_i) g_n\)

where \(k_i\) are stage vectors computed via the tableau coefficients. This ensures \(g_{n+1} \in G\) whenever \(g_n \in G\).

Module Attributes

`EULER`	Forward Euler method (1st-order, 1 stage).
`CG2`	Second-order Crouch-Grossmann method (2nd-order, 2 stages).
`CG3`	Third-order Crouch-Grossmann method (3rd-order, 3 stages).

Functions

`cg_stage`(y0, vect, manifold_types, ai, step_size)	Compute intermediate state for Crouch-Grossmann stage.
`crouch_grossmann`(vector_field, y0, args, t0, ...)	Integrate ODE using Crouch-Grossmann method with custom differentiation.
`crouch_grossmann_step`(manifold_types, ...)	Execute single step of Crouch-Grossmann integration method.

Classes

`ButcherTableau`	Butcher tableau defining coefficients for Runge-Kutta integration schemes.
`Euclidean`	Euclidean space with standard vector addition.
`ManifoldType`	Base class defining manifold structure and operations for ODE integration.
`Matrix`	Matrix Lie groups using exponential map integration.
`MatrixDeriv`	Tangent and cotangent spaces for matrix Lie groups.
`Unitary`	Unitary groups U(N) or SU(N) with optional projection to manifold.
`UnitaryDeriv`	Tangent/cotangent spaces for unitary groups with Lie algebra projection.

class bijx.cg.ButcherTableau[source]¶

Bases: Pytree

Butcher tableau defining coefficients for Runge-Kutta integration schemes.

Encodes the coefficient structure for explicit Runge-Kutta methods in the standard Butcher tableau format. Used to define Crouch-Grossmann schemes for integration on Lie groups.

The tableau has the structure:

c_1 | a_11  a_12  ...  a_1s
c_2 | a_21  a_22  ...  a_2s
 :  |  :     :    ⋱    :
c_s | a_s1  a_s2  ...  a_ss
----+----------------------
    | b_1   b_2   ...  b_s

For explicit methods: \(a_{ij} = 0\) for \(j \geq i\). Consistency requires: \(c_i = \sum_j a_{ij}\) and \(\sum_i b_i = 1\).

bijx.cg.EULER = ButcherTableau( stages=1, a=((0.0,),), b=(1.0,), c=(0.0,) )¶

Forward Euler method (1st-order, 1 stage).

The simplest integration scheme: \(y_{n+1} = y_n + h f(t_n, y_n)\).

For Lie groups: \(g_{n+1} = \exp(h A(t_n, g_n)) g_n\).

bijx.cg.CG2 = ButcherTableau( stages=2, a=((0.0, 0.0), (0.5, 0.0)), b=(0.0, 1.0), c=(0.0, 0.5) )¶

Second-order Crouch-Grossmann method (2nd-order, 2 stages).

A two-stage method achieving second-order accuracy for Lie group ODEs. This is the Lie group analogue of the classical midpoint rule.

Stages: 1. \(k_1 = A(t_n, g_n)\) 2. \(k_2 = A(t_n + h/2, \exp(h k_1/2) g_n)\)

Update: \(g_{n+1} = \exp(h k_2) g_n\)

bijx.cg.CG3 = ButcherTableau( stages=3, a=((0.0, 0.0, 0.0), (0.75, 0.0, 0.0), (0.5509259259259259, 0.1574074074074074, 0.0)), b=(0.2549019607843137, -0.6666666666666666, 1.411764705882353), c=(0.0, 0.75, 0.7083333333333334) )¶: Third-order Crouch-Grossmann method (3rd-order, 3 stages).

class bijx.cg.ManifoldType[source]¶

Bases: Pytree

Base class defining manifold structure and operations for ODE integration.

This abstract class specifies how to perform integration steps on different types of manifolds (Euclidean spaces, Lie groups, etc.). Subclasses define the specific operations needed for their manifold structure.

Key operations:

reduce: Accumulate increments using manifold-appropriate operations
post_stage/post_step: Optional projection back to manifold after updates
Adjoint methods: Handle cotangent space operations for differentiation

Common subclasses:

Euclidean: Standard vector spaces with addition
Matrix: Matrix Lie groups using exponential map
Unitary: Unitary groups with optional projection

derivative_type()[source]¶

Return the ManifoldType for tangent/cotangent spaces.

For most manifolds, derivatives live in a different space than the base manifold. This method specifies the manifold type for these spaces.

reduce(x, *deltas)[source]¶

Accumulate increments on the manifold.

Parameters:

x – Base point on manifold.
*deltas – Sequence of increments to accumulate.

Returns:

Updated point after accumulating all increments.

post_stage(x)[source]¶

Project state back to manifold after each internal stage.

Parameters:: x – State after stage update.
Returns:: Projected state (identity by default).

post_step(x)[source]¶

Project state back to manifold after each full integration step.

Parameters:: x – State after full step.
Returns:: Projected state (identity by default).

post_adjoint_vjp(x, adj_vec, v, adj)[source]¶

Adjust adjoint gradient after VJP computation.

Parameters:

x – Current state.
adj_vec – Adjoint vector from VJP.
v – Vector field value.
adj – Current adjoint state.

Returns:

Adjusted adjoint vector.

pre_adjoint_vjp(x, v)[source]¶

Transform vector field before adjoint VJP computation.

Parameters:

x – Current state.
v – Vector field value at identity/origin.

Returns:

Transformed vector (identity by default).

project_adjoint(x, adj)[source]¶

Project adjoint state to cotangent space.

Parameters:

x – Current state.
adj – Adjoint state.

Returns:

Projected adjoint (identity by default).

class bijx.cg.Euclidean[source]¶

Bases: ManifoldType

Euclidean space with standard vector addition.

Use for standard ODEs in flat space where derivatives are tangent vectors and integration uses standard addition. This is the default for most non-geometric problems.

derivative_type()[source]¶

Return the ManifoldType for tangent/cotangent spaces.

For most manifolds, derivatives live in a different space than the base manifold. This method specifies the manifold type for these spaces.

class bijx.cg.Matrix[source]¶

Bases: ManifoldType

Matrix Lie groups using exponential map integration.

Handles integration on matrix Lie groups (SO(N), SU(N), etc.) using the exponential map to ensure solutions remain on the manifold.

Parameters:

right_invariant – Use right-invariant vector fields (X_g = X_e @ g). Left-invariant not yet implemented.
transport_adjoint – If True, keep adjoint in dual tangent space at identity. If False, adjoint lives in ambient cotangent space.

derivative_type()[source]¶

Return the ManifoldType for tangent/cotangent spaces.

For most manifolds, derivatives live in a different space than the base manifold. This method specifies the manifold type for these spaces.

reduce(x, *deltas)[source]¶

Accumulate increments on the manifold.

Parameters:

x – Base point on manifold.
*deltas – Sequence of increments to accumulate.

Returns:

Updated point after accumulating all increments.

transport(x, v)[source]¶

Transport Lie algebra element to tangent space at group element.

Performs parallel transport of a Lie algebra element (tangent vector at identity) to the tangent space at an arbitrary group element.

For right-invariant vector fields: \(X_g = X_e \cdot g\) For left-invariant vector fields: \(X_g = g \cdot X_e\)

This implementation uses the right-invariant convention. The transport maps vectors from \(T_e G\) (Lie algebra) to \(T_g G\).

Parameters:

x – Group element.
v – Lie algebra element (vector at identity).

Returns:

Transported vector at x.

pre_adjoint_vjp(x, v)[source]¶

Transform vector field before adjoint VJP computation.

Parameters:

x – Current state.
v – Vector field value at identity/origin.

Returns:

Transformed vector (identity by default).

project_adjoint(x, adj)[source]¶

Project adjoint state to cotangent space.

Parameters:

x – Current state.
adj – Adjoint state.

Returns:

Projected adjoint (identity by default).

post_adjoint_vjp(x, adj_vec, v, adj)[source]¶

Adjust adjoint gradient after VJP computation.

Parameters:

x – Current state.
adj_vec – Adjoint vector from VJP.
v – Vector field value.
adj – Current adjoint state.

Returns:

Adjusted adjoint vector.

class bijx.cg.MatrixDeriv[source]¶

Bases: ManifoldType

Tangent and cotangent spaces for matrix Lie groups.

Used for both tangent and cotangent spaces since the projections for groups like SU(N) are the same in both cases.

Parameters:: right_invariant – Matches parent Matrix convention.

derivative_type()[source]¶

Return the ManifoldType for tangent/cotangent spaces.

For most manifolds, derivatives live in a different space than the base manifold. This method specifies the manifold type for these spaces.

class bijx.cg.UnitaryDeriv[source]¶

Bases: ManifoldType

Tangent/cotangent spaces for unitary groups with Lie algebra projection.

Projects vectors to the Lie algebra (anti-Hermitian matrices for U(N), traceless anti-Hermitian for SU(N)) at each stage or step.

Parameters:

project_stage – Project to Lie algebra after each internal stage.
project_step – Project to Lie algebra after each full step.
special – If True, enforce traceless condition (for SU(N)).

derivative_type()[source]¶

Return the ManifoldType for tangent/cotangent spaces.

For most manifolds, derivatives live in a different space than the base manifold. This method specifies the manifold type for these spaces.

project(x)[source]¶

post_stage(x)[source]¶

Project state back to manifold after each internal stage.

Parameters:: x – State after stage update.
Returns:: Projected state (identity by default).

post_step(x)[source]¶

Project state back to manifold after each full integration step.

Parameters:: x – State after full step.
Returns:: Projected state (identity by default).

class bijx.cg.Unitary[source]¶

Bases: Matrix

Unitary groups U(N) or SU(N) with optional projection to manifold.

Extends Matrix with projection operations that enforce unitarity constraints using polar decomposition after integration steps.

Parameters:

right_invariant – Use right-invariant vector fields.
transport_adjoint – Keep adjoint in tangent space at identity.
project_step – Project back to unitary manifold after each full step.
project_stage – Project back to unitary manifold after each internal stage.
special – If True, work with SU(N) (special unitary, det=1).
derivative – ManifoldType for tangent/cotangent spaces (default: UnitaryDeriv).

derivative_type()[source]¶

Return the ManifoldType for tangent/cotangent spaces.

For most manifolds, derivatives live in a different space than the base manifold. This method specifies the manifold type for these spaces.

project(m)[source]¶

post_stage(x)[source]¶

Project state back to manifold after each internal stage.

Parameters:: x – State after stage update.
Returns:: Projected state (identity by default).

post_step(x)[source]¶

Project state back to manifold after each full integration step.

Parameters:: x – State after full step.
Returns:: Projected state (identity by default).

bijx.cg.cg_stage(y0, vect, manifold_types, ai, step_size)[source]¶

Compute intermediate state for Crouch-Grossmann stage.

Combines previous stage vectors according to the Butcher tableau coefficients.

The computation follows: \(g_i = \left(\prod_j \exp(h a_{ij} k_j)\right) g_0\)

where \(k_j\) are stage vectors and \(a_{ij}\) are tableau coefficients.

Parameters:

y0 – Initial state for the current step.
vect – List of stage vectors from previous stages.
ai – Row of Butcher tableau coefficients for current stage.
step_size – Integration step size \(h\).

Returns:

Intermediate state for evaluating the next stage vector.

bijx.cg.crouch_grossmann_step(manifold_types, tableau, vector_field, step_size, t, y0)[source]¶

Execute single step of Crouch-Grossmann integration method.

Performs one integration step using the specified Butcher tableau, computing all intermediate stages and the final update.

Parameters:

manifold_types – Pytree of ManifoldType instances defining manifold structure.
tableau – ButcherTableau defining the integration method.
vector_field – Function \((t, g) \mapsto A\) where \(A\) is in the Lie algebra (for Lie groups) or tangent vector (for Euclidean).
step_size – Integration step size \(h\).
t – Current time.
y0 – Current state (group element or vector).

Returns:

Updated state after one integration step.

Important

The vector field must return values in the Lie algebra for Lie group components, enabling the exponential map to produce valid group elements.

bijx.cg.crouch_grossmann(vector_field, y0, args, t0, t1, step_size, manifold_types=None, args_types=None, tableau=ButcherTableau(stages=1, a=((0.0,),), b=(1.0,), c=(0.0,)))[source]¶

Integrate ODE using Crouch-Grossmann method with custom differentiation.

Solves: \(\dot{g} = f(t, g, \text{args})\) from \(t_0\) to \(t_1\)

For Lie groups: \(\dot{g} = A(t, g) g\) where \(A(t, g) \in \mathfrak{g}\) For Euclidean: \(\dot{y} = f(t, y)\) (standard ODE)

Parameters:

vector_field – Function \((t, g, \text{args}) \mapsto A\) where \(A\) is in the Lie algebra for Lie group components, or tangent vector for Euclidean components.
y0 – Initial condition (group element or vector).
args – Additional parameters passed to vector_field.
t0 – Initial time.
t1 – Final time.
step_size – Integration step size (positive for forward integration).
manifold_types – ManifoldType or pytree of ManifoldTypes specifying the manifold structure of y0. Defaults to Unitary().
args_types – ManifoldType or pytree of ManifoldTypes for args. Defaults to Euclidean().
tableau – ButcherTableau defining the integration method (default: EULER).

Returns:

Solution at time \(t_1\).

Example

>>> def rigid_body_eqs(t, R, omega):
...     # R(t) ∈ SO(3), omega is angular velocity
...     Omega = skew_symmetric(omega)
...     return Omega  # Lie algebra element
>>> R0 = jnp.eye(3)  # Initial orientation
>>> omega = jnp.array([1.0, 0.0, 0.0])  # Rotation about x-axis
>>> R_final = crouch_grossmann(
...     rigid_body_eqs, R0, omega, 0.0, 1.0, 0.01,
...     manifold_types=Matrix(), args_types=Euclidean(), tableau=CG2
... )

Important

The vector field must return values in the Lie algebra for Lie group components, enabling the exponential map to produce valid group elements.