Setting Up a Python Environment for HEP Machine Learning

Setting Up a Python Environment for HEP Machine Learning

Getting your Python environment right is the unglamorous prerequisite that determines whether the rest of your analysis goes smoothly or turns into a debugging marathon. If you've ever spent an afternoon watching ROOT and TensorFlow fight over library versions, you already know the pain. This guide walks you through the three main options—conda, venv, and containers—so you can pick the one that fits your workflow and get on with the actual machine learning.

Why Environment Management Matters in HEP

Particle physics software stacks are notoriously tangled. ROOT carries its own Python bindings (PyROOT), which are sensitive to the Python version they were compiled against. ML frameworks like PyTorch, TensorFlow, and JAX each have their own CUDA dependencies. Without isolation, a single pip install can silently break a working setup. A clean Python ROOT machine learning setup keeps these layers separate, reproducible, and shareable with collaborators.

Conda is the most forgiving choice for a conda environment particle physics workflow because it manages non-Python dependencies—compiled C++ libraries, CUDA toolkits, ROOT itself—alongside Python packages. This is the main reason it outperforms plain pip for HEP work.

Step 1: Install Miniforge

Prefer Miniforge over Anaconda. It defaults to the community-maintained conda-forge channel, which is where the HEP-specific packages live, and it's lighter weight.

# Linux/macOS
wget https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-x86_64.sh
bash Miniforge3-Linux-x86_64.sh

Follow the prompts, let it initialize your shell, then restart your terminal.

Step 2: Create a Dedicated Environment

Never work in base. Give your environment a descriptive name.

conda create -n hepml python=3.11
conda activate hepml

Step 3: Install ROOT via conda-forge

conda install -c conda-forge root

This pulls a pre-compiled ROOT build that is already matched to the Python version in your environment. Test it immediately:

import ROOT
ROOT.gSystem.Load("libRIO")

Step 4: Add ML Libraries

conda install -c conda-forge numpy scipy matplotlib scikit-learn
pip install torch torchvision  # or tensorflow; follow the project's own GPU install guide

Install PyTorch or TensorFlow via pip after conda, following the official GPU instructions for your CUDA version. Mixing conda and pip is safe when pip comes last.

Step 5: HEP-Specific Python Packages

pip install uproot awkward hist vector coffea

uproot lets you read and write ROOT files without a ROOT installation, which is useful for lightweight inference pipelines. awkward handles jagged arrays (variable-length jet constituents, for instance) in a NumPy-like style.

Option 2: venv + pip

If you're working on a cluster where conda isn't available or where a system ROOT is already maintained by your group, a plain virtual environment is perfectly reasonable.

python3 -m venv ~/envs/hepml
source ~/envs/hepml/bin/activate
pip install --upgrade pip
pip install numpy scipy scikit-learn matplotlib uproot awkward

The limitation is that you cannot install ROOT this way; you must load it from the system or a CVMFS release. For a python setup HEP machine learning workflow on CERN's lxplus or a similar cluster, this pattern is common and stable.

Option 3: Containers (Docker / Apptainer)

Containers are the most reproducible option and the right choice for production pipelines, CI/CD, or sharing a fully frozen environment with a paper. The HEP community maintains images through the hepstore and ML4HEP projects on Docker Hub.

docker pull hepstore/hep-ml
docker run -it --gpus all -v $(pwd):/work hepstore/hep-ml bash

On clusters that don't allow Docker, use Apptainer (formerly Singularity), which can convert Docker images directly. Containers shine when you need the exact same environment on your laptop, on the grid, and six months from now when a reviewer asks you to reproduce a plot.

Keeping Your Environment Sane

A few habits that save time:

  • Pin your environment. Export a snapshot with conda env export > environment.yml or pip freeze > requirements.txt and commit it alongside your analysis code.
  • One environment per project. Sharing an environment across analyses is how conflicts start.
  • Check ROOT–Python compatibility early. Run a trivial PyROOT script on day one, not after you've written a hundred lines of analysis code.
  • Use Jupyter with a kernel per environment. Register your conda env as a kernel: python -m ipykernel install --user --name hepml.

Going Further

Environment setup is the foundation, but the interesting work starts once you're inside a working shell. If you want a structured path from data loading through model training and evaluation, the complete HEP ML course walks through each stage with realistic particle physics examples. You can also start with the free Module 1 to get a feel for the pedagogy before committing.

For most analysts starting out, conda with conda-forge is the path of least resistance—get that working first, then layer in containers once your pipeline matures.

Want to go deeper?

Machine Learning for High Energy Physics: The Complete Course takes you from first principles to a defensible result in 6 structured modules. $97, 30-day guarantee.

See the course →

Not ready yet? Grab Module 1 free →