diffjeom

Differential geometry powered by jax. Given a metric function, this package lets you compute:

• Christoffel symbols of the second kind (get_christoffel2);
• The Riemann curvature tensor (get_riemann);
• The Ricci tensor (get_ricci_tensor);
• The Ricci scalar/scalar curvature (get_ricci_scalar).

These work with the usual jax transformations like jit and vmap.

The functions check_christoffel2_sym, check_riemann_sym, check_diff_bianchi and check_ricci_tensor_sym are provided to check identities for these objects.

Getting started

Install with

pip install diffjeom

This will install jax and jaxlib if you don't already have them. The tests require numpy and can be run with pytest. Check them out for some usage examples.