LUME-Model Tutorial: Impact-T + Distgen¶
This notebook demonstrates the use of the LUMEModel objects included in lume-impact.
LUMEModel (from lume-base) is a standard interface for simulation tools which allows them to be connected to control systems software (such as EPICS) for developing virtual accelerator software.
In particular, a model (named model) has the following:
model.supported_variables: Exposes the attributes that can be read / written to asVariableobjects.model.get(names: list[str]): Get the values of the variables.model.set(values: dict[str, Any]): Set the values of the variables (and update the state of the simulation)
Note that LUMEModel does not replace the raw Impact object, but rather is intended to provide a common interface for listing what can be changed and read in the simulation and how to do it.
In lume-impact, the LUMEModel is extended with an "actions framework" which provides a convenient way of connecting the Variable objects to how they interact with the simulation tool (Impact here).
As a note, this framework is intended to be upstreamed to lume-base and is subject to change slightly.
lume-impact includes tools that allow the construction of a LUMEModel from an existing Impact object in a configurable way.
That is, instead of having to manually define the included variables and how they talk to the lattice, the tool allows users to parse their Impact-T simulation and expose all of the elements as variables.
In this notebook, we'll build a model, inspect it, configure it, extend it with custom "actions", and finish with a small optimizer-style scan that touches only the generic get/set interface.
1. Setup¶
We use the LCLS injector template with a small particle count and space charge turned off so the whole notebook runs in a few seconds.
import numpy as np
import matplotlib.pyplot as plt
from collections import Counter
from impact import Impact
from distgen import Generator
from lume.variables import ScalarVariable, NDVariable
from impact.model.distgen.distgen_impact_model import LUMEDistgenImpactModel
from impact.model.config import (
VariableMappingConfig,
ElementsConfig,
StatsConfig,
RunInfoConfig,
)
from impact.model.actions import WritableImpactAction, StatAction
from impact.model.exceptions import ReadOnlyError
# Load the Distgen generator
gen = Generator("templates/lcls_injector/distgen.yaml")
gen["n_particle"] = 100
# Load the LCLS Impact-T lattice
impact = Impact("templates/lcls_injector/ImpactT.in")
impact.header["Np"] = gen["n_particle"]
# Turn space charge off so this runs quickly
impact.header["Bcurr"] = 0
impact.numprocs = 2
impact.run()
2. Build a default model¶
LUMEDistgenImpactModel.from_objects inspects both the generator and the Impact-T object and automatically creates control variables for distgen inputs, lattice element attributes, header settings, output statistics, run info, and particle groups. No configuration is required to get started.
model = LUMEDistgenImpactModel.from_objects(gen, impact)
len(model.supported_variables)
236
By default, variables are namespaced by a prefix that indicates where they come from.
prefixes = Counter(name.split(":")[0] for name in model.supported_variables)
prefixes
Counter({'ele': 164,
'header': 51,
'stat': 10,
'distgen': 6,
'particles': 3,
'run_info': 2})
Distgen inputs and a sample of the gun element attributes:
{
"distgen": [n for n in model.supported_variables if n.startswith("distgen:")],
"GUN": [n for n in model.supported_variables if n.startswith("ele:GUN:")],
}
{'distgen': ['distgen:n_particle',
'distgen:total_charge',
'distgen:start:cathode:MTE',
'distgen:r_dist:sigma_xy',
'distgen:t_dist:length',
'distgen:t_dist:ratio'],
'GUN': ['ele:GUN:rf_field_scale',
'ele:GUN:rf_frequency',
'ele:GUN:theta0_deg',
'ele:GUN:radius',
'ele:GUN:solenoid_field_scale',
'ele:GUN:x_offset',
'ele:GUN:y_offset',
'ele:GUN:x_rotation',
'ele:GUN:y_rotation',
'ele:GUN:z_rotation']}
Variable types¶
Each variable is a typed Variable from lume-base. The model uses three kinds:
ScalarVariable— single numbers (element attributes, header keys, distgen params)NDVariable— fixed-shape arrays (output statistics likesigma_xalongz)ParticleGroupVariable— openPMD particle groups
Outputs (stats, run info, non-initial particle groups) are marked read_only=True, so the model will refuse to write them.
Read-only NDVariable outputs (statistics) carry a shape and units.
[
(v.name, v.shape, v.unit)
for v in model.supported_variables.values()
if isinstance(v, NDVariable) and v.read_only
]
[('stat:mean_kinetic_energy', (969,), 'eV'),
('stat:mean_x', (969,), 'm'),
('stat:mean_y', (969,), 'm'),
('stat:mean_z', (969,), 'm'),
('stat:sigma_x', (969,), 'm'),
('stat:sigma_y', (969,), 'm'),
('stat:sigma_z', (969,), 'm'),
('stat:norm_emit_x', (969,), 'm'),
('stat:norm_emit_y', (969,), 'm'),
('stat:norm_emit_z', (969,), 'm')]
3. Reading and writing through get / set¶
get accepts a single name (returns the value) or a list (returns a dict). set writes the inputs, reruns Distgen and Impact-T, and refreshes the outputs.
def plot_beam(stats, label, color):
z = stats["stat:mean_z"]
plt.plot(z, 1e6 * stats["stat:sigma_x"], color=color, label=label)
plt.xlabel("s (m)")
plt.ylabel("RMS beam size $\\sigma_x$ (um)")
baseline = model.get(["stat:sigma_x", "stat:mean_z"])
plot_beam(baseline, "baseline", "C0")
plt.legend()
<matplotlib.legend.Legend at 0x7f77287b4aa0>
Change the main solenoid strength. This reruns Distgen + Impact-T automatically.
model.set({"ele:SOL1:solenoid_field_scale": 0.28})
changed = model.get(["stat:sigma_x", "stat:mean_z"])
plot_beam(baseline, "SOL1 = 0.2457 (baseline)", "C0")
plot_beam(changed, "SOL1 = 0.28", "C1")
plt.legend()
<matplotlib.legend.Legend at 0x7f76b5494140>
Outputs are read-only: attempting to set one is rejected before anything runs.
try:
model.set({"stat:sigma_x": np.zeros(10)})
rejection = "no error raised"
except (ReadOnlyError, ValueError) as exc:
rejection = exc
rejection
ValueError("Variable 'stat:sigma_x' is read-only. Cannot be set.")
# Restore the baseline before moving on
model.set({"ele:SOL1:solenoid_field_scale": 0.2457})
4. Configuring variable generation¶
The defaults are driven by pydantic config objects, so you control exactly which variables are created and how they are named. Here, you can rename simulation elements to machine names and choose your own naming pattern (these configs are designed to be driven from YAML).
Key settings in VariableMappingConfig:
- Set any category (
header,elements,stats,run_info,particles) toNoneto drop it entirely. patternis an f-string used to build variable names.control_to_tool_namemaps your chosen control name to the underlying Impact-T element name.
Expose Impact element QE01 under a machine-style name, rename the stat/run_info prefixes, and drop the header variables entirely.
imp_config = VariableMappingConfig(
header=None,
elements=ElementsConfig(
pattern="{name}:{attrib}",
control_to_tool_name={"QUAD:IN20:525": "QE01"},
),
stats=StatsConfig(pattern="STATS:{name}"),
run_info=RunInfoConfig(pattern="RUNINFO:{key}"),
)
configured = LUMEDistgenImpactModel.from_objects(gen, impact, impact_config=imp_config)
# QE01 now appears under the control-system name; header variables are gone.
{
"header_dropped": not any(
n.startswith("header:") for n in configured.supported_variables
),
"IN20_vars": [n for n in configured.supported_variables if "IN20" in n],
}
{'header_dropped': True,
'IN20_vars': ['QUAD:IN20:525:b1_gradient',
'QUAD:IN20:525:L_effective',
'QUAD:IN20:525:radius',
'QUAD:IN20:525:x_offset',
'QUAD:IN20:525:y_offset',
'QUAD:IN20:525:x_rotation',
'QUAD:IN20:525:y_rotation',
'QUAD:IN20:525:z_rotation']}
Stats are renamed by the new pattern.
[n for n in configured.supported_variables if n.startswith("STATS:")][:6]
['STATS:mean_kinetic_energy', 'STATS:mean_x', 'STATS:mean_y', 'STATS:mean_z', 'STATS:sigma_x', 'STATS:sigma_y']
5. Extending the model with custom actions¶
Every variable is backed by an action: a small pydantic object pairing a Variable with the _get/_set methods that read and write it on the simulator.
Note: the actions framework is intended to be moved into the lume-base package and may change slightly as it is developed.
Two base classes:
ImpactAction— read-only; its variable must beread_only=TrueWritableImpactAction— supports bothgetandset
Because actions are pydantic objects, validation happens at construction time.
Validation example: a read-only StatAction refuses a writable variable.
try:
StatAction(
stat_name="sigma_x",
var=NDVariable(name="bad", shape=(10,), default_value=None, read_only=False),
)
construction = "no error raised"
except ReadOnlyError as exc:
construction = exc
construction
impact.model.exceptions.ReadOnlyError('StatAction requires a read-only variable')
A custom writable action: expose a quad's integrated gradient (B-field integral) in control-system units while storing the raw gradient in Impact-T.
class IntegratedGradientAction(WritableImpactAction):
# Narrowing the `var` type gives pydantic validation for this action.
var: ScalarVariable
# Extra data the action needs is stored as plain pydantic fields (serializable).
ele_name: str
def _get(self, impact):
ele = impact.ele[self.ele_name]
return ele["b1_gradient"] * ele["L_effective"]
def _set(self, impact, value):
ele = impact.ele[self.ele_name]
impact.ele[self.ele_name]["b1_gradient"] = value / ele["L_effective"]
model.register_action(
IntegratedGradientAction(
var=ScalarVariable(name="QE01:GLINT", unit="T"),
ele_name="QE01",
)
)
"QE01:GLINT" in model.supported_variables
True
The new variable behaves like any other: read it, write it, and the model reruns.
before = model.get("QE01:GLINT")
model.set({"QE01:GLINT": 0.05})
after = model.get("QE01:GLINT")
{"before": before, "after": after}
{'before': 0.020217600168702478, 'after': 0.05}
6. Interacting with Impact Model as a Control System¶
The point of the standardized interface is that downstream code never needs to know it is driving Impact-T. The scan below could just as easily be an Xopt optimizer or an EPICS server. It touches only model.get and model.set.
Here we scan the solenoid strength and record the final transverse beam size, looking for the setting that focuses the beam.
def final_beam_size(model, knob, value):
"""Generic objective: set one knob, return the final RMS x size."""
model.set({knob: value})
# Output stat arrays are padded to a fixed length, so take the last finite value.
sigma_x = model.get("stat:sigma_x")
return float(sigma_x[np.isfinite(sigma_x)][-1])
knob = "ele:SOL1:solenoid_field_scale"
scan = np.linspace(0.18, 0.22, 9)
results = [final_beam_size(model, knob, v) for v in scan]
best = scan[int(np.argmin(results))]
best
np.float64(0.2)
plt.plot(scan, 1e6 * np.array(results), "o-")
plt.axvline(best, color="C1", ls="--", label=f"min @ {best:.4f}")
plt.xlabel("SOL1 solenoid_field_scale")
plt.ylabel("final $\\sigma_x$ (um)")
plt.legend()
<matplotlib.legend.Legend at 0x7f76b41e3d10>
Summary¶
LUMEDistgenImpactModel.from_objectsbuilds a fully working model from a generator and an Impact-T object with zero configuration.get/setprovide a tool-agnostic interface; onesetorchestrates Distgen -> Impact-T and refreshes outputs.VariableMappingConfigcontrols which variables exist and how they are named. It defines the boundary between control-system names/units and the simulator.- Custom actions extend the model with new, validated, serializable interactions without touching the model internals.
- Any generic consumer (ie optimizer, surrogate, or control server) can drive the simulation through
get/setalone.