Dynamics API

SolverInfo

Bases: NamedTuple

Diagnostics measured after a fixed nonlinear-solver iteration budget.

Attributes:

Name	Type	Description
residual_norm	Array	Euclidean norm of the discrete Moser–Veselov residual after the final Newton iteration. It is scalar for solve_F_with_info and has shape (steps,) when returned by simulate_free_rigid_body.
converged	jax.Array, dtype bool	Whether residual_norm is no greater than the requested tolerance. The flag reports the final residual; it does not imply that the fixed iteration loop terminated early.

solve_F

solve_F(pi, J, h, newton_iters=8)

Solve one Moser–Veselov step and return only its relative rotation.

Parameters:

Name	Type	Description	Default
pi	(Array, shape(3))	Body-frame angular momentum at the start of the step.	required
J	(Array, shape(3, 3))	Symmetric positive-definite body inertia tensor.	required
h	float or Array	Timestep.	required
newton_iters	int	Fixed Newton iteration budget. The default is 8.	8

Returns:

Type	Description
(Array, shape(3, 3))	Relative rotation for the step.

Notes

This convenience function discards the nonlinear-solver diagnostics. Use solve_F_with_info whenever convergence must be checked.

solve_F_with_info

solve_F_with_info(pi, J, h, newton_iters=8, tolerance=1e-10)

Solve the discrete Moser–Veselov equation for one relative rotation.

Given body angular momentum \(\boldsymbol{\pi}_k\), inertia \(J\), and timestep \(h\), this function finds \(F_k\in SO(3)\) satisfying

\[ F_k J_d - J_d F_k^T = h\,\widehat{\boldsymbol{\pi}_k}, \]

where \(J_d=\tfrac{1}{2}\operatorname{tr}(J)I-J\). The unknown is parameterized as \(F_k=\operatorname{Exp}(\mathbf{g})\) and solved with a fixed number of Newton iterations. Each iteration evaluates a small set of deterministic damping factors and retains the candidate with the lowest residual norm.

Parameters:

Name	Type	Description	Default
pi	(Array, shape(3))	Body-frame angular momentum at the start of the step.	required
J	(Array, shape(3, 3))	Symmetric positive-definite body inertia tensor. It may be non-diagonal.	required
h	float or Array	Timestep.	required
newton_iters	int	Number of Newton iterations. The value is static under JIT compilation. The default is 8.	8
tolerance	float or Array	Maximum final residual norm accepted as converged. The default is 1e-10.	1e-10

Returns:

Name	Type	Description
F	(Array, shape(3, 3))	Relative rotation for the step.
info	SolverInfo	Final residual norm and convergence flag.

Notes

Every call consumes the complete newton_iters budget. Callers should inspect info.converged before trusting the returned step; a failed solve does not raise an exception or stop a surrounding simulation.

rigid_body_step

rigid_body_step(R: ndarray, pi: ndarray, J: ndarray, torque_k: ndarray, torque_next: ndarray, dt: float, newton_iters: int = 8, tolerance: float = 1e-10) -> tuple[jnp.ndarray, jnp.ndarray, SolverInfo]

Advance one forced rigid-body step on \(SO(3)\).

The endpoint body torques enter as two half-step discrete impulses. First,

\[ \boldsymbol{\pi}_{k+\frac12} = \boldsymbol{\pi}_k + \frac{h}{2}\boldsymbol{\tau}_k. \]

The relative rotation \(F_k\) solves

\[ F_k J_d - J_d F_k^T = h\,\widehat{\boldsymbol{\pi}}_{k+\frac12}, \qquad J_d = \frac{1}{2}\operatorname{tr}(J)I-J. \]

The new state is

\[ R_{k+1}=R_kF_k, \qquad \boldsymbol{\pi}_{k+1} =F_k^T\boldsymbol{\pi}_{k+\frac12} +\frac{h}{2}\boldsymbol{\tau}_{k+1}. \]

Parameters:

Name	Type	Description	Default
R	(Array, shape(3, 3))	Body-to-spatial attitude at the start of the step.	required
pi	(Array, shape(3))	Body-frame angular momentum at the start of the step.	required
J	(Array, shape(3, 3))	Symmetric positive-definite body inertia tensor.	required
torque_k	(Array, shape(3))	Body-frame torque at the start of the step.	required
torque_next	(Array, shape(3))	Body-frame torque at the end of the step.	required
dt	float or Array	Timestep \(h\).	required
newton_iters	int	Newton iterations used to solve for \(F_k\). The default is 8.	8
tolerance	float or Array	Residual threshold used to form the convergence flag. The default is 1e-10.	1e-10

Returns:

