Integrating ML Into the Coffea Analysis Framework

October 7, 2026 · 4 min read

If you have spent time building a coffea-based analysis, you know how satisfying it is to process millions of events in a few minutes with clean, readable code. Then a collaborator hands you a trained classifier and says "just add it to the processor." Suddenly the elegant columnar workflow feels fragile. This guide walks you through exactly how to plug a trained model into a coffea processor without losing the columnar structure you worked hard to build.

Why Columnar Analysis and ML Fit Together Naturally

Columnar analysis treats each observable as a named array across all events — momentum, isolation, b-tag score — rather than looping event by event. ML models, at inference time, do exactly the same thing: they expect a 2D array of shape (n_events, n_features) and return a 1D score array. The conceptual alignment is nearly perfect. The main friction is plumbing — getting your coffea arrays into the model and your scores back out cleanly.

Understanding that mapping is the core skill behind coffea ML integration, and once you see it clearly, the implementation becomes straightforward.

If you want deeper background on how columnar analysis ML HEP workflows are structured before diving into the code, the full HEP ML course covers the complete pipeline from raw NanoAOD to trained models.

Step 1 — Load the Model Once, Outside the Processor

The most common mistake is loading the model inside process(). This re-deserialises the model for every chunk, which is slow and wasteful. Instead, load it in __init__ and store it as an instance attribute.

import awkward as ak
import numpy as np
import uproot
import pickle

class MyProcessor(processor.ProcessorABC):
    def __init__(self, model_path):
        with open(model_path, "rb") as f:
            self.clf = pickle.load(f)   # or joblib.load, torch.load, etc.

If your model is large or you are running on a distributed executor like Dask or Parsl, serialising the processor object and shipping it to workers can be slow. One clean solution is to pass the model path and let each worker load locally. Profile before optimising — on many clusters the default approach works fine.

Step 2 — Build a Feature Matrix Inside process()

Inside process() you have access to your event arrays. Select your working point (e.g., events with at least two jets passing quality cuts), then stack the observables into a 2D NumPy array.

def process(self, events):
    sel = events.Jet.pt[:, 0] > 30   # example selection
    jets = events[sel]

    # Build feature matrix — order must match training
    X = np.column_stack([
        ak.to_numpy(jets.Jet.pt[:, 0]),
        ak.to_numpy(jets.Jet.eta[:, 0]),
        ak.to_numpy(jets.MET.pt),
    ])

The call to ak.to_numpy() is the boundary between the awkward-array world and the NumPy world your model expects. Keep that boundary explicit and in one place — it makes debugging much easier.

A note on variable-length arrays: if your features involve per-jet quantities aggregated over a variable number of jets (sum, max, count), compute those aggregations first inside awkward, then convert to NumPy. Trying to feed ragged arrays directly into a scikit-learn or PyTorch model will raise an error.

Step 3 — Run Inference and Return the Score

Once you have X, inference is a single line:

    scores = self.clf.predict_proba(X)[:, 1]   # signal probability column

For PyTorch or JAX models, wrap the call in the appropriate no-gradient context (torch.no_grad()) before converting to NumPy. Keep inference on CPU unless you have confirmed that the overhead of moving data to GPU per chunk is worth it — for many HEP-scale classifiers it is not.

Now attach the scores back to your events so downstream cuts and histograms work normally:

    jets["clf_score"] = ak.Array(scores)

From here you can fill histograms, apply a score threshold as a selection stage, or write the scores to a Parquet file for later use.

Step 4 — Validate That the Columnar Workflow Is Intact

After integration, run a small sanity check before submitting to the full dataset. Process a single file with IterativeExecutor and verify:

• The output accumulator has the expected event counts.
• Score distributions look sensible (not all zeros, not all ones).
• No silent NaN propagation from missing values in the feature matrix.

This is the equivalent of checking a chi-square residual before trusting a fit — a quick look now saves hours of debugging later. If you want a structured checklist for coffea framework machine learning validation, the complete HEP ML course walks through debugging inference pipelines step by step.

Common Pitfalls

• Feature order mismatch — the model expects columns in training order. Document this explicitly in your processor.
• Selection mismatch — if you apply different cuts at inference time than at training time, the score distribution will be misleading.
• Chunk boundary effects — normalisation steps (like StandardScaler) must be fit on training data and applied frozen at inference; never refit per chunk.

The key insight is simple: columnar analysis already speaks the language ML models understand, so integration is mostly a matter of keeping your array conversions explicit and loading your model efficiently. Build that habit once, and every future model slots in cleanly.