Extend a Simulation with Python

Overview

ImpactX’s Python interface lets you weave custom Python code into a running simulation. Through this interface, you can access and modify simulation data, e.g., particle phase-space coordinates, the reference particle, and space-charge fields, as the simulation is tracking the beam. This facilitates a wide range of workflows, including:

In-situ diagnostics that compute quantities of interest at each turn or element without writing files to disk.

Dynamic lattice manipulation, e.g., changing element strengths between turns, ramping voltages, or feedback systems.

Custom beam optics elements, e.g., a non-linear kick or a tabulated map, plugged in via the impactx.elements.Programmable element.

AI/ML surrogate models built with PyTorch or TensorFlow to emulate detailed transport in a single element.

Because ImpactX exposes particles and fields without copies through pyAMReX, per-step Python callbacks have very low overhead and stay GPU-resident when ImpactX runs on GPUs (cupy, PyTorch, etc. can operate directly on the device data).

How to run a simulation with Python extensions

Install ImpactX with the Python interface: when compiling from source, this means -DImpactX_PYTHON=ON (see building). Pre-built ImpactX packages also include this interface.

Write a Python script that drives the simulation, registering callbacks or programmable elements before calling track_particles(). A minimal template looks like:

from impactx import ImpactX, distribution, elements

sim = ImpactX()
sim.space_charge = False
sim.init_grids()

# ... set up reference particle, distribution, and lattice ...

# register tracking hooks or programmable elements here:
def after_each_period(sim):
    turn = sim.tracking_period
    # e.g. compute a custom diagnostic from sim.beam

sim.hook["after_period"] = after_each_period

# advance the simulation
sim.track_particles()
sim.finalize()

Run the script like any Python script, optionally under MPI:
```
mpirun -np <n_ranks> python <script>.py
```

Callback Functions

ImpactX provides two complementary callback surfaces.

Tracking hooks

The hook mapping registers user-defined functions that are invoked at well-defined points of the tracking loop. The supported locations are:

"before_period" — before each lattice period (e.g. turn or channel period)
"after_period" — after each lattice period
"before_element" — before each lattice element is entered
"after_element" — after each lattice element is exited
"before_slice" — before each element slice (and before that slice’s space-charge solve, when enabled; to read freshly-computed space-charge fields, prefer "after_element")

Each hook receives the ImpactX instance and may inspect or modify simulation state through it (e.g., sim.beam, sim.lattice, sim.rho, sim.phi, …). Use tracking_step, tracking_period, and tracking_element to identify where in the loop the hook was called. For instance, you can return if the current element does not match a type or name.

def after_each_turn(sim):
    turn = sim.tracking_period
    element = sim.tracking_element
    df = sim.beam.to_df(local=True)
    if df is not None:
        print(f"turn {turn} at element={element.name}: "
              f"{len(df)} local particles, "
              f"<x>={df['position_x'].mean():.3e} m")

sim.hook["after_period"] = after_each_turn

Full example: Acceleration by RF Cavities.

Programmable lattice element

For workflows that replace a beamline element entirely (rather than observing the beam between elements), use the impactx.elements.Programmable element. It exposes three property hooks:

push: push the whole particle container.
ref_particle: push only the reference particle.
beam_particles: called once per particle tile (high-performance); the reference particle is updated before this callback fires.

Typical usage (per-tile push) looks like this:

from impactx import elements

def my_drift(pge, pti, refpart):
    """Push the beam particles as a drift over one slice."""
    soa = pti.soa().to_xp()   # NumPy (CPU) or CuPy (GPU)
    x = soa.real["position_x"]
    y = soa.real["position_y"]
    t = soa.real["position_t"]
    # ...

    slice_ds = pge.ds / pge.nslice
    betgam2 = refpart.pt**2 - 1.0

    x[:] += slice_ds * px[:]
    y[:] += slice_ds * py[:]
    t[:] += (slice_ds / betgam2) * pt[:]


def my_ref_drift(pge, refpart):
    """Push the reference particle over one slice."""
    x = refpart.x
    px = refpart.px
    s = refpart.s

    slice_ds = pge.ds / pge.nslice

    # ...
    refpart.x = x + 1.23
    refpart.s = s + slice_ds


dr = elements.Programmable(name="my_drift")
dr.ds = 0.25
dr.nslice = 5
dr.beam_particles = lambda pti, refpart: my_drift(dr, pti, refpart)
dr.ref_particle = lambda refpart: my_ref_drift(dr, refpart)

sim.lattice.append(pge)

A complete script is provided in this FODO Cell example. For a more complex example, see our ML surrogate element example.

Accessing simulation data through Python

While the simulation is running, the Python code (i.e., the code inside callback functions) has read and write access to the beam particles and the space-charge fields. The two sections below describe the specific Python syntax for each.

See also Data Analysis for file and streaming-based analysis of the openPMD output written by BeamMonitor.