jax_rb.manifolds

Global Manifold

Base class for manifold in global embedded coordinates

class jax_rb.manifolds.global_manifold.GlobalManifold

A manifold \(\mathcal{M}\) embedded in a vector space \(\mathcal{E}\) .

g_metric(x, omg)

the metric operator g, which is symmetric. The corresponding metric is \(\langle \omega, g(x)\omega \rangle_{\mathcal{E}}\) .

gamma(x, xi, eta)

Christoffel function. Symmetric for two tangent vectors xi, eta. The corresponding Levi-Civita connection is \(\nabla_{\mathtt{X}}\mathtt{Y} = \mathrm{D}_{\mathtt{X}}\mathtt{Y} + \Gamma(x; \mathtt{X}, \mathtt{Y})\) for two vector fields \(\mathtt{X}, \mathtt{Y}\).

inner(x, a, b)

Riemannian inner product.

Parameters:

• a -- a vector in ambient space,

• b -- a vector in ambient space,

Returns:

the inner product of a and b using the metric \(\mathsf{g}\) .

inv_g_metric(x, omg)

inverse of the metric operator g.

ito_drift(x)

Ito Brownian drift as an ambient vector.

laplace_beltrami(x, egradx, ehessvp)

Laplace Beltrami operator. This works in for vector and matrices. For a specific manifold, this may be simplified. We assume f is a scalar function in a tubular neighborhood of the manifold.

Parameters:

• x -- a point on the manifold,

• egradx -- is the Euclidean gradient of \(f\) , a matrix of the same shape with x,

• ehessvp -- is the Euclidean Hessian Productof \(f\) , a linear operator on \(\mathcal{E}\) ,

Returns:

the value of the Laplace Beltrami operator of \(f\) .

name()

name of the manifold.

proj(x, omg)

Metric compatible projection :param x: a point on the manifold, :param omg: a vector on the ambient manifold \(\mathcal{E}\) , :returns: a point the tangent space at x.

pseudo_transport(x, y, v)

an approximate parallel transport from x to y

Parameters:

• x -- a point on the manifold,

• y -- a point on the manifold,

• v -- a tangent vector at x,

Returns:

a tangent vector at y.

rand_ambient(key)

Random ambient vector.

rand_point(key)

A random point on the manifold.

rand_vec(key, x)

Random tangent vector at x.

retract(x, v)

Second order retraction

Parameters:

• x -- a point on the manifold,

• v -- a tangent vector at x,

Returns:

a point on the manifold.

sigma(x, dw)

Sigma map to generate Brownian motion.

Parameters:

• x -- a point on the manifold,

• dw -- a point on the ambient space,

Returns:

apoint on the ambient space such that \(\Pi(x) \sigma(x) \sigma^{\mathsf{T}}(x)\mathsf{g}^{-1}(x) dw = \Pi(x)dw\)

Sphere

Sphere of constant curvature

class jax_rb.manifolds.sphere.Sphere(n, r)

The sphere \(x_0^2+x_1^2+\cdots+x_d^2 = r^2\) with sectional curvature \(\frac{1}{r^2}\).

Parameters:

• n -- dimension of \(\mathcal{E}\)

• r -- the radius.

Symmetric Positive Definite Matrix Manifold

Symmetric Positive Definite Matrix Manifold \(\mathrm{S}^+(p)\) of positive definite matrices of shape \(p\times p\) .

class jax_rb.manifolds.spd.PositiveDefiniteManifold(p)

The manifold of positive definite matrices of size \(p\) .

Stiefel Manifold

Stiefel manifold \(\mathrm{St}(n, p, \alpha_0, \alpha_1)\) with metric defined by two parameters.

class jax_rb.manifolds.stiefel.RealStiefelAlpha(shape, alpha)

