diff --git a/.github/workflows/test-bmad.yml b/.github/workflows/test-bmad.yml index e80d702..02df149 100644 --- a/.github/workflows/test-bmad.yml +++ b/.github/workflows/test-bmad.yml @@ -1,7 +1,7 @@ name: Test bmad extra # Verifies the [bmad] optional extra installs and runs correctly. -# TestCUHXRBmad runs; cheetah, surrogate, and staged-model tests skip. +# BMAD-backed tests (including FACET2) run; cheetah and surrogate tests skip. on: push: @@ -33,6 +33,12 @@ jobs: repository: slaclab/lcls-lattice ref: temp_remove_fixers path: lcls-lattice + - name: Checkout facet2-lattice + uses: actions/checkout@v4 + with: + repository: slaclab/facet2-lattice + path: facet2-lattice + ref: temp_fix_fixer - name: Set up conda environment uses: conda-incubator/setup-miniconda@v3 @@ -58,7 +64,8 @@ jobs: - name: Verify bmad import works run: python -c "from lume_bmad.model import LUMEBmadModel; print('lume_bmad imported OK')" - - name: Run tests (TestCUHXRBmad runs; cheetah/surrogate/staged tests skip) + - name: Run tests (BMAD-backed tests including FACET2) env: LCLS_LATTICE: ${{ github.workspace }}/lcls-lattice + FACET2_LATTICE: ${{ github.workspace }}/facet2-lattice run: pytest virtual_accelerator/tests -v --tb=short diff --git a/.github/workflows/test-cheetah.yml b/.github/workflows/test-cheetah.yml index 03c6284..9d8eb56 100644 --- a/.github/workflows/test-cheetah.yml +++ b/.github/workflows/test-cheetah.yml @@ -25,7 +25,12 @@ jobs: repository: slaclab/lcls-lattice ref: temp_remove_fixers path: lcls-lattice - + - name: Checkout facet2-lattice + uses: actions/checkout@v4 + with: + repository: slaclab/facet2-lattice + path: facet2-lattice + ref: temp_fix_fixer - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v5 with: @@ -46,4 +51,5 @@ jobs: - name: Run tests (TestCUHXRCheetah runs; bmad/staged tests skip) env: LCLS_LATTICE: ${{ github.workspace }}/lcls-lattice + FACET2_LATTICE: ${{ github.workspace }}/facet2-lattice run: pytest virtual_accelerator/tests -v --tb=short diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 6d73550..03366e3 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -34,6 +34,12 @@ jobs: # temp change branch to address bmad bug TODO: should remove when bug is resolved ref: temp_remove_fixers path: lcls-lattice + - name: Checkout facet2-lattice + uses: actions/checkout@v4 + with: + repository: slaclab/facet2-lattice + path: facet2-lattice + ref: temp_fix_fixer - name: Set up conda environment uses: conda-incubator/setup-miniconda@v3 with: @@ -53,4 +59,5 @@ jobs: - name: Run tests env: LCLS_LATTICE: ${{ github.workspace }}/lcls-lattice + FACET2_LATTICE: ${{ github.workspace }}/facet2-lattice run: pytest virtual_accelerator/tests diff --git a/README.md b/README.md index 11f3a21..ed74775 100644 --- a/README.md +++ b/README.md @@ -23,9 +23,11 @@ pip install .[all] | Model / Factory Function | Optional dependency key(s) | Notes | | --- | --- | --- | | `get_cu_hxr_bmad_model` | `bmad` | Requires BMAD/PyTAO backend. | +| `get_facet_bmad_model` | `bmad` | FACET-II BMAD model; requires `FACET2_LATTICE`. | | `get_cu_hxr_cheetah_model` | `cheetah` | Requires Cheetah backend. | | `get_sc_diag0_cheetah_model` | `cheetah` | Requires Cheetah backend. | | `InjectorSurrogate` | `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. | | `virtual_accelerator.models.runners` CLI | `pva` (+ model backend key) | Runner requires `pva`; selected model backend must also be installed. | @@ -33,8 +35,10 @@ The package now lazily imports backend-specific dependencies. If you call a mode whose optional dependency is not installed, you will get an actionable error with the matching extra to install. -Creating the model instances requires the `$LCLS_LATTICE` environment variable to be set to a location containing the -contents of the lcls-lattice repo https://github.com/slaclab/lcls-lattice. +Creating model instances requires the `$LCLS_LATTICE` environment variable for LCLS-based models and +`$FACET2_LATTICE` for FACET-II models; each should point to a location containing the +contents of the lcls-lattice repo https://github.com/slaclab/lcls-lattice or the facet2-lattice +repo https://github.com/slaclab/facet2-lattice. #### Note diff --git a/examples/facet2_model_example.ipynb b/examples/facet2_model_example.ipynb new file mode 100644 index 0000000..c684787 --- /dev/null +++ b/examples/facet2_model_example.ipynb @@ -0,0 +1,113 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "3ecbbf9b", + "metadata": {}, + "outputs": [], + "source": [ + "from virtual_accelerator.models.facet2 import get_facet_bmad_model\n", + "\n", + "model = get_facet_bmad_model(\n", + " track_beam=True,\n", + " start_element=\"L0AFEND\",\n", + " end_element=\"PR10711\",\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2f5797bb", + "metadata": {}, + "outputs": [], + "source": [ + "variable_names = list(model.supported_variables.keys())\n", + "print(f\"Supported variables: {len(variable_names)}\")\n", + "print(\"Example variable names:\")\n", + "print(variable_names[:10])" + ] + }, + { + "cell_type": "markdown", + "id": "278c17bf", + "metadata": {}, + "source": [ + "## Predictions at elements" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6519be58", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "\n", + "out = model.get([\"s_ele\", \"a.beta\", \"b.beta\", \"e_tot\"])\n", + "plt.plot(out[\"s_ele\"], out[\"e_tot\"] / 1e6, label=\"b\")\n", + "\n", + "plt.figure()\n", + "plt.plot(out[\"s_ele\"], out[\"a.beta\"], label=\"a\")\n", + "plt.plot(out[\"s_ele\"], out[\"b.beta\"], label=\"b\")\n", + "plt.legend()" + ] + }, + { + "cell_type": "markdown", + "id": "ac1b293a", + "metadata": {}, + "source": [ + "## Predictions using comb" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0992b2d7", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "\n", + "out = model.get([\"s\", \"x.beta\", \"y.beta\"])\n", + "\n", + "plt.figure()\n", + "plt.plot(out[\"s\"], out[\"x.beta\"], label=\"x\")\n", + "plt.plot(out[\"s\"], out[\"y.beta\"], label=\"y\")\n", + "plt.legend()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "17fa073a", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "lume", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.12" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/pyproject.toml b/pyproject.toml index 5fbe38d..b7c1f65 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,7 +53,8 @@ pva = [ surrogate = [ "torch", "lume-torch @ git+https://github.com/lume-science/lume-torch", - "lume-cheetah @ git+https://github.com/lume-science/lume-cheetah" + "lume-cheetah @ git+https://github.com/lume-science/lume-cheetah", + "facet2_inj_ml_model @ git+https://github.com/slaclab/facet2_inj_ml_model", ] all = [ "torch", @@ -62,6 +63,7 @@ all = [ "lume-impact @ git+https://github.com/ChristopherMayes/lume-impact", "lume-torch @ git+https://github.com/lume-science/lume-torch", "lume-pva @ git+https://github.com/lume-science/lume-pva", + "facet2_inj_ml_model @ git+https://github.com/slaclab/facet2_inj_ml_model", ] dev = [ "pytest", diff --git a/virtual_accelerator/beams/2024-10-22_oneBunch.h5 b/virtual_accelerator/beams/2024-10-22_oneBunch.h5 new file mode 100644 index 0000000..851d6ed Binary files /dev/null and b/virtual_accelerator/beams/2024-10-22_oneBunch.h5 differ diff --git a/virtual_accelerator/beams/PR10241_impact_10000.h5 b/virtual_accelerator/beams/PR10241_impact_10000.h5 new file mode 100644 index 0000000..c3fb493 Binary files /dev/null and b/virtual_accelerator/beams/PR10241_impact_10000.h5 differ diff --git a/virtual_accelerator/beams/PR10241_impact_100000.h5 b/virtual_accelerator/beams/PR10241_impact_100000.h5 new file mode 100644 index 0000000..501d7d5 Binary files /dev/null and b/virtual_accelerator/beams/PR10241_impact_100000.h5 differ diff --git a/virtual_accelerator/bmad/factory.py b/virtual_accelerator/bmad/factory.py index a9df7a8..9849e82 100644 --- a/virtual_accelerator/bmad/factory.py +++ b/virtual_accelerator/bmad/factory.py @@ -34,6 +34,7 @@ def build_bmad_model( end_element: str, track_beam: bool, custom_beam_path: str | None, + custom_tao_commands: list[str] | None = None, ): """Build a lattice-specific LUMEBmadModel from a shared implementation.""" @@ -62,6 +63,10 @@ def build_bmad_model( init_file = os.path.join(lattice_root, spec.tao_init_relpath) tao = Tao(f"-init {init_file} -noplot -slice_lattice {start_element}:{end_element}") + if custom_tao_commands is not None: + for cmd in custom_tao_commands: + tao.cmd(cmd) + database_path = os.path.join(lattice_root, spec.database_relpath) control_name_to_element_name = get_epics_to_name_or_overlay_mapping( database_path, diff --git a/virtual_accelerator/models/facet2.py b/virtual_accelerator/models/facet2.py index 7e00734..0306dcb 100644 --- a/virtual_accelerator/models/facet2.py +++ b/virtual_accelerator/models/facet2.py @@ -1,16 +1,18 @@ +import os + from virtual_accelerator.bmad.factory import BmadModelSpec, build_bmad_model def get_facet_bmad_model( - start_element="PR10241", end_element="END", track_beam=False, custom_beam_path=None + start_element="L0AFEND", end_element="END", track_beam=False, custom_beam_path=None ): """ - Get the LUMEBmadModel for the FACET-II lattice from PRR10241 to END. + Get the LUMEBmadModel for the FACET-II lattice from L0AFEND to END. Parameters ------------- start_element: str, optional - The starting element for the model. Default is "PR10241". + The starting element for the model. Default is "L0AFEND". end_element: str, optional The ending element for the model. Default is "END". track_beam: bool, optional @@ -32,8 +34,8 @@ def get_facet_bmad_model( mapping_beampath=None, screens=("PR10571", "PR10711"), profmon_config_filename="facet2_profmon_info.yaml", - default_beam_relpath="bmad/bmad_set_beam2000_pg", - default_track_start="PR10241", + default_beam_relpath="beams/2024-10-22_oneBunch.h5", + default_track_start="L0AFEND", ) return build_bmad_model( spec=spec, @@ -41,4 +43,45 @@ def get_facet_bmad_model( end_element=end_element, track_beam=track_beam, custom_beam_path=custom_beam_path, + custom_tao_commands=["set bmad_com absolute_time_tracking=true"], + ) + + +def get_facet_staged_model(n_particles=10000, surrogate_inputs="machine", **kwargs): + """ + Get the StagedModel for the FACET-II lattice from PRR10241 to END, with an injector surrogate model. + + Parameters + ------------- + n_particles: int, optional + Number of particles to simulate in the surrogate model. Default is 1000. + surrogate_inputs: str, optional + Input for the surrogate model either "machine" or "sim". Default is "machine". + **kwargs: + Keyword arguments to be passed to the bmad LUMEModel instance as needed. + + Returns + ------- + StagedModel + Instance of the StagedModel for the FACET-II lattice. + """ + from facet2_inj_ml_model import load_model + from virtual_accelerator.surrogates.injector_surrogate import BeamOutputModel + from virtual_accelerator.models.facet2 import get_facet_bmad_model + from lume.staged_model import StagedModel + + injector_surrogate = BeamOutputModel( + load_model(surrogate_inputs), n_particles=n_particles ) + + # need to provide a beam distribution file to initialize the bmad model + fname = os.getcwd() + "/input_beam.h5" + injector_surrogate.final_particles.write(fname) + + facet_bmad_model = get_facet_bmad_model( + start_element="PR10241", track_beam=True, custom_beam_path=fname, **kwargs + ) + + staged_model = StagedModel([injector_surrogate, facet_bmad_model]) + + return staged_model diff --git a/virtual_accelerator/surrogates/injector_surrogate.py b/virtual_accelerator/surrogates/injector_surrogate.py index 562a2d1..a9fc3f0 100644 --- a/virtual_accelerator/surrogates/injector_surrogate.py +++ b/virtual_accelerator/surrogates/injector_surrogate.py @@ -117,7 +117,12 @@ class BeamOutputModel(LUMEModel, FinalParticlesMixIn): """ def __init__( - self, surrogate: TorchModel, n_particles: int = 10000, p0c: float = 1e8 + self, + surrogate: TorchModel, + n_particles: int = 10000, + p0c: float = 1e8, + t0: float = 0.0, + z0: float = 0.0, ) -> None: """ Initialize wrapper with surrogate model and internal cache copy. @@ -130,12 +135,18 @@ def __init__( The number of particles to generate in the output beam distribution (default: 10000). p0c: float, optional The reference momentum in eV/c to use for generating the output beam distribution (default: 1e8). + t0: float, optional + The reference time in seconds to use for generating the output beam distribution (default: 0.0). + z0: float, optional + The reference position in meters to use for generating the output beam distribution (default: 0.0). """ super().__init__() self.surrogate = LUMETorchModel(surrogate) self.n_particles = n_particles self.p0c = p0c + self.t0 = t0 + self.z0 = z0 self._cache: dict[str, Any] = {"output_beam": None} self.set({}) # Initializing with defaults of NN model self.update_state() @@ -174,14 +185,14 @@ def update_state(self): data = { "x": _tensor_to_numpy(particles[:, 0]), "y": _tensor_to_numpy(particles[:, 2]), - "z": _tensor_to_numpy(particles[:, 4]), + "t": _tensor_to_numpy(particles[:, 4] + self.t0), "px": _tensor_to_numpy(particles[:, 1]), "py": _tensor_to_numpy(particles[:, 3]), - "pz": _tensor_to_numpy(particles[:, 5]), - "t": 0.0, + "pz": _tensor_to_numpy(particles[:, 5] + self.p0c), + "z": self.z0, "weight": _tensor_to_numpy( torch.ones(self.n_particles) - ), # need to make at least 1d and negate + ), # need to make at least 1d "status": _tensor_to_numpy( torch.ones(self.n_particles, dtype=torch.int32) ), # need int @@ -191,7 +202,8 @@ def update_state(self): self._cache["output_beam"] = particle_group @property - def final_particles(self): + def final_particles(self) -> beamphysics.ParticleGroup: + """Return the final particle distribution as an openPMD ParticleGroup.""" return self._cache["output_beam"] diff --git a/virtual_accelerator/tests/test_facet.py b/virtual_accelerator/tests/test_facet.py index 206e236..f8a2c33 100644 --- a/virtual_accelerator/tests/test_facet.py +++ b/virtual_accelerator/tests/test_facet.py @@ -1,4 +1,5 @@ import os +from pathlib import Path import pytest from virtual_accelerator.tests._bmad_model_test_utils import ( @@ -10,7 +11,9 @@ assert_magnet_pvs_match_tao_lattice, assert_screen_image_pvs_in_supported_variables, ) -from virtual_accelerator.models.facet2 import get_facet_bmad_model +from virtual_accelerator.models.facet2 import ( + get_facet_bmad_model, +) HAS_FACET_LATTICE = bool(os.environ.get("FACET2_LATTICE")) @@ -24,6 +27,16 @@ class TestFACET2Bmad: def test_initialization(self): assert_bmad_model_initialization(get_facet_bmad_model) + def test_short_tracking(self): + get_facet_bmad_model( + track_beam=True, + start_element="PR10241", + end_element="PR10571", + custom_beam_path=os.path.join( + Path(__file__).parent, "../beams", "2024-10-22_oneBunch.h5" + ), + ) + def test_twiss(self): assert_bmad_model_twiss_outputs(get_facet_bmad_model) diff --git a/virtual_accelerator/tests/test_staged_model.py b/virtual_accelerator/tests/test_staged_model.py index 846822b..de129ba 100644 --- a/virtual_accelerator/tests/test_staged_model.py +++ b/virtual_accelerator/tests/test_staged_model.py @@ -10,12 +10,17 @@ "lume_torch", reason="requires lume-torch: pip install virtual-accelerator[surrogate]", ) +pytest.importorskip( + "facet2_inj_ml_model", + reason="requires facet2_inj_ml_model: pip install virtual-accelerator[surrogate]", +) from virtual_accelerator.models.staged_model import ( # noqa: E402 StagedModel, get_cu_hxr_staged_model, ) from virtual_accelerator.models.cu_hxr import get_cu_hxr_bmad_model # noqa: E402 +from virtual_accelerator.models.facet2 import get_facet_staged_model # noqa: E402 from virtual_accelerator.surrogates.injector_surrogate import InjectorSurrogate # noqa: E402 TEST_BEAM_PATH = os.path.join(Path(__file__).parent, "../bmad", "test_beam") @@ -33,6 +38,8 @@ def _has_module(name: str) -> bool: "lume_bmad", "cheetah", "lume_cheetah", + "lume_torch", + "facet2_inj_ml_model", ) ), reason="requires staged-model optional dependencies", @@ -162,3 +169,7 @@ def test_multiple_observable_get(self, staged_model): result = staged_model.get(["a.beta", "b.beta"]) assert "a.beta" in result assert "b.beta" in result + + def test_facet_model(self): + staged_model = get_facet_staged_model(end_element="PR10711") + staged_model.get(list(staged_model.supported_variables.keys()))