Facet dev

.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

.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

.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

README.md

@@ -23,18 +23,22 @@ 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. |
 
 The package now lazily imports backend-specific dependencies. If you call a model
 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

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
+}

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",

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,

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,13 +34,54 @@ 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,
         start_element=start_element,
         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

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"]