Code and Functionality Map: ParticleFlow

This document provides a high-level overview of the repository’s structure and the functionality of its core components. To run the python code in the right environment, use uv run python3 ....

1. Configuration & Specification¶

The project uses a hierarchical configuration system.

• particleflow_spec.yaml: The single source of truth for the entire project. It defines machine-specific paths (sites), data production scenarios, and model architectures.

• mlpf/conf.py: Defines the Pydantic models for the configuration, ensuring type safety, path resolution, and model-type schemas including attention, gnnlsh, litept, hept, and heptv2.

• mlpf/pipeline.py: Implements hierarchical configuration resolution: base defaults in mlpf/conf.py are overridden by scenario-specific values in particleflow_spec.yaml, which can further be overridden via command-line arguments (e.g., --model.num_convs 6).

• configs/: Site-specific Pixi environment configurations (local/, lxplus/, tallinn/) and reusable training scenarios/platform profiles under configs/training/. The root pixi.toml is a symlink into this directory.

• pixi.toml / pixi.lock / uv.lock / uv.singularity: Project environment, container definitions, and task management. Defines common tasks like gen, post, train, and validation.

• envs/: Isolated virtual environment specifications (e.g., ort-cpu, ort-gpu) for specific runtimes like ONNX.

• validation_cms.yaml / validation_key4hep.yaml: Specification files for validation scenarios.

2. Workflow Orchestration¶

Complex data production and training pipelines are managed using Snakemake or site-specific shell scripts.

• mlpf/snakemake/: Contains scripts to generate Snakemake workflows.

  • produce_snakemake.py: Generates workflows for generation, postprocessing, TFDS creation, and training.

  • produce_cms_validation_snakemake.py: Orchestrates validation workflows specifically for CMS.

  • produce_validation_snakemake.py: Orchestrates validation workflows for Key4Hep detectors (CLD, CLIC).

• mlpf/pipeline.py: The main CLI for training, testing, LR range tests (lr-find), and hyperparameter optimization. Supports standard and Ray-based execution.

• scripts/training/run_scenario.py: Resolves a generic scientific training scenario against a hardware profile, validates comparison invariants and global batch size (one value, or one per train group), and launches reproducible seeded jobs.

• scripts/training/compare_scenario.py: Tabulates and plots the (per-group) validation losses of the runs of one scenario.

3. Data Production & Preprocessing¶

• mlpf/data/: Simulator-specific code for generating and preprocessing data. ColliderML: colliderml/clustering.py (cluster_event dispatcher + the properties every clusterer enforces; algorithms in colliderml/radius_graph.py (radius-graph bfs/union_find) and colliderml/clue.py (CLUE), shared helpers in colliderml/cluster_common.py), colliderml/cluster_eval.py (decodes/caches ColliderML events and their per-hit truth, scores clustering settings and the v20 Pandora clusters with cluster_metrics.py). Shared across detectors: cluster_metrics.py (truth-based clustering metrics, detector- and clusterer-agnostic).

  • cms/: Scripts for CMSSW-based generation (genjob_pu.sh), postprocessing (postprocessing2.py), validation (valjob.sh, valjob_data.sh), and plotting (plot_cms.py).

  • key4hep/: Scripts for Key4Hep-based generation (gen/) and postprocessing (postprocessing.py).

• mlpf/heptfds/: TFDS (TensorFlow Datasets) builders for CMS, CLD, CLIC, IDEA, MAIA, and ColliderML datasets, including support for both cluster-based and raw hits-based data (cld_pf_edm4hep_hits, clic_pf_edm4hep_hits). Shared EDM4Hep utilities in edm4hep_utils/. ColliderML support: colliderml_pf/ttbar.py (clustered view), colliderml_hits/{ttbar,ttbar_pu200}.py (hits view)

  • shared conversion in colliderml_utils/utils.py reading the MLPF-format parquet written by mlpf/data/colliderml/postprocessing.py.

4. Machine Learning Core (mlpf/model/)¶

• mlpf.py: Implementation of the MLPF model, featuring multi-head attention and configurable sub-networks. Supports fused attention, simplified math attention for ONNX export, and HEPT/HEPTv2/LitePT layers that consume detector coordinates.

• gnnlsh.py: GNN layers with Locality-Sensitive Hashing (LSH) for scalable graph processing, including multi-hash binning, attention kernels, RMSNorm/SwiGLU blocks, and inter-bin attention.

• hept.py: Implementation of the Hashing-based Efficient Particle Transformer (HEPT).

• heptv2.py: HEPTv2 point encoder layer with hash-bucketed attention, learned positional encoding over eta/phi coordinates, Qwen-style RMSNorm/MLP blocks, and environment switches for attention/hash implementation variants.

• HEPTv2 activation checkpointing: Optional per-layer training recomputation in MLPF._run_backbone and a batch-384 comparison recipe (historical filename batch512). See heptv2_checkpointing.md for local memory measurements and target-platform validation.

• CLD hit-model evidence: Current recommendation, completed and partial ablations, and historical metric/dataset caveats: cld_hit_results.md.

• Tracker momentum features: Optional conformal curvature/quality, longitudinal moments and bounded cross-surface doublets for hit inputs. See tracker_momentum_features.md for definitions, ablations and measured cost.

• litept.py: Integration for the LitePT (Lightweight Point Transformer) architecture.

• PFDataset.py: Advanced data loading logic and multi-file handling. Training mixes one loader per train group with WeightedMixIterator (smooth weighted round-robin over the groups’ optional weights; each group cycles through its own epochs); validation interleaves its groups with InterleavedIterator.

