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..e4fed54 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/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", diff --git a/virtual_accelerator/models/cu_hxr_zfel.py b/virtual_accelerator/models/cu_hxr_zfel.py new file mode 100644 index 0000000..ddcabfe --- /dev/null +++ b/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() diff --git a/virtual_accelerator/models/runners.py b/virtual_accelerator/models/runners.py index fc978d6..ac8fd63 100644 --- a/virtual_accelerator/models/runners.py +++ b/virtual_accelerator/models/runners.py @@ -1,20 +1,23 @@ import argparse -from virtual_accelerator.models.facet2 import get_facet_staged_model from virtual_accelerator.utils.optional_dependencies import import_optional_symbol import logging def main(): - parser = argparse.ArgumentParser( - description="Run the CU HXR model with BMAD backend" - ) - choices = ["cu_hxr_bmad", "cu_hxr_staged", "facet_bmad", "facet_staged"] + 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, - 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,25 +50,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_model + + model = get_cu_hxr_zfel_model() + 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 new file mode 100644 index 0000000..8d79390 --- /dev/null +++ b/virtual_accelerator/tests/test_cu_hxr_zfel.py @@ -0,0 +1,166 @@ +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, + 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"] == "USEG:UNDH:1450:KAct" + + assert config["variables"]["DSKAct_14"]["pv"] == "USEG:UNDH:1450:DSKAct" + + assert config["variables"]["pulse_energy"]["pv"] == "ZFEL:PULSE_ENERGY" + + assert config["variables"]["pulse_intensity_mean"]["pv"] == "GDET:FEE1:361:ENRC" + + assert ( + config["variables"]["pulse_intensity_p80"]["pv"] == "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"] == "USEG:UNDH:1450:KAct" + + assert runner.config["variables"]["pulse_energy"]["pv"] == "ZFEL:PULSE_ENERGY" diff --git a/virtual_accelerator/tests/test_zfel_mapping.py b/virtual_accelerator/tests/test_zfel_mapping.py new file mode 100644 index 0000000..ea8a881 --- /dev/null +++ b/virtual_accelerator/tests/test_zfel_mapping.py @@ -0,0 +1,66 @@ +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) diff --git a/virtual_accelerator/tests/test_zfel_model.py b/virtual_accelerator/tests/test_zfel_model.py new file mode 100644 index 0000000..bb4d547 --- /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 ZFELBackend + + +def test_zfel_model_initialization(): + model = ZFELBackend() + + 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 = ZFELBackend() + + 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 = ZFELBackend() + + 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, + ) diff --git a/virtual_accelerator/zfel/__init__.py b/virtual_accelerator/zfel/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/virtual_accelerator/zfel/model.py b/virtual_accelerator/zfel/model.py new file mode 100644 index 0000000..1f3c3ea --- /dev/null +++ b/virtual_accelerator/zfel/model.py @@ -0,0 +1,215 @@ +from typing import Any + +import numpy as np + +from zfel import sase1d + +from virtual_accelerator.zfel.undulator_mapping import ( + MAGNETIC_LENGTH_M, + N_ACTIVE_SEGMENTS, + build_hxr_mapping, +) + +C_LIGHT = 299_792_458.0 + + +class ZFELBackend: + """ + 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() + + # 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() + + 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() diff --git a/virtual_accelerator/zfel/undulator_mapping.py b/virtual_accelerator/zfel/undulator_mapping.py new file mode 100644 index 0000000..254f69c --- /dev/null +++ b/virtual_accelerator/zfel/undulator_mapping.py @@ -0,0 +1,219 @@ +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, got {kact_array.size}." + ) + + if dskact_array.size != N_ACTIVE_SEGMENTS: + raise ValueError( + f"dskact must contain {N_ACTIVE_SEGMENTS} values, 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, + }