Made BPM readings more realistic (#36)

* Introduced charge scaling of BPM noise
        When averaged, orbit is also weighted by charge

* Function returns also SUM signal

* Option to calculate trajectory measurement errors
        Weighted RMS when all shots reached given BPM

* Optimised memory when fraction of BPMs is used

pySC/core/beam.py

@@ -11,7 +11,7 @@
 from numpy import ndarray
 
 from pySC.core.simulated_commissioning import SimulatedCommissioning
-from pySC.core.constants import TRACK_ORB, TRACK_PORB
+from pySC.core.constants import TRACK_ORB, TRACK_PORB, TRACK_TBT
 from pySC.utils.sc_tools import SCrandnc
 from pySC.utils.at_wrapper import atgetfieldvalues, atpass, findorbit6, findspos
 import warnings
@@ -21,7 +21,7 @@
 LOGGER = logging_tools.get_logger(__name__)
 
 
-def bpm_reading(SC: SimulatedCommissioning, bpm_ords: ndarray = None) -> ndarray:
+def bpm_reading(SC: SimulatedCommissioning, bpm_ords: ndarray = None, calculate_errors: bool = False) -> Tuple:
     """
     Calculates BPM readings with current injection setup `SC.INJ` and included all BPM uncertainties.
     Included uncertainties are offsets, rolls, calibration errors, and position noise.
@@ -33,25 +33,46 @@ def bpm_reading(SC: SimulatedCommissioning, bpm_ords: ndarray = None) -> ndarray
         SC: SimulatedCommissioning instance
         bpm_ords: array of element indices of registered BPMs for which to calculate readings
             (for convenience, otherwise `SC.ORD.BPM` is used)
+        calculate_errors: If true orbit errors are calculated
 
     Returns:
         Array of horizontal and vertical BPM readings (2, T x B) for T turns and B BPMs
+        Array of fractional transmission (eq. BPM sum signal) (2, T x B) for T turns and B BPMs
+        Optionally, in TBT mode and calculate_errors is True:
+            Array of measured errors for horizontal and vertical BPM readings (2, T x B) for T turns and B BPMs
     """
-    all_bpm_orbits_4d = np.full((2, len(SC.ORD.BPM), SC.INJ.nTurns, SC.INJ.nShots), np.nan)
+    bpm_inds = np.arange(len(SC.ORD.BPM), dtype=int) if bpm_ords is None else _only_registered_bpms(SC, bpm_ords)
+    bpm_orbits_4d = np.full((2, len(bpm_inds), SC.INJ.nTurns, SC.INJ.nShots), np.nan)
+    bpm_sums_4d = np.full((2, len(bpm_inds), SC.INJ.nTurns, SC.INJ.nShots), np.nan)
     for shot_num in range(SC.INJ.nShots):
-        tracking_4d = _tracking(SC, SC.ORD.BPM)
-        all_bpm_orbits_4d[:, :, :, shot_num] = _real_bpm_reading(SC, tracking_4d)
+        tracking_4d = _tracking(SC, SC.ORD.BPM[bpm_inds])
+        bpm_orbits_4d[:, :, :, shot_num], bpm_sums_4d[:, :, :, shot_num] = _real_bpm_reading(SC, tracking_4d, bpm_inds)
+
+    # mean_bpm_orbits_3d is 3D (dim, BPM, turn)
+    mean_bpm_orbits_3d = np.average(np.ma.array(bpm_orbits_4d, mask=np.isnan(bpm_orbits_4d)),
+                                    weights=np.ma.array(bpm_sums_4d, mask=np.isnan(bpm_sums_4d)), axis=3)
+    # averaging "charge" also when the beam did not reach the location
+    mean_bpm_sums_3d = np.nansum(bpm_sums_4d, axis=3) / SC.INJ.nShots
 
-    mean_bpm_orbits_3d = np.nanmean(all_bpm_orbits_4d, axis=3)  # mean_bpm_orbits_3d is 3D (dim, BPM, turn)
     if SC.plot:
-        _plot_bpm_reading(SC, mean_bpm_orbits_3d)
+        _plot_bpm_reading(SC, mean_bpm_orbits_3d, bpm_inds)
     if SC.INJ.trackMode == TRACK_PORB:   # ORB averaged over low amount of turns
-        mean_bpm_orbits_3d = np.nanmean(mean_bpm_orbits_3d, axis=2, keepdims=True)
-    if bpm_ords is not None:
-        ind = _only_registered_bpms(SC, bpm_ords)
-        mean_bpm_orbits_3d = mean_bpm_orbits_3d[:, ind, :]
-    # Organising the array 2 x (nturns x nbpms) sorted by "arrival time"
-    return _reshape_3d_to_matlab_like_2d(mean_bpm_orbits_3d)
+        mean_bpm_orbits_3d = np.average(np.ma.array(mean_bpm_orbits_3d, mask=np.isnan(mean_bpm_orbits_3d)),
+                                        weights=np.ma.array(mean_bpm_sums_3d, mask=np.isnan(mean_bpm_sums_3d)), axis=2,
+                                        keepdims=True)
+        mean_bpm_sums_3d = np.nansum(mean_bpm_sums_3d, axis=2, keepdims=True) / SC.INJ.nTurns
+    if calculate_errors and SC.INJ.trackMode == TRACK_TBT:
+        bpm_orbits_4d[np.sum(np.isnan(bpm_orbits_4d), axis=3) > 0, :] = np.nan
+        squared_orbit_diffs = np.square(bpm_orbits_4d - mean_bpm_orbits_3d)
+        err_bpm_orbits_3d = np.sqrt(np.average(np.ma.array(squared_orbit_diffs), mask=np.isnan(bpm_orbits_4d),
+                                   weights=np.ma.array(bpm_sums_4d, mask=np.isnan(bpm_orbits_4d)), axis=3))
+        # Organising the array 2 x (nturns x nbpms) sorted by "arrival time"
+        # TODO keep in 3D when the response matrices are changed
+        return (_reshape_3d_to_matlab_like_2d(mean_bpm_orbits_3d),
+                _reshape_3d_to_matlab_like_2d(mean_bpm_sums_3d),
+                _reshape_3d_to_matlab_like_2d(err_bpm_orbits_3d))
+    return (_reshape_3d_to_matlab_like_2d(mean_bpm_orbits_3d),
+            _reshape_3d_to_matlab_like_2d(mean_bpm_sums_3d))
 
 
 def all_elements_reading(SC: SimulatedCommissioning) -> Tuple[ndarray, ndarray]:
@@ -75,13 +96,13 @@ def all_elements_reading(SC: SimulatedCommissioning) -> Tuple[ndarray, ndarray]:
     all_bpm_orbits_4d = np.full((2, len(SC.ORD.BPM), SC.INJ.nTurns, SC.INJ.nShots), np.nan)
     for shot_num in range(SC.INJ.nShots):
         tracking_4d = _tracking(SC, np.arange(n_refs))
-        all_bpm_orbits_4d[:, :, :, shot_num] = _real_bpm_reading(SC, tracking_4d[:, :, SC.ORD.BPM, :])
+        all_bpm_orbits_4d[:, :, :, shot_num] = _real_bpm_reading(SC, tracking_4d[:, :, SC.ORD.BPM, :])[0]
         tracking_4d[:, np.isnan(tracking_4d[0, :])] = np.nan
         all_readings_5d[:, :, :, :, shot_num] = tracking_4d[:, :, :, :]
 
     mean_bpm_orbits_3d = np.nanmean(all_bpm_orbits_4d, axis=3)  # mean_bpm_orbits_3d is 3D (dim, BPM, turn)
     if SC.plot:
-        _plot_bpm_reading(SC, mean_bpm_orbits_3d, all_readings_5d)
+        _plot_bpm_reading(SC, mean_bpm_orbits_3d, None, all_readings_5d)
     return _reshape_3d_to_matlab_like_2d(mean_bpm_orbits_3d), all_readings_5d
 
 
@@ -159,21 +180,25 @@ def plot_transmission(ax, fraction_survived, n_turns, beam_lost_at):
     return ax
 
 
-def _real_bpm_reading(SC, track_mat):  # track_mat should be only x,y over all particles only at BPM positions
-    nBpms, nTurns = track_mat.shape[2:]
-    bpm_noise = np.transpose(atgetfieldvalues(SC.RING, SC.ORD.BPM, ('NoiseCO' if SC.INJ.trackMode == 'ORB' else "Noise")))
-    bpm_noise = bpm_noise[:, :, np.newaxis] * SCrandnc(2, (2, nBpms, nTurns))
-    bpm_offset = np.transpose(atgetfieldvalues(SC.RING, SC.ORD.BPM, 'Offset') + atgetfieldvalues(SC.RING, SC.ORD.BPM, 'SupportOffset'))
-    bpm_cal_error = np.transpose(atgetfieldvalues(SC.RING, SC.ORD.BPM, 'CalError'))
-    bpm_roll = np.squeeze(atgetfieldvalues(SC.RING, SC.ORD.BPM, 'Roll') + atgetfieldvalues(SC.RING, SC.ORD.BPM, 'SupportRoll'))
-    bpm_sum_error = np.transpose(atgetfieldvalues(SC.RING, SC.ORD.BPM, 'SumError'))[:, np.newaxis] * SCrandnc(2, (nBpms, nTurns))
+def _real_bpm_reading(SC, track_mat, bpm_inds=None):  # track_mat should be only x,y over all particles only at BPM positions
+    n_bpms, nTurns = track_mat.shape[2:]
+    bpm_ords = SC.ORD.BPM if bpm_inds is None else SC.ORD.BPM[bpm_inds]
+    bpm_noise = np.transpose(atgetfieldvalues(SC.RING, bpm_ords, ('NoiseCO' if SC.INJ.trackMode == 'ORB' else "Noise")))
+    bpm_noise = bpm_noise[:, :, np.newaxis] * SCrandnc(2, (2, n_bpms, nTurns))
+    bpm_offset = np.transpose(atgetfieldvalues(SC.RING, bpm_ords, 'Offset') + atgetfieldvalues(SC.RING, bpm_ords, 'SupportOffset'))
+    bpm_cal_error = np.transpose(atgetfieldvalues(SC.RING, bpm_ords, 'CalError'))
+    bpm_roll = np.squeeze(atgetfieldvalues(SC.RING, bpm_ords, 'Roll') + atgetfieldvalues(SC.RING, bpm_ords, 'SupportRoll'), axis=1)
+    bpm_sum_error = np.transpose(atgetfieldvalues(SC.RING, bpm_ords, 'SumError'))[:, np.newaxis] * SCrandnc(2, (n_bpms, nTurns))
     # averaging the X and Y positions at BPMs over particles
     mean_orbit = np.nanmean(track_mat, axis=1)
-    beam_lost = np.nonzero(np.mean(np.isnan(track_mat[0, :, :, :]), axis=0) * (1 + bpm_sum_error) > SC.INJ.beamLostAt)
+    transmission = np.mean(~np.isnan(track_mat[0, :, :, :]), axis=0) * (1 + bpm_sum_error)
+    beam_lost = np.nonzero(transmission < 1 - SC.INJ.beamLostAt)
     mean_orbit[:, beam_lost[0], beam_lost[1]] = np.nan
+    transmission[beam_lost[0], beam_lost[1]] = np.nan
 
     rolled_mean_orbit = np.einsum("ijk,jkl->ikl", _rotation_matrix(bpm_roll), mean_orbit)
-    return (rolled_mean_orbit - bpm_offset[:, :, np.newaxis]) * (1 + bpm_cal_error[:, :, np.newaxis]) + bpm_noise
+    return ((rolled_mean_orbit - bpm_offset[:, :, np.newaxis]) * (1 + bpm_cal_error[:, :, np.newaxis]) +
+            bpm_noise / transmission), np.tile(transmission, (2, 1, 1))
 
 
 def _rotation_matrix(a):
@@ -186,14 +211,19 @@ def _tracking(SC: SimulatedCommissioning, refs: ndarray) -> ndarray:
     #  lattice_pass output:            (6, N, R, T) coordinates of N particles at R reference points for T turns.
     #  findorbit second output value:  (R, 6) closed orbit vector at each specified location
     if SC.INJ.trackMode == TRACK_ORB:
-        return np.transpose(findorbit6(SC.RING, refs, keep_lattice=False)[1])[[0, 2], :].reshape(2, 1, len(refs), 1)
-    return atpass(SC.RING, generate_bunches(SC), SC.INJ.nTurns, refs, keep_lattice=False)[[0, 2], :, :, :]
+        pos = np.transpose(findorbit6(SC.RING, refs, keep_lattice=False)[1])[[0, 2], :].reshape(2, 1, len(refs), 1)
+    else:
+        pos = atpass(SC.RING, generate_bunches(SC), SC.INJ.nTurns, refs, keep_lattice=False)[[0, 2], :, :, :]
+    pos[1, np.isnan(pos[0, :, :, :])] = np.nan
+    return pos
 
 
 def _only_registered_bpms(SC: SimulatedCommissioning, bpm_ords: ndarray) -> ndarray:
     ind = np.where(np.isin(SC.ORD.BPM, bpm_ords))[0]
     if len(ind) != len(bpm_ords):
         LOGGER.warning('Not all specified ordinates are registered BPMs.')
+    if not len(ind):
+        raise ValueError("No registered BPMs specified.")
     return ind
 
 
@@ -202,17 +232,17 @@ def _reshape_3d_to_matlab_like_2d(mean_bpm_orbits_3d: ndarray) -> ndarray:
     return np.transpose(mean_bpm_orbits_3d, axes=(0, 2, 1)).reshape((2, np.prod(mean_bpm_orbits_3d.shape[1:])))
 
 
-def _plot_bpm_reading(SC, bpm_orbits_3d, all_readings_5d=None):
+def _plot_bpm_reading(SC, bpm_orbits_3d, bpm_inds=None, all_readings_5d=None):
     ap_ords, apers = _get_ring_aperture(SC)
     fig, ax = plt.subplots(num=1, nrows=2, ncols=1, sharex="all", figsize=(8, 6), dpi=100, facecolor="w")
     s_pos = findspos(SC.RING)
     circumference = s_pos[-1]
-
+    bpms = SC.ORD.BPM if bpm_inds is None else SC.ORD.BPM[bpm_inds]
     if all_readings_5d is not None:
         x = np.ravel(np.arange(SC.INJ.nTurns)[:, np.newaxis] * circumference + s_pos)
         ax = _plot_all_trajectories(ax, x, all_readings_5d)
     for n_dim in range(2):
-        x = np.ravel(np.arange(SC.INJ.nTurns)[:, np.newaxis] * circumference + s_pos[SC.ORD.BPM])
+        x = np.ravel(np.arange(SC.INJ.nTurns)[:, np.newaxis] * circumference + s_pos[bpms])
         y = 1E3 * np.ravel(bpm_orbits_3d[n_dim, :, :].T)
         ax[n_dim].plot(x, y, 'ro', label="BPM reading")
         if len(ap_ords):

pySC/correction/injection_fit.py

@@ -14,7 +14,7 @@ def fit_injection_drift(SC: SimulatedCommissioning, n_dims: ndarray = np.array([
     # uses SC.INJ.nShots
     s_pos = findspos(SC.RING)
     bpm_inds = np.arange(2, dtype=int) + len(SC.ORD.BPM) - 1
-    bref = bpm_reading(SC)[:, bpm_inds]
+    bref = bpm_reading(SC)[0][:, bpm_inds]
     s_bpm = np.array([s_pos[SC.ORD.BPM[-1]] - s_pos[-1], s_pos[SC.ORD.BPM[0]]])
     delta_z0 = _fit_bpm_data(s_bpm, bref)
     if np.sum(np.isnan(delta_z0)):
@@ -37,7 +37,7 @@ def fit_injection_drift(SC: SimulatedCommissioning, n_dims: ndarray = np.array([
 def fit_injection_trajectory(SC: SimulatedCommissioning, bpm_inds: ndarray = np.array([0, 1, 2]), plot: bool = False):
     # uses SC.INJ.nShots
     s_pos = findspos(SC.RING)
-    bref = bpm_reading(SC)[:, bpm_inds]
+    bref = bpm_reading(SC)[0][:, bpm_inds]
     delta_z0 = np.zeros(6)
     delta_z0[0:4] = -fmin(_merit_function, np.zeros(4), args=(SC, bref, SC.ORD.BPM[bpm_inds]))
     if np.sum(np.isnan(delta_z0)):
@@ -47,7 +47,7 @@ def fit_injection_trajectory(SC: SimulatedCommissioning, bpm_inds: ndarray = np.
     if plot:
         fig, ax = plt.subplots(nrows=2, num=342, sharex="all")
         SC.INJ.Z0 += delta_z0
-        bnew = bpm_reading(SC)[:, bpm_inds]
+        bnew = bpm_reading(SC)[0][:, bpm_inds]
         SC.INJ.Z0 -= delta_z0
         s_bpm = s_pos[SC.ORD.BPM[bpm_inds]]
         for n_dim in range(2):

pySC/correction/orbit_trajectory.py

@@ -192,7 +192,7 @@ def _check_ords(SC, Mplus, reference, BPMords, CMords):
 
 
 def _bpm_reading_and_logging(SC, BPMords, ind_history=None, orb_history=None):
-    bpm_readings = bpm_reading(SC, bpm_ords=BPMords)
+    bpm_readings = bpm_reading(SC, bpm_ords=BPMords)[0]
     bpms_reached = ~np.isnan(bpm_readings[0])
     if ind_history is None or orb_history is None:
         return bpm_readings, [np.sum(bpms_reached)], [np.sqrt(np.mean(np.square(bpm_readings[:, bpms_reached]), axis=1))]

pySC/correction/rf.py

@@ -116,7 +116,7 @@ def _check_data(test_vec, bpm_shift, bpm_shift_err, min_num_points=6):
 
 
 def _get_tbt_energy_shift(SC, bpm_ords):
-    bpm_readings = bpm_reading(SC, bpm_ords)
+    bpm_readings = bpm_reading(SC, bpm_ords)[0]
     x_reading = np.reshape(bpm_readings[0, :], (SC.INJ.nTurns, len(bpm_ords)))
     mean_tbt = np.mean(x_reading - x_reading[0, :], axis=1)
     mean_tbt_err = np.std(x_reading, axis=1) / np.sqrt(x_reading.shape[1])

pySC/example.py

@@ -142,11 +142,11 @@ def _marker(name):
     # RF cavity correction
     SC.INJ.nTurns = 5
     for nIter in range(2):
-        SC = SCsynchPhaseCorrection(SC, nSteps=25, plotResults=False, plotProgress=False)
+        SC = correct_rf_phase(SC, n_steps=25, plot_results=False, plot_progress=False)
 
-        SC = SCsynchEnergyCorrection(SC, f_range=40E3 * np.array([-1, 1]),  # Frequency range [kHz]
-                                         nSteps=15,  # Number of frequency steps
-                                         plotResults=False, plotProgress=False)
+        SC = correct_rf_frequency(SC, f_range=40E3 * np.array([-1, 1]),  # Frequency range [kHz]
+                                         n_steps=15,  # Number of frequency steps
+                                         plot_results=False, plot_progress=False)
 
     # Plot phasespace after RF correction
     plot_phase_space(SC, nParticles=10, nTurns=100)
@@ -168,8 +168,8 @@ def _marker(name):
             CUR = SCfeedbackRun(SC, MinvCO, target=0, maxsteps=50, scaleDisp=1E8)
         except RuntimeError:
             break
-        B0rms = np.sqrt(np.mean(np.square(bpm_reading(SC)), axis=1))
-        Brms = np.sqrt(np.mean(np.square(bpm_reading(CUR)), axis=1))
+        B0rms = np.sqrt(np.mean(np.square(bpm_reading(SC)[0]), axis=1))
+        Brms = np.sqrt(np.mean(np.square(bpm_reading(CUR)[0]), axis=1))
         if np.mean(B0rms) < np.mean(Brms):
             break
         SC = CUR

pySC/lattice_properties/response_measurement.py

@@ -22,7 +22,7 @@ def response_matrix(SC, amp, bpm_ords, cm_ords, mode='fixedKick', n_steps=2, fit
     rm = np.full((2 * SC.INJ.nTurns * bpm_ords.shape[0], n_hcm + n_vcm), np.nan)
     error = np.full((2 * SC.INJ.nTurns * bpm_ords.shape[0], n_hcm + n_vcm), np.nan)
     cm_steps = [np.zeros((n_steps, n_hcm)), np.zeros((n_steps, n_vcm))]
-    bref = np.ravel(bpm_reading(SC, bpm_ords=bpm_ords))
+    bref = np.ravel(bpm_reading(SC, bpm_ords=bpm_ords)[0])
     if SC.INJ.trackMode == 'ORB' and np.sum(np.isnan(bref)):
         raise ValueError('No closed orbit found.')
     i = 0
@@ -38,7 +38,7 @@ def response_matrix(SC, amp, bpm_ords, cm_ords, mode='fixedKick', n_steps=2, fit
                     if cm_step_vec[n_step] != 0 and cm_step_vec[n_step] != max_step:
                         SC = set_cm_setpoints(SC, cm_ords[n_dim][nCM], cmstart[nCM] + cm_step_vec[n_step], skewness=bool(n_dim))
                         real_cm_setpoint[n_step] = get_cm_setpoints(SC, cm_ords[n_dim][nCM], skewness=bool(n_dim))
-                        gradient[n_step, :] = np.ravel(bpm_reading(SC, bpm_ords=bpm_ords)) - bref
+                        gradient[n_step, :] = np.ravel(bpm_reading(SC, bpm_ords=bpm_ords)[0]) - bref
                 dCM = real_cm_setpoint - cmstart[nCM]
                 cm_steps[n_dim][:, nCM] = dCM
                 rm[:, i] = gradient / dCM
@@ -59,10 +59,10 @@ def response_matrix(SC, amp, bpm_ords, cm_ords, mode='fixedKick', n_steps=2, fit
 
 def dispersion(SC, rf_step, bpm_ords=None, cav_ords=None, n_steps=2):
     bpm_ords, cav_ords = _check_ords(SC, bpm_ords, cav_ords)
-    bref = np.ravel(bpm_reading(SC, bpm_ords=bpm_ords))
+    bref = np.ravel(bpm_reading(SC, bpm_ords=bpm_ords)[0])
     if n_steps == 2:
         SC = set_cavity_setpoints(SC, cav_ords, rf_step, 'Frequency', 'add')
-        B = np.ravel(bpm_reading(SC, bpm_ords=bpm_ords))
+        B = np.ravel(bpm_reading(SC, bpm_ords=bpm_ords)[0])
         SC = set_cavity_setpoints(SC, cav_ords, -rf_step, 'Frequency', 'add')
         return (B - bref) / rf_step
     rf_steps = np.zeros((len(cav_ords), n_steps))
@@ -72,7 +72,7 @@ def dispersion(SC, rf_step, bpm_ords=None, cav_ords=None, n_steps=2):
     rf0 = atgetfieldvalues(SC.RING, cav_ords, "FrequencySetPoint")
     for nStep in range(n_steps):
         SC = set_cavity_setpoints(SC, cav_ords, rf_steps[:, nStep], 'Frequency', 'abs')
-        dB[nStep, :] = np.ravel(bpm_reading(SC, bpm_ords=bpm_ords)) - bref
+        dB[nStep, :] = np.ravel(bpm_reading(SC, bpm_ords=bpm_ords)[0]) - bref
     SC = set_cavity_setpoints(SC, cav_ords, rf0, 'Frequency', 'abs')
     return np.linalg.lstsq(np.linspace(-rf_step, rf_step, n_steps), dB)[0]
 
@@ -103,7 +103,7 @@ def _kick_amplitude(SC, bref, bpm_ords, cm_ord, amp, skewness: bool, mode):
             LOGGER.debug(
                 f'Insufficient beam reach ({max_pos:d}/{max_pos_ref:d}). '
                 f'cm_step reduced to {1E6 * max_step:.1f}urad.')
-    return max_step, np.ravel(bpm_reading(SC, bpm_ords=bpm_ords)) - bref
+    return max_step, np.ravel(bpm_reading(SC, bpm_ords=bpm_ords)[0]) - bref
 
 
 def _try_setpoint(SC, bpm_ords, cm_ord, cmstart, max_step, skewness):
@@ -113,4 +113,4 @@ def _try_setpoint(SC, bpm_ords, cm_ord, cmstart, max_step, skewness):
         LOGGER.debug('CM  clipped. Using different CM direction.')
         max_step *= -1
         SC = set_cm_setpoints(SC, cm_ord, cmstart + max_step, skewness)
-    return SC, max_step, np.ravel(bpm_reading(SC, bpm_ords=bpm_ords))
+    return SC, max_step, np.ravel(bpm_reading(SC, bpm_ords=bpm_ords)[0])

pySC/utils/matlab_index.py

@@ -110,7 +110,7 @@ def SCgetBeamTransmission(SC: SimulatedCommissioning, /, *, nParticles: int = No
 
 
 def SCgetBPMreading(SC, /, *, BPMords=None):
-    return bpm_reading(SC, bpm_ords=BPMords)
+    return bpm_reading(SC, bpm_ords=BPMords)[0]
 
 
 def SCgetCMSetPoints(SC: SimulatedCommissioning, CMords: ndarray, nDim: int) -> ndarray:
@@ -200,7 +200,7 @@ def SCparticlesIn3D(*args):
 def SCplotBPMreading(SC, B=None, T=None):
     init_plot = SC.plot
     SC.plot = True
-    bpm_reading(SC)
+    _ = bpm_reading(SC)[0]
     SC.plot = init_plot
     return SC
 

tests/test_beam.py

@@ -32,18 +32,16 @@ def test_all_reading_basic(sc):
 
 
 def test_bpm_reading_basic(sc):
-    bpm_readings = bpm_reading(sc)
+    bpm_readings = bpm_reading(sc)[0]
     assert bpm_readings.shape == (2, sc.INJ.nTurns * sc.ORD.BPM.shape[0])
-    bpm_readings2 = bpm_reading(sc)
+    bpm_readings2 = bpm_reading(sc)[0]
     assert not np.allclose(bpm_readings2, bpm_readings)
-    bpm_readings = bpm_reading(sc, np.arange(3))
-    assert bpm_readings.shape == (2, 0)
     sc.INJ.trackMode = "PORB"
     sc.INJ.nTurns = 3
-    bpm_readings = bpm_reading(sc, np.arange(18))
+    bpm_readings = bpm_reading(sc, np.arange(18))[0]
     assert bpm_readings.shape == (2, 1)
     sc.plot = True
-    bpm_readings = bpm_reading(sc, 20 * np.ones(2))
+    bpm_readings = bpm_reading(sc, 20 * np.ones(2))[0]
     assert bpm_readings.shape == (2, 1)
 
 

tests/test_example.py

@@ -107,8 +107,8 @@ def test_example(at_lattice):
             cur = SCfeedbackRun(sc, minv_co, target=0, maxsteps=50, scaleDisp=1E8)
         except RuntimeError:
             break
-        B0rms = np.sqrt(np.mean(np.square(bpm_reading(sc)), axis=1))
-        Brms = np.sqrt(np.mean(np.square(bpm_reading(cur)), axis=1))
+        B0rms = np.sqrt(np.mean(np.square(bpm_reading(sc)[0]), axis=1))
+        Brms = np.sqrt(np.mean(np.square(bpm_reading(cur)[0]), axis=1))
         if np.mean(B0rms) < np.mean(Brms):
             break
         sc = cur