diff --git a/pySC/configuration/bpm_system_conf.py b/pySC/configuration/bpm_system_conf.py index f9a5452..2cc0a83 100644 --- a/pySC/configuration/bpm_system_conf.py +++ b/pySC/configuration/bpm_system_conf.py @@ -88,6 +88,29 @@ def configure_bpms(SC: SimulatedCommissioning) -> None: SC.bpm_system.initialize_empty_arrays() + # Dead / wrong-polarity BPM error model, applied per BPM category. + SC.bpm_system.dead = np.zeros(len(bpms_indices), dtype=bool) + + for bpm_category_name, cat_conf in bpms_conf.items(): + cat_indices = np.array([i for i, c in enumerate(bpms_categories) if c == bpm_category_name]) + if len(cat_indices) == 0: + continue + + fraction_dead = cat_conf.get('fraction_dead', 0.0) + if fraction_dead > 0: + n_dead = max(1, round(len(cat_indices) * fraction_dead)) + dead_idx = SC.rng.choice(len(cat_indices), size=n_dead, replace=False) + SC.bpm_system.dead[cat_indices[dead_idx]] = True + SC.tuning.bad_bpms.extend(cat_indices[dead_idx].tolist()) + + fraction_polarity = cat_conf.get('fraction_wrong_polarity', 0.0) + if fraction_polarity > 0: + n_flip = max(1, round(len(cat_indices) * fraction_polarity)) + flip_x = cat_indices[SC.rng.choice(len(cat_indices), size=n_flip, replace=False)] + flip_y = cat_indices[SC.rng.choice(len(cat_indices), size=n_flip, replace=False)] + SC.bpm_system.polarity_x[flip_x] = -1.0 + SC.bpm_system.polarity_y[flip_y] = -1.0 + for index, bpm_category in zip(bpms_indices, bpms_categories): generate_element_misalignments(SC, index, bpms_conf[bpm_category]) SC.bpm_system.update_rot_matrices() diff --git a/pySC/core/bpm_system.py b/pySC/core/bpm_system.py index 6d8e80c..918420f 100644 --- a/pySC/core/bpm_system.py +++ b/pySC/core/bpm_system.py @@ -17,7 +17,7 @@ def _rotation_matrix(a): 'reference_x', 'reference_y'] # These fields are initialized to ones (not zeros) — multiplicative corrections -BPM_FIELDS_TO_INITIALISE_ONES = ['gain_corrections_x', 'gain_corrections_y'] +BPM_FIELDS_TO_INITIALISE_ONES = ['gain_corrections_x', 'gain_corrections_y', 'polarity_x', 'polarity_y'] class BPMSystem(BaseModel, extra='forbid'): indices: list[int] = [] @@ -42,6 +42,10 @@ class BPMSystem(BaseModel, extra='forbid'): gain_corrections_x: NPARRAY = np.array([]) gain_corrections_y: NPARRAY = np.array([]) + dead: Optional[NPARRAY] = None + polarity_x: Optional[NPARRAY] = None + polarity_y: Optional[NPARRAY] = None + transmission_threshold: float = 0.4 _parent: Optional["SimulatedCommissioning"] = PrivateAttr(default=None) @@ -64,9 +68,21 @@ def initialize_empty_arrays(self): if not len(getattr(self, field)): # array is empty setattr(self, field, np.zeros(nbpm, dtype=float)) for field in BPM_FIELDS_TO_INITIALISE_ONES: - if not len(getattr(self, field)): # array is empty + value = getattr(self, field) + if value is None or not len(value): # optional/unset or empty array setattr(self, field, np.ones(nbpm, dtype=float)) + def _overwrite_dead_bpms(self, fake_x, fake_y, noise_x, noise_y): + """Overwrite dead-BPM readings with amplified noise, but only where the + beam actually reached the BPM. Dead BPMs downstream of a beam-loss point + keep their NaN reading instead of fabricating a non-NaN value that would + inflate transmission/reach metrics. Modifies fake_x/fake_y in place.""" + if self.dead is not None and self.dead.any(): + dead_alive_x = self.dead & ~np.isnan(fake_x) + dead_alive_y = self.dead & ~np.isnan(fake_y) + fake_x[dead_alive_x] = noise_x[dead_alive_x] * 10 + fake_y[dead_alive_y] = noise_y[dead_alive_y] * 10 + def update_rot_matrices(self): self._rot_matrices = _rotation_matrix(self.rolls) @@ -111,10 +127,13 @@ def capture_orbit(self, bba=True, subtract_reference=True, use_design=False) -> noise_x = self._parent.rng.normal(scale=self.noise_co_x) noise_y = self._parent.rng.normal(scale=self.noise_co_y) - fake_orbit_x = (rotated_orbit[0] - self.offsets_x) * (1 + self.calibration_errors_x) + noise_x - fake_orbit_y = (rotated_orbit[1] - self.offsets_y) * (1 + self.calibration_errors_y) + noise_y + pol_x = self.polarity_x if self.polarity_x is not None else 1.0 + pol_y = self.polarity_y if self.polarity_y is not None else 1.0 + fake_orbit_x = (rotated_orbit[0] - self.offsets_x) * (1 + self.calibration_errors_x) * pol_x + noise_x + fake_orbit_y = (rotated_orbit[1] - self.offsets_y) * (1 + self.calibration_errors_y) * pol_y + noise_y fake_orbit_x *= self.gain_corrections_x fake_orbit_y *= self.gain_corrections_y + self._overwrite_dead_bpms(fake_orbit_x, fake_orbit_y, noise_x, noise_y) if bba: # Apply BBA offsets @@ -154,6 +173,8 @@ def capture_injection(self, n_turns=1, bba=True, subtract_reference=True, use_de fake_trajectory_x_tbt = np.zeros([len(self.indices), n_turns]) fake_trajectory_y_tbt = np.zeros([len(self.indices), n_turns]) + pol_x = self.polarity_x if self.polarity_x is not None else 1.0 + pol_y = self.polarity_y if self.polarity_y is not None else 1.0 for n in range(n_turns): one_trajectory = trajectory[:, :, n] @@ -162,10 +183,11 @@ def capture_injection(self, n_turns=1, bba=True, subtract_reference=True, use_de noise_x = self._parent.rng.normal(scale=self.noise_tbt_x) noise_y = self._parent.rng.normal(scale=self.noise_tbt_y) - fake_trajectory_x = (rotated_trajectory[0] - self.offsets_x) * (1 + self.calibration_errors_x) + noise_x - fake_trajectory_y = (rotated_trajectory[1] - self.offsets_y) * (1 + self.calibration_errors_y) + noise_y + fake_trajectory_x = (rotated_trajectory[0] - self.offsets_x) * (1 + self.calibration_errors_x) * pol_x + noise_x + fake_trajectory_y = (rotated_trajectory[1] - self.offsets_y) * (1 + self.calibration_errors_y) * pol_y + noise_y fake_trajectory_x *= self.gain_corrections_x fake_trajectory_y *= self.gain_corrections_y + self._overwrite_dead_bpms(fake_trajectory_x, fake_trajectory_y, noise_x, noise_y) if bba: # Apply BBA offsets @@ -211,6 +233,8 @@ def capture_kick(self, n_turns=1, kick_px=0, kick_py=0, bba=True, subtract_refer fake_trajectory_x_tbt = np.zeros([len(self.indices), n_turns]) fake_trajectory_y_tbt = np.zeros([len(self.indices), n_turns]) + pol_x = self.polarity_x if self.polarity_x is not None else 1.0 + pol_y = self.polarity_y if self.polarity_y is not None else 1.0 for n in range(n_turns): one_trajectory = trajectory[:, :, n] @@ -219,10 +243,11 @@ def capture_kick(self, n_turns=1, kick_px=0, kick_py=0, bba=True, subtract_refer noise_x = self._parent.rng.normal(scale=self.noise_tbt_x) noise_y = self._parent.rng.normal(scale=self.noise_tbt_y) - fake_trajectory_x = (rotated_trajectory[0] - self.offsets_x) * (1 + self.calibration_errors_x) + noise_x - fake_trajectory_y = (rotated_trajectory[1] - self.offsets_y) * (1 + self.calibration_errors_y) + noise_y + fake_trajectory_x = (rotated_trajectory[0] - self.offsets_x) * (1 + self.calibration_errors_x) * pol_x + noise_x + fake_trajectory_y = (rotated_trajectory[1] - self.offsets_y) * (1 + self.calibration_errors_y) * pol_y + noise_y fake_trajectory_x *= self.gain_corrections_x fake_trajectory_y *= self.gain_corrections_y + self._overwrite_dead_bpms(fake_trajectory_x, fake_trajectory_y, noise_x, noise_y) if bba: # Apply BBA offsets diff --git a/pySC/core/rng.py b/pySC/core/rng.py index fe329ee..c66d15e 100644 --- a/pySC/core/rng.py +++ b/pySC/core/rng.py @@ -46,5 +46,8 @@ def normal(self, loc: float = 0, scale: float = 1, size: Optional[int] = None) - def uniform(self, low: float = 0, high: float = 1, size: Optional[int] = None) -> Union[float, np.ndarray]: return low + self._rng.random(size=size) * (high - low) + def choice(self, a, size=None, replace=True): + return self._rng.choice(a, size=size, replace=replace) + def randomize_rng(self) -> None: self._rng = default_rng() diff --git a/tests/configuration/test_bpm_system_conf.py b/tests/configuration/test_bpm_system_conf.py index 81988b1..123f119 100644 --- a/tests/configuration/test_bpm_system_conf.py +++ b/tests/configuration/test_bpm_system_conf.py @@ -190,3 +190,88 @@ def test_configure_bpms_multi_category_noise(hmba_lattice_file): f"BPM {name} should have cat_b orbit noise" assert SC.bpm_system.noise_tbt_x[i] == pytest.approx(5e-3), \ f"BPM {name} should have cat_b tbt noise" + + +# --------------------------------------------------------------------------- +# dead BPMs / wrong polarity (fraction_dead, fraction_wrong_polarity) +# --------------------------------------------------------------------------- + +# Single "standard" category matching all 10 HMBA BPMs. With n = 10: +# fraction_dead = 0.2 -> n_dead = max(1, round(10 * 0.2)) = 2 +# fraction_wrong_polarity = 0.3 -> n_flip = max(1, round(10 * 0.3)) = 3 +_FRACTION_DEAD = 0.2 +_FRACTION_POLARITY = 0.3 +_SEED = 42 + + +def _make_sc_with_dead_and_polarity(hmba_lattice_file): + """Fresh SC whose single BPM category configures dead + wrong-polarity fractions.""" + lattice = ATLattice(lattice_file=hmba_lattice_file, naming="FamName") + config = { + "error_table": {}, + "bpms": { + "standard": { + "regex": "^BPM", + "fraction_dead": _FRACTION_DEAD, + "fraction_wrong_polarity": _FRACTION_POLARITY, + }, + }, + } + return SimulatedCommissioning(lattice=lattice, configuration=config, seed=_SEED) + + +@pytest.mark.slow +@pytest.mark.regression +def test_configure_bpms_dead_count_and_membership(hmba_lattice_file): + """Dead-BPM count matches the formula and every dead BPM is in bad_bpms. + + Draw-order robust: asserts only on the total dead count and on set + membership of dead indices within SC.tuning.bad_bpms — never on which + specific indices were drawn. + """ + SC = _make_sc_with_dead_and_polarity(hmba_lattice_file) + configure_bpms(SC) + + n = len(SC.bpm_system.indices) + assert n == 10, "Expected all 10 HMBA BPMs in the single 'standard' category" + + expected_n_dead = max(1, round(n * _FRACTION_DEAD)) + dead_positions = np.flatnonzero(SC.bpm_system.dead) + assert dead_positions.size == expected_n_dead, \ + f"Expected {expected_n_dead} dead BPMs, got {dead_positions.size}" + + # Every dead BPM position must be recorded in tuning.bad_bpms (set membership). + bad_bpms = set(SC.tuning.bad_bpms) + for pos in dead_positions.tolist(): + assert pos in bad_bpms, f"Dead BPM position {pos} missing from tuning.bad_bpms" + + +@pytest.mark.slow +@pytest.mark.regression +def test_configure_bpms_wrong_polarity_counts(hmba_lattice_file): + """Exactly n_flip polarity_x and n_flip polarity_y entries are -1.0; rest +1.0. + + Draw-order robust: asserts only on the number of flipped (-1.0) entries and + that all remaining entries are +1.0 — never on which indices were flipped. + """ + SC = _make_sc_with_dead_and_polarity(hmba_lattice_file) + configure_bpms(SC) + + n = len(SC.bpm_system.indices) + expected_n_flip = max(1, round(n * _FRACTION_POLARITY)) + + polarity_x = np.asarray(SC.bpm_system.polarity_x) + polarity_y = np.asarray(SC.bpm_system.polarity_y) + + # Every entry is either +1.0 (unflipped) or -1.0 (flipped) — nothing else. + assert set(np.unique(polarity_x).tolist()) <= {-1.0, 1.0} + assert set(np.unique(polarity_y).tolist()) <= {-1.0, 1.0} + + assert np.count_nonzero(polarity_x == -1.0) == expected_n_flip, \ + f"Expected {expected_n_flip} flipped polarity_x entries" + assert np.count_nonzero(polarity_y == -1.0) == expected_n_flip, \ + f"Expected {expected_n_flip} flipped polarity_y entries" + + # All non-flipped entries equal +1.0. + assert np.count_nonzero(polarity_x == 1.0) == n - expected_n_flip + assert np.count_nonzero(polarity_y == 1.0) == n - expected_n_flip diff --git a/tests/core/test_bpm_system.py b/tests/core/test_bpm_system.py index 7f9d398..d7e09d5 100644 --- a/tests/core/test_bpm_system.py +++ b/tests/core/test_bpm_system.py @@ -4,6 +4,7 @@ from math import pi from pySC.core.bpm_system import BPMSystem, _rotation_matrix +from pySC.core.rng import RNG # --------------------------------------------------------------------------- @@ -352,3 +353,142 @@ def test_capture_injection_design_mode(sc): x_design2, y_design2 = bpm.capture_injection(n_turns=1, use_design=True) np.testing.assert_array_equal(x_design, x_design2) np.testing.assert_array_equal(y_design, y_design2) + + +# --------------------------------------------------------------------------- +# Dead-BPM alive-mask guard +# +# Regression coverage for the guard present in every capture method: +# +# if self.dead is not None and self.dead.any(): +# dead_alive_x = self.dead & ~np.isnan(fake_x) +# dead_alive_y = self.dead & ~np.isnan(fake_y) +# fake_x[dead_alive_x] = noise_x[dead_alive_x] * 10 +# fake_y[dead_alive_y] = noise_y[dead_alive_y] * 10 +# +# A dead BPM downstream of a beam-loss point already reads NaN (the beam never +# arrived) and MUST stay NaN — it must not be overwritten with fabricated +# ``noise * 10``, which would inflate transmission/reach metrics. A dead BPM +# upstream of the loss (where the beam did arrive) must still emit the +# amplified dead-BPM noise. +# --------------------------------------------------------------------------- + + +def test_capture_injection_dead_bpm_downstream_of_loss_stays_nan(sc, monkeypatch): + """Tracking path: a downstream dead BPM past a beam-loss point stays NaN. + + ``capture_injection`` obtains the mean trajectory from + ``lattice.track_mean`` and then applies the dead-BPM guard. On real beam + loss, ``track_mean`` returns NaN for every BPM whose transmission falls + below the threshold — i.e. every BPM downstream of the loss point. The + HMBA test lattice has no physical apertures, so a partial-NaN trajectory + cannot be produced by pushing amplitude; instead we craft the pre-noise + trajectory that ``track_mean`` hands back (finite up to the loss point, + NaN after it), which is exactly the shape real loss produces. The guard + under test lives in ``capture_injection``, and this exercises it on a + genuine partial-NaN trajectory. + """ + bpm = sc.bpm_system + n = len(bpm.indices) + + # Non-zero TbT noise so the amplified dead-BPM value is finite and non-zero. + bpm.noise_tbt_x = np.full(n, 1e-6) + bpm.noise_tbt_y = np.full(n, 1e-6) + + # Beam is lost at loss_idx: BPMs [0, loss_idx) see the beam (finite), + # BPMs [loss_idx, n) are downstream and read NaN. + loss_idx = n // 2 + 1 # 6 for the 10-BPM HMBA lattice + up_idx = loss_idx - 2 # dead BPM upstream of loss (finite reading) + down_idx = loss_idx + 2 # dead BPM downstream of loss (NaN reading) + pre_noise_val = 1.0 # absurd 1 m value: proves the overwrite fired + + def fake_track_mean(self, bunch, indices=None, n_turns=1, use_design=False, + coordinates=None, transmission_threshold=0): + # (coords, n_bpms, n_turns); NaN downstream of the loss point. + traj = np.full((2, n, n_turns), pre_noise_val) + traj[:, loss_idx:, :] = np.nan + transmission = np.zeros(n_turns) + return traj, transmission + + monkeypatch.setattr(type(sc.lattice), "track_mean", fake_track_mean) + + dead = np.zeros(n, dtype=bool) + dead[up_idx] = True + dead[down_idx] = True + bpm.dead = dead + + # Reproduce the guard's noise draw: capture_injection first calls + # generate_bunch (which consumes RNG via the per-shot injection jitter), + # then draws noise_x, then noise_y — mirror that exact order. + sc.rng = RNG(seed=2024) + x, y = bpm.capture_injection(n_turns=1, bba=False, subtract_reference=False) + + sc.rng = RNG(seed=2024) + sc.injection.generate_bunch() + noise_x = sc.rng.normal(scale=bpm.noise_tbt_x) + noise_y = sc.rng.normal(scale=bpm.noise_tbt_y) + + # Downstream dead BPM: beam never arrived -> stays NaN, NOT fabricated. + assert np.isnan(x[down_idx, 0]) + assert np.isnan(y[down_idx, 0]) + + # Upstream dead BPM: beam arrived -> emits amplified noise (noise * 10). + assert np.isfinite(x[up_idx, 0]) + assert np.isfinite(y[up_idx, 0]) + assert x[up_idx, 0] == noise_x[up_idx] * 10 + assert y[up_idx, 0] == noise_y[up_idx] * 10 + # The 1 m pre-noise value was overwritten by the (tiny) amplified noise. + assert abs(x[up_idx, 0]) < 1e-3 + + +def test_capture_orbit_dead_bpm_guard_nan_index_stays_nan(sc, monkeypatch): + """Guard unit test for ``capture_orbit``: a NaN orbit index stays NaN. + + NOTE: this is a UNIT test of the dead-BPM guard, not a physical + closed-orbit claim. ``capture_orbit`` sources its orbit from + ``at.find_orbit`` (via ``lattice.get_orbit``), which is all-or-nothing — + a real closed orbit is never partially NaN. We therefore monkeypatch + ``get_orbit`` to inject a NaN at one index purely to drive the guard, and + assert the guard keeps that index NaN while a finite index still receives + amplified dead-BPM noise. + """ + bpm = sc.bpm_system + n = len(bpm.indices) + + bpm.noise_co_x = np.full(n, 1e-6) + bpm.noise_co_y = np.full(n, 1e-6) + + nan_idx = n - 2 # "downstream" index whose crafted orbit reads NaN + finite_idx = 2 # index with a finite orbit reading + orbit_val = 5e-4 # distinctive finite pre-noise orbit value + + def fake_get_orbit(self, indices=None, use_design=False): + # get_orbit returns shape (2, n_bpms): (x, y) per BPM. + orbit = np.full((2, n), orbit_val) + orbit[:, nan_idx] = np.nan + return orbit + + monkeypatch.setattr(type(sc.lattice), "get_orbit", fake_get_orbit) + + dead = np.zeros(n, dtype=bool) + dead[nan_idx] = True + dead[finite_idx] = True + bpm.dead = dead + + # capture_orbit draws noise_x then noise_y (no bunch generation). + sc.rng = RNG(seed=2024) + fake_x, fake_y = bpm.capture_orbit(bba=False, subtract_reference=False) + + sc.rng = RNG(seed=2024) + noise_x = sc.rng.normal(scale=bpm.noise_co_x) + noise_y = sc.rng.normal(scale=bpm.noise_co_y) + + # Dead BPM at the NaN index: stays NaN, not fabricated. + assert np.isnan(fake_x[nan_idx]) + assert np.isnan(fake_y[nan_idx]) + + # Dead BPM at a finite index: overwritten with amplified noise (noise * 10). + assert np.isfinite(fake_x[finite_idx]) + assert np.isfinite(fake_y[finite_idx]) + assert fake_x[finite_idx] == noise_x[finite_idx] * 10 + assert fake_y[finite_idx] == noise_y[finite_idx] * 10