Add covariance_matrix output to InjectorSurrogate

virtual_accelerator/surrogates/covariance.py

@@ -0,0 +1,41 @@
+from typing import Any, Mapping
+
+import numpy as np
+from scipy import constants
+
+
+def compute_covariance_matrix(state: Mapping[str, Any], energy: float) -> np.ndarray:
+    """Compute a diagonal 6x6 covariance matrix from scalar beam parameters.
+
+    The matrix is in trace space (x, x', y, y', z, z') with no off-diagonal terms.
+
+    Parameters
+    ----------
+    state : dict
+        Model output state containing XRMS, YRMS, sigma_z, norm_emit_x, norm_emit_y.
+    energy : float
+        Beam energy in eV, used to convert normalized emittance to geometric.
+
+    Returns
+    -------
+    np.ndarray
+        6x6 diagonal covariance matrix.
+    """
+    sigma_x = state["OTRS:IN20:571:XRMS"] * 1e-6  # microns -> meters
+    sigma_y = state["OTRS:IN20:571:YRMS"] * 1e-6
+    sigma_z = state["sigma_z"] * 1e-6
+
+    relativistic_gamma = energy / (
+        constants.value("electron mass energy equivalent in MeV") * 1e6
+    )
+    emit_x = state["norm_emit_x"] / relativistic_gamma  # geometric emittance
+    emit_y = state["norm_emit_y"] / relativistic_gamma
+
+    cov = np.zeros((6, 6))
+    cov[0, 0] = sigma_x**2
+    cov[2, 2] = sigma_y**2
+    cov[1, 1] = emit_x / cov[0, 0]
+    cov[3, 3] = emit_y / cov[2, 2]
+    cov[4, 4] = sigma_z**2
+    # cov[5, 5] is left as zero — energy spread not available from model
+    return cov

virtual_accelerator/surrogates/injector_surrogate.py

@@ -14,6 +14,8 @@
 import beamphysics
 from cheetah import ParticleBeam
 
+from virtual_accelerator.surrogates.covariance import compute_covariance_matrix
+
 
 OTR2_BEAM_ENERGY = 135.0e6  # eV
 
@@ -248,18 +250,28 @@ def _find_config(cls) -> Path:
             "'git subtree add --prefix subtrees/lcls_cu_injector_ml_model <remote> <ref>'."
         )
 
-    def __init__(self, n_particles: int = 10000) -> None:
+    def __init__(
+        self, n_particles: int = 10000, energy: float = OTR2_BEAM_ENERGY
+    ) -> None:
         """Initialize surrogate model and internal cache copy.
 
         Resource paths inside ``model_config.yaml`` are relative to the
         submodule directory.  A temporary config file with those paths
         rewritten to absolute paths is passed to ``TorchModel`` so that
         initialization succeeds regardless of the current working directory.
+
+        Parameters
+        ----------
+        n_particles : int, optional
+            Number of particles for beam distribution (default: 10000).
+        energy : float, optional
+            Beam energy in eV for covariance matrix calculation (default: OTR2_BEAM_ENERGY).
         """
         super().__init__()
         tm = self._load_torch_model()
         self.model = LUMETorchModel(tm)
         self.n_particles = n_particles
+        self.energy = energy
         self._cache: dict[str, Any] = {}
         self.set({})  # Initializing with defaults of NN model
         self.update_state()
@@ -326,6 +338,9 @@ def update_state(self):
         self._cache = {k: _to_python_scalar(v, k) for k, v in self._cache.items()}
         beam = create_beam_distribution_from_state(self._cache, self.n_particles)
         self._cache["output_beam"] = to_openpmd_particlegroup(beam)
+        self._cache["covariance_matrix"] = compute_covariance_matrix(
+            self._cache, self.energy
+        )
 
     @property
     def final_particles(self):