Test for use of keep_lattice in find_orbit6().

pyat/at/physics/orbit.py

@@ -278,6 +278,10 @@ def _orbit6(ring, cavpts=None, guess=None, keep_lattice=False, **kwargs):
 
     if guess is None:
         _, dt = get_timelag_fromU0(ring, method=method, cavpts=cavpts)
+        # Getting timelag by tracking uses a different lattice,
+        # so we cannot now use the same one again.
+        if method is ELossMethod.TRACKING:
+            keep_lattice = False
         ref_in = numpy.zeros((6,), order='F')
         ref_in[5] = -dt
     else:

pyat/setup.py

@@ -133,7 +133,7 @@ def integrator_ext(pass_method):
 
 setup(
     name='accelerator-toolbox',
-    version='0.2.0',
+    version='0.2.1',
     description='Accelerator Toolbox',
     long_description=long_description,
     author='The AT collaboration',

pyat/test/test_physics.py

@@ -47,6 +47,12 @@ def test_find_orbit4_result_unchanged_by_atpass(dba_lattice):
     assert_close(orbit[:4], orbit_copy[:4], atol=1e-12)
 
 
+def test_find_orbit4_produces_same_result_with_keep_lattice_True(dba_lattice):
+    orbit0, _ = physics.find_orbit4(dba_lattice)
+    orbit1, _ = physics.find_orbit4(dba_lattice, keep_lattice=True)
+    assert_close(orbit0, orbit1, rtol=0, atol=1e-12)
+
+
 def test_find_orbit4_with_two_refpts_with_and_without_guess(dba_lattice):
     expected = numpy.array(
         [[8.148212e-6, 1.0993354e-5, 0, 0, DP, 2.963929e-6],
@@ -125,6 +131,20 @@ def test_find_orbit6(hmba_lattice):
     assert_close(orbit6, orbit6_MATLAB, rtol=0, atol=1e-12)
 
 
+def test_find_orbit6_produces_same_result_with_keep_lattice_True(hmba_lattice):
+    hmba_lattice = hmba_lattice.radiation_on(quadrupole_pass=None, copy=True)
+    orbit0, _ = physics.find_orbit6(hmba_lattice)
+    # Technicality - the default arguments to find_orbit6 mean that
+    # keep_lattice argument is always false.
+    orbit1, _ = physics.find_orbit6(hmba_lattice, keep_lattice=True)
+    # With INTEGRAL keep_lattice does take effect.
+    orbit2, _ = physics.find_orbit6(
+        hmba_lattice, keep_lattice=True, method=physics.ELossMethod.INTEGRAL
+    )
+    assert_close(orbit0, orbit1, rtol=0, atol=1e-12)
+    assert_close(orbit0, orbit2, rtol=0, atol=1e-12)
+
+
 def test_find_orbit6_raises_AtError_if_there_is_no_cavity(dba_lattice):
     with pytest.raises(at.lattice.utils.AtError):
         physics.find_orbit6(dba_lattice)