Name	Type	Description
R_next	(Array, shape(3, 3))	Body-to-spatial attitude at the end of the step.
pi_next	(Array, shape(3))	Body-frame angular momentum at the end of the step.
solver_info	SolverInfo	Residual norm and convergence flag for the nonlinear solve.

simulate_free_rigid_body

simulate_free_rigid_body(R0: ndarray, pi0: ndarray, J: ndarray, dt: float | ndarray, steps: int, newton_iters: int = 8, tolerance: float = 1e-10) -> tuple[jnp.ndarray, jnp.ndarray, SolverInfo]

Simulate torque-free rigid-body motion on \(SO(3)\).

The state is \((R_k,\boldsymbol{\pi}_k)\), where \(R_k\) maps body coordinates to spatial coordinates and \(\boldsymbol{\pi}_k\) is body angular momentum. For each transition, solve_F_with_info computes a relative rotation \(F_k\), followed by

\[ R_{k+1} = R_k F_k, \qquad \boldsymbol{\pi}_{k+1} = F_k^T\boldsymbol{\pi}_k. \]

This discrete update preserves spatial angular momentum algebraically:

\[ R_{k+1}\boldsymbol{\pi}_{k+1} = R_k\boldsymbol{\pi}_k. \]

Parameters:

Name	Type	Description	Default
R0	(Array, shape(3, 3))	Initial body-to-spatial attitude. The caller is responsible for providing a proper rotation.	required
pi0	(Array, shape(3))	Initial body-frame angular momentum.	required
J	(Array, shape(3, 3))	Symmetric positive-definite body inertia tensor. It may be non-diagonal.	required
dt	(float or Array, shape(steps))	Positive simulation timestep. A scalar applies the same interval to every transition; an array supplies one interval per transition.	required
steps	int	Number of discrete transitions. This value is static under JIT compilation.	required
newton_iters	int	Newton iterations used for every transition. The default is 8 and the value is static under JIT compilation.	8
tolerance	float or Array	Residual threshold used to form each convergence flag. The default is 1e-10.	1e-10

Returns:

Name	Type	Description
Rs	(Array, shape(steps + 1, 3, 3))	Attitude history including R0.
pis	(Array, shape(steps + 1, 3))	Body-angular-momentum history including pi0.
solver_info	SolverInfo	Arrays of residual norms and convergence flags with shape (steps,); there is one diagnostic entry per transition.

Notes

The function returns trajectories even when one or more nonlinear solves fail the requested tolerance. Always inspect solver_info.converged before using the result. The complete simulation is compatible with jax.jit and automatic differentiation.

simulate_rigid_body

simulate_rigid_body(R0: ndarray, pi0: ndarray, J: ndarray, body_torques: ndarray, dt: float | ndarray, newton_iters: int = 8, tolerance: float = 1e-10) -> tuple[jnp.ndarray, jnp.ndarray, SolverInfo]

Simulate a rigid body driven by prescribed body-frame torques.

The state at node \(k\) is \((R_k,\boldsymbol{\pi}_k)\), where \(R_k\) maps body coordinates to spatial coordinates and \(\boldsymbol{\pi}_k\) is body angular momentum. The supplied torque \(\boldsymbol{\tau}_k\) is also expressed in body coordinates.

For a timestep \(h\), define the half-step momentum

\[ \boldsymbol{\pi}_{k+\frac12} = \boldsymbol{\pi}_k + \frac{h}{2}\boldsymbol{\tau}_k. \]

The forced Lie-group variational update finds \(F_k\in SO(3)\) from

\[ F_k J_d - J_d F_k^T = h\,\widehat{\boldsymbol{\pi}}_{k+\frac12}, \qquad J_d = \frac{1}{2}\operatorname{tr}(J)I-J. \]

The state then advances according to

\[ R_{k+1} = R_k F_k, \]

\[ \boldsymbol{\pi}_{k+1} = F_k^T\boldsymbol{\pi}_{k+\frac12} + \frac{h}{2}\boldsymbol{\tau}_{k+1}. \]

The two half-step torque terms are the discrete forces at the ends of the interval. Consequently, a trajectory with steps transitions requires steps + 1 torque samples. Setting every torque to zero recovers simulate_free_rigid_body.

Parameters:

