File size: 3,634 Bytes
213d4a7 ac1fa8d 213d4a7 ac1fa8d 1128d9c ac1fa8d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | ---
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}
}
```
|