Reducing the amount of single element NumPy arrays

    also using lattice_track name

pySC/core/beam.py

@@ -185,7 +185,7 @@ def _real_bpm_reading(SC, track_mat, bpm_inds=None):  # track_mat should be only
     bpm_noise = bpm_noise[:, :, np.newaxis] * sc_tools.randnc(2, (2, n_bpms, nTurns))
     bpm_offset = np.transpose(at_wrapper.atgetfieldvalues(SC.RING, bpm_ords, 'Offset') + at_wrapper.atgetfieldvalues(SC.RING, bpm_ords, 'SupportOffset'))
     bpm_cal_error = np.transpose(at_wrapper.atgetfieldvalues(SC.RING, bpm_ords, 'CalError'))
-    bpm_roll = np.squeeze(at_wrapper.atgetfieldvalues(SC.RING, bpm_ords, 'Roll') + at_wrapper.atgetfieldvalues(SC.RING, bpm_ords, 'SupportRoll'), axis=1)
+    bpm_roll = np.ravel(at_wrapper.atgetfieldvalues(SC.RING, bpm_ords, 'Roll')) + np.ravel(at_wrapper.atgetfieldvalues(SC.RING, bpm_ords, 'SupportRoll'))
     bpm_sum_error = np.transpose(at_wrapper.atgetfieldvalues(SC.RING, bpm_ords, 'SumError'))[:, np.newaxis] * sc_tools.randnc(2, (n_bpms, nTurns))
     # averaging the X and Y positions at BPMs over particles
     mean_orbit = np.nanmean(track_mat, axis=1)
@@ -211,7 +211,7 @@ def _tracking(SC: SimulatedCommissioning, refs: ndarray) -> ndarray:
     if SC.INJ.trackMode == TRACK_ORB:
         pos = np.transpose(at_wrapper.findorbit6(SC.RING, refs, keep_lattice=False)[1])[[0, 2], :].reshape(2, 1, len(refs), 1)
     else:
-        pos = at_wrapper.atpass(SC.RING, generate_bunches(SC), SC.INJ.nTurns, refs, keep_lattice=False)[[0, 2], :, :, :]
+        pos = at_wrapper.lattice_track(SC.RING, generate_bunches(SC), SC.INJ.nTurns, refs, keep_lattice=False)[[0, 2], :, :, :]
     pos[1, np.isnan(pos[0, :, :, :])] = np.nan
     return pos
 

pySC/core/simulated_commissioning.py

@@ -560,7 +560,7 @@ def apply_errors(self, nsigmas: float = 2):
         self.INJ.randomInjectionZ = 1 * self.SIG.randomInjectionZ
         # Circumference
         if 'Circumference' in self.SIG.keys():
-            circScaling = 1 + self.SIG.Circumference * sc_tools.randnc(nsigmas, (1, ))[0]
+            circScaling = 1 + self.SIG.Circumference * sc_tools.randnc(nsigmas, ())
             self.RING = sc_tools.scale_circumference(self.RING, circScaling, 'rel')
             LOGGER.info('Circumference error applied.')
         # Misalignments
@@ -696,7 +696,7 @@ def update_supports(self, offset_bpms: bool = True, offset_magnets: bool = True)
                 for i, ind in enumerate(self.ORD.BPM):
                     setattr(self.RING[ind], "SupportOffset", offsets[0:2, i])  # No longitudinal BPM offsets implemented
                     setattr(self.RING[ind], "SupportRoll",
-                            np.array([rolls[0, i]]))  # BPM pitch and yaw angles not  implemented
+                            rolls[0, i])  # BPM pitch and yaw angles not  implemented
             else:
                 LOGGER.warning('SC: No BPMs have been registered!')
 

pySC/correction/bba.py

@@ -94,7 +94,7 @@ def fake_bba(SC, bpm_ords, mag_ords, errors=None, fake_offset=None):
         fake_bpm_offset = (SC.RING[mag_ords[inds[0], inds[1]]].MagnetOffset[inds[0]]
                            + SC.RING[mag_ords[inds[0], inds[1]]].SupportOffset[inds[0]]
                            - SC.RING[bpm_ords[inds[0], inds[1]]].SupportOffset[inds[0]]
-                           + fake_offset[inds[0]] * sc_tools.randnc())
+                           + fake_offset[inds[0]] * sc_tools.randnc(2, ()))
         if not np.isnan(fake_bpm_offset):
             SC.RING[bpm_ords[inds[0], inds[1]]].Offset[inds[0]] = fake_bpm_offset
         else:

pySC/correction/injection_fit.py

@@ -68,7 +68,7 @@ def _fit_bpm_data(s_bpm, bref):
 
 
 def _merit_function(x, SC, bref, ords_used):
-    t = at_wrapper.atpass(SC.IDEALRING, np.concatenate((x, np.zeros(2))), 1, ords_used, keep_lattice=False)[[0, 2], 0, :, 0]
+    t = at_wrapper.lattice_track(SC.IDEALRING, np.concatenate((x, np.zeros(2))), 1, ords_used, keep_lattice=False)[[0, 2], 0, :, 0]
     return np.sqrt(np.mean(bref - t) ** 2)
 
 

pySC/lattice_properties/apertures.py

@@ -26,7 +26,7 @@ def SCdynamicAperture(RING, dE, bounds=np.array([0, 1e-3]), nturns=1000, thetas=
     for cntt in range(len(thetas)):  # Loop over angles
         theta = thetas[cntt]
         limits = inibounds
-        at_wrapper.atpass(RING, np.full(6, np.nan), 1, [1])  # Fake Track to initialize lattice
+        at_wrapper.lattice_track(RING, np.full(6, np.nan), 1, [1])  # Fake Track to initialize lattice
         scales = 0
         while scales < 16:
             if _check_bounds(RING, ZCO, nturns, theta, limits):
@@ -102,7 +102,7 @@ def _check_bounds(RING, ZCO, nturns, theta, boundsIn):
     Zmax = ZCO[:]
     Zmax[0] = boundsIn[1] * np.cos(theta)
     Zmax[2] = boundsIn[1] * np.sin(theta)
-    ROUT = at_wrapper.atpass(RING, [Zmin, Zmax], nturns, len(RING), keep_lattice=True)  # Track
+    ROUT = at_wrapper.lattice_track(RING, [Zmin, Zmax], nturns, len(RING), keep_lattice=True)  # Track
     RLAST = ROUT[:, len(ROUT) - 1:len(ROUT)]  # Get positions after last turn
     return ~np.isnan(RLAST[0, 0]) and np.isnan(RLAST[0, 1])
 
@@ -114,7 +114,7 @@ def _refine_bounds(RING, ZCO, nturns, theta, boundsIn):
     Z = ZCO[:]
     Z[0] = rmid * np.cos(theta)
     Z[2] = rmid * np.sin(theta)
-    ROUT = at_wrapper.atpass(RING, Z, nturns, len(RING), keep_lattice=True)  # Track
+    ROUT = at_wrapper.lattice_track(RING, Z, nturns, len(RING), keep_lattice=True)  # Track
     RLAST = ROUT[:, len(ROUT) - 1]  # Get positions after last turn
     return [rmin, rmid] if np.isnan(RLAST[0]) else [rmid, rmax]  # Midpoint is outside or inside DA
 
@@ -146,7 +146,7 @@ def refine_bounds(local_bounds, RING, ZCO, nturns):
     dmean = np.mean(local_bounds)
     Z0 = ZCO
     Z0[4] = Z0[4] + dmean
-    ROUT = at_wrapper.atpass(RING, Z0, nturns, len(RING))
+    ROUT = at_wrapper.lattice_track(RING, Z0, nturns, len(RING))
     if np.isnan(ROUT[0]):  # Particle dead :(
         local_bounds[1] = dmean  # Set abs-upper bound to midpoint
     else:  # Particle alive :)
@@ -157,7 +157,7 @@ def refine_bounds(local_bounds, RING, ZCO, nturns):
 def check_bounds(local_bounds, RING, ZCO, nturns):
     Z = np.array([ZCO, ZCO])
     Z[4, :] = Z[4, :] + local_bounds[:]
-    ROUT = at_wrapper.atpass(RING, Z, nturns, len(RING))
+    ROUT = at_wrapper.lattice_track(RING, Z, nturns, len(RING))
     if np.isnan(ROUT[0, 0]) and not np.isnan(ROUT[0, 1]):
         LOGGER.warning('Closer-to-momentum particle is unstable. This shouldnt be!')
     return not np.isnan(ROUT[0, 0]) and np.isnan(ROUT[0, 1])

pySC/lattice_properties/response_model.py

@@ -49,7 +49,7 @@ def SCgetModelRM(SC, BPMords, CMords, trackMode='TBT', Z0=np.zeros(6), nTurns=1,
 
     """
     LOGGER.info('Calculating model response matrix')
-    track_methods = dict(TBT=at_wrapper.atpass, ORB=orbpass)
+    track_methods = dict(TBT=at_wrapper.lattice_track, ORB=orbpass)
     if trackMode not in track_methods.keys():
         ValueError(f'Unknown track mode {trackMode}. Valid values are {track_methods.keys()}')
     ring = SC.IDEALRING.deepcopy() if useIdealRing else SCgetModelRING(SC)
@@ -121,7 +121,7 @@ def SCgetModelDispersion(SC, BPMords, CAVords, trackMode='ORB', Z0=np.zeros(6),
 
     """
     LOGGER.info('Calculating model dispersion')
-    track_methods = dict(TBT=at_wrapper.atpass, ORB=orbpass)
+    track_methods = dict(TBT=at_wrapper.lattice_track, ORB=orbpass)
     if trackMode not in track_methods.keys():
         ValueError(f'Unknown track mode {trackMode}. Valid values are {track_methods.keys()}')
     ring = SC.IDEALRING.deepcopy() if useIdealRing else SCgetModelRING(SC)

pySC/plotting/plot_phase_space.py

@@ -10,7 +10,7 @@ def plot_phase_space(SC, ords=np.zeros(1, dtype=int), custom_bunch=None, nPartic
     init_font = plt.rcParams["font.size"]
     plt.rcParams.update({'font.size': 18})
     z_in, n_particles, n_turns = _check_input(SC, custom_bunch, nParticles, nTurns)
-    T = at_wrapper.atpass(SC.RING, z_in, n_turns, ords, keep_lattice=False)
+    T = at_wrapper.lattice_track(SC.RING, z_in, n_turns, ords, keep_lattice=False)
     T[:, np.isnan(T[0, :])] = np.nan
     label_str = [r'$\Delta x$ [$\mu$m]', r"$\Delta x'$ [$\mu$rad]", r'$\Delta y$ [$\mu$m]', r"$\Delta y'$ [$\mu$rad]",
                  r'$\Delta S$ [m]', r'$\delta E$ $[\%]$']

pySC/plotting/plot_support.py

@@ -64,9 +64,9 @@ def plot_support(SC: SimulatedCommissioning, font_size: int = 8, x_lim: Tuple[fl
     # Longitudinal offsets and Pitch and Yaw angles not supported for BPMs
     pad_off, pad_roll = ((0, 0), (0, 1)), ((0, 0), (0, 2))
     off_bpm = np.pad(at_wrapper.atgetfieldvalues(SC.RING, SC.ORD.BPM, "Offset"), pad_off)
-    roll_bpm = np.pad(at_wrapper.atgetfieldvalues(SC.RING, SC.ORD.BPM, "Roll"), pad_roll)
+    roll_bpm = np.pad(at_wrapper.atgetfieldvalues(SC.RING, SC.ORD.BPM, "Roll")[:, np.newaxis], pad_roll)
     off_bpm_support = np.pad(at_wrapper.atgetfieldvalues(SC.RING, SC.ORD.BPM, "SupportOffset"), pad_off)
-    roll_bpm_support = np.pad(at_wrapper.atgetfieldvalues(SC.RING, SC.ORD.BPM, "SupportRoll"), pad_roll)
+    roll_bpm_support = np.pad(at_wrapper.atgetfieldvalues(SC.RING, SC.ORD.BPM, "SupportRoll")[:, np.newaxis], pad_roll)
 
     # create figure
     fig, ax = plt.subplots(nrows=9, ncols=2, num=1213, sharex="all", figsize=(10, 15))

pySC/utils/at_wrapper.py

@@ -6,7 +6,7 @@
 which are not member functions of used ``pyAT`` objects.
 This is due to observed side effects, such as modification of input parameters.
 
-Tracking fuctions ``latice_pass``, ``find_orbit4``, ``find_orbit6``
+Tracking fuctions ``latice_track``, ``find_orbit4``, ``find_orbit6``
 index the result the same way as ``get_s_pos``,
 i.e. 0 means entrance of the first element, len(elements) means end of the last element.
 Function ``get_value_refpts`` indexes elements as usual.
@@ -21,14 +21,14 @@
 from at import Lattice
 
 
-def atpass(ring: Lattice, init_pos: ndarray, nturns: int, refpts: ndarray, keep_lattice: bool = False):
+def lattice_track(ring: Lattice, init_pos: ndarray, nturns: int, refpts: ndarray, keep_lattice: bool = False):
     return at.lattice_track(lattice=ring.copy(), r_in=init_pos.copy(), nturns=nturns, refpts=refpts,
-                           keep_lattice=keep_lattice)[0]
+                           keep_lattice=keep_lattice, in_place=False)[0]
 
 
 def patpass(ring: Lattice, init_pos: ndarray, nturns: int, refpts: ndarray, keep_lattice: bool = False):
     return at.lattice_track(lattice=ring.copy(), r_in=init_pos.copy(), nturns=nturns, refpts=refpts,
-                      keep_lattice=keep_lattice, use_mp=True)[0]
+                      keep_lattice=keep_lattice, in_place=False, use_mp=True)[0]
 
 
 def atgetfieldvalues(ring: Lattice, refpts: ndarray, attrname: str, index: int = None):

pySC/utils/sc_tools.py

@@ -2,18 +2,20 @@
 
 import numpy as np
 from matplotlib import pyplot as plt
-from numpy import ndarray
+from numpy import ndarray, float64
 from at import Lattice
+from typing import Union
 
 from pySC.utils import at_wrapper, logging_tools
 
 
 LOGGER = logging_tools.get_logger(__name__)
 
 
-def randnc(cut_off: float = 2, shape: tuple = (1,)) -> ndarray:
+def randnc(cut_off: float = 2, shape: tuple = (1,)) -> Union[ndarray, float64]:
     """
     Generates an array of random number(s) from normal distribution with a cut-off.
+    If shape = () a single float value is returned.
 
     Parameters
     ----------
@@ -24,15 +26,19 @@ def randnc(cut_off: float = 2, shape: tuple = (1,)) -> ndarray:
 
     Returns
     -------
-    out : ndarray
-        The output array.
+    out : Union[ndarray, float64]
+        The output array (If shape = () a single float value)
     """
+    if np.any([s < 0 for s in shape]):
+        raise ValueError('The shape of the output array must be positive.')
     out_shape = (1,) if np.sum(shape) < 1 else shape
     out = np.random.randn(np.prod(out_shape))
     outindex = np.abs(out) > np.abs(cut_off)
     while np.sum(outindex):
         out[outindex] = np.random.randn(np.sum(outindex))
         outindex = np.abs(out) > np.abs(cut_off)
+    if np.sum(shape) < 1:
+        return out[0]
     return out.reshape(out_shape)
 
 

tests/test_at_wrapper.py

@@ -5,7 +5,7 @@
 from numpy.testing import assert_equal
 import at
 from at import Lattice
-from pySC.utils.at_wrapper import findspos, atgetfieldvalues, atpass, patpass, findorbit6, findorbit4
+from pySC.utils.at_wrapper import findspos, atgetfieldvalues, lattice_track, patpass, findorbit6, findorbit4
 
 
 
@@ -30,7 +30,7 @@ def test_atpass(at_lattice):
     indices = np.arange(11, 450, 22, dtype=int)
     initial_pos = np.random.randn(6)
     copy_initial_pos = copy.deepcopy(initial_pos)
-    tracking = atpass(at_lattice, initial_pos, 3, indices,)
+    tracking = lattice_track(at_lattice, initial_pos, 3, indices, )
     assert tracking.shape == (6, 1, 20, 3)
     assert_equal(initial_pos, copy_initial_pos)
     assert at_lattice.__repr__() == lattice_copy.__repr__()

tests/test_loco.py

@@ -30,7 +30,7 @@ def test_loco_hmba(at_ring):
     cav_ords = sc_tools.ords_from_regex(sc.RING, 'RFC')
     quads_ords = [sc_tools.ords_from_regex(sc.RING, 'QF'), sc_tools.ords_from_regex(sc.RING, 'QD')]
 
-    CMstep = np.array([1.e-4])  # correctors change [rad]
+    CMstep = 1.e-4  # correctors change [rad]
     dk = 1.e-4  # quads change
     RFstep = 1e3