Clean code and update docstrings

sample_config.toml

@@ -11,20 +11,21 @@ description = "Sample configuration for FieldFlow training"
 # Model architecture parameters
 data_size = 2  # Dimensionality of the input data (2D for x,y coordinates)
 exact_logp = true  # Use exact log probability computation (more accurate but slower)
-width_size = 48  # Width of neural network hidden layers
-depth = 3  # Number of hidden layers in the neural network
+width_size = 256  # Width of neural network hidden layers
+depth = 16  # Number of hidden layers in the neural network
+scalar = true # Use scalar potential model
 
 # ODE solver settings - these control the accuracy and efficiency of the flow
-use_pid_controller = true  # Use adaptive PIDController (recommended) vs constant step size
+use_pid_controller = false  # Use adaptive PIDController (recommended) vs constant step size
 rtol = 1e-3  # Relative tolerance for PIDController (smaller = more accurate, slower)
 atol = 1e-6  # Absolute tolerance for PIDController (smaller = more accurate, slower)
-dtmax = 5.0  # Maximum step size for PIDController
+dtmax = 2.0  # Maximum step size for PIDController
 dtmin = 0.05  # Minimum step size for PIDController
 
 # Time integration parameters
 t0 = 0.0  # Starting time for ODE integration
 extract_t1 = 10.0  # End time for extract phase
-dt0 = 1.0  # Initial time step size
+dt0 = 0.5  # Initial time step size
 
 [training]
 # Multi-GPU Training Support:
@@ -39,45 +40,45 @@ dt0 = 1.0  # Initial time step size
 
 # Training process parameters
 seed = 42  # Random seed for reproducibility
-learning_rate = 2e-3  # Initial learning rate (will be scheduled during training)
+learning_rate = 1e-5  # Initial learning rate (will be scheduled during training)
 weight_decay = 1e-4  # L2 regularization parameter
-epochs = 100  # Number of training epochs
+epochs = 300  # Number of training epochs
 enable_scheduler = true  # Enable learning rate scheduling for training from scratch
                         # When false, uses constant LR = learning_rate * 0.01
                         # Set to false when loading pretrained models to continue training
 
 # Data and batching parameters
-batch_size = 2048  # Training batch size (adjust based on GPU memory)
+batch_size = 1024  # Training batch size (adjust based on GPU memory)
 num_devices = 1  # Number of GPUs for data parallelization
                  # Examples: 1 (single GPU), 2 (dual GPU), 4 (quad GPU), 8 (octa GPU)
                  # Note: batch_size should be divisible by num_devices for optimal performance
 n_samples = 16  # Number of samples per instance for likelihood estimation
-n_train = 200000  # Size of training set
-n_test = 20000  # Size of test/validation set
+n_train = 65536  # Size of training set
+n_test = 4096  # Size of test/validation set
 
 # Training strategy parameters
 use_best = true  # Use the best model based on validation loss (recommended)
 curl_loss_multiplier = 1000.0  # Weight for curl penalty (encourages curl-free fields)
 z_scale = 5.0  # Scaling factor for z dimension coordinates
-multisteps_every_k = 4  # Gradient accumulation steps for MultiSteps optimizer
+multisteps_every_k = 2  # Gradient accumulation steps for MultiSteps optimizer
 
 [experiment]
 # Physical experimental setup parameters
-tpc_height = 148.6515  # Height of the TPC in cm (for filtering z coordinates)
-tpc_r = 66.4  # Radius of the TPC in cm (for boundary constraints)
+tpc_height = 259.92  # Height of the TPC in cm (for filtering z coordinates)
+tpc_r = 129.96  # Radius of the TPC in cm (for boundary constraints)
 
 [posrec]
 # Position reconstruction flow model parameters
 # These should match the pretrained position reconstruction model
-flow_layers = 5  # Number of coupling layers in the flow
+flow_layers = 6  # Number of coupling layers in the flow
 nn_width = 128  # Width of neural networks in coupling layers
-nn_depth = 3  # Depth of neural networks in coupling layers
+nn_depth = 6  # Depth of neural networks in coupling layers
 invert_bool = false  # Whether to invert the flow (should match pretrained model)
-cond_dim = 494  # Conditioning dimension (should match hit pattern size)
+cond_dim = 860  # Conditioning dimension (should match hit pattern size)
 
 # Spline transformation parameters
 spline_knots = 5  # Number of knots for rational quadratic splines
 spline_interval = 5.0  # Interval for spline transformations
 
 # Coordinate transformation parameters
