Traditional High-Energy Physics (HEP) tracking architectures rely on Graph Neural Networks (GNNs), which suffer from a fatal
This repository is heavily modularized to reflect the physical production pipeline for the ALICE O2 Data Processing Layer.
Uses Conformal Mapping (
A highly optimized, batched Runge-Kutta spatial integrator. It integrates purely over Transverse Arc Length (
The active neural component embedded inside the ODE. A zero-initialized FP8 Multi-Layer Perceptron (MLP) that evaluates the kinematic state at each layer boundary and predicts the momentum decay. By wrapping this in a Softplus constraint and scaling by
The C++ bridge. Proves the device-agnostic execution of the Python prototype. It wraps the exported static ONNX computational graph in an asynchronous std::future, allocating pinned host memory via #ifdef USE_NVIDIA (CUDA) and #elif defined(USE_AMD) (ROCm) to guarantee non-blocking, zero-copy inference inside the O2 framework.
Strict dependency versions are required to ensure deterministic execution across O2 hardware farms.
pip install -r requirements.txtTo prove the necessity of the Neural ODE, you can evaluate the pure RK4 vacuum solver without ML friction. In the master evaluation script, the MLP can be toggled via the zero_forward injection.
python evaluate_pipeline.pyExpected Output: You will see the Z-axis Vacuum Math Error hit ~83 cm for low-momentum loopers, proving that mathematical geometry alone cannot account for ionization energy loss.
The main pipeline performs end-to-end backpropagation through the batched RK4 stepper. It utilizes zero-initialization, gradient clipping (max norm 1.0), and
python evaluate_pipeline.pyExpected Output: The network will aggressively pull the longitudinal tracking error down to the millimeter scale (e.g., ~0.31 cm).
Running the training pipeline automatically traces the PyTorch model and compiles it down to a branchless, 0-dynamic-loop graph for hardware inference.
python train_pipeline.pyExpected Output: Generates helical_flow_tracker_fp8.onnx and outputs the exact Min/Max FP8 normalization constants to be hardcoded into the C++ HAL.