From 9ef8dd26ffae9749f4646547765d9f8206cd5522 Mon Sep 17 00:00:00 2001 From: hong274SLAC Date: Tue, 11 Aug 2026 17:10:05 -0700 Subject: [PATCH 01/11] Add ZFEL backend and CU HXR model Add the ZFEL LUME backend, HXR undulator mapping, and CU HXR scalar model wrapper. --- virtual_accelerator/models/cu_hxr_zfel.py | 240 +++++++++++++++ virtual_accelerator/zfel/__init__.py | 1 + virtual_accelerator/zfel/model.py | 288 ++++++++++++++++++ virtual_accelerator/zfel/undulator_mapping.py | 262 ++++++++++++++++ 4 files changed, 791 insertions(+) create mode 100644 virtual_accelerator/models/cu_hxr_zfel.py create mode 100644 virtual_accelerator/zfel/__init__.py create mode 100644 virtual_accelerator/zfel/model.py create mode 100644 virtual_accelerator/zfel/undulator_mapping.py diff --git a/virtual_accelerator/models/cu_hxr_zfel.py b/virtual_accelerator/models/cu_hxr_zfel.py new file mode 100644 index 0000000..2e163ca --- /dev/null +++ b/virtual_accelerator/models/cu_hxr_zfel.py @@ -0,0 +1,240 @@ +from typing import Any + +import numpy as np + +from lume.model import LUMEModel +from lume.variables import ScalarVariable + +from virtual_accelerator.zfel.undulator_mapping import HXR_CELLS +from virtual_accelerator.zfel.model import ZFELModel + + +class ZFELPVModel(LUMEModel): + """ + Scalar EPICS-facing wrapper around ZFELModel. + + EPICS-facing controls: + KAct_14, DSKAct_14, ..., KAct_47, DSKAct_47 + + Internal physics model: + Kact[32], DSKact[32] -> zfel + """ + + def __init__(self): + self._backend = ZFELModel() + + backend_state = self._backend.get([ + "Kact", + "DSKact", + "power_max", + "exit_power", + "pulse_energy", + ]) + + kact = np.asarray( + backend_state["Kact"], + dtype=float, + ) + + dskact = np.asarray( + backend_state["DSKact"], + dtype=float, + ) + + self._state: dict[str, Any] = {} + self._variables = {} + + # ------------------------------------------------------ + # Scalar, real-machine-like undulator controls + # ------------------------------------------------------ + + for index, cell in enumerate(HXR_CELLS): + kact_name = f"KAct_{cell}" + dskact_name = f"DSKAct_{cell}" + + self._state[kact_name] = float(kact[index]) + self._state[dskact_name] = float(dskact[index]) + + self._variables[kact_name] = ScalarVariable( + name=kact_name, + default_value=float(kact[index]), + value_range=(0.0, 5.0), + unit="dimensionless", + read_only=False, + ) + + self._variables[dskact_name] = ScalarVariable( + name=dskact_name, + default_value=float(dskact[index]), + value_range=(0.0, 5.0), + unit="dimensionless", + read_only=False, + ) + + # ------------------------------------------------------ + # Read-only FEL diagnostics + # ------------------------------------------------------ + + self._variables.update({ + "power_max": ScalarVariable( + name="power_max", + default_value=0.0, + unit="W", + read_only=True, + ), + "exit_power": ScalarVariable( + name="exit_power", + default_value=0.0, + unit="W", + read_only=True, + ), + "pulse_energy": ScalarVariable( + name="pulse_energy", + default_value=0.0, + unit="J", + read_only=True, + ), + "pulse_intensity_mean": ScalarVariable( + name="pulse_intensity_mean", + default_value=0.0, + unit="J", + read_only=True, + ), + "pulse_intensity_p80": ScalarVariable( + name="pulse_intensity_p80", + default_value=0.0, + unit="J", + read_only=True, + ), + "pulse_intensity_std_relative": ScalarVariable( + name="pulse_intensity_std_relative", + default_value=0.0, + unit="dimensionless", + read_only=True, + ), + }) + + self._sync_from_backend(include_controls=True) + + @property + def supported_variables(self): + return self._variables + + def _get(self, names: list[str]) -> dict[str, Any]: + return { + name: self._state[name] + for name in names + } + + def _set(self, values: dict[str, Any]) -> None: + """ + Update any scalar KAct/DSKAct values, then run zfel once. + + Lume-PVA batches multiple PV writes before calling this, + according to Runner update_rate. + """ + + if not values: + self._sync_from_backend( + include_controls=True + ) + return + + for name, value in values.items(): + self._state[name] = float(value) + + kact = np.asarray( + [ + self._state[f"KAct_{cell}"] + for cell in HXR_CELLS + ], + dtype=float, + ) + + dskact = np.asarray( + [ + self._state[f"DSKAct_{cell}"] + for cell in HXR_CELLS + ], + dtype=float, + ) + + self._backend.set({ + "Kact": kact, + "DSKact": dskact, + }) + + self._sync_from_backend( + include_controls=True + ) + + def _sync_from_backend( + self, + *, + include_controls: bool, + ) -> None: + backend_state = self._backend.get([ + "Kact", + "DSKact", + "power_max", + "exit_power", + "pulse_energy", + ]) + + if include_controls: + kact = np.asarray( + backend_state["Kact"], + dtype=float, + ) + + dskact = np.asarray( + backend_state["DSKact"], + dtype=float, + ) + + for index, cell in enumerate(HXR_CELLS): + self._state[f"KAct_{cell}"] = float( + kact[index] + ) + + self._state[f"DSKAct_{cell}"] = float( + dskact[index] + ) + + pulse_energy = float( + backend_state["pulse_energy"] + ) + + self._state["power_max"] = float( + backend_state["power_max"] + ) + + self._state["exit_power"] = float( + backend_state["exit_power"] + ) + + self._state["pulse_energy"] = pulse_energy + + # Fixed-seed deterministic model for now. + self._state["pulse_intensity_mean"] = ( + pulse_energy + ) + self._state["pulse_intensity_p80"] = ( + pulse_energy + ) + self._state[ + "pulse_intensity_std_relative" + ] = 0.0 + + def reset(self) -> None: + self._backend.reset() + + self._sync_from_backend( + include_controls=True + ) + +def get_cu_hxr_zfel_model() -> ZFELPVModel: + """ + Construct the CU HXR ZFEL virtual-accelerator model. + """ + return ZFELPVModel() \ No newline at end of file diff --git a/virtual_accelerator/zfel/__init__.py b/virtual_accelerator/zfel/__init__.py new file mode 100644 index 0000000..b320b98 --- /dev/null +++ b/virtual_accelerator/zfel/__init__.py @@ -0,0 +1 @@ +from virtual_accelerator.zfel.model import ZFELModel \ No newline at end of file diff --git a/virtual_accelerator/zfel/model.py b/virtual_accelerator/zfel/model.py new file mode 100644 index 0000000..d3ef616 --- /dev/null +++ b/virtual_accelerator/zfel/model.py @@ -0,0 +1,288 @@ +from typing import Any + +import numpy as np + +from lume.model import LUMEModel +from lume.variables import NDVariable, ScalarVariable +from zfel import sase1d + +C_LIGHT = 299_792_458.0 + +from virtual_accelerator.zfel.undulator_mapping import ( + MAGNETIC_LENGTH_M, + N_ACTIVE_SEGMENTS, + build_hxr_mapping, +) + + +class ZFELModel(LUMEModel): + """ + LUMEModel wrapper around zfel using real-machine-like + HXR Kact and DSKact controls. + + Flow + ---- + Kact / DSKact + -> realistic HXR machine mapping + -> effective zfel magnetic coordinate + -> zfel simulation + -> cached FEL diagnostics + """ + + ZFEL_Z_STEPS = 320 + INITIAL_K = 3.5 + + def __init__(self): + # ---------------------------------------------------------- + # Baseline zfel beam inputs + # + # These are still the generic zfel test parameters. + # Only the undulator representation is being upgraded in + # Phase 4D. + # ---------------------------------------------------------- + + self._sase_input = { + "npart": 512, + "s_steps": 200, + "z_steps": self.ZFEL_Z_STEPS, + "energy": 4313.34e6, + "eSpread": 0.0, + "emitN": 1.2e-6, + "currentMax": 3400, + "beta": 26, + "unduPeriod": 0.03, + "unduK": self.INITIAL_K, + "unduL": MAGNETIC_LENGTH_M, + "radWavelength": None, + "random_seed": 31, + "particle_position": None, + "hist_rule": "square-root", + "iopt": "sase", + "P0": 0, + } + + # ---------------------------------------------------------- + # Initial virtual-machine controls + # + # Constant K is used only to validate the new Kact/DSKact + # control path. + # ---------------------------------------------------------- + + initial_kact = np.full( + N_ACTIVE_SEGMENTS, + self.INITIAL_K, + dtype=float, + ) + + initial_dskact = np.full( + N_ACTIVE_SEGMENTS, + self.INITIAL_K, + dtype=float, + ) + + self._initial_state = { + "Kact": initial_kact, + "DSKact": initial_dskact, + "unduK": np.full( + self.ZFEL_Z_STEPS, + self.INITIAL_K, + dtype=float, + ), + "power_max": 0.0, + "exit_power": 0.0, + "pulse_energy": 0.0, + } + + self._state = self._copy_initial_state() + + # ---------------------------------------------------------- + # LUMEModel variable definitions + # ---------------------------------------------------------- + + self._variables = { + # Writable virtual-machine controls + "Kact": NDVariable( + name="Kact", + shape=(N_ACTIVE_SEGMENTS,), + default_value=initial_kact.copy(), + unit="dimensionless", + read_only=False, + ), + "DSKact": NDVariable( + name="DSKact", + shape=(N_ACTIVE_SEGMENTS,), + default_value=initial_dskact.copy(), + unit="dimensionless", + read_only=False, + ), + + # Read-only mapped zfel K profile + "unduK": NDVariable( + name="unduK", + shape=(self.ZFEL_Z_STEPS,), + default_value=np.full( + self.ZFEL_Z_STEPS, + self.INITIAL_K, + dtype=float, + ), + unit="dimensionless", + read_only=True, + ), + + # Read-only FEL diagnostics + "power_max": ScalarVariable( + name="power_max", + default_value=0.0, + unit="W", + read_only=True, + ), + "exit_power": ScalarVariable( + name="exit_power", + default_value=0.0, + unit="W", + read_only=True, + ), + "pulse_energy": ScalarVariable( + name="pulse_energy", + default_value=0.0, + unit="J", + read_only=True, + ), + } + + # Store the latest complete HXR mapping for debugging, + # plotting, and later Badger integration. + self._mapping = None + + # Run the initial constant-K case. + self._run_simulation() + + @property + def supported_variables(self): + return self._variables + + def _copy_initial_state(self) -> dict[str, Any]: + """ + Return an independent copy of the initial machine state. + """ + + return { + "Kact": self._initial_state["Kact"].copy(), + "DSKact": self._initial_state["DSKact"].copy(), + "unduK": self._initial_state["unduK"].copy(), + "power_max": float(self._initial_state["power_max"]), + "exit_power": float(self._initial_state["exit_power"]), + "pulse_energy": float(self._initial_state["pulse_energy"]), + } + + def _get(self, names: list[str]) -> dict[str, Any]: + """ + Return cached state only. + + Calling get() does not rerun zfel. + """ + + out = {} + + for name in names: + value = self._state[name] + + if isinstance(value, np.ndarray): + out[name] = value.copy() + else: + out[name] = value + + return out + + def _set(self, values: dict[str, Any]) -> None: + """ + Update Kact and/or DSKact, then run zfel once. + """ + + if "Kact" in values: + self._state["Kact"] = np.asarray( + values["Kact"], + dtype=float, + ).copy() + + if "DSKact" in values: + self._state["DSKact"] = np.asarray( + values["DSKact"], + dtype=float, + ).copy() + + self._run_simulation() + + def _run_simulation(self) -> None: + """ + Map the virtual HXR machine state into zfel and run once. + """ + + mapping = build_hxr_mapping( + self._state["Kact"], + self._state["DSKact"], + z_steps=self.ZFEL_Z_STEPS, + ) + + self._mapping = mapping + + k_zfel = np.asarray( + mapping["k_zfel"], + dtype=float, + ) + + self._state["unduK"] = k_zfel.copy() + + sase_input = self._sase_input.copy() + sase_input["unduK"] = k_zfel + + output = sase1d.sase(sase_input) + + power_z = np.asarray( + output["power_z"], + dtype=float, + ) + + self._state["power_max"] = float( + np.max(power_z) + ) + + self._state["exit_power"] = float( + power_z[-1] + ) + power_s = np.asarray( + output["power_s"], + dtype=float, + ) + + s_m = np.asarray( + output["s"], + dtype=float, + ) + + if s_m.size < 2: + raise ValueError( + "zfel must return at least two longitudinal s points." + ) + + ds_m = float( + np.mean(np.diff(s_m)) + ) + + # Each power_s sample represents one longitudinal slice. + # Convert ds [m] to dt [s] using dt = ds/c. + pulse_energy_j = float( + np.sum(power_s[-1, :]) + * ds_m + / C_LIGHT + ) + + self._state["pulse_energy"] = pulse_energy_j + + def reset(self) -> None: + """ + Restore the initial 32-segment virtual-machine state. + """ + + self._state = self._copy_initial_state() + self._run_simulation() \ No newline at end of file diff --git a/virtual_accelerator/zfel/undulator_mapping.py b/virtual_accelerator/zfel/undulator_mapping.py new file mode 100644 index 0000000..df0c5ad --- /dev/null +++ b/virtual_accelerator/zfel/undulator_mapping.py @@ -0,0 +1,262 @@ +from typing import Sequence + +import numpy as np + + +# ------------------------------------------------------------------ +# LCLS HXR machine geometry +# ------------------------------------------------------------------ + +# Physical HXR cell slots covered by the present taper environment. +HXR_SLOTS: tuple[int, ...] = tuple(range(14, 48)) + +# Active HXR undulator cells. +# This matches the current real-machine Badger environment. +HXR_CELLS: tuple[int, ...] = ( + tuple(range(14, 21)) + + tuple(range(22, 28)) + + tuple(range(29, 48)) +) + +# Cells omitted from the active-undulator list. +INACTIVE_HXR_CELLS: tuple[int, ...] = (21, 28) + +FIRST_HXR_CELL = HXR_SLOTS[0] +LAST_HXR_CELL = HXR_SLOTS[-1] + +CELL_LENGTH_M = 4.4 +UNDULATOR_LENGTH_M = 3.4 +INTERSPACE_LENGTH_M = 1.0 + +N_ACTIVE_SEGMENTS = len(HXR_CELLS) + +PHYSICAL_SPAN_M = len(HXR_SLOTS) * CELL_LENGTH_M +MAGNETIC_LENGTH_M = N_ACTIVE_SEGMENTS * UNDULATOR_LENGTH_M + + +def _validate_endpoints( + kact: Sequence[float], + dskact: Sequence[float], +) -> tuple[np.ndarray, np.ndarray]: + """ + Validate and return the per-segment K endpoints. + """ + + kact_array = np.asarray(kact, dtype=float).reshape(-1) + dskact_array = np.asarray(dskact, dtype=float).reshape(-1) + + if kact_array.size != N_ACTIVE_SEGMENTS: + raise ValueError( + f"kact must contain {N_ACTIVE_SEGMENTS} values, " + f"got {kact_array.size}." + ) + + if dskact_array.size != N_ACTIVE_SEGMENTS: + raise ValueError( + f"dskact must contain {N_ACTIVE_SEGMENTS} values, " + f"got {dskact_array.size}." + ) + + if not np.all(np.isfinite(kact_array)): + raise ValueError("kact contains non-finite values.") + + if not np.all(np.isfinite(dskact_array)): + raise ValueError("dskact contains non-finite values.") + + return kact_array, dskact_array + + +def build_hxr_mapping( + kact: Sequence[float], + dskact: Sequence[float], + *, + z_steps: int = 320, + physical_points_per_cell: int = 100, +) -> dict[str, np.ndarray | float]: + """ + Build both the physical HXR map and the effective zfel K profile. + + Parameters + ---------- + kact + K at the upstream end of each active HXR segment. + + dskact + K at the downstream end of each active HXR segment. + + z_steps + Number of longitudinal integration steps for zfel. + + The default is 320: + 32 segments x 10 zfel steps per segment. + + physical_points_per_cell + Plotting resolution for the real 149.6 m machine layout. + + Returns + ------- + dict + Physical machine representation: + z_physical_m + k_physical + physical_cell + + zfel representation: + z_zfel_m + k_zfel + zfel_cell + + Geometry: + physical_span_m + magnetic_length_m + """ + + kact_array, dskact_array = _validate_endpoints( + kact, + dskact, + ) + + z_steps = int(z_steps) + physical_points_per_cell = int(physical_points_per_cell) + + if z_steps < N_ACTIVE_SEGMENTS: + raise ValueError( + "z_steps must be at least the number of active segments." + ) + + if physical_points_per_cell < 2: + raise ValueError( + "physical_points_per_cell must be at least 2." + ) + + # ============================================================== + # 1. PHYSICAL MACHINE COORDINATE + # + # Includes: + # - 3.4 m active undulators + # - 1.0 m interspaces + # - inactive cells 21 and 28 + # + # NaN is used outside active magnetic sections so plotting does + # not falsely imply that the drift/interspace has an undulator K. + # ============================================================== + + n_physical_points = ( + len(HXR_SLOTS) * physical_points_per_cell + ) + + dz_physical = PHYSICAL_SPAN_M / n_physical_points + + z_physical_m = ( + np.arange(n_physical_points, dtype=float) + 0.5 + ) * dz_physical + + physical_cell = ( + FIRST_HXR_CELL + + np.floor( + z_physical_m / CELL_LENGTH_M + ).astype(int) + ) + + cell_start_m = ( + physical_cell - FIRST_HXR_CELL + ) * CELL_LENGTH_M + + local_z_m = z_physical_m - cell_start_m + + # NaN means: no active undulator field represented here. + k_physical = np.full( + n_physical_points, + np.nan, + dtype=float, + ) + + for segment_index, cell in enumerate(HXR_CELLS): + + mask = ( + (physical_cell == cell) + & (local_z_m < UNDULATOR_LENGTH_M) + ) + + fraction = ( + local_z_m[mask] / UNDULATOR_LENGTH_M + ) + + k_physical[mask] = ( + kact_array[segment_index] + + fraction + * ( + dskact_array[segment_index] + - kact_array[segment_index] + ) + ) + + # ============================================================== + # 2. ZFEL MAGNETIC COORDINATE + # + # Current zfel does not explicitly model the machine drifts, + # interspaces, or chicanes. + # + # We therefore concatenate only the 32 active 3.4 m magnetic + # segments: + # + # total magnetic length = 32 x 3.4 m = 108.8 m + # + # Each zfel K value represents the center of one integration step. + # ============================================================== + + z_edges_m = np.linspace( + 0.0, + MAGNETIC_LENGTH_M, + z_steps + 1, + dtype=float, + ) + + z_zfel_m = 0.5 * ( + z_edges_m[:-1] + z_edges_m[1:] + ) + + segment_index = np.floor( + z_zfel_m / UNDULATOR_LENGTH_M + ).astype(int) + + segment_index = np.clip( + segment_index, + 0, + N_ACTIVE_SEGMENTS - 1, + ) + + segment_start_m = ( + segment_index * UNDULATOR_LENGTH_M + ) + + local_fraction = ( + z_zfel_m - segment_start_m + ) / UNDULATOR_LENGTH_M + + k_zfel = ( + kact_array[segment_index] + + local_fraction + * ( + dskact_array[segment_index] + - kact_array[segment_index] + ) + ) + + hxr_cells_array = np.asarray( + HXR_CELLS, + dtype=int, + ) + + zfel_cell = hxr_cells_array[segment_index] + + return { + "z_physical_m": z_physical_m, + "k_physical": k_physical, + "physical_cell": physical_cell, + "z_zfel_m": z_zfel_m, + "k_zfel": k_zfel, + "zfel_cell": zfel_cell, + "physical_span_m": PHYSICAL_SPAN_M, + "magnetic_length_m": MAGNETIC_LENGTH_M, + } \ No newline at end of file From d14e3d88ca28bb94e5fffedddb41f12a0e5048eb Mon Sep 17 00:00:00 2001 From: hong274SLAC Date: Tue, 11 Aug 2026 22:07:32 -0700 Subject: [PATCH 02/11] Integrate CU HXR ZFEL runner Add dedicated CU HXR ZFEL runner configuration and machine-style PV mapping. Lazy-load model backends so optional dependencies are only imported when selected. --- virtual_accelerator/models/cu_hxr_zfel.py | 88 ++++++++++++++++++++++- virtual_accelerator/models/runners.py | 23 +++--- 2 files changed, 99 insertions(+), 12 deletions(-) diff --git a/virtual_accelerator/models/cu_hxr_zfel.py b/virtual_accelerator/models/cu_hxr_zfel.py index 2e163ca..d7cc874 100644 --- a/virtual_accelerator/models/cu_hxr_zfel.py +++ b/virtual_accelerator/models/cu_hxr_zfel.py @@ -237,4 +237,90 @@ def get_cu_hxr_zfel_model() -> ZFELPVModel: """ Construct the CU HXR ZFEL virtual-accelerator model. """ - return ZFELPVModel() \ No newline at end of file + return ZFELPVModel() + +def build_cu_hxr_zfel_runner_config( + runner_cls, + model, + *, + prefix: str = "VA:", + protocols: tuple[str, ...] = ("ca", "pva"), + update_rate: float = 0.5, +): + """ + Build the lume-pva Runner configuration for the CU HXR ZFEL model. + """ + + config = runner_cls.generate_config( + model=model, + prefix="", + ) + + # PV names below already include the requested prefix. + config["prefix"] = "" + config["protocol"] = list(protocols) + config["update_rate"] = update_rate + + def va_pv(name: str) -> str: + return f"{prefix}{name}" + + for cell in HXR_CELLS: + config["variables"][f"KAct_{cell}"]["pv"] = va_pv( + f"USEG:UNDH:{cell}50:KAct" + ) + + config["variables"][f"DSKAct_{cell}"]["pv"] = va_pv( + f"USEG:UNDH:{cell}50:DSKAct" + ) + + config["variables"]["power_max"]["pv"] = va_pv( + "ZFEL:POWER_MAX" + ) + + config["variables"]["exit_power"]["pv"] = va_pv( + "ZFEL:EXIT_POWER" + ) + + config["variables"]["pulse_energy"]["pv"] = va_pv( + "ZFEL:PULSE_ENERGY" + ) + + config["variables"]["pulse_intensity_mean"]["pv"] = va_pv( + "GDET:FEE1:361:ENRC" + ) + + config["variables"]["pulse_intensity_p80"]["pv"] = va_pv( + "GDET:FEE1:361:ENRCHSTCUHBR" + ) + + config["variables"]["pulse_intensity_std_relative"]["pv"] = va_pv( + "ZFEL:PULSE_INTENSITY_STD_REL" + ) + + return config + +def get_cu_hxr_zfel_runner( + runner_cls, + *, + prefix: str = "VA:", + protocols: tuple[str, ...] = ("ca", "pva"), + update_rate: float = 0.5, +): + """ + Construct a lume-pva Runner for the CU HXR ZFEL virtual accelerator. + """ + + model = get_cu_hxr_zfel_model() + + config = build_cu_hxr_zfel_runner_config( + runner_cls, + model, + prefix=prefix, + protocols=protocols, + update_rate=update_rate, + ) + + return runner_cls( + model=model, + config=config, + ) \ No newline at end of file diff --git a/virtual_accelerator/models/runners.py b/virtual_accelerator/models/runners.py index fc978d6..04fdfb4 100644 --- a/virtual_accelerator/models/runners.py +++ b/virtual_accelerator/models/runners.py @@ -1,6 +1,5 @@ import argparse -from virtual_accelerator.models.facet2 import get_facet_staged_model from virtual_accelerator.utils.optional_dependencies import import_optional_symbol import logging @@ -8,13 +7,13 @@ def main(): parser = argparse.ArgumentParser( - description="Run the CU HXR model with BMAD backend" + description="Run a virtual accelerator model" ) - choices = ["cu_hxr_bmad", "cu_hxr_staged", "facet_bmad", "facet_staged"] + choices = ["cu_hxr_bmad", "cu_hxr_staged", "facet_bmad", "facet_staged", "cu_hxr_zfel"] parser.add_argument( "model", choices=choices, - help="Model backend to run (cu_hxr_bmad, cu_hxr_staged, facet_bmad, or facet_staged)", + help="Model backend to run (cu_hxr_bmad, cu_hxr_staged, facet_bmad, facet_staged, or cu_hxr_zfel)", ) parser.add_argument( "--end-element", @@ -47,30 +46,32 @@ def main(): extra="pva", ) - from virtual_accelerator.models.cu_hxr import ( - get_cu_hxr_bmad_model, - get_cu_hxr_staged_model, - ) - from virtual_accelerator.models.facet2 import get_facet_bmad_model - # Get the appropriate model based on user input if args.model == "cu_hxr_bmad": + from virtual_accelerator.models.cu_hxr import (get_cu_hxr_bmad_model) model = get_cu_hxr_bmad_model(end_element=args.end_element, track_beam=True) elif args.model == "cu_hxr_staged": + from virtual_accelerator.models.cu_hxr import (get_cu_hxr_staged_model) model = get_cu_hxr_staged_model( end_element=args.end_element, n_particles=args.n_particles ) elif args.model == "facet_bmad": + from virtual_accelerator.models.facet2 import (get_facet_bmad_model) model = get_facet_bmad_model(end_element=args.end_element, track_beam=True) elif args.model == "facet_staged": + from virtual_accelerator.models.facet2 import (get_facet_staged_model) model = get_facet_staged_model( end_element=args.end_element, n_particles=args.n_particles ) + elif args.model == "cu_hxr_zfel": + from virtual_accelerator.models.cu_hxr_zfel import ( get_cu_hxr_zfel_runner,) + runner = get_cu_hxr_zfel_runner(Runner) else: raise ValueError(f"Invalid model choice. Please choose one of {choices}.") # Run the model - runner = Runner(model) + if args.model != "cu_hxr_zfel": + runner = Runner(model) runner.run() From 22cddc2e9ffbbae74b409b4f8bf87e057da9ac91 Mon Sep 17 00:00:00 2001 From: hong274SLAC Date: Tue, 11 Aug 2026 23:49:06 -0700 Subject: [PATCH 03/11] Add tests for CU HXR ZFEL virtual accelerator Add unit tests for HXR-to-ZFEL mapping, ZFEL model behavior, CU HXR scalar controls, and runner PV configuration. --- virtual_accelerator/tests/test_cu_hxr_zfel.py | 186 ++++++++++++++++++ .../tests/test_zfel_mapping.py | 68 +++++++ virtual_accelerator/tests/test_zfel_model.py | 143 ++++++++++++++ 3 files changed, 397 insertions(+) create mode 100644 virtual_accelerator/tests/test_cu_hxr_zfel.py create mode 100644 virtual_accelerator/tests/test_zfel_mapping.py create mode 100644 virtual_accelerator/tests/test_zfel_model.py diff --git a/virtual_accelerator/tests/test_cu_hxr_zfel.py b/virtual_accelerator/tests/test_cu_hxr_zfel.py new file mode 100644 index 0000000..d9ed431 --- /dev/null +++ b/virtual_accelerator/tests/test_cu_hxr_zfel.py @@ -0,0 +1,186 @@ +import numpy as np +import pytest + + +pytest.importorskip("zfel") + +from virtual_accelerator.models.cu_hxr_zfel import ( + build_cu_hxr_zfel_runner_config, + get_cu_hxr_zfel_model, + get_cu_hxr_zfel_runner, +) +from virtual_accelerator.zfel.undulator_mapping import HXR_CELLS + + +class DummyRunner: + """ + Minimal stand-in for lume-pva Runner. + + This allows runner configuration to be tested without + starting an EPICS server. + """ + + def __init__(self, model, config=None): + self.model = model + self.config = config + + @staticmethod + def generate_config(model, prefix=""): + return { + "prefix": prefix, + "protocol": ["pva"], + "update_rate": 0.1, + "variables": { + name: { + "pv": f"{prefix}{name}", + } + for name in model.supported_variables + }, + } + + +def test_cu_hxr_zfel_model_variables(): + model = get_cu_hxr_zfel_model() + + supported = model.supported_variables + + for cell in HXR_CELLS: + assert f"KAct_{cell}" in supported + assert f"DSKAct_{cell}" in supported + + expected_diagnostics = { + "power_max", + "exit_power", + "pulse_energy", + "pulse_intensity_mean", + "pulse_intensity_p80", + "pulse_intensity_std_relative", + } + + assert expected_diagnostics.issubset(supported) + + state = model.get( + [ + "KAct_14", + "DSKAct_14", + "pulse_energy", + "pulse_intensity_mean", + "pulse_intensity_p80", + "pulse_intensity_std_relative", + ] + ) + + assert np.isclose(state["KAct_14"], 3.5) + assert np.isclose(state["DSKAct_14"], 3.5) + + assert state["pulse_energy"] > 0.0 + + assert np.isclose( + state["pulse_intensity_mean"], + state["pulse_energy"], + ) + + assert np.isclose( + state["pulse_intensity_p80"], + state["pulse_energy"], + ) + + assert state["pulse_intensity_std_relative"] == 0.0 + + +def test_scalar_kact_write_updates_zfel_backend(): + model = get_cu_hxr_zfel_model() + + baseline = model.get( + [ + "KAct_47", + "pulse_energy", + ] + ) + + target_k = baseline["KAct_47"] - 0.02 + + model.set( + { + "KAct_47": target_k, + } + ) + + changed = model.get( + [ + "KAct_47", + "pulse_energy", + ] + ) + + assert np.isclose( + changed["KAct_47"], + target_k, + ) + + assert not np.isclose( + changed["pulse_energy"], + baseline["pulse_energy"], + rtol=1e-8, + atol=0.0, + ) + + +def test_runner_config_uses_machine_style_pvs(): + model = get_cu_hxr_zfel_model() + + config = build_cu_hxr_zfel_runner_config( + DummyRunner, + model, + prefix="VA:", + protocols=("ca",), + update_rate=0.5, + ) + + assert config["prefix"] == "" + assert config["protocol"] == ["ca"] + assert config["update_rate"] == 0.5 + + assert ( + config["variables"]["KAct_14"]["pv"] + == "VA:USEG:UNDH:1450:KAct" + ) + + assert ( + config["variables"]["DSKAct_14"]["pv"] + == "VA:USEG:UNDH:1450:DSKAct" + ) + + assert ( + config["variables"]["pulse_energy"]["pv"] + == "VA:ZFEL:PULSE_ENERGY" + ) + + assert ( + config["variables"]["pulse_intensity_mean"]["pv"] + == "VA:GDET:FEE1:361:ENRC" + ) + + assert ( + config["variables"]["pulse_intensity_p80"]["pv"] + == "VA:GDET:FEE1:361:ENRCHSTCUHBR" + ) + + +def test_zfel_runner_factory_returns_configured_runner(): + runner = get_cu_hxr_zfel_runner( + DummyRunner, + protocols=("ca",), + ) + + assert isinstance(runner, DummyRunner) + + assert ( + runner.config["variables"]["KAct_14"]["pv"] + == "VA:USEG:UNDH:1450:KAct" + ) + + assert ( + runner.config["variables"]["pulse_energy"]["pv"] + == "VA:ZFEL:PULSE_ENERGY" + ) \ No newline at end of file diff --git a/virtual_accelerator/tests/test_zfel_mapping.py b/virtual_accelerator/tests/test_zfel_mapping.py new file mode 100644 index 0000000..2a24212 --- /dev/null +++ b/virtual_accelerator/tests/test_zfel_mapping.py @@ -0,0 +1,68 @@ +import numpy as np + +from virtual_accelerator.zfel.undulator_mapping import ( + HXR_CELLS, + build_hxr_mapping, +) + + +def test_hxr_cells_have_expected_layout(): + expected_cells = tuple( + list(range(14, 21)) + + list(range(22, 28)) + + list(range(29, 48)) + ) + + assert HXR_CELLS == expected_cells + assert len(HXR_CELLS) == 32 + + # Bypass-chicane locations are not active undulator segments. + assert 21 not in HXR_CELLS + assert 28 not in HXR_CELLS + + +def test_constant_k_maps_to_expected_zfel_profile(): + kact = np.full(len(HXR_CELLS), 3.5) + dskact = np.full(len(HXR_CELLS), 3.5) + + mapping = build_hxr_mapping( + kact, + dskact, + z_steps=320, + ) + + k_zfel = np.asarray(mapping["k_zfel"]) + + assert k_zfel.shape == (320,) + assert np.isfinite(k_zfel).all() + assert np.allclose(k_zfel, 3.5) + + +def test_changing_segment_changes_zfel_profile(): + kact = np.full(len(HXR_CELLS), 3.5) + dskact = np.full(len(HXR_CELLS), 3.5) + + baseline = build_hxr_mapping( + kact, + dskact, + z_steps=320, + ) + + changed_kact = kact.copy() + changed_dskact = dskact.copy() + + changed_kact[-1] -= 0.02 + changed_dskact[-1] -= 0.02 + + changed = build_hxr_mapping( + changed_kact, + changed_dskact, + z_steps=320, + ) + + baseline_k = np.asarray(baseline["k_zfel"]) + changed_k = np.asarray(changed["k_zfel"]) + + assert baseline_k.shape == changed_k.shape + assert not np.allclose(changed_k, baseline_k) + assert np.min(changed_k) < np.min(baseline_k) \ No newline at end of file diff --git a/virtual_accelerator/tests/test_zfel_model.py b/virtual_accelerator/tests/test_zfel_model.py new file mode 100644 index 0000000..c013699 --- /dev/null +++ b/virtual_accelerator/tests/test_zfel_model.py @@ -0,0 +1,143 @@ +import numpy as np +import pytest + + +pytest.importorskip("zfel") + +from virtual_accelerator.zfel.model import ZFELModel + + +def test_zfel_model_initialization(): + model = ZFELModel() + + state = model.get( + [ + "Kact", + "DSKact", + "unduK", + "power_max", + "exit_power", + "pulse_energy", + ] + ) + + assert np.asarray(state["Kact"]).shape == (32,) + assert np.asarray(state["DSKact"]).shape == (32,) + assert np.asarray(state["unduK"]).shape == (320,) + + assert np.isfinite(state["power_max"]) + assert np.isfinite(state["exit_power"]) + assert np.isfinite(state["pulse_energy"]) + + assert state["pulse_energy"] > 0.0 + + +def test_setting_k_changes_zfel_result(): + model = ZFELModel() + + baseline = model.get( + [ + "Kact", + "DSKact", + "unduK", + "pulse_energy", + ] + ) + + changed_kact = np.asarray( + baseline["Kact"], + dtype=float, + ).copy() + + changed_kact[-1] -= 0.02 + + model.set( + { + "Kact": changed_kact, + "DSKact": baseline["DSKact"], + } + ) + + changed = model.get( + [ + "Kact", + "unduK", + "pulse_energy", + ] + ) + + assert np.isclose( + changed["Kact"][-1], + changed_kact[-1], + ) + + assert not np.allclose( + changed["unduK"], + baseline["unduK"], + ) + + assert not np.isclose( + changed["pulse_energy"], + baseline["pulse_energy"], + rtol=1e-8, + atol=0.0, + ) + + +def test_reset_restores_initial_state(): + model = ZFELModel() + + baseline = model.get( + [ + "Kact", + "DSKact", + "unduK", + "pulse_energy", + ] + ) + + changed_kact = np.asarray( + baseline["Kact"], + dtype=float, + ).copy() + + changed_kact[-1] -= 0.02 + + model.set( + { + "Kact": changed_kact, + } + ) + + model.reset() + + reset_state = model.get( + [ + "Kact", + "DSKact", + "unduK", + "pulse_energy", + ] + ) + + assert np.allclose( + reset_state["Kact"], + baseline["Kact"], + ) + + assert np.allclose( + reset_state["DSKact"], + baseline["DSKact"], + ) + + assert np.allclose( + reset_state["unduK"], + baseline["unduK"], + ) + + assert np.isclose( + reset_state["pulse_energy"], + baseline["pulse_energy"], + rtol=1e-10, + atol=0.0, + ) \ No newline at end of file From 6a5b0287adb797d2bbf127b1d0cbed3f04920a6b Mon Sep 17 00:00:00 2001 From: hong274SLAC Date: Fri, 14 Aug 2026 14:38:31 -0700 Subject: [PATCH 04/11] Add ZFEL optional dependency and documentation Add the ZFEL optional dependency pinned to the NumPy-compatible fork commit and document installation and CU HXR ZFEL runner usage. --- README.md | 9 +++++++++ pyproject.toml | 4 ++++ 2 files changed, 13 insertions(+) diff --git a/README.md b/README.md index 1f53375..33d73cb 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,7 @@ Lastly, install backend-specific extras depending on which simulation types you pip install .[bmad] pip install .[cheetah] pip install .[impact] +pip install .[zfel] pip install .[pva] pip install .[surrogate] pip install .[all] @@ -40,6 +41,7 @@ Optional Dependency Keys by Model: | `get_cu_hxr_injector_surrogate_model` | `surrogate` | Uses torch surrogate + cheetah particles. | | `get_facet_staged_model` | `surrogate`, `bmad` | FACET-II staged model (injector surrogate + FACET-II BMAD). | | `get_cu_hxr_staged_model` | `surrogate`, `bmad` | Stages `InjectorSurrogate` + CU HXR BMAD model. | +| `get_cu_hxr_zfel_model` | `zfel` | CU HXR taper model using the 1D ZFEL backend. | | `virtual_accelerator.models.runners` CLI | `pva` (+ model backend key) | Runner requires `pva`; selected model backend must also be installed. | The package now lazily imports backend-specific dependencies. If you call a model @@ -61,6 +63,13 @@ For example: python virtual_accelerator/models/runners.py cu_hxr_bmad --end-element OTR4 ``` +CU HXR ZFEL runner serves machine-style PVs with the VA: prefix, for example: +``` +python -m virtual_accelerator.models.runners cu_hxr_zfel +VA:USEG:UNDH:1450:KAct +VA:ZFEL:PULSE_ENERGY +``` + For more info, run: ``` python virtual_accelerator/models/runners.py -h diff --git a/pyproject.toml b/pyproject.toml index 25dd995..4c83df3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,6 +47,9 @@ cheetah = [ impact = [ "lume-impact @ git+https://github.com/ChristopherMayes/lume-impact" ] +zfel = [ + "zfel @ git+https://github.com/hong274SLAC/zfel.git@801d00bb24a867fbaae132a44b791a3dcc0343cf" +] pva = [ "lume-pva @ git+https://github.com/lume-science/lume-pva" ] @@ -68,6 +71,7 @@ all = [ "distgen", "facet2-inj-ml-model @ git+https://github.com/slaclab/facet2_inj_ml_model", "lcls-cu-inj-model @ git+https://github.com/slaclab/lcls_cu_injector_ml_model.git", + "zfel @ git+https://github.com/hong274SLAC/zfel.git@801d00bb24a867fbaae132a44b791a3dcc0343cf", ] dev = [ "pytest", From e72131611bd1ec2e48e563a92403857c52bea194 Mon Sep 17 00:00:00 2001 From: hong274SLAC Date: Fri, 14 Aug 2026 15:16:30 -0700 Subject: [PATCH 05/11] Apply pre-commit formatting and lint fixes Apply repository formatting, end-of-file, and Ruff lint fixes to the ZFEL backend and tests. --- virtual_accelerator/models/cu_hxr_zfel.py | 204 ++++++++---------- virtual_accelerator/models/runners.py | 29 ++- virtual_accelerator/tests/test_cu_hxr_zfel.py | 30 +-- .../tests/test_zfel_mapping.py | 6 +- virtual_accelerator/tests/test_zfel_model.py | 2 +- virtual_accelerator/zfel/__init__.py | 1 - virtual_accelerator/zfel/model.py | 30 +-- virtual_accelerator/zfel/undulator_mapping.py | 83 ++----- 8 files changed, 144 insertions(+), 241 deletions(-) diff --git a/virtual_accelerator/models/cu_hxr_zfel.py b/virtual_accelerator/models/cu_hxr_zfel.py index d7cc874..b1639e9 100644 --- a/virtual_accelerator/models/cu_hxr_zfel.py +++ b/virtual_accelerator/models/cu_hxr_zfel.py @@ -23,13 +23,15 @@ class ZFELPVModel(LUMEModel): def __init__(self): self._backend = ZFELModel() - backend_state = self._backend.get([ - "Kact", - "DSKact", - "power_max", - "exit_power", - "pulse_energy", - ]) + backend_state = self._backend.get( + [ + "Kact", + "DSKact", + "power_max", + "exit_power", + "pulse_energy", + ] + ) kact = np.asarray( backend_state["Kact"], @@ -75,44 +77,46 @@ def __init__(self): # Read-only FEL diagnostics # ------------------------------------------------------ - self._variables.update({ - "power_max": ScalarVariable( - name="power_max", - default_value=0.0, - unit="W", - read_only=True, - ), - "exit_power": ScalarVariable( - name="exit_power", - default_value=0.0, - unit="W", - read_only=True, - ), - "pulse_energy": ScalarVariable( - name="pulse_energy", - default_value=0.0, - unit="J", - read_only=True, - ), - "pulse_intensity_mean": ScalarVariable( - name="pulse_intensity_mean", - default_value=0.0, - unit="J", - read_only=True, - ), - "pulse_intensity_p80": ScalarVariable( - name="pulse_intensity_p80", - default_value=0.0, - unit="J", - read_only=True, - ), - "pulse_intensity_std_relative": ScalarVariable( - name="pulse_intensity_std_relative", - default_value=0.0, - unit="dimensionless", - read_only=True, - ), - }) + self._variables.update( + { + "power_max": ScalarVariable( + name="power_max", + default_value=0.0, + unit="W", + read_only=True, + ), + "exit_power": ScalarVariable( + name="exit_power", + default_value=0.0, + unit="W", + read_only=True, + ), + "pulse_energy": ScalarVariable( + name="pulse_energy", + default_value=0.0, + unit="J", + read_only=True, + ), + "pulse_intensity_mean": ScalarVariable( + name="pulse_intensity_mean", + default_value=0.0, + unit="J", + read_only=True, + ), + "pulse_intensity_p80": ScalarVariable( + name="pulse_intensity_p80", + default_value=0.0, + unit="J", + read_only=True, + ), + "pulse_intensity_std_relative": ScalarVariable( + name="pulse_intensity_std_relative", + default_value=0.0, + unit="dimensionless", + read_only=True, + ), + } + ) self._sync_from_backend(include_controls=True) @@ -121,10 +125,7 @@ def supported_variables(self): return self._variables def _get(self, names: list[str]) -> dict[str, Any]: - return { - name: self._state[name] - for name in names - } + return {name: self._state[name] for name in names} def _set(self, values: dict[str, Any]) -> None: """ @@ -135,51 +136,45 @@ def _set(self, values: dict[str, Any]) -> None: """ if not values: - self._sync_from_backend( - include_controls=True - ) + self._sync_from_backend(include_controls=True) return for name, value in values.items(): self._state[name] = float(value) kact = np.asarray( - [ - self._state[f"KAct_{cell}"] - for cell in HXR_CELLS - ], + [self._state[f"KAct_{cell}"] for cell in HXR_CELLS], dtype=float, ) dskact = np.asarray( - [ - self._state[f"DSKAct_{cell}"] - for cell in HXR_CELLS - ], + [self._state[f"DSKAct_{cell}"] for cell in HXR_CELLS], dtype=float, ) - self._backend.set({ - "Kact": kact, - "DSKact": dskact, - }) - - self._sync_from_backend( - include_controls=True + self._backend.set( + { + "Kact": kact, + "DSKact": dskact, + } ) + self._sync_from_backend(include_controls=True) + def _sync_from_backend( self, *, include_controls: bool, ) -> None: - backend_state = self._backend.get([ - "Kact", - "DSKact", - "power_max", - "exit_power", - "pulse_energy", - ]) + backend_state = self._backend.get( + [ + "Kact", + "DSKact", + "power_max", + "exit_power", + "pulse_energy", + ] + ) if include_controls: kact = np.asarray( @@ -193,45 +188,28 @@ def _sync_from_backend( ) for index, cell in enumerate(HXR_CELLS): - self._state[f"KAct_{cell}"] = float( - kact[index] - ) + self._state[f"KAct_{cell}"] = float(kact[index]) - self._state[f"DSKAct_{cell}"] = float( - dskact[index] - ) + self._state[f"DSKAct_{cell}"] = float(dskact[index]) - pulse_energy = float( - backend_state["pulse_energy"] - ) + pulse_energy = float(backend_state["pulse_energy"]) - self._state["power_max"] = float( - backend_state["power_max"] - ) + self._state["power_max"] = float(backend_state["power_max"]) - self._state["exit_power"] = float( - backend_state["exit_power"] - ) + self._state["exit_power"] = float(backend_state["exit_power"]) self._state["pulse_energy"] = pulse_energy # Fixed-seed deterministic model for now. - self._state["pulse_intensity_mean"] = ( - pulse_energy - ) - self._state["pulse_intensity_p80"] = ( - pulse_energy - ) - self._state[ - "pulse_intensity_std_relative" - ] = 0.0 + self._state["pulse_intensity_mean"] = pulse_energy + self._state["pulse_intensity_p80"] = pulse_energy + self._state["pulse_intensity_std_relative"] = 0.0 def reset(self) -> None: self._backend.reset() - self._sync_from_backend( - include_controls=True - ) + self._sync_from_backend(include_controls=True) + def get_cu_hxr_zfel_model() -> ZFELPVModel: """ @@ -239,6 +217,7 @@ def get_cu_hxr_zfel_model() -> ZFELPVModel: """ return ZFELPVModel() + def build_cu_hxr_zfel_runner_config( runner_cls, model, @@ -265,29 +244,19 @@ def va_pv(name: str) -> str: return f"{prefix}{name}" for cell in HXR_CELLS: - config["variables"][f"KAct_{cell}"]["pv"] = va_pv( - f"USEG:UNDH:{cell}50:KAct" - ) + config["variables"][f"KAct_{cell}"]["pv"] = va_pv(f"USEG:UNDH:{cell}50:KAct") config["variables"][f"DSKAct_{cell}"]["pv"] = va_pv( f"USEG:UNDH:{cell}50:DSKAct" ) - config["variables"]["power_max"]["pv"] = va_pv( - "ZFEL:POWER_MAX" - ) + config["variables"]["power_max"]["pv"] = va_pv("ZFEL:POWER_MAX") - config["variables"]["exit_power"]["pv"] = va_pv( - "ZFEL:EXIT_POWER" - ) + config["variables"]["exit_power"]["pv"] = va_pv("ZFEL:EXIT_POWER") - config["variables"]["pulse_energy"]["pv"] = va_pv( - "ZFEL:PULSE_ENERGY" - ) + config["variables"]["pulse_energy"]["pv"] = va_pv("ZFEL:PULSE_ENERGY") - config["variables"]["pulse_intensity_mean"]["pv"] = va_pv( - "GDET:FEE1:361:ENRC" - ) + config["variables"]["pulse_intensity_mean"]["pv"] = va_pv("GDET:FEE1:361:ENRC") config["variables"]["pulse_intensity_p80"]["pv"] = va_pv( "GDET:FEE1:361:ENRCHSTCUHBR" @@ -299,6 +268,7 @@ def va_pv(name: str) -> str: return config + def get_cu_hxr_zfel_runner( runner_cls, *, @@ -323,4 +293,4 @@ def get_cu_hxr_zfel_runner( return runner_cls( model=model, config=config, - ) \ No newline at end of file + ) diff --git a/virtual_accelerator/models/runners.py b/virtual_accelerator/models/runners.py index 04fdfb4..95f0fd2 100644 --- a/virtual_accelerator/models/runners.py +++ b/virtual_accelerator/models/runners.py @@ -6,10 +6,14 @@ def main(): - parser = argparse.ArgumentParser( - description="Run a virtual accelerator model" - ) - choices = ["cu_hxr_bmad", "cu_hxr_staged", "facet_bmad", "facet_staged", "cu_hxr_zfel"] + parser = argparse.ArgumentParser(description="Run a virtual accelerator model") + choices = [ + "cu_hxr_bmad", + "cu_hxr_staged", + "facet_bmad", + "facet_staged", + "cu_hxr_zfel", + ] parser.add_argument( "model", choices=choices, @@ -48,23 +52,30 @@ def main(): # Get the appropriate model based on user input if args.model == "cu_hxr_bmad": - from virtual_accelerator.models.cu_hxr import (get_cu_hxr_bmad_model) + from virtual_accelerator.models.cu_hxr import get_cu_hxr_bmad_model + model = get_cu_hxr_bmad_model(end_element=args.end_element, track_beam=True) elif args.model == "cu_hxr_staged": - from virtual_accelerator.models.cu_hxr import (get_cu_hxr_staged_model) + from virtual_accelerator.models.cu_hxr import get_cu_hxr_staged_model + model = get_cu_hxr_staged_model( end_element=args.end_element, n_particles=args.n_particles ) elif args.model == "facet_bmad": - from virtual_accelerator.models.facet2 import (get_facet_bmad_model) + from virtual_accelerator.models.facet2 import get_facet_bmad_model + model = get_facet_bmad_model(end_element=args.end_element, track_beam=True) elif args.model == "facet_staged": - from virtual_accelerator.models.facet2 import (get_facet_staged_model) + from virtual_accelerator.models.facet2 import get_facet_staged_model + model = get_facet_staged_model( end_element=args.end_element, n_particles=args.n_particles ) elif args.model == "cu_hxr_zfel": - from virtual_accelerator.models.cu_hxr_zfel import ( get_cu_hxr_zfel_runner,) + from virtual_accelerator.models.cu_hxr_zfel import ( + get_cu_hxr_zfel_runner, + ) + runner = get_cu_hxr_zfel_runner(Runner) else: raise ValueError(f"Invalid model choice. Please choose one of {choices}.") diff --git a/virtual_accelerator/tests/test_cu_hxr_zfel.py b/virtual_accelerator/tests/test_cu_hxr_zfel.py index d9ed431..230f0f2 100644 --- a/virtual_accelerator/tests/test_cu_hxr_zfel.py +++ b/virtual_accelerator/tests/test_cu_hxr_zfel.py @@ -141,25 +141,13 @@ def test_runner_config_uses_machine_style_pvs(): assert config["protocol"] == ["ca"] assert config["update_rate"] == 0.5 - assert ( - config["variables"]["KAct_14"]["pv"] - == "VA:USEG:UNDH:1450:KAct" - ) + assert config["variables"]["KAct_14"]["pv"] == "VA:USEG:UNDH:1450:KAct" - assert ( - config["variables"]["DSKAct_14"]["pv"] - == "VA:USEG:UNDH:1450:DSKAct" - ) + assert config["variables"]["DSKAct_14"]["pv"] == "VA:USEG:UNDH:1450:DSKAct" - assert ( - config["variables"]["pulse_energy"]["pv"] - == "VA:ZFEL:PULSE_ENERGY" - ) + assert config["variables"]["pulse_energy"]["pv"] == "VA:ZFEL:PULSE_ENERGY" - assert ( - config["variables"]["pulse_intensity_mean"]["pv"] - == "VA:GDET:FEE1:361:ENRC" - ) + assert config["variables"]["pulse_intensity_mean"]["pv"] == "VA:GDET:FEE1:361:ENRC" assert ( config["variables"]["pulse_intensity_p80"]["pv"] @@ -175,12 +163,6 @@ def test_zfel_runner_factory_returns_configured_runner(): assert isinstance(runner, DummyRunner) - assert ( - runner.config["variables"]["KAct_14"]["pv"] - == "VA:USEG:UNDH:1450:KAct" - ) + assert runner.config["variables"]["KAct_14"]["pv"] == "VA:USEG:UNDH:1450:KAct" - assert ( - runner.config["variables"]["pulse_energy"]["pv"] - == "VA:ZFEL:PULSE_ENERGY" - ) \ No newline at end of file + assert runner.config["variables"]["pulse_energy"]["pv"] == "VA:ZFEL:PULSE_ENERGY" diff --git a/virtual_accelerator/tests/test_zfel_mapping.py b/virtual_accelerator/tests/test_zfel_mapping.py index 2a24212..ea8a881 100644 --- a/virtual_accelerator/tests/test_zfel_mapping.py +++ b/virtual_accelerator/tests/test_zfel_mapping.py @@ -8,9 +8,7 @@ def test_hxr_cells_have_expected_layout(): expected_cells = tuple( - list(range(14, 21)) - + list(range(22, 28)) - + list(range(29, 48)) + list(range(14, 21)) + list(range(22, 28)) + list(range(29, 48)) ) assert HXR_CELLS == expected_cells @@ -65,4 +63,4 @@ def test_changing_segment_changes_zfel_profile(): assert baseline_k.shape == changed_k.shape assert not np.allclose(changed_k, baseline_k) - assert np.min(changed_k) < np.min(baseline_k) \ No newline at end of file + assert np.min(changed_k) < np.min(baseline_k) diff --git a/virtual_accelerator/tests/test_zfel_model.py b/virtual_accelerator/tests/test_zfel_model.py index c013699..1b996ee 100644 --- a/virtual_accelerator/tests/test_zfel_model.py +++ b/virtual_accelerator/tests/test_zfel_model.py @@ -140,4 +140,4 @@ def test_reset_restores_initial_state(): baseline["pulse_energy"], rtol=1e-10, atol=0.0, - ) \ No newline at end of file + ) diff --git a/virtual_accelerator/zfel/__init__.py b/virtual_accelerator/zfel/__init__.py index b320b98..e69de29 100644 --- a/virtual_accelerator/zfel/__init__.py +++ b/virtual_accelerator/zfel/__init__.py @@ -1 +0,0 @@ -from virtual_accelerator.zfel.model import ZFELModel \ No newline at end of file diff --git a/virtual_accelerator/zfel/model.py b/virtual_accelerator/zfel/model.py index d3ef616..5918363 100644 --- a/virtual_accelerator/zfel/model.py +++ b/virtual_accelerator/zfel/model.py @@ -6,14 +6,14 @@ from lume.variables import NDVariable, ScalarVariable from zfel import sase1d -C_LIGHT = 299_792_458.0 - from virtual_accelerator.zfel.undulator_mapping import ( MAGNETIC_LENGTH_M, N_ACTIVE_SEGMENTS, build_hxr_mapping, ) +C_LIGHT = 299_792_458.0 + class ZFELModel(LUMEModel): """ @@ -115,7 +115,6 @@ def __init__(self): unit="dimensionless", read_only=False, ), - # Read-only mapped zfel K profile "unduK": NDVariable( name="unduK", @@ -128,7 +127,6 @@ def __init__(self): unit="dimensionless", read_only=True, ), - # Read-only FEL diagnostics "power_max": ScalarVariable( name="power_max", @@ -243,13 +241,9 @@ def _run_simulation(self) -> None: dtype=float, ) - self._state["power_max"] = float( - np.max(power_z) - ) + self._state["power_max"] = float(np.max(power_z)) - self._state["exit_power"] = float( - power_z[-1] - ) + self._state["exit_power"] = float(power_z[-1]) power_s = np.asarray( output["power_s"], dtype=float, @@ -261,21 +255,13 @@ def _run_simulation(self) -> None: ) if s_m.size < 2: - raise ValueError( - "zfel must return at least two longitudinal s points." - ) + raise ValueError("zfel must return at least two longitudinal s points.") - ds_m = float( - np.mean(np.diff(s_m)) - ) + ds_m = float(np.mean(np.diff(s_m))) # Each power_s sample represents one longitudinal slice. # Convert ds [m] to dt [s] using dt = ds/c. - pulse_energy_j = float( - np.sum(power_s[-1, :]) - * ds_m - / C_LIGHT - ) + pulse_energy_j = float(np.sum(power_s[-1, :]) * ds_m / C_LIGHT) self._state["pulse_energy"] = pulse_energy_j @@ -285,4 +271,4 @@ def reset(self) -> None: """ self._state = self._copy_initial_state() - self._run_simulation() \ No newline at end of file + self._run_simulation() diff --git a/virtual_accelerator/zfel/undulator_mapping.py b/virtual_accelerator/zfel/undulator_mapping.py index df0c5ad..254f69c 100644 --- a/virtual_accelerator/zfel/undulator_mapping.py +++ b/virtual_accelerator/zfel/undulator_mapping.py @@ -13,9 +13,7 @@ # Active HXR undulator cells. # This matches the current real-machine Badger environment. HXR_CELLS: tuple[int, ...] = ( - tuple(range(14, 21)) - + tuple(range(22, 28)) - + tuple(range(29, 48)) + tuple(range(14, 21)) + tuple(range(22, 28)) + tuple(range(29, 48)) ) # Cells omitted from the active-undulator list. @@ -47,14 +45,12 @@ def _validate_endpoints( if kact_array.size != N_ACTIVE_SEGMENTS: raise ValueError( - f"kact must contain {N_ACTIVE_SEGMENTS} values, " - f"got {kact_array.size}." + f"kact must contain {N_ACTIVE_SEGMENTS} values, got {kact_array.size}." ) if dskact_array.size != N_ACTIVE_SEGMENTS: raise ValueError( - f"dskact must contain {N_ACTIVE_SEGMENTS} values, " - f"got {dskact_array.size}." + f"dskact must contain {N_ACTIVE_SEGMENTS} values, got {dskact_array.size}." ) if not np.all(np.isfinite(kact_array)): @@ -120,14 +116,10 @@ def build_hxr_mapping( physical_points_per_cell = int(physical_points_per_cell) if z_steps < N_ACTIVE_SEGMENTS: - raise ValueError( - "z_steps must be at least the number of active segments." - ) + raise ValueError("z_steps must be at least the number of active segments.") if physical_points_per_cell < 2: - raise ValueError( - "physical_points_per_cell must be at least 2." - ) + raise ValueError("physical_points_per_cell must be at least 2.") # ============================================================== # 1. PHYSICAL MACHINE COORDINATE @@ -141,26 +133,15 @@ def build_hxr_mapping( # not falsely imply that the drift/interspace has an undulator K. # ============================================================== - n_physical_points = ( - len(HXR_SLOTS) * physical_points_per_cell - ) + n_physical_points = len(HXR_SLOTS) * physical_points_per_cell dz_physical = PHYSICAL_SPAN_M / n_physical_points - z_physical_m = ( - np.arange(n_physical_points, dtype=float) + 0.5 - ) * dz_physical + z_physical_m = (np.arange(n_physical_points, dtype=float) + 0.5) * dz_physical - physical_cell = ( - FIRST_HXR_CELL - + np.floor( - z_physical_m / CELL_LENGTH_M - ).astype(int) - ) + physical_cell = FIRST_HXR_CELL + np.floor(z_physical_m / CELL_LENGTH_M).astype(int) - cell_start_m = ( - physical_cell - FIRST_HXR_CELL - ) * CELL_LENGTH_M + cell_start_m = (physical_cell - FIRST_HXR_CELL) * CELL_LENGTH_M local_z_m = z_physical_m - cell_start_m @@ -172,23 +153,12 @@ def build_hxr_mapping( ) for segment_index, cell in enumerate(HXR_CELLS): + mask = (physical_cell == cell) & (local_z_m < UNDULATOR_LENGTH_M) - mask = ( - (physical_cell == cell) - & (local_z_m < UNDULATOR_LENGTH_M) - ) - - fraction = ( - local_z_m[mask] / UNDULATOR_LENGTH_M - ) + fraction = local_z_m[mask] / UNDULATOR_LENGTH_M - k_physical[mask] = ( - kact_array[segment_index] - + fraction - * ( - dskact_array[segment_index] - - kact_array[segment_index] - ) + k_physical[mask] = kact_array[segment_index] + fraction * ( + dskact_array[segment_index] - kact_array[segment_index] ) # ============================================================== @@ -212,13 +182,9 @@ def build_hxr_mapping( dtype=float, ) - z_zfel_m = 0.5 * ( - z_edges_m[:-1] + z_edges_m[1:] - ) + z_zfel_m = 0.5 * (z_edges_m[:-1] + z_edges_m[1:]) - segment_index = np.floor( - z_zfel_m / UNDULATOR_LENGTH_M - ).astype(int) + segment_index = np.floor(z_zfel_m / UNDULATOR_LENGTH_M).astype(int) segment_index = np.clip( segment_index, @@ -226,21 +192,12 @@ def build_hxr_mapping( N_ACTIVE_SEGMENTS - 1, ) - segment_start_m = ( - segment_index * UNDULATOR_LENGTH_M - ) + segment_start_m = segment_index * UNDULATOR_LENGTH_M - local_fraction = ( - z_zfel_m - segment_start_m - ) / UNDULATOR_LENGTH_M + local_fraction = (z_zfel_m - segment_start_m) / UNDULATOR_LENGTH_M - k_zfel = ( - kact_array[segment_index] - + local_fraction - * ( - dskact_array[segment_index] - - kact_array[segment_index] - ) + k_zfel = kact_array[segment_index] + local_fraction * ( + dskact_array[segment_index] - kact_array[segment_index] ) hxr_cells_array = np.asarray( @@ -259,4 +216,4 @@ def build_hxr_mapping( "zfel_cell": zfel_cell, "physical_span_m": PHYSICAL_SPAN_M, "magnetic_length_m": MAGNETIC_LENGTH_M, - } \ No newline at end of file + } From 32d83457baef01edd4548ca288e224e93cb42181 Mon Sep 17 00:00:00 2001 From: hong274SLAC Date: Thu, 20 Aug 2026 14:43:22 -0700 Subject: [PATCH 06/11] Simplify ZFEL model layering Refactor the ZFEL physics layer into a plain backend and keep the CU HXR wrapper as the single LUMEModel. Also standardize KAct/DSKAct naming. --- virtual_accelerator/models/cu_hxr_zfel.py | 26 +++--- virtual_accelerator/tests/test_zfel_model.py | 50 +++++------ virtual_accelerator/zfel/model.py | 91 ++++---------------- 3 files changed, 54 insertions(+), 113 deletions(-) diff --git a/virtual_accelerator/models/cu_hxr_zfel.py b/virtual_accelerator/models/cu_hxr_zfel.py index b1639e9..4c2d5c3 100644 --- a/virtual_accelerator/models/cu_hxr_zfel.py +++ b/virtual_accelerator/models/cu_hxr_zfel.py @@ -6,7 +6,7 @@ from lume.variables import ScalarVariable from virtual_accelerator.zfel.undulator_mapping import HXR_CELLS -from virtual_accelerator.zfel.model import ZFELModel +from virtual_accelerator.zfel.model import ZFELBackend class ZFELPVModel(LUMEModel): @@ -17,16 +17,16 @@ class ZFELPVModel(LUMEModel): KAct_14, DSKAct_14, ..., KAct_47, DSKAct_47 Internal physics model: - Kact[32], DSKact[32] -> zfel + KAct[32], DSKact[32] -> zfel """ def __init__(self): - self._backend = ZFELModel() + self._backend = ZFELBackend() backend_state = self._backend.get( [ - "Kact", - "DSKact", + "KAct", + "DSKAct", "power_max", "exit_power", "pulse_energy", @@ -34,12 +34,12 @@ def __init__(self): ) kact = np.asarray( - backend_state["Kact"], + backend_state["KAct"], dtype=float, ) dskact = np.asarray( - backend_state["DSKact"], + backend_state["DSKAct"], dtype=float, ) @@ -154,8 +154,8 @@ def _set(self, values: dict[str, Any]) -> None: self._backend.set( { - "Kact": kact, - "DSKact": dskact, + "KAct": kact, + "DSKAct": dskact, } ) @@ -168,8 +168,8 @@ def _sync_from_backend( ) -> None: backend_state = self._backend.get( [ - "Kact", - "DSKact", + "KAct", + "DSKAct", "power_max", "exit_power", "pulse_energy", @@ -178,12 +178,12 @@ def _sync_from_backend( if include_controls: kact = np.asarray( - backend_state["Kact"], + backend_state["KAct"], dtype=float, ) dskact = np.asarray( - backend_state["DSKact"], + backend_state["DSKAct"], dtype=float, ) diff --git a/virtual_accelerator/tests/test_zfel_model.py b/virtual_accelerator/tests/test_zfel_model.py index 1b996ee..bb4d547 100644 --- a/virtual_accelerator/tests/test_zfel_model.py +++ b/virtual_accelerator/tests/test_zfel_model.py @@ -4,16 +4,16 @@ pytest.importorskip("zfel") -from virtual_accelerator.zfel.model import ZFELModel +from virtual_accelerator.zfel.model import ZFELBackend def test_zfel_model_initialization(): - model = ZFELModel() + model = ZFELBackend() state = model.get( [ - "Kact", - "DSKact", + "KAct", + "DSKAct", "unduK", "power_max", "exit_power", @@ -21,8 +21,8 @@ def test_zfel_model_initialization(): ] ) - assert np.asarray(state["Kact"]).shape == (32,) - assert np.asarray(state["DSKact"]).shape == (32,) + assert np.asarray(state["KAct"]).shape == (32,) + assert np.asarray(state["DSKAct"]).shape == (32,) assert np.asarray(state["unduK"]).shape == (320,) assert np.isfinite(state["power_max"]) @@ -33,19 +33,19 @@ def test_zfel_model_initialization(): def test_setting_k_changes_zfel_result(): - model = ZFELModel() + model = ZFELBackend() baseline = model.get( [ - "Kact", - "DSKact", + "KAct", + "DSKAct", "unduK", "pulse_energy", ] ) changed_kact = np.asarray( - baseline["Kact"], + baseline["KAct"], dtype=float, ).copy() @@ -53,21 +53,21 @@ def test_setting_k_changes_zfel_result(): model.set( { - "Kact": changed_kact, - "DSKact": baseline["DSKact"], + "KAct": changed_kact, + "DSKAct": baseline["DSKAct"], } ) changed = model.get( [ - "Kact", + "KAct", "unduK", "pulse_energy", ] ) assert np.isclose( - changed["Kact"][-1], + changed["KAct"][-1], changed_kact[-1], ) @@ -85,19 +85,19 @@ def test_setting_k_changes_zfel_result(): def test_reset_restores_initial_state(): - model = ZFELModel() + model = ZFELBackend() baseline = model.get( [ - "Kact", - "DSKact", + "KAct", + "DSKAct", "unduK", "pulse_energy", ] ) changed_kact = np.asarray( - baseline["Kact"], + baseline["KAct"], dtype=float, ).copy() @@ -105,7 +105,7 @@ def test_reset_restores_initial_state(): model.set( { - "Kact": changed_kact, + "KAct": changed_kact, } ) @@ -113,21 +113,21 @@ def test_reset_restores_initial_state(): reset_state = model.get( [ - "Kact", - "DSKact", + "KAct", + "DSKAct", "unduK", "pulse_energy", ] ) assert np.allclose( - reset_state["Kact"], - baseline["Kact"], + reset_state["KAct"], + baseline["KAct"], ) assert np.allclose( - reset_state["DSKact"], - baseline["DSKact"], + reset_state["DSKAct"], + baseline["DSKAct"], ) assert np.allclose( diff --git a/virtual_accelerator/zfel/model.py b/virtual_accelerator/zfel/model.py index 5918363..1f3c3ea 100644 --- a/virtual_accelerator/zfel/model.py +++ b/virtual_accelerator/zfel/model.py @@ -2,8 +2,6 @@ import numpy as np -from lume.model import LUMEModel -from lume.variables import NDVariable, ScalarVariable from zfel import sase1d from virtual_accelerator.zfel.undulator_mapping import ( @@ -15,7 +13,7 @@ C_LIGHT = 299_792_458.0 -class ZFELModel(LUMEModel): +class ZFELBackend: """ LUMEModel wrapper around zfel using real-machine-like HXR Kact and DSKact controls. @@ -81,8 +79,8 @@ def __init__(self): ) self._initial_state = { - "Kact": initial_kact, - "DSKact": initial_dskact, + "KAct": initial_kact, + "DSKAct": initial_dskact, "unduK": np.full( self.ZFEL_Z_STEPS, self.INITIAL_K, @@ -95,59 +93,6 @@ def __init__(self): self._state = self._copy_initial_state() - # ---------------------------------------------------------- - # LUMEModel variable definitions - # ---------------------------------------------------------- - - self._variables = { - # Writable virtual-machine controls - "Kact": NDVariable( - name="Kact", - shape=(N_ACTIVE_SEGMENTS,), - default_value=initial_kact.copy(), - unit="dimensionless", - read_only=False, - ), - "DSKact": NDVariable( - name="DSKact", - shape=(N_ACTIVE_SEGMENTS,), - default_value=initial_dskact.copy(), - unit="dimensionless", - read_only=False, - ), - # Read-only mapped zfel K profile - "unduK": NDVariable( - name="unduK", - shape=(self.ZFEL_Z_STEPS,), - default_value=np.full( - self.ZFEL_Z_STEPS, - self.INITIAL_K, - dtype=float, - ), - unit="dimensionless", - read_only=True, - ), - # Read-only FEL diagnostics - "power_max": ScalarVariable( - name="power_max", - default_value=0.0, - unit="W", - read_only=True, - ), - "exit_power": ScalarVariable( - name="exit_power", - default_value=0.0, - unit="W", - read_only=True, - ), - "pulse_energy": ScalarVariable( - name="pulse_energy", - default_value=0.0, - unit="J", - read_only=True, - ), - } - # Store the latest complete HXR mapping for debugging, # plotting, and later Badger integration. self._mapping = None @@ -155,25 +100,21 @@ def __init__(self): # Run the initial constant-K case. self._run_simulation() - @property - def supported_variables(self): - return self._variables - def _copy_initial_state(self) -> dict[str, Any]: """ Return an independent copy of the initial machine state. """ return { - "Kact": self._initial_state["Kact"].copy(), - "DSKact": self._initial_state["DSKact"].copy(), + "KAct": self._initial_state["KAct"].copy(), + "DSKAct": self._initial_state["DSKAct"].copy(), "unduK": self._initial_state["unduK"].copy(), "power_max": float(self._initial_state["power_max"]), "exit_power": float(self._initial_state["exit_power"]), "pulse_energy": float(self._initial_state["pulse_energy"]), } - def _get(self, names: list[str]) -> dict[str, Any]: + def get(self, names: list[str]) -> dict[str, Any]: """ Return cached state only. @@ -192,20 +133,20 @@ def _get(self, names: list[str]) -> dict[str, Any]: return out - def _set(self, values: dict[str, Any]) -> None: + def set(self, values: dict[str, Any]) -> None: """ - Update Kact and/or DSKact, then run zfel once. + Update KAct and/or DSKAct, then run zfel once. """ - if "Kact" in values: - self._state["Kact"] = np.asarray( - values["Kact"], + if "KAct" in values: + self._state["KAct"] = np.asarray( + values["KAct"], dtype=float, ).copy() - if "DSKact" in values: - self._state["DSKact"] = np.asarray( - values["DSKact"], + if "DSKAct" in values: + self._state["DSKAct"] = np.asarray( + values["DSKAct"], dtype=float, ).copy() @@ -217,8 +158,8 @@ def _run_simulation(self) -> None: """ mapping = build_hxr_mapping( - self._state["Kact"], - self._state["DSKact"], + self._state["KAct"], + self._state["DSKAct"], z_steps=self.ZFEL_Z_STEPS, ) From 890ef6a860251bf87c6cc9701663f0427b6b735c Mon Sep 17 00:00:00 2001 From: hong274SLAC Date: Thu, 20 Aug 2026 14:59:13 -0700 Subject: [PATCH 07/11] Use upstream ZFEL dependency --- pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 4c83df3..e4fed54 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,7 +48,7 @@ impact = [ "lume-impact @ git+https://github.com/ChristopherMayes/lume-impact" ] zfel = [ - "zfel @ git+https://github.com/hong274SLAC/zfel.git@801d00bb24a867fbaae132a44b791a3dcc0343cf" + "zfel @ git+https://github.com/slaclab/zfel.git" ] pva = [ "lume-pva @ git+https://github.com/lume-science/lume-pva" @@ -71,7 +71,7 @@ all = [ "distgen", "facet2-inj-ml-model @ git+https://github.com/slaclab/facet2_inj_ml_model", "lcls-cu-inj-model @ git+https://github.com/slaclab/lcls_cu_injector_ml_model.git", - "zfel @ git+https://github.com/hong274SLAC/zfel.git@801d00bb24a867fbaae132a44b791a3dcc0343cf", + "zfel @ git+https://github.com/slaclab/zfel.git", ] dev = [ "pytest", From 6ee0ef79bb741e957c359445f9d20ccd791100bc Mon Sep 17 00:00:00 2001 From: hong274SLAC Date: Mon, 31 Aug 2026 11:56:02 -0700 Subject: [PATCH 08/11] Remove VA prefix from CU HXR ZFEL PVs --- virtual_accelerator/models/cu_hxr_zfel.py | 27 ++++++------------- virtual_accelerator/tests/test_cu_hxr_zfel.py | 16 +++++------ 2 files changed, 15 insertions(+), 28 deletions(-) diff --git a/virtual_accelerator/models/cu_hxr_zfel.py b/virtual_accelerator/models/cu_hxr_zfel.py index 4c2d5c3..7fb43dd 100644 --- a/virtual_accelerator/models/cu_hxr_zfel.py +++ b/virtual_accelerator/models/cu_hxr_zfel.py @@ -222,7 +222,6 @@ def build_cu_hxr_zfel_runner_config( runner_cls, model, *, - prefix: str = "VA:", protocols: tuple[str, ...] = ("ca", "pva"), update_rate: float = 0.5, ): @@ -240,31 +239,23 @@ def build_cu_hxr_zfel_runner_config( config["protocol"] = list(protocols) config["update_rate"] = update_rate - def va_pv(name: str) -> str: - return f"{prefix}{name}" for cell in HXR_CELLS: - config["variables"][f"KAct_{cell}"]["pv"] = va_pv(f"USEG:UNDH:{cell}50:KAct") + config["variables"][f"KAct_{cell}"]["pv"] = f"USEG:UNDH:{cell}50:KAct" - config["variables"][f"DSKAct_{cell}"]["pv"] = va_pv( - f"USEG:UNDH:{cell}50:DSKAct" - ) + config["variables"][f"DSKAct_{cell}"]["pv"] = f"USEG:UNDH:{cell}50:DSKAct" - config["variables"]["power_max"]["pv"] = va_pv("ZFEL:POWER_MAX") + config["variables"]["power_max"]["pv"] = "ZFEL:POWER_MAX" - config["variables"]["exit_power"]["pv"] = va_pv("ZFEL:EXIT_POWER") + config["variables"]["exit_power"]["pv"] = "ZFEL:EXIT_POWER" - config["variables"]["pulse_energy"]["pv"] = va_pv("ZFEL:PULSE_ENERGY") + config["variables"]["pulse_energy"]["pv"] = "ZFEL:PULSE_ENERGY" - config["variables"]["pulse_intensity_mean"]["pv"] = va_pv("GDET:FEE1:361:ENRC") + config["variables"]["pulse_intensity_mean"]["pv"] = "GDET:FEE1:361:ENRC" - config["variables"]["pulse_intensity_p80"]["pv"] = va_pv( - "GDET:FEE1:361:ENRCHSTCUHBR" - ) + config["variables"]["pulse_intensity_p80"]["pv"] = "GDET:FEE1:361:ENRCHSTCUHBR" - config["variables"]["pulse_intensity_std_relative"]["pv"] = va_pv( - "ZFEL:PULSE_INTENSITY_STD_REL" - ) + config["variables"]["pulse_intensity_std_relative"]["pv"] = ("ZFEL:PULSE_INTENSITY_STD_REL") return config @@ -272,7 +263,6 @@ def va_pv(name: str) -> str: def get_cu_hxr_zfel_runner( runner_cls, *, - prefix: str = "VA:", protocols: tuple[str, ...] = ("ca", "pva"), update_rate: float = 0.5, ): @@ -285,7 +275,6 @@ def get_cu_hxr_zfel_runner( config = build_cu_hxr_zfel_runner_config( runner_cls, model, - prefix=prefix, protocols=protocols, update_rate=update_rate, ) diff --git a/virtual_accelerator/tests/test_cu_hxr_zfel.py b/virtual_accelerator/tests/test_cu_hxr_zfel.py index 230f0f2..8d79390 100644 --- a/virtual_accelerator/tests/test_cu_hxr_zfel.py +++ b/virtual_accelerator/tests/test_cu_hxr_zfel.py @@ -132,7 +132,6 @@ def test_runner_config_uses_machine_style_pvs(): config = build_cu_hxr_zfel_runner_config( DummyRunner, model, - prefix="VA:", protocols=("ca",), update_rate=0.5, ) @@ -141,17 +140,16 @@ def test_runner_config_uses_machine_style_pvs(): assert config["protocol"] == ["ca"] assert config["update_rate"] == 0.5 - assert config["variables"]["KAct_14"]["pv"] == "VA:USEG:UNDH:1450:KAct" + assert config["variables"]["KAct_14"]["pv"] == "USEG:UNDH:1450:KAct" - assert config["variables"]["DSKAct_14"]["pv"] == "VA:USEG:UNDH:1450:DSKAct" + assert config["variables"]["DSKAct_14"]["pv"] == "USEG:UNDH:1450:DSKAct" - assert config["variables"]["pulse_energy"]["pv"] == "VA:ZFEL:PULSE_ENERGY" + assert config["variables"]["pulse_energy"]["pv"] == "ZFEL:PULSE_ENERGY" - assert config["variables"]["pulse_intensity_mean"]["pv"] == "VA:GDET:FEE1:361:ENRC" + assert config["variables"]["pulse_intensity_mean"]["pv"] == "GDET:FEE1:361:ENRC" assert ( - config["variables"]["pulse_intensity_p80"]["pv"] - == "VA:GDET:FEE1:361:ENRCHSTCUHBR" + config["variables"]["pulse_intensity_p80"]["pv"] == "GDET:FEE1:361:ENRCHSTCUHBR" ) @@ -163,6 +161,6 @@ def test_zfel_runner_factory_returns_configured_runner(): assert isinstance(runner, DummyRunner) - assert runner.config["variables"]["KAct_14"]["pv"] == "VA:USEG:UNDH:1450:KAct" + assert runner.config["variables"]["KAct_14"]["pv"] == "USEG:UNDH:1450:KAct" - assert runner.config["variables"]["pulse_energy"]["pv"] == "VA:ZFEL:PULSE_ENERGY" + assert runner.config["variables"]["pulse_energy"]["pv"] == "ZFEL:PULSE_ENERGY" From 71ae70d55e86947f52bfdf3b90cfa0da54f0e619 Mon Sep 17 00:00:00 2001 From: hong274SLAC Date: Thu, 10 Sep 2026 11:34:14 -0700 Subject: [PATCH 09/11] Added eval id Added VA model-evaluation synchronization --- virtual_accelerator/models/cu_hxr_zfel.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/virtual_accelerator/models/cu_hxr_zfel.py b/virtual_accelerator/models/cu_hxr_zfel.py index 7fb43dd..af388ca 100644 --- a/virtual_accelerator/models/cu_hxr_zfel.py +++ b/virtual_accelerator/models/cu_hxr_zfel.py @@ -45,6 +45,7 @@ def __init__(self): self._state: dict[str, Any] = {} self._variables = {} + self._state["model_eval_id"] = 0.0 # ------------------------------------------------------ # Scalar, real-machine-like undulator controls @@ -115,6 +116,12 @@ def __init__(self): unit="dimensionless", read_only=True, ), + "model_eval_id": ScalarVariable( + name="model_eval_id", + default_value=0.0, + unit="", + read_only=True, + ), } ) @@ -161,6 +168,8 @@ def _set(self, values: dict[str, Any]) -> None: self._sync_from_backend(include_controls=True) + self._state["model_eval_id"] += 1.0 + def _sync_from_backend( self, *, @@ -257,6 +266,7 @@ def build_cu_hxr_zfel_runner_config( config["variables"]["pulse_intensity_std_relative"]["pv"] = ("ZFEL:PULSE_INTENSITY_STD_REL") + config["variables"]["model_eval_id"]["pv"] = ("ZFEL:MODEL_EVAL_ID") return config From 300d6069f7c8a3aa8e75d6125ca79ddb78f547bd Mon Sep 17 00:00:00 2001 From: hong274SLAC Date: Fri, 11 Sep 2026 15:36:34 -0700 Subject: [PATCH 10/11] update zfel runner construction return the zfel model from cu_hxr_zfel and construct runner in the shared runners path. preserve the mode-specific PV config and machine-style PV names --- virtual_accelerator/models/cu_hxr_zfel.py | 25 ----------------------- virtual_accelerator/models/runners.py | 14 +++++++++---- 2 files changed, 10 insertions(+), 29 deletions(-) diff --git a/virtual_accelerator/models/cu_hxr_zfel.py b/virtual_accelerator/models/cu_hxr_zfel.py index af388ca..2895e2e 100644 --- a/virtual_accelerator/models/cu_hxr_zfel.py +++ b/virtual_accelerator/models/cu_hxr_zfel.py @@ -268,28 +268,3 @@ def build_cu_hxr_zfel_runner_config( config["variables"]["model_eval_id"]["pv"] = ("ZFEL:MODEL_EVAL_ID") return config - - -def get_cu_hxr_zfel_runner( - runner_cls, - *, - protocols: tuple[str, ...] = ("ca", "pva"), - update_rate: float = 0.5, -): - """ - Construct a lume-pva Runner for the CU HXR ZFEL virtual accelerator. - """ - - model = get_cu_hxr_zfel_model() - - config = build_cu_hxr_zfel_runner_config( - runner_cls, - model, - protocols=protocols, - update_rate=update_rate, - ) - - return runner_cls( - model=model, - config=config, - ) diff --git a/virtual_accelerator/models/runners.py b/virtual_accelerator/models/runners.py index 95f0fd2..51a2c7c 100644 --- a/virtual_accelerator/models/runners.py +++ b/virtual_accelerator/models/runners.py @@ -51,6 +51,7 @@ def main(): ) # Get the appropriate model based on user input + runner_kwargs = {} if args.model == "cu_hxr_bmad": from virtual_accelerator.models.cu_hxr import get_cu_hxr_bmad_model @@ -73,16 +74,21 @@ def main(): ) elif args.model == "cu_hxr_zfel": from virtual_accelerator.models.cu_hxr_zfel import ( - get_cu_hxr_zfel_runner, + get_cu_hxr_zfel_model, + build_cu_hxr_zfel_runner_config, + ) + + model = get_cu_hxr_zfel_model() + runner_kwargs["config"] = build_cu_hxr_zfel_runner_config( + Runner, + model, ) - runner = get_cu_hxr_zfel_runner(Runner) else: raise ValueError(f"Invalid model choice. Please choose one of {choices}.") # Run the model - if args.model != "cu_hxr_zfel": - runner = Runner(model) + runner = Runner(model=model, **runner_kwargs) runner.run() From 91d127511a8bf6c89089114c64eda408c11478a1 Mon Sep 17 00:00:00 2001 From: hong274SLAC Date: Fri, 11 Sep 2026 16:15:03 -0700 Subject: [PATCH 11/11] Expose machine PV names in ZFEL model add machine PV aliases to ZFELPVModel as suggested, while keeping runners construction same as before --- virtual_accelerator/models/cu_hxr_zfel.py | 136 ++++++++++++++-------- virtual_accelerator/models/runners.py | 12 +- 2 files changed, 92 insertions(+), 56 deletions(-) diff --git a/virtual_accelerator/models/cu_hxr_zfel.py b/virtual_accelerator/models/cu_hxr_zfel.py index 2895e2e..ddcabfe 100644 --- a/virtual_accelerator/models/cu_hxr_zfel.py +++ b/virtual_accelerator/models/cu_hxr_zfel.py @@ -44,6 +44,7 @@ def __init__(self): ) self._state: dict[str, Any] = {} + self._pv_aliases = {} self._variables = {} self._state["model_eval_id"] = 0.0 @@ -58,6 +59,12 @@ def __init__(self): self._state[kact_name] = float(kact[index]) self._state[dskact_name] = float(dskact[index]) + kact_pv = f"USEG:UNDH:{cell}50:KAct" + dskact_pv = f"USEG:UNDH:{cell}50:DSKAct" + + self._pv_aliases[kact_pv] = kact_name + self._pv_aliases[dskact_pv] = dskact_name + self._variables[kact_name] = ScalarVariable( name=kact_name, default_value=float(kact[index]), @@ -74,6 +81,22 @@ def __init__(self): read_only=False, ) + self._variables[kact_pv] = ScalarVariable( + name=kact_pv, + default_value=float(kact[index]), + value_range=(0.0, 5.0), + unit="dimensionless", + read_only=False, + ) + + self._variables[dskact_pv] = ScalarVariable( + name=dskact_pv, + default_value=float(dskact[index]), + value_range=(0.0, 5.0), + unit="dimensionless", + read_only=False, + ) + # ------------------------------------------------------ # Read-only FEL diagnostics # ------------------------------------------------------ @@ -127,12 +150,75 @@ def __init__(self): self._sync_from_backend(include_controls=True) + self._pv_aliases.update( + { + "GDET:FEE1:361:ENRC": "pulse_intensity_mean", + "GDET:FEE1:361:ENRCHSTCUHBR": "pulse_intensity_p80", + "ZFEL:POWER_MAX": "power_max", + "ZFEL:EXIT_POWER": "exit_power", + "ZFEL:PULSE_ENERGY": "pulse_energy", + "ZFEL:PULSE_INTENSITY_STD_REL": + "pulse_intensity_std_relative", + "ZFEL:MODEL_EVAL_ID": "model_eval_id", + } + ) + + self._variables.update( + { + "GDET:FEE1:361:ENRC": ScalarVariable( + name="GDET:FEE1:361:ENRC", + default_value=0.0, + unit="J", + read_only=True, + ), + "GDET:FEE1:361:ENRCHSTCUHBR": ScalarVariable( + name="GDET:FEE1:361:ENRCHSTCUHBR", + default_value=0.0, + unit="J", + read_only=True, + ), + "ZFEL:POWER_MAX": ScalarVariable( + name="ZFEL:POWER_MAX", + default_value=0.0, + unit="W", + read_only=True, + ), + "ZFEL:EXIT_POWER": ScalarVariable( + name="ZFEL:EXIT_POWER", + default_value=0.0, + unit="W", + read_only=True, + ), + "ZFEL:PULSE_ENERGY": ScalarVariable( + name="ZFEL:PULSE_ENERGY", + default_value=0.0, + unit="J", + read_only=True, + ), + "ZFEL:PULSE_INTENSITY_STD_REL": ScalarVariable( + name="ZFEL:PULSE_INTENSITY_STD_REL", + default_value=0.0, + unit="dimensionless", + read_only=True, + ), + "ZFEL:MODEL_EVAL_ID": ScalarVariable( + name="ZFEL:MODEL_EVAL_ID", + default_value=0.0, + unit="", + read_only=True, + ), + } + ) + @property def supported_variables(self): return self._variables - def _get(self, names: list[str]) -> dict[str, Any]: - return {name: self._state[name] for name in names} + def _get(self, names): + return { + name: self._state[self._pv_aliases.get(name, name)] + for name in names + } def _set(self, values: dict[str, Any]) -> None: """ @@ -147,7 +233,8 @@ def _set(self, values: dict[str, Any]) -> None: return for name, value in values.items(): - self._state[name] = float(value) + state_name = self._pv_aliases.get(name, name) + self._state[state_name] = float(value) kact = np.asarray( [self._state[f"KAct_{cell}"] for cell in HXR_CELLS], @@ -225,46 +312,3 @@ def get_cu_hxr_zfel_model() -> ZFELPVModel: Construct the CU HXR ZFEL virtual-accelerator model. """ return ZFELPVModel() - - -def build_cu_hxr_zfel_runner_config( - runner_cls, - model, - *, - protocols: tuple[str, ...] = ("ca", "pva"), - update_rate: float = 0.5, -): - """ - Build the lume-pva Runner configuration for the CU HXR ZFEL model. - """ - - config = runner_cls.generate_config( - model=model, - prefix="", - ) - - # PV names below already include the requested prefix. - config["prefix"] = "" - config["protocol"] = list(protocols) - config["update_rate"] = update_rate - - - for cell in HXR_CELLS: - config["variables"][f"KAct_{cell}"]["pv"] = f"USEG:UNDH:{cell}50:KAct" - - config["variables"][f"DSKAct_{cell}"]["pv"] = f"USEG:UNDH:{cell}50:DSKAct" - - config["variables"]["power_max"]["pv"] = "ZFEL:POWER_MAX" - - config["variables"]["exit_power"]["pv"] = "ZFEL:EXIT_POWER" - - config["variables"]["pulse_energy"]["pv"] = "ZFEL:PULSE_ENERGY" - - config["variables"]["pulse_intensity_mean"]["pv"] = "GDET:FEE1:361:ENRC" - - config["variables"]["pulse_intensity_p80"]["pv"] = "GDET:FEE1:361:ENRCHSTCUHBR" - - config["variables"]["pulse_intensity_std_relative"]["pv"] = ("ZFEL:PULSE_INTENSITY_STD_REL") - - config["variables"]["model_eval_id"]["pv"] = ("ZFEL:MODEL_EVAL_ID") - return config diff --git a/virtual_accelerator/models/runners.py b/virtual_accelerator/models/runners.py index 51a2c7c..ac8fd63 100644 --- a/virtual_accelerator/models/runners.py +++ b/virtual_accelerator/models/runners.py @@ -51,7 +51,6 @@ def main(): ) # Get the appropriate model based on user input - runner_kwargs = {} if args.model == "cu_hxr_bmad": from virtual_accelerator.models.cu_hxr import get_cu_hxr_bmad_model @@ -73,22 +72,15 @@ def main(): end_element=args.end_element, n_particles=args.n_particles ) elif args.model == "cu_hxr_zfel": - from virtual_accelerator.models.cu_hxr_zfel import ( - get_cu_hxr_zfel_model, - build_cu_hxr_zfel_runner_config, - ) + from virtual_accelerator.models.cu_hxr_zfel import get_cu_hxr_zfel_model model = get_cu_hxr_zfel_model() - runner_kwargs["config"] = build_cu_hxr_zfel_runner_config( - Runner, - model, - ) else: raise ValueError(f"Invalid model choice. Please choose one of {choices}.") # Run the model - runner = Runner(model=model, **runner_kwargs) + runner = Runner(model) runner.run()