-radius_buffer = 20.0  # Buffer for predictions beyond TPC radius (in cm)
+radius_buffer = 0.0  # Buffer for predictions beyond TPC radius (in cm)

src/fieldflow/__init__.py

@@ -1,8 +1,17 @@
-"""JAX-based normalizing flows for physical field modeling.
+"""JAX-based continuous normalizing flows for electric field modeling.
 
-This package provides tools for modeling physical fields using continuous
-normalizing flows, with a focus on electric field modeling for particle
-detectors.
+FieldFlow provides tools for modeling electric fields in dual-phase Time
+Projection Chambers (TPCs) using continuous normalizing flows (CNFs). The
+architecture mirrors the physical structure of dual-phase TPCs, with separate
+neural networks for the extraction field (z-independent distortions) and
+drift field (z-dependent distortions).
+
+The library supports two approaches for enforcing Maxwell's equations:
+
+- **Scalar potential method**: Models the field as the negative gradient of a
+  learned scalar potential, which is curl-free by construction.
+- **Vector field with curl loss**: Directly learns the vector field while
+  penalizing non-zero curl during training.
 """
 
 from fieldflow.config import Config, load_config

src/fieldflow/__main__.py

@@ -1,7 +1,12 @@
-"""Entry point for FieldFlow training.
+"""Command-line interface for FieldFlow training.
 
-This module provides a simple command-line interface for training FieldFlow
-models from configuration, with optional fine-tuning of pre-trained models.
+This module provides the CLI entry point for training CNF models to learn
+electric field distortions in dual-phase TPCs. Supports training from
+scratch or fine-tuning pretrained models.
+
+Usage:
+    python -m fieldflow config.toml
+    python -m fieldflow config.toml --pretrained model.eqx
 """
 
 import argparse
@@ -23,14 +28,18 @@
 
 
 def create_model_from_config(config, key):
-    """Create a CNF model from configuration.
+    """Create a CNF model from configuration parameters.
+
+    Initializes a ContinuousNormalizingFlow with architecture and ODE solver
+    settings from the config. Uses DriftFromPotential (scalar method) or
+    MLPFunc (vector method) based on config.model.scalar.
 
     Args:
-        config: Configuration object containing model parameters
-        key: JAX PRNG key for model initialization
+        config: Config object with model parameters.
+        key: JAX random key for parameter initialization.
 
     Returns:
-        Initialized ContinuousNormalizingFlow model
+        Initialized ContinuousNormalizingFlow model.
     """
     # Create step size controller based on config
     if config.model.use_pid_controller:

src/fieldflow/config.py

@@ -18,28 +18,31 @@
 
 @dataclass
 class ModelConfig:
-    """Configuration for continuous normalizing flow model architecture and
-    behavior.
+    """Configuration for continuous normalizing flow model architecture.
 
-    This class encapsulates the parameters used to define a Continuous
-    Normalizing Flow model, including neural network architecture, ODE
-    solver settings, and model-specific hyperparameters.
+    This class defines parameters for the CNF model that learns electric field
+    distortions in dual-phase TPCs. The model uses separate neural networks for
+    extraction (z-independent) and drift (z-dependent) field components.
 
     Attributes:
-        data_size: Dimensionality of the input data.
-        exact_logp: Whether to use exact log probability calculation.
-        width_size: Width of neural network layers.
-        depth: Depth of neural network.
-        scalar: Whether to use scalar field instead of vector field.
-        use_pid_controller: Whether to use PIDController instead of
-            ConstantStepSize.
-        rtol: Relative tolerance for PIDController.
-        atol: Absolute tolerance for PIDController.
-        dtmax: Maximum step size for PIDController.
-        dtmin: Minimum step size for PIDController.
-        t0: Starting time for ODE.
-        extract_t1: End time for extract phase.
-        dt0: Initial time step.
+        data_size: Dimensionality of the input data (default 2 for x,y).
+        exact_logp: If True, compute exact log probability using full Jacobian
+            trace. If False, use Hutchinson trace estimator (faster but
+            approximate).
+        width_size: Width of hidden layers in the neural networks.
+        depth: Number of hidden layers in the neural networks.
+        scalar: If True, use scalar potential method (curl-free by
+            construction). If False, use direct vector field with curl
+            penalty loss.
+        use_pid_controller: If True, use adaptive PID step size controller.
+            If False, use constant step size.
+        rtol: Relative tolerance for adaptive ODE solver.
+        atol: Absolute tolerance for adaptive ODE solver.
+        dtmax: Maximum step size for adaptive ODE solver.
+        dtmin: Minimum step size for adaptive ODE solver.
+        t0: Starting time for ODE integration.
+        extract_t1: End time for extraction phase ODE integration.
+        dt0: Initial time step for ODE solver.
     """
 
     data_size: int = 2
