|
|
---
|
|
|
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}
|
|
|
}
|
|
|
```
|
|
|
|
|
|
|
