# Source code for geomstats.geometry.sub_riemannian_metric

"""Implementation of sub-Riemannian geometry.

Lead author: Morten Pedersen.
"""

import geomstats.backend as gs

[docs]
class SubRiemannianMetric:
"""Class for Sub-Riemannian metrics.

This implementation assumes a bracket-generating distribution of constant dimension.

Only one of the arguments 'cometric_matrix' and 'frame' can be different from
None. If the frame is supplied, it is assumed orthonormal.

Parameters
----------
space : Manifold object
cometric_matrix : callable
Optional, default: 'None'

The cometric matrix as a function of a point.

You should pass :
base_point : array-like, shape=[..., dim]
It returns:
_ : array-like, shape=[..., dim, dim]

frame : callable
Optional, default: 'None'

Matrix representing the frame spanning the distribution,
as a function of a point.

You should pass:
point : array-like, shape=[..., dim]
It returns
_ : array-like, shape=[..., dim, dist_dim]
Frame field matrix. Each column is a vector field of the frame
spanning the distribution.
"""

def __init__(
self,
space,
cometric_matrix=None,
frame=None,
):
if not bool(cometric_matrix is not None) ^ bool(frame is not None):
raise ValueError(
"Either 'cometric_matrix' or 'frame' must be passed, and not both."
)

if space.point_ndim != 1:
raise ValueError("Can only handle 'vector' point types.")

self.cometric_matrix = cometric_matrix
self.frame = frame

self._def_type = "cometric" if self.frame is None else "frame"

[docs]
def sr_sharp(self, base_point, cotangent_vec):
r"""Compute sub-Riemannian sharp map.

This is the sub-Riemannian sharp map, mapping a covector at base_point to a
tangent vector in the distribution subspace at base_point. For an orthonormal
frame (F_i)_{i=1..dist_dim}, the sharp map is given by

.. math::
\sharp(q, p) = \sum_i^{dist_dim} p(F_i(q)) * F_i(q)

Parameters
----------
base_point : array-like, shape=[..., dim]
Point on the manifold.

cotangent_vec : array-like, shape=[..., dim]
Cotangent vector at base_point.

Returns
-------
sr_sharp : array-like, shape=[..., dim]
sub-Riemannian sharp of 'cotangent_vec' at 'base_point'
"""
if self.frame is None:
raise NotImplementedError(
"The sub-Riemannian sharp map is only"
" implemented when a frame is passed."
)

frame = self.frame(base_point)
coefs = gs.einsum("...i,...ij->...j", cotangent_vec, frame)
coefs_on_frame = gs.einsum("...j,...ij->...ij", coefs, frame)

return gs.einsum("...ij->...i", coefs_on_frame)

def _get_method(self, method_name):
return getattr(self, f"_{method_name}_{self._def_type}")

def _inner_coproduct_frame(self, cotangent_vec_a, cotangent_vec_b, base_point):
sharp = self.sr_sharp(base_point=base_point, cotangent_vec=cotangent_vec_b)
return gs.einsum("...i,...i -> ...", cotangent_vec_a, sharp)

def _inner_coproduct_cometric(self, cotangent_vec_a, cotangent_vec_b, base_point):
C_b = gs.einsum(
"...ij,...j->...i", self.cometric_matrix(base_point), cotangent_vec_b
)
return gs.einsum("...i,...i->...", cotangent_vec_a, C_b)

[docs]
def inner_coproduct(self, cotangent_vec_a, cotangent_vec_b, base_point):
"""Compute inner coproduct between two cotangent vectors at base point.

This is the inner product associated to the cometric matrix.

Parameters
----------
cotangent_vec_a : array-like, shape=[..., dim]
Cotangent vector at base_point.
cotangent_vec_b : array-like, shape=[..., dim]
Cotangent vector at base_point.
base_point : array-like, shape=[..., dim]
Point on the manifold.

Returns
-------
inner_coproduct : float
Inner coproduct between the two cotangent vectors.
"""
return self._get_method("inner_coproduct")(
cotangent_vec_a, cotangent_vec_b, base_point
)

def _hamiltonian_frame(self, state):
position, momentum = state

inner_products = gs.einsum("...i,...ij->...j", momentum, self.frame(position))

out = 0.5 * gs.einsum("...i,...i->...", inner_products, inner_products)
if not gs.is_array(out):
return gs.array(out)
return out

def _hamiltonian_cometric(self, state):
position, momentum = state
return 0.5 * self.inner_coproduct(momentum, momentum, position)

[docs]
def hamiltonian(self, state):
r"""Compute the hamiltonian energy associated to the cometric.

The Hamiltonian at state :math:(q, p) is defined by

.. math::
H(q, p) = \frac{1}{2} <p, p>_q

where :math:<\cdot, \cdot>_q is the cometric at :math:q.

Parameters
----------
state : array-like, shape=[[..., dim], [..., dim]]
The first array in 'state' contains positions, the second contains
covectors (momentums).

Returns
-------
energy :  float
Hamiltonian energy at state.
"""
return self._get_method("hamiltonian")(state)

[docs]
@staticmethod
def symp_grad(hamiltonian):
r"""Compute the symplectic gradient of the Hamiltonian.

Parameters
----------
hamiltonian : callable
The hamiltonian function from the tangent bundle to the reals.

Returns
-------
vector : array-like, shape=[, 2*dim]
The symplectic gradient of the Hamiltonian.
"""
value_and_grad = gs.autodiff.value_and_grad(
hamiltonian,
point_ndims=2,
)

