---
license: mit
library_name: flax
pipeline_tag: other
tags:
- neural-fields
- general-relativity
- physics
- jax
- flax
- siren
- wire
- mlp
---

# 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

```python
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](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}
}
```