Add CU HXR ZFEL virtual accelerator backend

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

pyproject.toml

@@ -47,6 +47,9 @@ cheetah = [
 impact = [
     "lume-impact @ git+https://github.com/ChristopherMayes/lume-impact"
 ]
+zfel = [
+    "zfel @ git+https://github.com/slaclab/zfel.git"
+]
 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/slaclab/zfel.git",
 ]
 dev = [
     "pytest",

virtual_accelerator/models/cu_hxr_zfel.py

@@ -0,0 +1,314 @@
+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 ZFELBackend
+
+
+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 = ZFELBackend()
+
+        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._pv_aliases = {}
+        self._variables = {}
+        self._state["model_eval_id"] = 0.0
+
+        # ------------------------------------------------------
+        # 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])
+
+            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]),
+                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,
+            )
+
+            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
+        # ------------------------------------------------------
+
+        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,
+                ),
+                "model_eval_id": ScalarVariable(
+                    name="model_eval_id",
+                    default_value=0.0,
+                    unit="",
+                    read_only=True,
+                ),
+            }
+        )
+
+        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):
+        return {
+            name: self._state[self._pv_aliases.get(name, 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():
+            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],
+            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)
+
+        self._state["model_eval_id"] += 1.0
+
+    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()