def vector(x):
"""Compute symplectic gradient at x."""
_, (h_q, h_p) = value_and_grad(x)
return gs.array([h_p, -h_q])

return vector

def _symp_euler_frame(self, hamiltonian, step_size):
def step(state):
position, momentum = state
dq = self.sr_sharp(base_point=position, cotangent_vec=momentum)
y = gs.array([position + step_size * dq, momentum])
_, dp = self.symp_grad(hamiltonian)(y)
return gs.array([position + step_size * dq, momentum + step_size * dp])

return step

def _symp_euler_cometric(self, hamiltonian, step_size):
def step(state):
r"""Compute an integration step from state."""
position, momentum = state
dq, _ = self.symp_grad(hamiltonian)(state)
y = gs.array([position + step_size * dq, momentum])
_, dp = self.symp_grad(hamiltonian)(y)
return gs.array([position + step_size * dq, momentum + step_size * dp])

return step

[docs]
def symp_euler(self, hamiltonian, step_size):
r"""Compute a function which calculates a step of symplectic euler integration.

Parameters
----------
hamiltonian : callable
The hamiltonian function from the tangent bundle to the reals.
step_size : float
Step size of the symplectic euler step

Returns
-------
step : callable
Given a state, 'step' returns the next symplectic euler step
"""
return self._get_method("symp_euler")(hamiltonian, step_size)

[docs]
@staticmethod
def iterate(func, n_steps):
r"""Construct a function which iterates another function n_steps times.

Parameters
----------
func : callable
A function which calculates the next step of a
sequence to be calculated.
n_steps : int
The number of times to iterate func.

Returns
-------
flow : callable
Given a state, 'flow' returns a sequence with n_steps
iterations of func.
"""

def flow(x):
r"""Compute flow starting at x."""
xs = [x]
for i in range(n_steps):
xs.append(func(xs[i]))
return gs.array(xs)

return flow

[docs]
def symp_flow(self, hamiltonian, end_time=1.0, n_steps=20):
r"""Compute the symplectic flow of the hamiltonian.

Parameters
----------
hamiltonian : callable
The hamiltonian function from the tangent bundle to the reals.
end_time : float
The last time point until which to calculate the flow.
n_steps : int
Number of integration steps.

Returns
-------
_ : array-like, shape[,n_steps]
Given a state, 'symp_flow' returns a
sequence with n_steps iterations of func.
"""
step = self.symp_euler
step_size = end_time / n_steps
return self.iterate(step(hamiltonian, step_size), n_steps)

[docs]
def exp(self, cotangent_vec, base_point, n_steps=20):
"""Exponential map associated to the cometric.

Exponential map at base_point of cotangent_vec computed by integration
of the Hamiltonian equation (initial value problem), using the cometric.
In the Riemannian case this yields the exponential associated to the
Levi-Civita connection of the metric (the inverse of the cometric).

Parameters
----------
cotangent_vec : array-like, shape=[..., dim]
Tangent vector at the base point.
base_point : array-like, shape=[..., dim]
Point on the manifold.
n_steps : int
Number of discrete time steps to take in the integration.
Optional, default: N_STEPS.

Returns
-------
exp : array-like, shape=[..., dim]
Point on the manifold.
"""
# TODO: add log (caveat in docstrings)
# TODO: add distance (based on log - also caveat)
base_point = gs.broadcast_to(base_point, cotangent_vec.shape)

# TODO: integrate; allow euler
initial_state = gs.stack([base_point, cotangent_vec])

flow = self.symp_flow(self.hamiltonian, n_steps=n_steps)

return flow(initial_state)[-1][0]

[docs]
def geodesic(self, initial_point, initial_cotangent_vec, n_steps=20):
"""Generate parameterized function for the normal geodesic curve.

Normal geodesic curve defined by an initial point and an initial
cotangent vector.

Parameters
----------
initial_point : array-like, shape=[..., dim]
Point on the manifold, initial point of the geodesic.
initial_cotangent_vec : array-like, shape=[..., dim]
Cotangent vector at base point, the initial speed of the geodesics.

Returns
-------
path : callable
Time parameterized normal geodesic curve. If a batch of initial
conditions is passed, the output array's first dimension
represents the different initial conditions, and the second
corresponds to time.
"""
if initial_cotangent_vec is None:
raise ValueError(
"Specify an initial cotangent " "vector to define the geodesic."
)

initial_point = gs.to_ndarray(initial_point, to_ndim=2)
initial_cotangent_vec = gs.to_ndarray(initial_cotangent_vec, to_ndim=2)

n_initial_conditions = initial_cotangent_vec.shape[0]

if n_initial_conditions > 1 and len(initial_point) == 1:
initial_point = gs.stack([initial_point[0]] * n_initial_conditions)

def path(t):
"""Generate parameterized function for geodesic curve.

Parameters
----------
t : array-like, shape=[n_points,]
Times at which to compute points of the geodesics.
"""
t = gs.array(t)
t = gs.cast(t, initial_cotangent_vec.dtype)
t = gs.to_ndarray(t, to_ndim=1)

cotangent_vecs = gs.einsum("i,...k->...ik", t, initial_cotangent_vec)

points_at_time_t = [
self.exp(tv, pt, n_steps=n_steps)
for tv, pt in zip(cotangent_vecs, initial_point)
]
points_at_time_t = gs.stack(points_at_time_t, axis=0)

return (
points_at_time_t[0] if n_initial_conditions == 1 else points_at_time_t
)

return path