@@ -63,7 +66,23 @@ class ModelConfig:
 
 @dataclass
 class PosRecFlowConfig:
-    """Configuration for Position Reconstruction Flow model."""
+    """Configuration for the position reconstruction normalizing flow.
+
+    The position reconstruction flow is a pretrained conditional normalizing
+    flow that maps detector hit patterns to (x, y) position distributions.
+    This flow provides the prior distribution for CNF training.
+
+    Attributes:
+        flow_layers: Number of coupling layers in the normalizing flow.
+        nn_width: Width of hidden layers in coupling layer neural networks.
+        nn_depth: Number of hidden layers in coupling layer neural networks.
+        invert_bool: Whether to invert the flow direction.
+        cond_dim: Dimension of the conditioning vector (hit pattern size).
+        spline_knots: Number of knots for rational quadratic spline bijections.
+        spline_interval: Interval parameter for spline transformations.
+        radius_buffer: Buffer beyond TPC radius for coordinate transformation,
+            allowing predictions slightly outside the physical boundary.
+    """
 
     # Neural network architecture
     flow_layers: int = 5
@@ -86,27 +105,31 @@ class PosRecFlowConfig:
 class TrainingConfig:
     """Configuration for training CNF models.
 
-    This class encapsulates the parameters used during the training process,
-    including optimization settings, data handling, and training strategies.
+    This class defines parameters for the training process including
+    optimization, batching, loss computation, and multi-GPU parallelization.
 
     Attributes:
-        seed: Random seed for training reproducibility.
-        learning_rate: Initial learning rate for the optimizer.
-        weight_decay: L2 regularization parameter.
+        seed: Random seed for reproducibility.
+        learning_rate: Initial learning rate for AdamW optimizer.
+        weight_decay: L2 regularization coefficient.
         epochs: Number of training epochs.
-        batch_size: Batch size for training.
-        enable_scheduler: Whether to enable learning rate scheduling.
-        epoch_start: First epoch to begin training at.
-        n_samples: Number of samples per instance for likelihood estimation.
-        n_train: Size of training set.
-        n_test: Size of test/validation set.
-        use_best: Whether to use the best model based on validation.
-        curl_loss_multiplier: Coefficient for curl loss component.
-        z_scale: Scaling factor for z dimension.
-        multisteps_every_k: Steps for MultiSteps optimizer.
-        num_devices: Number of devices to use for data parallelization.
-        save_iter: Every n number of epochs to save the model
-        save_file_name: Path and file name start (/path/to/model_name)
+        batch_size: Number of samples per training batch.
+        enable_scheduler: If True, use learning rate schedule that reduces
+            LR at epochs 20 and 70. If False, use constant learning rate.
+        epoch_start: Starting epoch number (useful for resuming training).
+        n_samples: Number of Monte Carlo samples per event for likelihood
+            estimation from the position reconstruction flow.
+        n_train: Number of training samples to use from the dataset.
+        n_test: Number of samples for validation/test set.
+        use_best: If True, return the model with lowest validation loss.
+        curl_loss_multiplier: Weight for curl penalty term (only used when
+            scalar=False in ModelConfig).
+        z_scale: Scaling factor to normalize z coordinates.
+        multisteps_every_k: Gradient accumulation steps before optimizer
+            update.
+        num_devices: Number of GPUs for data parallelization.
+        save_iter: Save model checkpoint every N epochs.
+        save_file_name: Base filename for saved model checkpoints.
     """
 
     # Training process parameters
@@ -138,14 +161,16 @@ class TrainingConfig:
 
 @dataclass
 class ExperimentConfig:
-    """Configuration for experimental parameters.
+    """Configuration for the physical TPC geometry.
 
-    This class encapsulates parameters that describe the physical
-    experimental setup and constraints.
+    This class defines physical parameters of the dual-phase Time Projection
+    Chamber that constrain the model.
 
     Attributes:
-        tpc_height: Height of the TPC for filtering z coordinates.
-        tpc_r: Radius of the TPC for boundary constraints.
+        tpc_height: Height of the TPC drift region in cm. Used to filter
+            events by z coordinate (keeping -tpc_height < z < 0).
+        tpc_r: Radius of the cylindrical TPC in cm. Used for boundary
+            constraints in the loss function.
     """
 
     tpc_height: float = 259.92