diff --git a/pySC/core/beam.py b/pySC/core/beam.py index 090427e..80d6aaa 100644 --- a/pySC/core/beam.py +++ b/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): diff --git a/pySC/correction/injection_fit.py b/pySC/correction/injection_fit.py index aeadef2..40e1602 100644 --- a/pySC/correction/injection_fit.py +++ b/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): diff --git a/pySC/correction/orbit_trajectory.py b/pySC/correction/orbit_trajectory.py index 2275cf1..8eee225 100644 --- a/pySC/correction/orbit_trajectory.py +++ b/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))] diff --git a/pySC/correction/rf.py b/pySC/correction/rf.py index 3197b53..015c01e 100644 --- a/pySC/correction/rf.py +++ b/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]) diff --git a/pySC/example.py b/pySC/example.py index 9bdb1c2..542d20b 100644 --- a/pySC/example.py +++ b/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 diff --git a/pySC/lattice_properties/response_measurement.py b/pySC/lattice_properties/response_measurement.py index 433aaa1..8bc0a20 100644 --- a/pySC/lattice_properties/response_measurement.py +++ b/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]) diff --git a/pySC/utils/matlab_index.py b/pySC/utils/matlab_index.py index c4e63db..656e42e 100644 --- a/pySC/utils/matlab_index.py +++ b/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 diff --git a/tests/test_beam.py b/tests/test_beam.py index 257a136..a1b9ce3 100644 --- a/tests/test_beam.py +++ b/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) diff --git a/tests/test_example.py b/tests/test_example.py index 66f1067..baee93a 100644 --- a/tests/test_example.py +++ b/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