The manifold \(Y^TY = I\) where \(Y\) is a matrix of size \(shape=n\times p\) with metric \(\lvert \omega\rvert^2_{\mathsf{g}} =\alpha_0 Tr(\omega^T\omega) +(\alpha_0-\alpha_1)Tr(\omega^TYY^T\omega)\) .

Parameters:

• shape -- tuple (n, p),

• alpha -- array of 2 positive numbers.

Grassmann Manifold

Grassmann manifold \(\mathrm{Gr}(n, p)\) of vector spaces of rank \(p\) in a \(n\) -dimension vector space.

class jax_rb.manifolds.grassmann.Grassmann(shape)

The lift of the Grassman manifold to to the Stiefel manifold \(Y^TY = I\) where \(Y\) is a matrix of \(shape=n\times p\) with metric \(\lvert \omega\rvert^2_{\mathsf{g}} = Tr(\omega^T\omega)\). The lift is with respect to the submersion defined by the relationship \(Y\sim YU\) for an orthogonal matrix \(U\).

Hypersurface with Diagonal constraints Manifold

Hypersurface with a constraint of the form \(\sum_i d_i x_i^p = 1\)

class jax_rb.manifolds.diag_hypersurface.DiagHypersurface(dvec, p)

Hypersurface of the form \(\sum_i d_ix_i^p = 1\).

Parameters:

• dvec -- vector \(d_i\) of coefficients. Sort dvec so dvec[-1] is positive.

• p -- \(p > 0\) is an integer, degree of the constraint.

Use embedded metric.

Matrix Lie Group Left Invariant Metric

Base class for matrix groups with left invariant metrics.

Required implementations

jax_rb.manifolds.matrix_left_invariant.MatrixLeftInvariant._lie_algebra_proj(self, omg)

The projection \(p_{\mathfrak{g}}\) at the identity.

jax_rb.manifolds.matrix_left_invariant.MatrixLeftInvariant._mat_apply(self, mat, omg)

Implementing the operator \(\mathcal{I}\) applied on omg in \(\mathcal{E}\).

Base class

class jax_rb.manifolds.matrix_left_invariant.MatrixLeftInvariant(p, g_mat)

Matrix group with left invariant metric.

Parameters:

• p -- the size of the matrix

• g_mat -- The matrix defining the inner product at the identity. Typically g_mat is of size \(\dim \mathrm{G}\) .

approx_nearest(q)

find point on the manifold that is nearest to q, same order as the nearest point.

g_metric(x, omg)

the metric operator g, which is symmetric. The corresponding metric is \(\langle \omega, g(x)\omega \rangle_{\mathcal{E}}\) .

gamma_ambient(x, xi, eta)

Christoffel function for ambient manifold.

id_drift

Ito drift at the identity.

inner(x, a, b)

Riemannian inner product.

Parameters:

• a -- a vector in ambient space,

• b -- a vector in ambient space,

Returns:

the inner product of a and b using the metric \(\mathsf{g}\) .

inv_g_metric(x, omg)

inverse of the metric operator g.

laplace_beltrami(x, egradx, ehessvp)

Laplace Beltrami operator. This works in for vector and matrices. For a specific manifold, this may be simplified. We assume f is a scalar function in a tubular neighborhood of the manifold.

Parameters:

• x -- a point on the manifold,

• egradx -- is the Euclidean gradient of \(f\) , a matrix of the same shape with x,

• ehessvp -- is the Euclidean Hessian Productof \(f\) , a linear operator on \(\mathcal{E}\) ,

Returns:

the value of the Laplace Beltrami operator of \(f\) .

left_invariant_vector_field(x, v)

map from a unit vector in the trace metric to a vector field with unit length in the left invariant metric.

name()

name of the manifold.

proj(x, omg)

Metric compatible projection :param x: a point on the manifold, :param omg: a vector on the ambient manifold \(\mathcal{E}\) , :returns: a point the tangent space at x.

pseudo_transport(x, y, v)

the easy one

rand_ambient(key)

Random ambient vector.