• training.py: Core training loop implementation, supporting DDP and Ray-based distributed training, TensorBoard loss logging (also per dataset group), and per-step data-load/model-forward timing.

• losses.py: Specialized loss functions for particle classification and energy regression.

• lr_finder.py: LR range test behind pipeline.py lr-find: exponential lr ramp after the task-weight calibration steps, smoothed loss, divergence stop, lr suggestions, lr_find/lr_find.{json,png} in the experiment dir.

• inference.py: Utilities for running model inference and generating predictions.

• plots.py: Confusion matrix logging and model performance plotting during training.

• utils.py: Model-level utility functions including target unpacking, object-condensation clustering helpers, learning rate schedules, and metric computation.

• tracker_momentum_readout.py: Optional tracker-only momentum residual and sparse differentiable hit-to-slot pooling; see the small-model ablation.

• Tracker grouping losses: Particle-balanced tracker CE and soft tracker overlap supervision; see the from-scratch small-model factorial.

• Tracker width scaling: Four fixed-depth widths and three seeds with tracker momentum readout and balanced CE; see the scaling recipe.

• scripts/studies/train_frozen_tracker_readout.py: Local head-only hard/soft/oracle momentum comparison with cached frozen-backbone summaries; see the experiment and resume instructions.

• distributed_ray.py: Integration with Ray for distributed training and HPO.

5. Validation, Plotting & Monitoring¶

• Local ttbar detector comparison: see the detector-comparison runbook for inputs, rerun commands, statistical definitions, cache behavior, and limitations.

• mlpf/standalone_eval/key4hep/: Standalone tools for evaluating model checkpoints and generating performance plots on ROOT files for Key4Hep detectors. Note the ColliderML validator is the shared tests/validate_parquet.py (it knows how to handle the colliderml detector name).

• scripts/studies/probe_calo_continuation.py: Cross-layer, distance-matched calorimeter continuation probes with equal particle weighting. See calo_continuation_probe.md.

• scripts/studies/probe_hit_instances.py: Frozen linear pair-affinity probes comparing raw hit features and HEPTv2 embeddings for separate tracker/calorimeter instance segmentation. See hit_instance_probe.md.

• mlpf/model/reconstruction_diagnostics.py: Optional validation diagnostics for hit-grouping purity/completeness and momentum errors by class and eta. Enable with validation_diagnostics_batches; definitions and output paths are in reconstruction_diagnostics.md.

• Charged fit correction: An opt-in charged-pT branch in set_decoder.py combines tracker_group_readout.py circle fits with a bounded learned correction; configuration and limits are in charged_fit_correction.md.

• mlpf/plotting/: Comprehensive suite of plotting tools.

  • plot_validation.py / plot_met_validation.py: Standard validation plots for jets and MET.

  • corrections.py: Derivation of jet energy corrections.

  • cmssw_validation_data.py: Validation scripts for CMS data.

  • cms_fwlite.py: CMS FWLite-based event analysis and plotting.

• scripts/cms-validate-onnx.py: Exports PyTorch models to ONNX (supporting FP32, FP16, Fused Flash Attention, HEPT, HEPTv2, and GNNLSH) and validates inference with per-configuration PyTorch/ONNX sessions.

• mlpf/model/onnx_attention.py, onnx_graph.py, onnx_runtime.py: Explicit Attention-24/decomposed export, padding and output rewrites, and provider selection/placement checks. See ONNX validation.

• scripts/export-checkpoint-onnx.py, qualify-heptv2-onnx.py, benchmark-model-export.py: Trained CMS and seeded-model reference bundles, runtime qualification, and full-model measurements; see Export qualification.

• scripts/plot-onnx-summary.py: Generates summary plots for ONNX validation results.

• mlpf/model/monitoring.py: System resource monitoring and logging.

6. Utilities & Miscellaneous¶

• mlpf/utils.py / mlpf/logger.py: Common utilities and centralized logging.

• mlpf/customizations.py: Config customization helpers for fast CI/test pipelines.

• mlpf/timing.py: Performance timing utilities.

• mlpf/optimizers/: Custom optimizers like LAMB (lamb.py).

• mlpf/standalone/: Standalone, hackable MLPF training and architecture search (train.py, eval.py, dsl.py, puppi.py, plot_evolution.py, run_evolution.py).

• mlpf/raytune/: Ray Tune integration for hyperparameter search (search_space.py, utils.py).

• mlpf/jet_utils.py: Jet clustering and matching logic.

• scripts/: Miscellaneous utility scripts.

  • benchmark.py: Benchmarks the forward and backward pass timings and peak memory usage for various model architectures.

  • tallinn/, lxplus/, flatiron/, lumi/, local/: Site-specific and local orchestration scripts for training and evaluation.

  • legacy/: Older site-specific scripts retained for previous cluster workflows.

  • upload_model_hf.py / upload_hf.py: Utilities for uploading experiment results and model checkpoints to HuggingFace Hub.

  • visualize_hits.py: Tool for visualizing detector hits and model embeddings using UMAP and 3D plotting.

  • local_test_cld.sh / local_test_cms.sh: Scripts for quick local verification of the pipeline.

  • fetch_test_data_cld.sh / fetch_test_data_cms.sh: Scripts to download test data for CLD and CMS.

• tests/: Unit tests and validation helpers for configuration, data loading, model forward passes, attention stability, HEPT/HEPTv2 layers, ONNX export, and Key4Hep assignment/ground-truth consistency.

7. Local ttbar Detector Comparison Report¶

The detector-comparison runbook contains the commands, local input inventory, cache behavior, statistical definitions, and limitations for the five-detector report and real ColliderML training check.