lume-pyat

pyAT-specific implementation of LUME model classes for virtual accelerators.

lume-pyat sits between lume-base and pyAT. You build the lattice, 4D or 6D; this package wraps it in a PyATSimulator, binds free-form variable names to element attributes — or to the ring energy — in native pyAT units, and exposes the result as a lume-base ActionModel. It is a sibling of lume-cheetah and follows the same conventions.

The package is facility-agnostic. It ships no lattice data and knows nothing about any particular machine's naming, units, or calibration. Mapping control-system addresses and engineering units onto simulator parameters is the job of the layer above — and PyATWritableScalarVariable and PyATLatticeScalarVariable are the seams where that layer plugs in.

Installation

pip install lume-pyat

Requires Python 3.10 or newer.

What it gives you

• Atomic multi-set. A batch of writes either lands completely or not at all. Everything the batch touches — every bound element, and for a lattice-level variable the ring energy and every cavity — is snapshotted first, the whole batch is applied, the orbit is solved once, and only then is anything committed. A failure restores every snapshot, so a rejected write is a complete no-op.
• One setpoint, several elements. A writable variable binds a list of elements, each with a weight: a write puts value × weight on every one of them, and a read comes from the first. A magnet sliced into several lattice elements, or a kick shared across them, is one variable.
• The ring energy is a variable too. PyATLatticeScalarVariable sets ring.energy and every cavity's Energy as one atomic write, with a hook for the layer above to rescale strengths in the same batch.
• Bindings are checked where you declare them. A variable naming an element the lattice does not have, a name two elements share, an attribute that does not exist, or an index past the end of one is rejected when the model is built — not at whichever later write happens to touch it.
• Instability is an exception, not a NaN. pyAT reports an unstable ring by returning non-finite values rather than raising — or, with a cavity in the ring, by returning a finite one-turn matrix whose eigenvalues have left the unit circle. solve_orbit checks the one-turn matrix's eigenvalues in all three planes and the closed orbit by value, and raises OrbitSolveError naming the unstable plane, so a caller cannot serve garbage as a monitor reading.
• 4D or 6D. A ring without a cavity is solved in 4D, a ring with one enabled in 6D, so RF-driven dispersive orbits solve as they should.
• One persistent lattice. It is mutated in place, never copied or rebuilt. Writes compose the way their physical counterparts do, and seeded misalignments survive every later write.
• Cheap imports. import lume_pyat pulls in neither pyAT nor lume-base nor HDF5 — resolution is lazy, so error paths and light entry points stay light.

Example

import at
import numpy as np

from lume_pyat import (
    LUMEPyATModel,
    PyATReadOnlyScalarVariable,
    PyATSimulator,
    PyATWritableScalarVariable,
)

# 1. Build a lattice. lume-pyat never builds one for you -- this is a small
#    FODO ring, but any 4D-canonical at.Lattice works.
N_CELLS = 8
elements = []
for cell in range(1, N_CELLS + 1):
    tag = f"{cell:02d}"
    elements += [
        at.Monitor(f"BPM_{tag}"),
        at.Drift("DRIFT", 0.4),
        at.Quadrupole(f"QUAD_F_{tag}", 0.3, 1.0),
        at.Drift("DRIFT", 0.4),
        at.Corrector(f"COR_H_{tag}", 0.0, [0.0, 0.0]),
        at.Dipole(f"BEND_{tag}", 1.0, 2 * np.pi / N_CELLS),
        at.Drift("DRIFT", 0.4),
        at.Quadrupole(f"QUAD_D_{tag}", 0.3, -1.0),
        at.Drift("DRIFT", 0.4),
    ]
ring = at.Lattice(elements, name="EXAMPLE", energy=1.0e9, periodicity=1)
ring.disable_6d()

# 2. Hand the lattice to a simulator. It is mutated in place, never copied.
simulator = PyATSimulator(ring)

# 3. Bind names to element attributes. Names are free-form -- use whatever your
#    own control system calls these. Values are in native pyAT units.
variables = [
    PyATWritableScalarVariable(
        name="corrector_kick",
        element_name="COR_H_01",
        attribute="KickAngle",
        index=0,
        default_value=0.0,
        unit="rad",
    ),
    *(
        PyATReadOnlyScalarVariable(
            name=f"bpm_{cell:02d}_x", element_name=f"BPM_{cell:02d}", axis="x", unit="m"
        )
        for cell in range(1, N_CELLS + 1)
    ),
]

# 4. The model solves the boot orbit once and caches it.
model = LUMEPyATModel(simulator=simulator, action_variables=variables)

bpm_names = [f"bpm_{cell:02d}_x" for cell in range(1, N_CELLS + 1)]
print(f"{len(model.supported_variables)} variables: 1 input, {len(bpm_names)} outputs")

# 5. Read: the orbit is flat and the corrector sits at its default.
print(f"before:  corrector = {model.get('corrector_kick')} rad")
print(f"         max |x|   = {max(abs(v) for v in model.get(bpm_names).values()):.3e} m")

# 6. Write: one batch, one solve, all or nothing.
model.set({"corrector_kick": 1.0e-4})
print(f"after:   corrector = {model.get('corrector_kick')} rad   <- retained exactly")
print(f"         max |x|   = {max(abs(v) for v in model.get(bpm_names).values()):.3e} m")

