removing patpass

pySC/core/beam.py

@@ -128,7 +128,8 @@ def beam_transmission(SC: SimulatedCommissioning, nParticles: int = None, nTurns
     if nTurns is None:
         nTurns = SC.INJ.nTurns
     LOGGER.debug(f'Calculating maximum beam transmission for {nParticles} particles and {nTurns} turns: ')
-    T = at_wrapper.patpass(SC.RING, generate_bunches(SC, nParticles=nParticles), nTurns, np.array([len(SC.RING)]), keep_lattice=False)
+    T = at_wrapper.lattice_track(SC.RING, generate_bunches(SC, nParticles=nParticles), nTurns, np.array([len(SC.RING)]),
+                                 keep_lattice=False, use_mp=True)
     fraction_survived = np.mean(~np.isnan(T[0, :, :, :]), axis=(0, 1))
     max_turns = np.sum(fraction_survived > 1 - SC.INJ.beamLostAt)
     if plot:

pySC/utils/at_wrapper.py

@@ -21,14 +21,10 @@
 from at import Lattice
 
 
-def lattice_track(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,
+                  use_mp: bool = False):
     return at.lattice_track(lattice=ring.copy(), r_in=init_pos.copy(), nturns=nturns, refpts=refpts,
-                           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, in_place=False, use_mp=True)[0]
+                           keep_lattice=keep_lattice, in_place=False, use_mp=use_mp)[0]
 
 
 def atgetfieldvalues(ring: Lattice, refpts: ndarray, attrname: str, index: int = None):

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, lattice_track, patpass, findorbit6, findorbit4
+from pySC.utils.at_wrapper import findspos, atgetfieldvalues, lattice_track, findorbit6, findorbit4
 
 
 
@@ -25,7 +25,7 @@ def test_atgetfieldvalues(at_lattice):
     assert_equal(indices, np.arange(11, 450, 22, dtype=int))
 
 
-def test_atpass(at_lattice):
+def test_lattice_track1(at_lattice):
     lattice_copy = copy.deepcopy(at_lattice)
     indices = np.arange(11, 450, 22, dtype=int)
     initial_pos = np.random.randn(6)
@@ -42,7 +42,7 @@ def test_patpass(at_lattice):
     indices = np.arange(11, 450, 22, dtype=int)
     initial_pos = np.random.randn(6)
     copy_initial_pos = copy.deepcopy(initial_pos)
-    tracking = patpass(at_lattice, initial_pos, 3, indices,)
+    tracking = lattice_track(at_lattice, initial_pos, 3, indices, use_mp=True)
     assert tracking.shape == (6, 1, 20, 3)
     assert_equal(initial_pos, copy_initial_pos)
     assert at_lattice.__repr__() == lattice_copy.__repr__()