Physics-informed continuous normalizing flows to learn the electric field in time projection chambers
FieldFlow is a JAX-based library implementing neural ODEs to learn the electric field transformation in time projection chambers directly from calibration data.
In a xenon TPC, ionization electrons drift through the detector under an applied electric field. Imperfections in this field distort the electrons' paths, so the positions measured at the top of the detector don't match where the interaction actually occurred. FieldFlow models the electric field and thereby can be used to correct these distortions.
The model enforces physical constraints (field conservativity) through its architecture to ensure a curl-free electric field. Compared to traditional field distortion corrections in xenon TPCs, the flow can achieve comparable accuracy in position reconstruction while requiring significantly fewer calibration events.
The library implements continuous normalizing flows with configurable ODE solvers, supporting both exact and approximate log probability computation. Multi-GPU training is supported with automatic data parallelization across devices.
Physics-informed continuous normalizing flows to learn the electric field within a time-projection chamber Ivy Li, Peter Gaemers, Juehang Qin, Naija Bruckner, Maris Arthurs, Maria Elena Monzani, Christopher Tunnell (2025) arXiv:2511.01897 (under review)
- Continuous normalizing flows with exact or approximate log probability computation
- Multi-GPU training with automatic data sharding across devices
- Configurable ODE solvers including adaptive PID controllers
- Fine-tuning support for transfer learning from pretrained models
- Position reconstruction integration with pretrained flow models
- JAX-based implementation for GPU acceleration and automatic differentiation
Solving the neural ODE: A continuous normalizing flow takes an interaction's transverse position (x, y) and time as input and outputs the dynamics of that position over drift time — where drift time maps directly to depth z inside the detector (z = drift velocity x time). Solving the neural ODE backward in time traces how a position gets distorted as electrons drift through the field. Inverting this transformation recovers the true interaction position.
Constraining the electric field: Rather than learning the electric field directly, the neural network parameterizes the negative gradient of a scalar potential. This architectural choice guarantees that the learned electric field is curl-free, as we make the assumption of electrostatic electric fields which do not significantly vary over time. Thus, unlike other physics-informed models which enforce physics constraints through a loss penalty, the physics here is baked into the model architecture instead.
Fig. 1: A position s inside the TPC evolves over drift time t, where depth corresponds to z(t) = vdt. The neural network learns a function gϕ approximating a scaled scalar potential, and outputs −∇sgϕ — the transverse gradient that governs how each position transforms as electrons drift. Because the dynamics are derived from a scalar potential, the resulting field is guaranteed to be curl-free.
Requires Python ≥3.10. Install from source:
git clone https://github.com/RiceAstroparticleLab/fieldflow.git
cd fieldflow
pip install -e .Train a new model:
python -m fieldflow config.tomlFine-tune a pretrained model:
python -m fieldflow config.toml --pretrained model.eqxSee the documentation (in progress) for example usage details. Configuration options are documented in sample_config.toml.
fieldflow/
├── src/fieldflow/ # Core library (flows, ODE solvers, training)
├── tests/ # Unit tests
├── docs/ # Documentation
├── sample_config.toml # Example training configuration
└── pyproject.toml # Package metadata
MIT License. See LICENSE for details.
Report issues at GitHub Issues.