# 7. A write that leaves the ring without a stable orbit is a complete no-op.
try:
    model.set({"corrector_kick": 1.0})
except Exception as exc:
    print(f"rejected: {type(exc).__name__} -- corrector still {model.get('corrector_kick')}")

# 8. reset() puts every input back to its default, in one atomic batch.
model.reset()
print(f"reset:   corrector = {model.get('corrector_kick')} rad")
print(f"         max |x|   = {max(abs(v) for v in model.get(bpm_names).values()):.3e} m")

assert model.get("corrector_kick") == 0.0
assert max(abs(v) for v in model.get(bpm_names).values()) < 1e-12

Which prints:

9 variables: 1 input, 8 outputs
before:  corrector = 0.0 rad
         max |x|   = 0.000e+00 m
after:   corrector = 0.0001 rad   <- retained exactly
         max |x|   = 8.155e-04 m
rejected: OrbitSolveError -- corrector still 0.0001
reset:   corrector = 0.0 rad
         max |x|   = 0.000e+00 m

The example code above is extracted and executed on every CI run, so it cannot drift from the package. (The sample output is illustrative — the last digits depend on the platform's linear algebra.)

Adding your own units

PyATWritableScalarVariable writes native pyAT quantities. To control an element in your own units, subclass it and put the conversion in _set and _get:

class CurrentVariable(PyATWritableScalarVariable):
    amps_per_unit: float

    def _set(self, simulator, value):
        super()._set(simulator, value / self.amps_per_unit)

    def _get(self, simulator):
        return super()._get(simulator) * self.amps_per_unit

The conversion belongs to your layer, not to this package. Everything else — the atomicity, the solve, the caching — keeps working unchanged.

Several elements from one setpoint

The element_name= / attribute= / index= form above is shorthand for a single ElementBinding with unit weight. Give bindings= instead to drive several elements from one variable. A write of v puts v × weight on every bound element; a read returns the first element's value divided by its weight.

from lume_pyat import ElementBinding, PyATWritableScalarVariable

# A corrector split over three lattice slices: each slice carries a third of
# the kick, and the variable reads back the whole kick from the first slice.
split_kick = PyATWritableScalarVariable(
    name="corrector_kick",
    bindings=[
        ElementBinding(element_name=name, attribute="KickAngle", index=0, weight=1 / 3)
        for name in ("COR_H_01A", "COR_H_01B", "COR_H_01C")
    ],
    default_value=0.0,
    unit="rad",
)

# A quadrupole sliced into two elements that share one strength.
sliced_quad = PyATWritableScalarVariable(
    name="quad_strength",
    bindings=[
        ElementBinding(element_name="QUAD_01A", attribute="K"),
        ElementBinding(element_name="QUAD_01B", attribute="K"),
    ],
    default_value=1.0,
)

Every binding is checked when the model is built, and a failed batch restores every bound element.

The ring energy

PyATLatticeScalarVariable binds the lattice as a whole rather than an element. Its value is ring.energy in eV; a write sets that and the Energy of every at.RFCavity, and a failed batch restores all of them.

Changing the energy does not change the normalised strengths pyAT stores. A layer that holds magnet currents fixed across an energy change — so that K falls as 1/E — does the rescaling itself in the _after_write hook, which runs inside the same batch as the energy write, and declares what it touches in snapshot_targets so the rescale rolls back with the energy:

class EnergyFromDipoleCurrent(PyATLatticeScalarVariable):
    ev_per_amp: float
    reference_energy: float  # the energy the reference strengths belong to, in eV
    reference_strengths: dict[str, float]  # FamName -> K at that energy, current held

    def _set(self, simulator, amps):
        super()._set(simulator, amps * self.ev_per_amp)

    def _get(self, simulator):
        return super()._get(simulator) / self.ev_per_amp

    def _after_write(self, ring, energy):
        for name, k_reference in self.reference_strengths.items():
            index = ring.get_uint32_index(name)[0]
            ring[index].K = k_reference * self.reference_energy / energy

    def snapshot_targets(self, simulator):
        return [
            *super().snapshot_targets(simulator),
            *(
                (simulator.element_index(name), "K")
                for name in self.reference_strengths
            ),
        ]

4D and 6D rings

solve_orbit follows the ring: a lattice without an enabled cavity is solved with find_orbit4, one with a cavity enabled (ring.enable_6d(at.RFCavity)) with find_orbit6, so an RF frequency change produces the dispersive orbit it should. In both cases the stability guard is the same: every eigenvalue of the find_m66 one-turn matrix must satisfy |eigenvalue| ≤ 1 + 1e-6, in all three planes.

Rings with radiation enabled are out of scope. The guards still run — a damped map passes them — but nothing in this package accounts for the energy loss per turn, and the orbit such a ring solves to is not the one its variables describe. Enable the cavity only, and leave radiation off.

API

What raises what

Everything that can be settled when a variable is declared is settled then, so the write path raises only about the physics.

Duplicate names are only a problem for names you address. A lattice whose drifts all share one name — most of them — is fine.

Development

uv sync --extra dev
uv run pytest lume_pyat/tests
uv run pre-commit run --all-files

License

Apache-2.0. See LICENSE.