EinFields: Neural Implicit Representations for General Relativity

EinFields is a JAX/Flax-based library for learning neural implicit representations of spacetime metrics in General Relativity. The repository provides pre-trained models on various neural network architectures: MLP, SIREN, WIRE.

Overview

This model HF repository contains plenty of pre-trained models. The reason for so many is that geodesics require different training domains, although it would've been possible and cleaner to enlarge this domain and have only one model.

Watch for metric_type in the architecture.yml as it says if the model is trained on the full metric or distortion. If distortion, then you can combine with the Minkowski metrics from our Github repo to get the full one.

Also, very important to have a look at train_data.yml to see what was the training grid on which the model was trained.

Repository Structure

EinFields/
├── flax_models/          # Core model implementations
│   ├── __init__.py       # Model factory and utilities
│   ├── activations.py    # Activation functions
│   ├── mlp.py            # Multi-Layer Perceptron
│   ├── siren.py          # SIREN architecture
│   └── wire.py           # WIRE
├── gw/                   # GW metrics
│   └── cartesian/
│       ├── silu/        # SILU activation models
│       ├── siren/       # SIREN models
│       └── wire/        # WIRE models
├── schwarzschild/        # Schwarzschild black hole models
│   └── spherical/
│       ├── close_event_horizon/
│       └── perihelion/
├── kerr/                 # Kerr black hole models
│   ├── boyer_lindquist/
└── └── kerr_schild_cartesian/

Each model directory contains:

architecture.yml: Model configuration
params.msgpack: Model parameters
train_data.yml: Training grid info

Loading Models

from huggingface_hub import hf_hub_download, snapshot_download
import os

# First option: get the full repository
repo_path = snapshot_download(repo_id="AndreiB137/EinFields")
# or clone the repository if you prefer

# Second option: get only flax_models and the model file you want

flax_models_folder = snapshot_download(
    repo_id="AndreiB137/EinFields",
    allow_patterns="flax_models/*"
)

model_folder = snapshot_download(
    repo_id="AndreiB137/EinFields",
    allow_patterns="kerr/boyer_lindquist/prograde/*"
)

# Then move the content in flax_models_folder to a directory where you are working with flax_models folder name. Afterwads:

from flax_models import load_metric_from_model

# Example
# `load_metric_from_model` returns directly the metric tensor function
metric_fn = load_metric_from_model("/your_path_to_model_folder/kerr/boyer_lindquist/zackiger")

# Now is ready to be used.

Code

The associated code and framework for EinFields is available on GitHub: https://github.com/AndreiB137/EinFields

Citation

Paper: https://arxiv.org/abs/2507.11589

@article{
      title={EINSTEIN FIELDS: A NEURAL PERSPECTIVE TO COMPUTATIONAL GENERAL RELATIVITY},
      author={Cranganore, Bodnar and Berzins},
      year={2025},
      eprint={2507.11589},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}