https://github.com/danielkelshaw/riemax
Riemannian geometry in JAX
https://github.com/danielkelshaw/riemax
deep-learning differential-geometry geodesics jax riemannian-geometry
Last synced: 6 months ago
JSON representation
Riemannian geometry in JAX
- Host: GitHub
- URL: https://github.com/danielkelshaw/riemax
- Owner: danielkelshaw
- License: mit
- Created: 2023-10-17T21:07:35.000Z (over 2 years ago)
- Default Branch: main
- Last Pushed: 2024-08-21T00:53:16.000Z (almost 2 years ago)
- Last Synced: 2024-12-06T21:45:11.053Z (over 1 year ago)
- Topics: deep-learning, differential-geometry, geodesics, jax, riemannian-geometry
- Language: Python
- Homepage: https://riemax.danielkelshaw.com
- Size: 2.32 MB
- Stars: 13
- Watchers: 1
- Forks: 2
- Open Issues: 0
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
# Riemax: Riemannian geometry for JAX
> [!IMPORTANT]
> This library is very much a work in progress, as is the documentation. Naturally, the contents of this library is subject to change depending on choice of research avenue. At the moment, there exists a solid basis on which to build, but the API may change to improve accessiblity for new users. If there is anything you would like to see in the library, or anything you think that needs to change -- please do let me know.
Riemax is a [JAX] library for Riemannian geometry, providing the ability to define induced metrics and operate on manifolds directly. This includes functionality, such as:
- Computing geometric quantities on manifolds.
- Defining operators on manifolds.
- Tools for geometric statistics.
## Installation
```bash
pip install git+https://github.com/danielkelshaw/riemax
```
Requires Python 3.12+ and JAX 0.4.13+.
## Quick Example
Manifolds can be defined by a function transformation, $\iota: M \rightarrow N$:
```python
import riemax as rx
fn_transformation = rx.fn_transformations.fn_peaks
manifold = rx.Manifold.from_fn_transformation(fn_transformation)
```
and can be used to compute properties on the manifold:
```python
import jax
import jax.numpy as jnp
p = jnp.array([0.0, 0.0])
metric = manifold.metric_tensor(p)
christoffel = manifold.sk_christoffel(p)
ricci_scalar = manifold.ricci_scalar(p)
```
We can also define operators. Given a function $f: M \rightarrow \mathbb{R}$:
```python
from riemax.manifold import M, Rn
def f(p: M[jax.Array]) -> Rn[jax.Array]:
return jnp.einsum('i -> ', p ** 2)
fn_grad = manifold.grad(f)
# we can define the laplacian explicitly:
fn_laplacian = manifold.div(fn_grad)
# ... or from manifold:
fn_laplacian = manifold.laplace_beltrami(f)
```
We can exploit the ability to compute geometric quantities to compute the exponential map:
```python
dt = 1e-3
n_steps = int(1.0 // dt)
fn_exp_map = rx.manifold.maps.exponential_map_factory(
integrator=rx.numerical.integrators.odeint,
dt=dt,
metric=manifold.metric_tensor,
n_steps=n_steps
)
p_in_tm = rx.manifold.TangentSpace(p, jnp.array([1.0, 1.0]))
q_in_tm, trajectory = fn_exp_map(p_in_tm)
```
## Next Steps
The examples above show just a fraction of what is possible with Riemax. If this quick start has piqued your interest, please feel free to take a look at the examples to get a better feeling for what is possible with this library.
## Citation
If you found this library to be useful in academic work, then please cite:
```tex
@misc{kelshaw2023riemax
title = {{Riemax}: {R}iemannian geometry in {JAX} via automatic differentiation}
author = {Daniel Kelshaw}
year = {2023},
url = {https://github.com/danielkelshaw/riemax}
}
```
[JAX]: https://github.com/google/jax