kparasch · thellert · Jul 12, 2026 · kparasch · Jul 21, 2026 · kparasch
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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