From 64b83d350e0e53a2bc7229a56f70ed85734f8bda Mon Sep 17 00:00:00 2001 From: Gopika Bhardwaj Date: Thu, 21 May 2026 08:59:39 -0700 Subject: [PATCH 1/2] Add covariance_matrix output to InjectorSurrogate Derive a diagonal 6x6 covariance matrix from the scalar beam parameters (XRMS, YRMS, sigma_z, norm_emit_x, norm_emit_y) with no off-diagonal terms. Energy is a constructor parameter for normalized-to-geometric emittance conversion. --- .../surrogates/injector_surrogate.py | 52 ++++++++++++++++++- 1 file changed, 51 insertions(+), 1 deletion(-) diff --git a/virtual_accelerator/surrogates/injector_surrogate.py b/virtual_accelerator/surrogates/injector_surrogate.py index 562a2d1..d778181 100644 --- a/virtual_accelerator/surrogates/injector_surrogate.py +++ b/virtual_accelerator/surrogates/injector_surrogate.py @@ -79,6 +79,43 @@ def to_openpmd_particlegroup(beam) -> Any: return particle_group +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 + + def create_beam_distribution_from_state(state: Mapping[str, Any], n_particles: int): sigma_x = torch.tensor(state["OTRS:IN20:571:XRMS"] * 1e-6) sigma_y = torch.tensor(state["OTRS:IN20:571:YRMS"] * 1e-6) @@ -248,18 +285,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 +373,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): From 8ab072bbd7f239d2d0f9bbe6f44b1833b635f0f3 Mon Sep 17 00:00:00 2001 From: Gopika Bhardwaj Date: Thu, 21 May 2026 09:14:09 -0700 Subject: [PATCH 2/2] Moving to utility function --- virtual_accelerator/surrogates/covariance.py | 41 +++++++++++++++++++ .../surrogates/injector_surrogate.py | 39 +----------------- 2 files changed, 43 insertions(+), 37 deletions(-) create mode 100644 virtual_accelerator/surrogates/covariance.py 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 d778181..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 @@ -79,43 +81,6 @@ def to_openpmd_particlegroup(beam) -> Any: return particle_group -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 - - def create_beam_distribution_from_state(state: Mapping[str, Any], n_particles: int): sigma_x = torch.tensor(state["OTRS:IN20:571:XRMS"] * 1e-6) sigma_y = torch.tensor(state["OTRS:IN20:571:YRMS"] * 1e-6)