rand_point(key)

A random point on the manifold.

rand_vec(key, x)

Random tangent vector at x.

retract(x, v)

Second order retraction

Parameters:

• x -- a point on the manifold,

• v -- a tangent vector at x,

Returns:

a point on the manifold.

sigma(x, dw)

sigma, to generate the Brownian motion.

sigma_id(dw)

sigma, to generate the Brownian motion at the identity.

stratonovich_drift(x)

Stratonovich drift.

v0

Stratonovich drift at the identity.

\(\mathrm{GL}^+(n)\) Generalized Linear Group Positive Determinant

\(GL^+\): Positive Component of the Generalized Linear group with left-invariant metric.

class jax_rb.manifolds.glp_left_invariant.GLpLeftInvariant(p, g_mat)

\(GL^+\) with left invariant metric defined by g_mat.

Parameters:

• p -- the size of the matrix

• g_mat -- The matrix defining the inner product at the identity. g_mat is in \(\mathbb{R}^{p^2\times p^2}\) .

\(\mathrm{Aff}^+(n)\) Affine Linear Group Positive Determinant

\(Aff^+\): Positive Component of the Affine group with left invariant metric.

class jax_rb.manifolds.affine_left_invariant.AffineLeftInvariant(n, g_mat)

Group of affine tranformations of \(\mathbb{R}^{n}\), represented by a pair \((A, v)\in GL^+(n)\times \mathbb{R}^{n}\) with action \((A, v).w = Aw + v\) for \(w\in\mathbb{R}^{n}\) .

Alternatively, it is represented as a matrix \(\begin{bmatrix} A & v \\ 0 & 1 \end{bmatrix}\in GL(n+1)\).

Parameters:

• n -- size of A

• g_mat -- a positive definite matrix in \(\mathbb{R}^{n(n+1)\times n(n+1)}\) defining the metric at \(I_{n+1}\)

\(\mathrm{SL}(n)\) Special Linear Group

\(SL\): special linear group of matrices of determinant 1.

class jax_rb.manifolds.sl_left_invariant.SLLeftInvariant(n, g_mat)

\(SL(n)\) Special linear group, of size \(n\) with left invariant metric defined by g_mat.

Parameters:

• n -- size of the matrix

• g_mat -- a positive definite matrix in \(\mathbb{R}^{(n^2-1)\times (n^2-1)}\) defining the metric at \(I_{n}\)

\(\mathrm{SO}(n)\) Special Orthogonal Group

\(SO\): Special Orthogonal group.

class jax_rb.manifolds.so_left_invariant.SOLeftInvariant(n, g_mat)

The group \(SO(n)\) of orthogonal matrices \(U\in R^{n\times n}\) of determinant 1 with metric \(Tr(\omega^TU^T\mathcal{I}(U^T\omega))\) where \(\mathcal{I}\) is the metric defined by g_mat, a matrix of size \(\mathbb{R}^{\frac{n(n-1)}{2}\times \frac{n(n-1)}{2}}\).

\(\mathrm{SE}(n)\) Special Euclidean Group

\(SE\): Special Euclidean group.

class jax_rb.manifolds.se_left_invariant.SELeftInvariant(n, g_mat)

Specian Euclidean group of Euclidean transformation of \(\mathbb{R}^{n}\), represented by a pair \((A, v)\in SO^(n)\times \mathbb{R}^{n}\) with action \((A, v).w = Aw + v\) for \(w\in\mathbb{R}^{n}\) .

Alternatively, it is represented as a matrix \(\begin{bmatrix} A & v \\ 0 & 1 \end{bmatrix}\in GL(n+1)\) where \(A\in SO(n)\).

Parameters:

• n -- size of A

• g_mat -- a positive definite matrix in \(\mathbb{R}^{\frac{n(n+1)}{2}\times\frac{n(n+1)}{2}}\) defining the metric at \(I_{n+1}\)