diff --git a/virtual_accelerator/surrogates/covariance.py b/virtual_accelerator/surrogates/covariance.py new file mode 100644 index 0000000..fb0459c --- /dev/null +++ b/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 diff --git a/virtual_accelerator/surrogates/injector_surrogate.py b/virtual_accelerator/surrogates/injector_surrogate.py index 562a2d1..ac5436c 100644 --- a/virtual_accelerator/surrogates/injector_surrogate.py +++ b/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 '." ) - 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):