Name	Type	Description	Default
R0	(Array, shape(3, 3))	Initial body-to-spatial attitude. The caller is responsible for providing a proper rotation.	required
pi0	(Array, shape(3))	Initial body-frame angular momentum.	required
J	(Array, shape(3, 3))	Symmetric positive-definite body inertia tensor. It may be non-diagonal.	required
body_torques	(Array, shape(steps + 1, 3))	Prescribed body-frame torque at every state node, including both the initial and final nodes. The number of transitions is inferred from this leading dimension.	required
dt	(float or Array, shape(steps))	Positive simulation timestep. A scalar applies the same interval to every transition; an array supplies one interval per transition.	required
newton_iters	int	Newton iterations used for every transition. The default is 8 and the value is static under JIT compilation.	8
tolerance	float or Array	Residual threshold used to form each convergence flag. The default is 1e-10.	1e-10

Returns:

Name	Type	Description
Rs	(Array, shape(steps + 1, 3, 3))	Attitude history including R0.
pis	(Array, shape(steps + 1, 3))	Body-angular-momentum history including pi0.
solver_info	SolverInfo	Arrays of residual norms and convergence flags with shape (steps,); there is one diagnostic entry per transition.

Notes

The nonlinear equation is the same Moser--Veselov solve used by simulate_free_rigid_body, with the first half of the discrete torque impulse included in its momentum argument. As in the free-body simulator, the returned trajectory must not be trusted unless every entry of solver_info.converged is true.

References

T. Lee, N. H. McClamroch, and M. Leok, "A Lie Group Variational Integrator for the Attitude Dynamics of a Rigid Body with Applications to the 3D Pendulum," Proceedings of the 2005 IEEE Conference on Control Applications, pp. 962--967, 2005. doi:10.1109/CCA.2005.1507254.

energy

energy(pi: ndarray, J: ndarray) -> jnp.ndarray

Compute rotational kinetic energy from body angular momentum.

\[ E(\boldsymbol{\pi},J) = \frac{1}{2}\boldsymbol{\pi}^{T}J^{-1}\boldsymbol{\pi}. \]

Parameters:

Name	Type	Description	Default
pi	(Array, shape(3))	Body-frame angular momentum.	required
J	(Array, shape(3, 3))	Body inertia tensor.	required

Returns:

Type	Description
(Array, scalar)	Rotational kinetic energy.

Notes

J is applied through a linear solve rather than an explicit inverse. This function operates on one state; use jax.vmap for a trajectory or batch.

spatial_momentum

spatial_momentum(R: ndarray, pi: ndarray) -> jnp.ndarray

Express body angular momentum in spatial coordinates.

For a body-to-spatial attitude \(R\) and body angular momentum \(\boldsymbol{\pi}\), the spatial angular momentum is

\[ \mathbf{L} = R\boldsymbol{\pi}. \]

Parameters:

Name	Type	Description	Default
R	(Array, shape(3, 3))	Body-to-spatial rotation matrix.	required
pi	(Array, shape(3))	Body-frame angular momentum.	required

Returns:

Type	Description
(Array, shape(3))	Angular momentum expressed in the spatial frame.

Notes

This function operates on one state; use jax.vmap for a trajectory or batch.

ortho_error

ortho_error(R: ndarray) -> jnp.ndarray

Measure violation of the rotation-matrix orthogonality condition.

\[ e_{\mathrm{orth}}(R) = \lVert R^T R-I\rVert_F. \]

Parameters:

Name	Type	Description	Default
R	(Array, shape(3, 3))	Matrix to evaluate.	required

Returns:

Type	Description
(Array, scalar)	Frobenius norm of the orthogonality residual. The value is zero for an exactly orthogonal matrix.

Notes

Orthogonality alone does not distinguish rotations from reflections. Pair this diagnostic with determinant_error when testing membership in \(SO(3)\). Use jax.vmap for a trajectory or batch.

determinant_error

determinant_error(R: ndarray) -> jnp.ndarray

Measure violation of the proper-rotation determinant condition.

\[ e_{\det}(R) = |\det(R)-1|. \]

Parameters:

Name	Type	Description	Default
R	(Array, shape(3, 3))	Matrix to evaluate.	required

Returns:

Type	Description
(Array, scalar)	Absolute determinant error. The value is zero when det(R) is exactly +1.

Notes

A unit determinant alone does not imply orthogonality. Pair this diagnostic with ortho_error when testing membership in \(SO(3)\). Use jax.vmap for a trajectory or batch.

discrete_inertia

discrete_inertia(J: ndarray) -> jnp.ndarray

Construct the discrete inertia used by the Moser–Veselov equation.

The discrete inertia associated with a physical inertia tensor \(J\) is

\[ J_d = \frac{1}{2}\operatorname{tr}(J)I - J. \]

The definition is coordinate-independent; J need not be diagonal or expressed in principal axes.

Parameters:

Name	Type	Description	Default
J	(Array, shape(3, 3))	Body inertia tensor.	required

Returns:

Type	Description
(Array, shape(3, 3))	Discrete inertia with the same dtype as J.