Add documentation

.github/workflows/docs.yml

@@ -0,0 +1,41 @@
+name: Deploy Documentation
+
+concurrency:
+  group: pages
+  cancel-in-progress: false
+
+on:
+  push:
+    branches:
+      - main 
+
+permissions:
+  contents: read
+  pages: write
+  id-token: write
+
+jobs:
+  deploy:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.x'
+
+      - name: Install Zensical
+        run: pip install zensical
+
+      - name: Build documentation
+        run: zensical build --clean
+
+      - name: Upload Pages artifact
+        uses: actions/upload-pages-artifact@v4
+        with:
+          path: site
+
+      - name: Deploy to GitHub Pages
+        uses: actions/deploy-pages@v4

README.md

@@ -2,10 +2,45 @@
 
 # FieldFlow
 
-JAX-based continuous normalizing flows for electric field modeling in Time Projection Chambers (TPCs). FieldFlow uses neural ODEs to learn the mapping between hit patterns and electric field configurations for astrophysical particle detection experiments.
+**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.
 
+## Paper
+
+> **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](https://arxiv.org/abs/2511.01897) *(under review)*
+
+## Key Features
+
+- **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
+
+## How It Works
+
+**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.
+
+<p align="center">
+  <img src="docs/assets/model_architecture.png" alt="FieldFlow model architecture" width="700">
+</p>
+
+<p align="center"><em>
+Fig. 1: A position <b>s</b> inside the TPC evolves over drift time t, where depth corresponds to z(t) = v<sub>d</sub>t. The neural network learns a function g<sub>ϕ</sub> approximating a scaled scalar potential, and outputs −∇<sub>s</sub>g<sub>ϕ</sub> — 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.
+</em></p>
+
 ## Installation
 
 Requires Python ≥3.10. Install from source:
@@ -16,7 +51,7 @@ cd fieldflow
 pip install -e .
 ```
 
-## Quick Usage
+## Quick Start
 
 Train a new model:
 ```bash
@@ -28,16 +63,18 @@ Fine-tune a pretrained model:
 python -m fieldflow config.toml --pretrained model.eqx
 ```
 
-See [USAGE.md](USAGE.md) for detailed usage instructions and [sample_config.toml](sample_config.toml) for configuration options.
+See the [documentation](https://riceastroparticlelab.github.io/fieldflow/) (in progress) for example usage details. Configuration options are documented in [`sample_config.toml`](sample_config.toml).
 
-## Key Features
+## Project Structure
 
-- **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
+```
+fieldflow/
+├── src/fieldflow/       # Core library (flows, ODE solvers, training)
+├── tests/               # Unit tests
+├── docs/                # Documentation
+├── sample_config.toml   # Example training configuration
+└── pyproject.toml       # Package metadata
+```
 
 ## License
 

docs/index.md

@@ -0,0 +1,8 @@
+# FieldFlow Documentation
+
+JAX-based continuous normalizing flows for electric field modeling in Time Projection Chambers (TPCs).
+
+## Quick Links
+
+- [Usage Guide](USAGE.md)
+- [GitHub Repository](https://github.com/RiceAstroparticleLab/fieldflow)

mkdocs.yml

@@ -0,0 +1,12 @@
+site_name: FieldFlow Documentation
+site_description: JAX-based continuous normalizing flows for electric field modeling
+repo_url: https://github.com/RiceAstroparticleLab/fieldflow
+
+docs_dir: docs  # This tells Zensical to use the 'docs' folder
+
+nav:
+  - Home: index.md
+  - Usage: USAGE.md
+
+theme:
+  name: material