diff --git a/pyat/test/test_lattice_object.py b/pyat/test/test_lattice_object.py index 2dbe64a84..336ba1347 100644 --- a/pyat/test/test_lattice_object.py +++ b/pyat/test/test_lattice_object.py @@ -1,53 +1,63 @@ import numpy -from numpy.testing import assert_allclose, assert_equal import pytest +from numpy.testing import assert_allclose, assert_equal + from at import elements from at.lattice import Lattice, AtWarning, AtError def test_lattice_creation_gets_attributes_from_arguments(): - lat = Lattice(name='lattice', energy=3.e+6, periodicity=32, an_attr=12) + lat = Lattice(name="lattice", energy=3.0e6, periodicity=32, an_attr=12) assert len(lat) == 0 - assert lat.name == 'lattice' - assert lat.energy == 3.e+6 + assert lat.name == "lattice" + assert lat.energy == 3.0e6 assert lat.periodicity == 32 assert lat.radiation is False - assert lat.particle.name == 'relativistic' + assert lat.particle.name == "relativistic" assert lat.an_attr == 12 def test_lattice_energy_radiation_periodicity(): - d = elements.Dipole('d1', 1, BendingAngle=numpy.pi/16, Energy=5.e+6, - PassMethod='BndMPoleSymplectic4RadPass') - lat = Lattice([d], name='lattice', energy=3.e+6) - assert lat.energy == 3.e+6 + d = elements.Dipole( + "d1", + 1, + BendingAngle=numpy.pi / 16, + Energy=5.0e6, + PassMethod="BndMPoleSymplectic4RadPass", + ) + lat = Lattice([d], name="lattice", energy=3.0e6) + assert lat.energy == 3.0e6 assert lat.periodicity == 32 assert lat.radiation is True assert lat.is_6d is True def test_lattice_voltage_harmonic_number(): - rf = elements.RFCavity('rf', 0, 0.2e6, 0.5e9, 5, 3.e6) - d = elements.Dipole('d1', 2.99792458, BendingAngle=numpy.pi/5) - lat = Lattice([rf, d], name='lattice') - assert lat.energy == 3.e+6 + rf = elements.RFCavity("rf", 0, 0.2e6, 0.5e9, 5, 3.0e6) + d = elements.Dipole("d1", 2.99792458, BendingAngle=numpy.pi / 5) + lat = Lattice([rf, d], name="lattice") + assert lat.energy == 3.0e6 assert lat.periodicity == 10 assert lat.rf_voltage == 2e6 - assert lat.revolution_frequency == 10.e6 + assert lat.revolution_frequency == 10.0e6 assert lat.harmonic_number == 50 assert lat.radiation is True def test_lattice_creation_from_lattice_inherits_attributes(): - d = elements.Dipole('d1', 1, BendingAngle=numpy.pi, Energy=5.e+6, - PassMethod='BndMPoleSymplectic4RadPass') - lat1 = Lattice([d], name='lattice', energy=3.e+6, periodicity=32, - an_attr=12) + d = elements.Dipole( + "d1", + 1, + BendingAngle=numpy.pi, + Energy=5.0e6, + PassMethod="BndMPoleSymplectic4RadPass", + ) + lat1 = Lattice([d], name="lattice", energy=3.0e6, periodicity=32, an_attr=12) lat2 = Lattice(lat1, another_attr=5) assert id(lat1) != id(lat2) assert len(lat2) == 1 - assert lat2.name == 'lattice' - assert lat2.energy == 3.e+6 + assert lat2.name == "lattice" + assert lat2.energy == 3.0e6 assert lat2.periodicity == 32 assert lat2.is_6d is True assert lat2.another_attr == 5 @@ -56,30 +66,39 @@ def test_lattice_creation_from_lattice_inherits_attributes(): def test_lattice_energy_is_not_defined_raises_AtError(): - d = elements.Dipole('d1', 1, BendingAngle=numpy.pi) + d = elements.Dipole("d1", 1, BendingAngle=numpy.pi) with pytest.raises(AtError): Lattice([d]) def test_item_is_not_an_AT_element_warns_correctly(): with pytest.warns(AtWarning): - Lattice(['a'], energy=0, periodicity=1) + Lattice(["a"], energy=0, periodicity=1) def test_lattice_string_ordering(): - lat = Lattice([elements.Drift('D0', 1.0, attr1=numpy.array(0))], - name='lat', energy=5, periodicity=1, attr2=3) + lat = Lattice( + [elements.Drift("D0", 1.0, attr1=numpy.array(0))], + name="lat", + energy=5, + periodicity=1, + attr2=3, + ) latstr = str(lat) - assert latstr.startswith("Lattice(<1 elements>, name='lat', " - "energy=5, particle=Particle('relativistic'), " - "periodicity=1, beam_current=0.0, nbunch=1") + assert latstr.startswith( + "Lattice(<1 elements>, name='lat', " + "energy=5, particle=Particle('relativistic'), " + "periodicity=1, beam_current=0.0, nbunch=1" + ) assert latstr.endswith("attr2=3)") latrepr = repr(lat) - assert latrepr.startswith("Lattice([Drift('D0', 1.0, attr1=array(0))], " - "name='lat', " - "energy=5, particle=Particle('relativistic'), " - "periodicity=1, beam_current=0.0, nbunch=1") + assert latrepr.startswith( + "Lattice([Drift('D0', 1.0, attr1=array(0))], " + "name='lat', " + "energy=5, particle=Particle('relativistic'), " + "periodicity=1, beam_current=0.0, nbunch=1" + ) assert latrepr.endswith("attr2=3)") @@ -90,7 +109,7 @@ def test_getitem(simple_lattice, simple_ring): def test_setitem(simple_lattice): - new = elements.Monitor('M2') + new = elements.Monitor("M2") old = simple_lattice[5] simple_lattice[5] = new assert simple_lattice[5] != old @@ -102,7 +121,7 @@ def test_setitem(simple_lattice): def test_delitem(simple_lattice, simple_ring): - mon = elements.Monitor('M2') + mon = elements.Monitor("M2") simple_lattice.append(mon) assert len(simple_lattice) == 7 del simple_lattice[-1] @@ -128,34 +147,32 @@ def test_property_values_against_known(hmba_lattice): assert hmba_lattice.rf_voltage == 6000000 assert hmba_lattice.harmonic_number == 992 assert hmba_lattice.radiation is False - numpy.testing.assert_almost_equal(hmba_lattice.energy_loss, - 2526188.658758993, decimal=1) + numpy.testing.assert_almost_equal( + hmba_lattice.energy_loss, 2526188.658758993, decimal=1 + ) def test_radiation_change(hmba_lattice): - rfs = [elem for elem in hmba_lattice if isinstance(elem, - elements.RFCavity)] - dipoles = [elem for elem in hmba_lattice if isinstance(elem, - elements.Dipole)] - quads = [elem for elem in hmba_lattice if isinstance(elem, - elements.Quadrupole)] - hmba_lattice.radiation_on(None, 'pass2', 'auto') + rfs = [elem for elem in hmba_lattice if isinstance(elem, elements.RFCavity)] + dipoles = [elem for elem in hmba_lattice if isinstance(elem, elements.Dipole)] + quads = [elem for elem in hmba_lattice if isinstance(elem, elements.Quadrupole)] + hmba_lattice.radiation_on(None, "pass2", "auto") assert hmba_lattice.radiation is True assert hmba_lattice.has_cavity is False for elem in rfs: - assert elem.PassMethod == 'IdentityPass' + assert elem.PassMethod == "IdentityPass" for elem in dipoles: - assert elem.PassMethod == 'pass2' + assert elem.PassMethod == "pass2" for elem in quads: - assert elem.PassMethod == 'StrMPoleSymplectic4RadPass' - hmba_lattice.radiation_off(None, 'BndMPoleSymplectic4Pass', 'auto') + assert elem.PassMethod == "StrMPoleSymplectic4RadPass" + hmba_lattice.radiation_off(None, "BndMPoleSymplectic4Pass", "auto") assert hmba_lattice.radiation is False for elem in rfs: - assert elem.PassMethod == 'IdentityPass' + assert elem.PassMethod == "IdentityPass" for elem in dipoles: - assert elem.PassMethod == 'BndMPoleSymplectic4Pass' + assert elem.PassMethod == "BndMPoleSymplectic4Pass" for elem in quads: - assert elem.PassMethod == 'StrMPoleSymplectic4Pass' + assert elem.PassMethod == "StrMPoleSymplectic4Pass" def test_radiation_state_errors(hmba_lattice): @@ -173,35 +190,38 @@ def test_radiation_state_errors(hmba_lattice): hmba_lattice.disable_6d() hmba_lattice.get_mcf() -@pytest.mark.parametrize('ring', - [pytest.lazy_fixture('hmba_lattice')]) -def test_develop(ring): + +@pytest.mark.parametrize("ring", ["hmba_lattice"]) +def test_develop(request, ring): + ring = request.getfixturevalue(ring) newring = ring.develop() assert_allclose(ring.circumference, newring.circumference) assert_allclose(ring.rf_voltage, newring.rf_voltage) - assert_allclose(ring.revolution_frequency*ring.harmonic_number, - newring.revolution_frequency*newring.harmonic_number) + assert_allclose( + ring.revolution_frequency * ring.harmonic_number, + newring.revolution_frequency * newring.harmonic_number, + ) rp1 = ring.radiation_parameters() rp2 = newring.radiation_parameters() assert_allclose(rp1.fulltunes, rp2.fulltunes) assert_allclose(rp1.U0, rp2.U0) -@pytest.mark.parametrize('ring', - [pytest.lazy_fixture('hmba_lattice')]) -def test_operators(ring): + +@pytest.mark.parametrize("ring", ["hmba_lattice"]) +def test_operators(request, ring): + ring = request.getfixturevalue(ring) newring1 = ring + ring newring2 = ring.concatenate(ring, copy=True) - assert_equal(len(newring1), len(ring)*2) + assert_equal(len(newring1), len(ring) * 2) assert_equal(len(newring2), len(ring) * 2) newring1 += ring newring2.concatenate(ring, copy=False) - assert_equal(len(newring1), len(ring)*3) - assert_equal(len(newring2), len(ring)*3) + assert_equal(len(newring1), len(ring) * 3) + assert_equal(len(newring2), len(ring) * 3) newring1 = ring * 2 newring2 = ring.repeat(2) - assert_equal(len(newring1), len(ring)*2) + assert_equal(len(newring1), len(ring) * 2) assert_equal(len(newring2), len(ring) * 2) assert_equal(newring1.harmonic_number, ring.harmonic_number) newring1 = ring.reverse(copy=True) assert_equal(newring1[-1].FamName, ring[0].FamName) - diff --git a/pyat/test/test_lattice_utils.py b/pyat/test/test_lattice_utils.py index 7fcc72cfa..79c00636a 100644 --- a/pyat/test/test_lattice_utils.py +++ b/pyat/test/test_lattice_utils.py @@ -1,18 +1,22 @@ import numpy -from at.lattice import elements, uint32_refpts, bool_refpts, checkattr +import pytest + from at.lattice import checktype, get_cells, refpts_iterator, get_elements +from at.lattice import elements, uint32_refpts, bool_refpts, checkattr from at.lattice import get_s_pos, tilt_elem, shift_elem, set_tilt, set_shift -import pytest -@pytest.mark.parametrize('ref_in, expected', ( - [2, numpy.array([2], dtype=numpy.uint32)], - [-1, numpy.array([4], dtype=numpy.uint32)], - [7, numpy.array([2], dtype=numpy.uint32)], - [[1, 2, 3], numpy.array([1, 2, 3], dtype=numpy.uint32)], - [[0, 1, 2, 3, 4, 5], numpy.array([0, 1, 2, 3, 4, 5], dtype=numpy.uint32)], - [[0, 6, 2, -2, 4, 5], numpy.array([0, 1, 2, 3, 4, 5], dtype=numpy.uint32)], -)) +@pytest.mark.parametrize( + "ref_in, expected", + ( + [2, numpy.array([2], dtype=numpy.uint32)], + [-1, numpy.array([4], dtype=numpy.uint32)], + [7, numpy.array([2], dtype=numpy.uint32)], + [[1, 2, 3], numpy.array([1, 2, 3], dtype=numpy.uint32)], + [[0, 1, 2, 3, 4, 5], numpy.array([0, 1, 2, 3, 4, 5], dtype=numpy.uint32)], + [[0, 6, 2, -2, 4, 5], numpy.array([0, 1, 2, 3, 4, 5], dtype=numpy.uint32)], + ), +) def test_uint32_refpts_converts_numerical_inputs_correctly(ref_in, expected): numpy.testing.assert_equal(uint32_refpts(ref_in, 5), expected) # float indexes are deprecated @@ -22,17 +26,24 @@ def test_uint32_refpts_converts_numerical_inputs_correctly(ref_in, expected): numpy.testing.assert_equal(uint32_refpts(ref_in2, 5), expected) -@pytest.mark.parametrize('ref_in, expected', ( - [None, numpy.array([], dtype=numpy.uint32)], - [[], numpy.array([], dtype=numpy.uint32)], - [True, numpy.array([0], dtype=numpy.uint32)], - [False, numpy.array([], dtype=numpy.uint32)], - [[True, False, True], numpy.array([0, 2], dtype=numpy.uint32)], - [[False, False, False, False, False, False], - numpy.array([], dtype=numpy.uint32)], - [[True, True, True, True, True, True], - numpy.array([0, 1, 2, 3, 4, 5], dtype=numpy.uint32)] -)) +@pytest.mark.parametrize( + "ref_in, expected", + ( + [None, numpy.array([], dtype=numpy.uint32)], + [[], numpy.array([], dtype=numpy.uint32)], + [True, numpy.array([0], dtype=numpy.uint32)], + [False, numpy.array([], dtype=numpy.uint32)], + [[True, False, True], numpy.array([0, 2], dtype=numpy.uint32)], + [ + [False, False, False, False, False, False], + numpy.array([], dtype=numpy.uint32), + ], + [ + [True, True, True, True, True, True], + numpy.array([0, 1, 2, 3, 4, 5], dtype=numpy.uint32), + ], + ), +) def test_uint32_refpts_converts_other_input_types_correctly(ref_in, expected): numpy.testing.assert_equal(uint32_refpts(ref_in, 5), expected) @@ -41,62 +52,72 @@ def test_uint32_refpts_converts_other_input_types_correctly(ref_in, expected): # duplicate, wrap-around indexing misordered, wrap-around indexing duplicate, # -ve indexing and wrap-around indexing duplicate, -ve indexing and wrap-around # indexing misordered -@pytest.mark.parametrize('ref_in', ([0, 1, 2, 3], [2, 1], [0, 0], [-1, 0], - [0, -2], [3, 0], [1, 3], [-1, 3], [3, -2])) +@pytest.mark.parametrize( + "ref_in", + ([0, 1, 2, 3], [2, 1], [0, 0], [-1, 0], [0, -2], [3, 0], [1, 3], [-1, 3], [3, -2]), +) def test_uint32_refpts_throws_IndexError_correctly(ref_in): with pytest.raises(IndexError): uint32_refpts(ref_in, 2) -@pytest.mark.parametrize('ref_in, expected', ( - [[True], numpy.array([True, False, False, False], dtype=bool)], - [[], numpy.array([False, False, False, False], dtype=bool)], - [None, numpy.array([False, False, False, False], dtype=bool)], - [2, numpy.array([False, False, True, False], dtype=bool)], - [[1, 2], numpy.array([False, True, True, False], dtype=bool)], -)) +@pytest.mark.parametrize( + "ref_in, expected", + ( + [[True], numpy.array([True, False, False, False], dtype=bool)], + [[], numpy.array([False, False, False, False], dtype=bool)], + [None, numpy.array([False, False, False, False], dtype=bool)], + [2, numpy.array([False, False, True, False], dtype=bool)], + [[1, 2], numpy.array([False, True, True, False], dtype=bool)], + ), +) def test_bool_refpts(ref_in, expected): numpy.testing.assert_equal(bool_refpts(ref_in, 3), expected) -@pytest.mark.parametrize('ref_in', (10, range(100))) +@pytest.mark.parametrize("ref_in", (10, range(100))) def test_bool_refpts_throws_IndexErrorrs(ref_in): with pytest.raises(IndexError): bool_refpts(ref_in, 3) def test_checkattr(simple_ring): - assert checkattr('Length')(simple_ring[0]) is True - assert checkattr('not_an_attr')(simple_ring[0]) is False - assert (list(filter(checkattr('Length', 1), simple_ring)) == - [simple_ring[0], simple_ring[3], simple_ring[4], simple_ring[5]]) - assert list(filter(checkattr('Length', 2), simple_ring)) == [] - assert list(filter(checkattr('not_an_attr'), simple_ring)) == [] + assert checkattr("Length")(simple_ring[0]) is True + assert checkattr("not_an_attr")(simple_ring[0]) is False + assert list(filter(checkattr("Length", 1), simple_ring)) == [ + simple_ring[0], + simple_ring[3], + simple_ring[4], + simple_ring[5], + ] + assert list(filter(checkattr("Length", 2), simple_ring)) == [] + assert list(filter(checkattr("not_an_attr"), simple_ring)) == [] def test_checktype(simple_ring): assert checktype(elements.Drift)(simple_ring[0]) is True assert checktype(elements.Marker)(simple_ring[0]) is False - assert (list(filter(checktype(elements.Drift), simple_ring)) == - [simple_ring[0], simple_ring[3], simple_ring[4], simple_ring[5]]) + assert list(filter(checktype(elements.Drift), simple_ring)) == [ + simple_ring[0], + simple_ring[3], + simple_ring[4], + simple_ring[5], + ] assert list(filter(checktype(elements.Monitor), simple_ring)) == [] def test_get_cells(simple_ring): a = numpy.array([True, True, True, True, True, True, False]) - numpy.testing.assert_equal(get_cells(simple_ring, checkattr('Length')), a) + numpy.testing.assert_equal(get_cells(simple_ring, checkattr("Length")), a) a = numpy.array([False, True, False, False, False, False, False]) - numpy.testing.assert_equal(get_cells(simple_ring, 'attr'), a) + numpy.testing.assert_equal(get_cells(simple_ring, "attr"), a) a = numpy.array([True, False, False, False, False, False, False]) - numpy.testing.assert_equal(get_cells(simple_ring, 'FamName', 'D1'), a) + numpy.testing.assert_equal(get_cells(simple_ring, "FamName", "D1"), a) def test_refpts_iterator(simple_ring): - assert (list(refpts_iterator(simple_ring, [0, 1, 2, 3, 4, 5])) == - simple_ring) - assert (list(refpts_iterator(simple_ring, - numpy.ones(6, dtype=bool))) == - simple_ring) + assert list(refpts_iterator(simple_ring, [0, 1, 2, 3, 4, 5])) == simple_ring + assert list(refpts_iterator(simple_ring, numpy.ones(6, dtype=bool))) == simple_ring assert list(refpts_iterator(simple_ring, [1])) == [simple_ring[1]] a = numpy.array([False, True, False, False, False, False]) assert list(refpts_iterator(simple_ring, a)) == [simple_ring[1]] @@ -104,20 +125,17 @@ def test_refpts_iterator(simple_ring): def test_get_elements(hmba_lattice): # test FamName direct match - assert get_elements(hmba_lattice, 'BPM_06') == [hmba_lattice[65]] + assert get_elements(hmba_lattice, "BPM_06") == [hmba_lattice[65]] # test FamName wildcard matching - assert get_elements(hmba_lattice, 'QD2?') == list(hmba_lattice[9, 113]) - assert get_elements(hmba_lattice, 'QD3*') == list(hmba_lattice[19, 105]) - assert get_elements(hmba_lattice, 'S*H2B') == list([hmba_lattice[55]]) - assert get_elements(hmba_lattice, '*C_1') == list(hmba_lattice[59, 60]) - assert get_elements(hmba_lattice, - 'DR_2[1-3]') == list(hmba_lattice[54, 56, 58]) - assert get_elements(hmba_lattice, - 'DR_2[!1-7]') == list(hmba_lattice[52, 78, 80]) + assert get_elements(hmba_lattice, "QD2?") == list(hmba_lattice[9, 113]) + assert get_elements(hmba_lattice, "QD3*") == list(hmba_lattice[19, 105]) + assert get_elements(hmba_lattice, "S*H2B") == list([hmba_lattice[55]]) + assert get_elements(hmba_lattice, "*C_1") == list(hmba_lattice[59, 60]) + assert get_elements(hmba_lattice, "DR_2[1-3]") == list(hmba_lattice[54, 56, 58]) + assert get_elements(hmba_lattice, "DR_2[!1-7]") == list(hmba_lattice[52, 78, 80]) # test element instance - marker = elements.Marker('M1') - assert get_elements(hmba_lattice, - marker) == list(hmba_lattice[1, 12, 61, 67, 73]) + marker = elements.Marker("M1") + assert get_elements(hmba_lattice, marker) == list(hmba_lattice[1, 12, 61, 67, 73]) # test element type assert get_elements(hmba_lattice, elements.RFCavity) == [hmba_lattice[0]] # test invalid key raises TypeError @@ -130,35 +148,37 @@ def test_get_s_pos_returns_zero_for_empty_lattice(): def test_get_s_pos_returns_length_for_lattice_with_one_element(): - e = elements.LongElement('e', 0.1) + e = elements.LongElement("e", 0.1) assert get_s_pos([e], [1]) == numpy.array([0.1]) def test_get_s_pos_returns_all_pts_for_lat_with_2_elements(): - e = elements.LongElement('e', 0.1) - f = elements.LongElement('e', 2.1) + e = elements.LongElement("e", 0.1) + f = elements.LongElement("e", 2.1) numpy.testing.assert_equal(get_s_pos([e, f]), numpy.array([0, 0.1, 2.2])) def test_get_s_pos_returns_all_pts_for_lat_with_2_elements_using_int_refpts(): - e = elements.LongElement('e', 0.1) - f = elements.LongElement('e', 2.1) + e = elements.LongElement("e", 0.1) + f = elements.LongElement("e", 2.1) lat = [e, f] - numpy.testing.assert_equal(get_s_pos(lat, range(len(lat) + 1)), - numpy.array([0, 0.1, 2.2])) + numpy.testing.assert_equal( + get_s_pos(lat, range(len(lat) + 1)), numpy.array([0, 0.1, 2.2]) + ) def test_get_s_pos_returns_all_pts_for_lat_with_2_elements_using_bool_refpts(): - e = elements.LongElement('e', 0.1) - f = elements.LongElement('e', 2.1) + e = elements.LongElement("e", 0.1) + f = elements.LongElement("e", 2.1) lat = [e, f] - numpy.testing.assert_equal(get_s_pos(lat, numpy.ones(3, dtype=bool)), - numpy.array([0, 0.1, 2.2])) + numpy.testing.assert_equal( + get_s_pos(lat, numpy.ones(3, dtype=bool)), numpy.array([0, 0.1, 2.2]) + ) def test_tilt_elem(simple_ring): - tilt_elem(simple_ring[0], (numpy.pi/4)) - v = 1/2**0.5 + tilt_elem(simple_ring[0], (numpy.pi / 4)) + v = 1 / 2**0.5 a = numpy.diag([v, v, v, v, 1.0, 1.0]) a[0, 2], a[1, 3], a[2, 0], a[3, 1] = v, v, -v, -v numpy.testing.assert_allclose(simple_ring[0].R1, a) @@ -173,8 +193,8 @@ def test_shift_elem(simple_ring): def test_set_tilt(simple_ring): - set_tilt(simple_ring, [(numpy.pi/4), 0, 0, 0, (numpy.pi/4), 0]) - v = 1/2**0.5 + set_tilt(simple_ring, [(numpy.pi / 4), 0, 0, 0, (numpy.pi / 4), 0]) + v = 1 / 2**0.5 a = numpy.diag([v, v, v, v, 1.0, 1.0]) a[0, 2], a[1, 3], a[2, 0], a[3, 1] = v, v, -v, -v numpy.testing.assert_allclose(simple_ring[0].R1, a) @@ -182,21 +202,24 @@ def test_set_tilt(simple_ring): numpy.testing.assert_allclose(simple_ring[0].R1, a) numpy.testing.assert_allclose(simple_ring[0].R2, a.T) ring = [simple_ring[0]] - set_tilt(ring, (0)) + set_tilt(ring, 0) numpy.testing.assert_allclose(ring[0].R1, numpy.eye(6)) numpy.testing.assert_allclose(ring[0].R2, numpy.eye(6)) def test_set_shift(simple_ring): - set_shift(simple_ring, numpy.array([0., 0., 0., 1., 0., 0.5]), - numpy.array([0., 0., 0., 2., 0., 1.])) - a = numpy.array([0.5, 0., 1., 0., 0., 0.]) - numpy.testing.assert_equal(simple_ring[3].T1, -a*2) - numpy.testing.assert_equal(simple_ring[3].T2, a*2) + set_shift( + simple_ring, + numpy.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.5]), + numpy.array([0.0, 0.0, 0.0, 2.0, 0.0, 1.0]), + ) + a = numpy.array([0.5, 0.0, 1.0, 0.0, 0.0, 0.0]) + numpy.testing.assert_equal(simple_ring[3].T1, -a * 2) + numpy.testing.assert_equal(simple_ring[3].T2, a * 2) numpy.testing.assert_equal(simple_ring[5].T1, -a) numpy.testing.assert_equal(simple_ring[5].T2, a) ring = [simple_ring[3]] set_shift(ring, 3, 5) - a = numpy.array([3., 0., 5., 0., 0., 0.]) + a = numpy.array([3.0, 0.0, 5.0, 0.0, 0.0, 0.0]) numpy.testing.assert_equal(simple_ring[3].T1, -a) numpy.testing.assert_equal(simple_ring[3].T2, a) diff --git a/pyat/test/test_linopt6.py b/pyat/test/test_linopt6.py index c45b95f77..ba3c60862 100644 --- a/pyat/test/test_linopt6.py +++ b/pyat/test/test_linopt6.py @@ -1,29 +1,29 @@ -from at import linopt2, linopt4, linopt6, get_optics -from numpy.testing import assert_allclose as assert_close import pytest +from numpy.testing import assert_allclose as assert_close + +from at import linopt2, linopt4, linopt6, get_optics -@pytest.mark.parametrize('lattice', - [pytest.lazy_fixture('dba_lattice'), - pytest.lazy_fixture('hmba_lattice')]) -def test_linopt6_norad(lattice): +@pytest.mark.parametrize("lattice", ["dba_lattice", "hmba_lattice"]) +def test_linopt6_norad(request, lattice): """Compare the results of linopt2 and linopt6 in 4d""" + lattice = request.getfixturevalue(lattice) refpts = range(len(lattice) + 1) ld02, rd2, ld2 = linopt2(lattice, refpts, get_w=True) ld06, rd6, ld6 = linopt6(lattice, refpts, get_w=True) assert_close(rd2.tune, rd6.tune, atol=1e-10, rtol=0) assert_close(rd2.chromaticity, rd6.chromaticity, atol=1e-10, rtol=0) - for field in ['s_pos', 'closed_orbit', 'dispersion', - 'alpha', 'beta', 'mu']: + for field in ["s_pos", "closed_orbit", "dispersion", "alpha", "beta", "mu"]: assert_close(ld2[field], ld6[field], atol=1e-10, rtol=0, err_msg=field) assert_close(ld2.W, ld6.W, atol=1e-6, rtol=0) -@pytest.mark.parametrize('lattice', [pytest.lazy_fixture('hmba_lattice')]) -def test_linopt6_rad(lattice): +@pytest.mark.parametrize("lattice", ["hmba_lattice"]) +def test_linopt6_rad(request, lattice): """Compare the results with and without radiation""" - refpts = range(len(lattice)+1) + lattice = request.getfixturevalue(lattice) + refpts = range(len(lattice) + 1) # Turn cavity ON, without radiation radlattice = lattice.radiation_on(dipole_pass=None, copy=True) ld04, rd4, ld4 = linopt6(lattice, refpts, get_w=True) @@ -32,30 +32,29 @@ def test_linopt6_rad(lattice): assert_close(rd4.tune, rd6.tune[:2], atol=1e-10, rtol=0) assert_close(rd4.chromaticity, rd6.chromaticity[:2], atol=1e-8, rtol=0) - for field in ['s_pos', 'closed_orbit', 'dispersion', 'alpha', 'beta']: - assert_close(ld4[field], ld6[field], atol=1.e-8, rtol=0, err_msg=field) - assert_close(ld4.mu, ld6.mu[:, :2], atol=1.e-8, rtol=0) + for field in ["s_pos", "closed_orbit", "dispersion", "alpha", "beta"]: + assert_close(ld4[field], ld6[field], atol=1.0e-8, rtol=0, err_msg=field) + assert_close(ld4.mu, ld6.mu[:, :2], atol=1.0e-8, rtol=0) assert_close(ld4.W, ld6.W, atol=1e-6, rtol=2e-5) -@pytest.mark.parametrize('dp', (-0.005, 0.0, 0.005)) -@pytest.mark.parametrize('lattice', - [pytest.lazy_fixture('dba_lattice'), - pytest.lazy_fixture('hmba_lattice')]) -@pytest.mark.parametrize('method', [linopt2, linopt4, linopt6]) -def test_linopt6_line(lattice, dp, method): +@pytest.mark.parametrize("dp", (-0.005, 0.0, 0.005)) +@pytest.mark.parametrize("lattice", ["dba_lattice", "hmba_lattice"]) +@pytest.mark.parametrize("method", [linopt2, linopt4, linopt6]) +def test_linopt6_line(request, lattice, dp, method): """Run a ring as a transfer line and compare with the ring results""" - refpts = lattice.uint32_refpts(range(len(lattice)+1)) + lattice = request.getfixturevalue(lattice) + refpts = lattice.uint32_refpts(range(len(lattice) + 1)) ld04, rd4, ld4 = get_optics(lattice, refpts, dp=dp, method=method) - twin = dict(alpha=ld04.alpha, beta=ld04.beta, - dispersion=ld04.dispersion, closed_orbit=ld04.closed_orbit) + twin = dict( + alpha=ld04.alpha, + beta=ld04.beta, + dispersion=ld04.dispersion, + closed_orbit=ld04.closed_orbit, + ) # twin = ld04 - tr04, bd4, tr4 = get_optics(lattice, refpts, dp=dp, - twiss_in=twin, method=method) - for field in ['s_pos', 'closed_orbit', 'dispersion', - 'alpha', 'beta', 'mu']: - assert_close(ld04[field], tr04[field], atol=1.e-10, - rtol=0, err_msg=field) - assert_close(ld4[field], tr4[field], atol=1.e-8, - rtol=0, err_msg=field) + tr04, bd4, tr4 = get_optics(lattice, refpts, dp=dp, twiss_in=twin, method=method) + for field in ["s_pos", "closed_orbit", "dispersion", "alpha", "beta", "mu"]: + assert_close(ld04[field], tr04[field], atol=1.0e-10, rtol=0, err_msg=field) + assert_close(ld4[field], tr4[field], atol=1.0e-8, rtol=0, err_msg=field) diff --git a/pyat/test/test_track.py b/pyat/test/test_track.py index 6b8dcc21a..56f6de1a8 100644 --- a/pyat/test/test_track.py +++ b/pyat/test/test_track.py @@ -1,12 +1,13 @@ -from at import elements -from at import lattice_track -from at import lattice_pass, internal_lpass import numpy import pytest +from at import elements +from at import lattice_pass, internal_lpass +from at import lattice_track + -@pytest.mark.parametrize('func', (lattice_track, lattice_pass, internal_lpass)) -@pytest.mark.parametrize('input_dim', [(0,), (5,), (7,), (1, 1), (6, 1, 1)]) +@pytest.mark.parametrize("func", (lattice_track, lattice_pass, internal_lpass)) +@pytest.mark.parametrize("input_dim", [(0,), (5,), (7,), (1, 1), (6, 1, 1)]) def test_lattice_track_raises_AssertionError_if_rin_incorrect_shape(input_dim, func): rin = numpy.zeros(input_dim) lattice = [] @@ -14,16 +15,16 @@ def test_lattice_track_raises_AssertionError_if_rin_incorrect_shape(input_dim, f func(lattice, rin) -@pytest.mark.parametrize('func', (lattice_track, lattice_pass, internal_lpass)) +@pytest.mark.parametrize("func", (lattice_track, lattice_pass, internal_lpass)) def test_multiple_particles_lattice_track(func): - lattice = [elements.Drift('Drift', 1.0)] + lattice = [elements.Drift("Drift", 1.0)] rin = numpy.zeros((6, 2)) rin[0, 0] = 1e-6 # particle one offset in x rin[2, 1] = 1e-6 # particle two offset in y r_original = numpy.copy(rin) r_out = func(lattice, rin, nturns=2) if isinstance(r_out, tuple): - r_out, *_ = r_out + r_out, *_ = r_out # particle position is not changed passing through the drift numpy.testing.assert_equal(r_original[:, 0], r_out[:, 0, 0, 0]) numpy.testing.assert_equal(r_original[:, 0], r_out[:, 0, 0, 1]) @@ -31,19 +32,20 @@ def test_multiple_particles_lattice_track(func): numpy.testing.assert_equal(r_original[:, 1], r_out[:, 1, 0, 1]) -@pytest.mark.parametrize('func', (lattice_track, lattice_pass)) +@pytest.mark.parametrize("func", (lattice_track, lattice_pass)) def test_lattice_convert_to_list_if_incorrect_type(func): - lattice = numpy.array([elements.Drift('Drift', 1.0)]) + lattice = numpy.array([elements.Drift("Drift", 1.0)]) rin = numpy.zeros((6, 2)) rin[0, 0] = 1e-6 r_original = numpy.copy(rin) r_out = func(lattice, rin, 1) if isinstance(r_out, tuple): - r_out, *_ = r_out + r_out, *_ = r_out numpy.testing.assert_equal(r_original, r_out.reshape(6, 2)) + def test_in_place_flag_tracking(): - lattice = numpy.array([elements.Drift('Drift', 1.0)]) + lattice = numpy.array([elements.Drift("Drift", 1.0)]) rin = numpy.zeros((6, 1)) rin[1] = 1e-6 r_original = numpy.copy(rin) diff --git a/pyat/test_matlab/conftest.py b/pyat/test_matlab/conftest.py index 0cf98b068..dda44839c 100644 --- a/pyat/test_matlab/conftest.py +++ b/pyat/test_matlab/conftest.py @@ -2,9 +2,12 @@ conftest.py is a special pytest file that allows you to share fixtures between other modules. """ -import pytest + import os import sys + +import pytest + if sys.version_info.minor < 9: from importlib_resources import files, as_file else: @@ -13,14 +16,15 @@ import numpy import machine_data import at + try: from matlab.engine import connect_matlab, start_matlab, EngineError except ImportError: - print('Matlab comparison tests require Matlab Python Engine installed.') - print('Python will exit.') + print("Matlab comparison tests require Matlab Python Engine installed.") + print("Python will exit.") sys.exit() -ROOT_DIR = os.path.realpath(os.path.join(os.path.dirname(__file__), '../..')) +ROOT_DIR = os.path.realpath(os.path.join(os.path.dirname(__file__), "../..")) # noinspection PyUnusedLocal @@ -29,21 +33,21 @@ def pytest_report_header(config): a = platform.uname() sysinfo = "system: {} {}".format(a[0], a[2]) except AttributeError: - sysinfo = '' + sysinfo = "" numpyinfo = "numpy version: {}".format(numpy.__version__) return [sysinfo, numpyinfo] -@pytest.fixture(scope='session') +@pytest.fixture(scope="session") def engine(): try: - eng = connect_matlab('pytest') + eng = connect_matlab("pytest") # Keep the existing Matlab path except EngineError: eng = start_matlab() # Add the local AT path - eng.addpath(eng.genpath(os.path.join(ROOT_DIR, 'atintegrators/'))) - eng.addpath(eng.genpath(os.path.join(ROOT_DIR, 'atmat/'))) + eng.addpath(eng.genpath(os.path.join(ROOT_DIR, "atintegrators/"))) + eng.addpath(eng.genpath(os.path.join(ROOT_DIR, "atmat/"))) # Add the local test_matlab directory eng.addpath(os.path.dirname(__file__)) yield eng @@ -58,35 +62,35 @@ def _load_lattice(engine, name, key): # ----------------------- dba lattice ----------------------- -@pytest.fixture(scope='session') +@pytest.fixture(scope="session") def dba(engine): - return _load_lattice(engine, 'dba.mat', 'RING') + return _load_lattice(engine, "dba.mat", "RING") # ----------------------- hmba lattice ----------------------- -@pytest.fixture(scope='session') +@pytest.fixture(scope="session") def hmba(engine): - return _load_lattice(engine, 'hmba.mat', 'RING') + return _load_lattice(engine, "hmba.mat", "RING") # hmba lattice with cavities on -@pytest.fixture(scope='session') +@pytest.fixture(scope="session") def hmba_cav(engine): - pyl, myl, _ = _load_lattice(engine, 'hmba.mat', 'RING') - myl, radindex = engine.pyproxy('atradon', myl, 'BendPass', '', nargout=2) + pyl, myl, _ = _load_lattice(engine, "hmba.mat", "RING") + myl, radindex = engine.pyproxy("atradon", myl, "BendPass", "", nargout=2) return pyl.radiation_on(copy=True, dipole_pass=None), myl, radindex # hmba lattice with radiation on -@pytest.fixture(scope='session') +@pytest.fixture(scope="session") def hmba_rad(engine): - pyl, myl, _ = _load_lattice(engine, 'hmba.mat', 'RING') - myl, radindex = engine.pyproxy('atradon', myl, nargout=2) + pyl, myl, _ = _load_lattice(engine, "hmba.mat", "RING") + myl, radindex = engine.pyproxy("atradon", myl, nargout=2) return pyl.radiation_on(copy=True), myl, radindex # -------------------- lattice with errors -------------------- # Takes too much time for usual tests -@pytest.fixture(scope='session') +@pytest.fixture(scope="session") def err(engine): - return _load_lattice(engine, 'err.mat', 'RING') + return _load_lattice(engine, "err.mat", "RING") diff --git a/pyat/test_matlab/test_cmp_atpass.py b/pyat/test_matlab/test_cmp_atpass.py index 08d892676..9924bb90f 100644 --- a/pyat/test_matlab/test_cmp_atpass.py +++ b/pyat/test_matlab/test_cmp_atpass.py @@ -1,44 +1,53 @@ -import at -import numpy import matlab -from numpy.testing import assert_allclose as assert_close +import numpy as np import pytest +from numpy.testing import assert_allclose as assert_close +import at -@pytest.mark.parametrize('lattices', - [pytest.lazy_fixture('dba'), - pytest.lazy_fixture('hmba'), - pytest.lazy_fixture('hmba_cav'), - pytest.lazy_fixture('hmba_rad')]) -def test_atpass(engine, lattices): - py_lattice, ml_lattice, _ = lattices - xy_step = 1.e-8 + +@pytest.mark.parametrize( + "lattices", + [ + "dba", + "hmba", + "hmba_cav", + "hmba_rad", + ], +) +def test_atpass(engine, request, lattices): + py_lattice, ml_lattice, _ = request.getfixturevalue(lattices) + xy_step = 1.0e-8 scaling = [xy_step, xy_step, xy_step, xy_step, xy_step, xy_step] ml_rin = engine.diag(matlab.double(scaling)) ml_rout = engine.atpass(ml_lattice, ml_rin, 1, 1) - py_rin = numpy.asfortranarray(numpy.diag(scaling)) + py_rin = np.asfortranarray(np.diag(scaling)) at.lattice_track(py_lattice, py_rin, 1, in_place=True) - assert_close(py_rin, ml_rout, rtol=0, atol=1.e-30) - - -@pytest.mark.parametrize('lattices', - [pytest.lazy_fixture('dba'), - pytest.lazy_fixture('hmba'), - pytest.lazy_fixture('hmba_cav'), - pytest.lazy_fixture('hmba_rad')]) -def test_linepass(engine, lattices): - py_lattice, ml_lattice, _ = lattices - xy_step = 1.e-8 + assert_close(py_rin, ml_rout, rtol=0, atol=1.0e-30) + + +@pytest.mark.parametrize( + "lattices", + [ + "dba", + "hmba", + "hmba_cav", + "hmba_rad", + ], +) +def test_linepass(engine, request, lattices): + py_lattice, ml_lattice, _ = request.getfixturevalue(lattices) + xy_step = 1.0e-8 scaling = [xy_step, xy_step, xy_step, xy_step, xy_step, xy_step] ml_rin = engine.diag(matlab.double(scaling)) - ml_rout = numpy.array(engine.linepass(ml_lattice, ml_rin)) + ml_rout = np.array(engine.linepass(ml_lattice, ml_rin)) - py_rin = numpy.asfortranarray(numpy.diag(scaling)) + py_rin = np.asfortranarray(np.diag(scaling)) py_rout, *_ = at.lattice_track(py_lattice, py_rin, refpts=len(py_lattice)) - py_rout = numpy.squeeze(py_rout) + py_rout = np.squeeze(py_rout) - assert_close(py_rout, ml_rout, rtol=0, atol=1.e-30) + assert_close(py_rout, ml_rout, rtol=0, atol=1.0e-30) diff --git a/pyat/test_matlab/test_cmp_avlinopt.py b/pyat/test_matlab/test_cmp_avlinopt.py index 294c649cc..400895c6d 100644 --- a/pyat/test_matlab/test_cmp_avlinopt.py +++ b/pyat/test_matlab/test_cmp_avlinopt.py @@ -1,33 +1,34 @@ -import at import matlab -import numpy -from numpy.testing import assert_allclose as assert_close +import numpy as np import pytest +from numpy.testing import assert_allclose as assert_close + +import at def _ml_refs(refpts, nelems): """Convert refpoints to Matlab""" uintrefs = at.uint32_refpts(refpts, nelems) - return matlab.double([float(ref+1) for ref in uintrefs]) + return matlab.double([float(ref + 1) for ref in uintrefs]) -@pytest.mark.parametrize('dp', (-0.01, 0.0, 0.01)) -@pytest.mark.parametrize('lattices', - [pytest.lazy_fixture('dba'), - pytest.lazy_fixture('hmba')]) -def test_avlinopt(engine, lattices, dp): - py_lattice, ml_lattice, _ = lattices +@pytest.mark.parametrize("dp", (-0.01, 0.0, 0.01)) +@pytest.mark.parametrize("lattices", ["dba", "hmba"]) +def test_avlinopt(engine, request, lattices, dp): + py_lattice, ml_lattice, _ = request.getfixturevalue(lattices) nelems = len(py_lattice) refpts = range(nelems) mlrefs = _ml_refs(refpts, nelems) # Python call - lindata, avebeta, avemu, avedisp, aves, tune, chrom = \ - at.avlinopt(py_lattice, dp, refpts=refpts) + lindata, avebeta, avemu, avedisp, aves, tune, chrom = at.avlinopt( + py_lattice, dp, refpts=refpts + ) # Matlab call - ml_data, ml_avebeta, ml_avemu, ml_avedisp, ml_tune, ml_chrom = \ - engine.pyproxy('atavedata', ml_lattice, dp, mlrefs, nargout=6) + ml_data, ml_avebeta, ml_avemu, ml_avedisp, ml_tune, ml_chrom = engine.pyproxy( + "atavedata", ml_lattice, dp, mlrefs, nargout=6 + ) # Comparison - assert_close(avebeta, numpy.asarray(ml_avebeta), rtol=1e-7) - assert_close(avemu, numpy.asarray(ml_avemu), rtol=0, atol=1e-7) - assert_close(avedisp, numpy.asarray(ml_avedisp), rtol=0, atol=1e-8) + assert_close(avebeta, np.asarray(ml_avebeta), rtol=1e-7) + assert_close(avemu, np.asarray(ml_avemu), rtol=0, atol=1e-7) + assert_close(avedisp, np.asarray(ml_avedisp), rtol=0, atol=1e-8) diff --git a/pyat/test_matlab/test_cmp_linopt6.py b/pyat/test_matlab/test_cmp_linopt6.py index 464c539bc..6d316ceec 100644 --- a/pyat/test_matlab/test_cmp_linopt6.py +++ b/pyat/test_matlab/test_cmp_linopt6.py @@ -1,97 +1,111 @@ -from at import linopt6 -import numpy import matlab -from numpy.testing import assert_allclose as assert_close +import numpy as np import pytest +from numpy.testing import assert_allclose as assert_close + +from at import linopt6 def _py_data(ml_data): """Convert a Matlab vector to a numpy vector""" - return numpy.squeeze(numpy.asarray(ml_data)) + return np.squeeze(np.asarray(ml_data)) def _ml_refs(ring, py_refs): """Convert refpoints to Matlab""" uintrefs = ring.uint32_refpts(py_refs) - return matlab.double([float(ref+1) for ref in uintrefs]) + return matlab.double([float(ref + 1) for ref in uintrefs]) -def _compare_4d(py_data, ml_data, fields, atol=1.e-12, rtol=1.e-7): - for (ml_key, py_key) in fields: +def _compare_4d(py_data, ml_data, fields, atol=1.0e-12, rtol=1.0e-7): + for ml_key, py_key in fields: ml_val = _py_data(ml_data[ml_key]) - if py_key == 'closed_orbit': - py_val = numpy.squeeze(py_data[py_key][:, :4]) + if py_key == "closed_orbit": + py_val = np.squeeze(py_data[py_key][:, :4]) else: - py_val = numpy.squeeze(py_data[py_key]) + py_val = np.squeeze(py_data[py_key]) assert_close(py_val, ml_val, atol=atol, rtol=rtol) -def _compare_6d(py_data, ml_data, fields, atol=1.e-12, rtol=1.e-7): - for (ml_key, py_key) in fields: +def _compare_6d(py_data, ml_data, fields, atol=1.0e-12, rtol=1.0e-7): + for ml_key, py_key in fields: ml_val = _py_data(ml_data[ml_key]) - py_val = numpy.squeeze(py_data[py_key]) + py_val = np.squeeze(py_data[py_key]) assert_close(py_val, ml_val, atol=atol, rtol=rtol) -@pytest.mark.parametrize('dp', (-0.01, 0.0, 0.01)) -@pytest.mark.parametrize('lattices', - [pytest.lazy_fixture('hmba'), - pytest.lazy_fixture('dba')]) -def test_4d_analysis(engine, lattices, dp): +@pytest.mark.parametrize("dp", (-0.01, 0.0, 0.01)) +@pytest.mark.parametrize("lattices", ["dba", "hmba"]) +def test_4d_analysis(engine, request, lattices, dp): """Compare linopt6 in 4D""" - py_lattice, ml_lattice, _ = lattices - fields = [('SPos', 's_pos'), ('ClosedOrbit', 'closed_orbit'), - ('M', 'M'), ('A', 'A'), ('Dispersion', 'dispersion'), - ('alpha', 'alpha'), ('beta', 'beta'), ('mu', 'mu'), - ] + py_lattice, ml_lattice, _ = request.getfixturevalue(lattices) + fields = [ + ("SPos", "s_pos"), + ("ClosedOrbit", "closed_orbit"), + ("M", "M"), + ("A", "A"), + ("Dispersion", "dispersion"), + ("alpha", "alpha"), + ("beta", "beta"), + ("mu", "mu"), + ] pypts = range(10) mlpts = _ml_refs(py_lattice, pypts) # python call - py_data0, py_ringdata, py_data = linopt6(py_lattice, refpts=pypts, - dp=dp, get_w=True) + py_data0, py_ringdata, py_data = linopt6( + py_lattice, refpts=pypts, dp=dp, get_w=True + ) # Matlab call - ml_ringdata, ml_data = engine.pyproxy('atlinopt6', ml_lattice, mlpts, - 'dp', dp, 'get_w', nargout=2) + ml_ringdata, ml_data = engine.pyproxy( + "atlinopt6", ml_lattice, mlpts, "dp", dp, "get_w", nargout=2 + ) # Comparison - assert_close(py_ringdata.tune, - _py_data(ml_ringdata['tune']), - atol=1.e-8, rtol=0) - assert_close(py_ringdata.chromaticity, - _py_data(ml_ringdata['chromaticity']), - atol=5.e-4, rtol=0) + assert_close(py_ringdata.tune, _py_data(ml_ringdata["tune"]), atol=1.0e-8, rtol=0) + assert_close( + py_ringdata.chromaticity, + _py_data(ml_ringdata["chromaticity"]), + atol=5.0e-4, + rtol=0, + ) - _compare_6d(py_data, ml_data, fields, atol=1.e-6, rtol=1e-8) - _compare_6d(py_data, ml_data, [('W', 'W')], atol=1.e-6, rtol=4.e-4) + _compare_6d(py_data, ml_data, fields, atol=1.0e-6, rtol=1e-8) + _compare_6d(py_data, ml_data, [("W", "W")], atol=1.0e-6, rtol=4.0e-4) -@pytest.mark.parametrize('lattices', - [pytest.lazy_fixture('hmba')]) -def test_6d_analysis(engine, lattices): +@pytest.mark.parametrize("lattices", ["hmba"]) +def test_6d_analysis(engine, request, lattices): """Compare linopt6 in 6D""" - py_lattice, ml_lattice, _ = lattices - fields = [('SPos', 's_pos'), ('ClosedOrbit', 'closed_orbit'), - ('M', 'M'), ('A', 'A'), ('Dispersion', 'dispersion'), - ('alpha', 'alpha'), ('beta', 'beta'), ('mu', 'mu'), - ('W', 'W') - ] - ml_lattice = engine.pyproxy('atradon', ml_lattice) + py_lattice, ml_lattice, _ = request.getfixturevalue(lattices) + fields = [ + ("SPos", "s_pos"), + ("ClosedOrbit", "closed_orbit"), + ("M", "M"), + ("A", "A"), + ("Dispersion", "dispersion"), + ("alpha", "alpha"), + ("beta", "beta"), + ("mu", "mu"), + ("W", "W"), + ] + ml_lattice = engine.pyproxy("atradon", ml_lattice) py_lattice = py_lattice.radiation_on(copy=True) pypts = range(10) mlpts = _ml_refs(py_lattice, pypts) # python call - py_data0, py_ringdata, py_data = linopt6(py_lattice, refpts=pypts, - get_w=True) + py_data0, py_ringdata, py_data = linopt6(py_lattice, refpts=pypts, get_w=True) # Matlab call - ml_ringdata, ml_data = engine.pyproxy('atlinopt6', ml_lattice, mlpts, - 'get_w', nargout=2) + ml_ringdata, ml_data = engine.pyproxy( + "atlinopt6", ml_lattice, mlpts, "get_w", nargout=2 + ) # Comparison - assert_close(py_ringdata.tune, - _py_data(ml_ringdata['tune']), - atol=1.e-12, rtol=0) - assert_close(py_ringdata.chromaticity, - _py_data(ml_ringdata['chromaticity']), - atol=1.e-5, rtol=0) - - _compare_6d(py_data, ml_data, fields, atol=1.e-6, rtol=1e-8) + assert_close(py_ringdata.tune, _py_data(ml_ringdata["tune"]), atol=1.0e-12, rtol=0) + assert_close( + py_ringdata.chromaticity, + _py_data(ml_ringdata["chromaticity"]), + atol=1.0e-5, + rtol=0, + ) + + _compare_6d(py_data, ml_data, fields, atol=1.0e-6, rtol=1e-8) diff --git a/pyat/test_matlab/test_cmp_physics.py b/pyat/test_matlab/test_cmp_physics.py index 5d2c3f6f6..699e9a150 100644 --- a/pyat/test_matlab/test_cmp_physics.py +++ b/pyat/test_matlab/test_cmp_physics.py @@ -1,69 +1,75 @@ from at import physics, get_s_pos, uint32_refpts from scipy.constants import speed_of_light import matlab -import numpy +import numpy as np from numpy.testing import assert_allclose as assert_close import pytest def _py_data(ml_data): """Convert a Matlab vector to a numpy vector""" - return numpy.squeeze(numpy.asarray(ml_data)) + return np.squeeze(np.asarray(ml_data)) def _ml_refs(refpts, nelems): """Convert refpoints to Matlab""" uintrefs = uint32_refpts(refpts, nelems) - return matlab.double([float(ref+1) for ref in uintrefs]) + return matlab.double([float(ref + 1) for ref in uintrefs]) -def _compare_physdata(py_data, ml_data, fields, atol=12, rtol=1.e-7): - for (ml_key, py_key) in fields: +def _compare_physdata(py_data, ml_data, fields, atol=12, rtol=1.0e-7): + for ml_key, py_key in fields: ml_val = _py_data(ml_data[ml_key]) - if py_key == 'closed_orbit': - py_val = numpy.squeeze(py_data[py_key][:, :4]) + if py_key == "closed_orbit": + py_val = np.squeeze(py_data[py_key][:, :4]) else: - py_val = numpy.squeeze(py_data[py_key]) + py_val = np.squeeze(py_data[py_key]) assert_close(py_val, ml_val, atol=atol, rtol=rtol) -@pytest.mark.parametrize('dp', (-0.01, 0.0, 0.01)) -@pytest.mark.parametrize('lattices', - [pytest.lazy_fixture('hmba'), - pytest.lazy_fixture('dba')]) -def test_linear_analysis(engine, lattices, dp): - py_lattice, ml_lattice, _ = lattices +@pytest.mark.parametrize("dp", (-0.01, 0.0, 0.01)) +@pytest.mark.parametrize("lattices", ["dba", "hmba"]) +def test_linear_analysis(engine, request, lattices, dp): + py_lattice, ml_lattice, _ = request.getfixturevalue(lattices) nelems = len(py_lattice) - fields = [('SPos', 's_pos'), - ('ClosedOrbit', 'closed_orbit'), - ('Dispersion', 'dispersion'), - ('alpha', 'alpha'), ('beta', 'beta'), - ('mu', 'mu'), ('M44', 'm44'), - ('A', 'A'), ('B', 'B'), ('C', 'C'), - ('gamma', 'gamma')] + fields = [ + ("SPos", "s_pos"), + ("ClosedOrbit", "closed_orbit"), + ("Dispersion", "dispersion"), + ("alpha", "alpha"), + ("beta", "beta"), + ("mu", "mu"), + ("M44", "m44"), + ("A", "A"), + ("B", "B"), + ("C", "C"), + ("gamma", "gamma"), + ] refpts = range(nelems + 1) - py_data0, py_tune, py_chrom, py_data = physics.linopt(py_lattice, dp, - refpts, True) + py_data0, py_tune, py_chrom, py_data = physics.linopt(py_lattice, dp, refpts, True) # Matlab call - ml_data, ml_tune, ml_chrom = engine.pyproxy('atlinopt', ml_lattice, dp, - _ml_refs(refpts, nelems), - nargout=3) - ml_data0 = engine.pyproxy('atlinopt', ml_lattice, dp, - _ml_refs(nelems, nelems), nargout=3)[0] + ml_data, ml_tune, ml_chrom = engine.pyproxy( + "atlinopt", ml_lattice, dp, _ml_refs(refpts, nelems), nargout=3 + ) + ml_data0 = engine.pyproxy( + "atlinopt", ml_lattice, dp, _ml_refs(nelems, nelems), nargout=3 + )[0] # Comparison - assert_close(py_tune, _py_data(ml_tune), atol=1.e-8, rtol=0) - assert_close(py_chrom, _py_data(ml_chrom), atol=3.e-3, rtol=0) - _compare_physdata(numpy.expand_dims(py_data0, 0), ml_data0, fields, - rtol=1e-8) + assert_close(py_tune, _py_data(ml_tune), atol=1.0e-8, rtol=0) + assert_close(py_chrom, _py_data(ml_chrom), atol=3.0e-3, rtol=0) + _compare_physdata(np.expand_dims(py_data0, 0), ml_data0, fields, rtol=1e-8) _compare_physdata(py_data, ml_data, fields, rtol=1e-8) -@pytest.mark.parametrize('lattices', - [pytest.lazy_fixture('hmba')]) -def test_ohmi_envelope(engine, lattices): - py_lattice, ml_lattice, _ = lattices - fields = [('beam66', 'r66'), ('beam44', 'r44'), ('emit66', 'emitXYZ'), - ('emit44', 'emitXY')] +@pytest.mark.parametrize("lattices", ["hmba"]) +def test_ohmi_envelope(engine, request, lattices): + py_lattice, ml_lattice, _ = request.getfixturevalue(lattices) + fields = [ + ("beam66", "r66"), + ("beam44", "r44"), + ("emit66", "emitXYZ"), + ("emit44", "emitXY"), + ] nelems = len(py_lattice) refpts = range(nelems + 1) @@ -71,86 +77,84 @@ def test_ohmi_envelope(engine, lattices): py_lattice = py_lattice.radiation_on(copy=True) py_emit0, py_beamdata, py_emit = py_lattice.ohmi_envelope(refpts) # Matlab call - ml_emit = engine.pyproxy('atx', ml_lattice, 0.0, _ml_refs(refpts, nelems)) - ml_emit0, ml_params = engine.pyproxy('atx', ml_lattice, 0.0, - _ml_refs(0, nelems), nargout=2) - revolution_period = get_s_pos(py_lattice, nelems) / speed_of_light + ml_emit = engine.pyproxy("atx", ml_lattice, 0.0, _ml_refs(refpts, nelems)) + ml_emit0, ml_params = engine.pyproxy( + "atx", ml_lattice, 0.0, _ml_refs(0, nelems), nargout=2 + ) + revolution_period = get_s_pos(py_lattice, nelems)[0] / speed_of_light damping_times = revolution_period / py_beamdata.damping_rates # Comparison - assert_close(damping_times, _py_data(ml_params['dampingtime']), - atol=1.e-6, rtol=1.e-3) - assert_close(py_beamdata.mode_emittances, _py_data(ml_emit0['modemit']), - atol=1.e-20, rtol=2.e-3) - _compare_physdata(numpy.expand_dims(py_emit0, 0), ml_emit0, fields) + assert_close( + damping_times, _py_data(ml_params["dampingtime"]), atol=1.0e-6, rtol=1.0e-3 + ) + assert_close( + py_beamdata.mode_emittances, + _py_data(ml_emit0["modemit"]), + atol=1.0e-20, + rtol=2.0e-3, + ) + _compare_physdata(np.expand_dims(py_emit0, 0), ml_emit0, fields) _compare_physdata(py_emit, ml_emit, fields) -@pytest.mark.parametrize('lattices', - [pytest.lazy_fixture('hmba_rad')]) -def test_quantdiff(engine, lattices): - py_lattice, ml_lattice, radindex = lattices +@pytest.mark.parametrize("lattices", ["hmba_rad"]) +def test_quantdiff(engine, request, lattices): + py_lattice, ml_lattice, radindex = request.getfixturevalue(lattices) # Python call dmat = physics.radiation.quantdiffmat(py_lattice) # Matlab call - dmat_ml = engine.pyproxy('quantumDiff', ml_lattice, radindex) + dmat_ml = engine.pyproxy("quantumDiff", ml_lattice, radindex) # Comparison assert_close(dmat, dmat_ml, rtol=1.0e-8, atol=1.0e-18) -@pytest.mark.parametrize('lattices', - [pytest.lazy_fixture('hmba')]) -def test_fastring(engine, lattices): - py_lattice, ml_lattice, _ = lattices +@pytest.mark.parametrize("lattices", ["hmba"]) +def test_fastring(engine, request, lattices): + py_lattice, ml_lattice, _ = request.getfixturevalue(lattices) # Python call ring, ringrad = physics.fastring.fast_ring(py_lattice) # Matlab call - ring_ml, ringrad_ml = engine.pyproxy('atfastring', - ml_lattice, nargout=2) + ring_ml, ringrad_ml = engine.pyproxy("atfastring", ml_lattice, nargout=2) # Comparison for r, rml in zip([ring, ringrad], [ring_ml, ringrad_ml]): - assert_close(r[0].Frequency, rml[1]['Frequency'], rtol=1.0e-20) - assert_close(r[0].Voltage, rml[1]['Voltage'], rtol=1.0e-20) - assert_close(r[1].I2, rml[2]['I2'], rtol=1.0e-20) - assert_close(r[1].Length, rml[2]['Length'], rtol=1.0e-20) - assert_close(r[1].M66, rml[2]['M66'], atol=1.0e-7) - assert_close(r[1].T1, numpy.squeeze(rml[2]['T1']), - rtol=1.0e-8, atol=1.0e-11) - assert_close(r[1].T2, numpy.squeeze(rml[2]['T2']), - rtol=1.0e-8, atol=1.0e-11) - assert_close(r[2].A1, rml[-1]['A1'], rtol=0.02) - assert_close(r[2].A2, rml[-1]['A2'], rtol=0.02) - assert_close(r[2].A3, rml[-1]['A3'], rtol=0.02) - assert_close(r[2].Alphax, rml[-1]['Alphax'], rtol=1.0e-5) - assert_close(r[2].Alphay, rml[-1]['Alphay'], rtol=1.0e-6) - assert_close(r[2].Betax, rml[-1]['Betax'], rtol=1.0e-10) - assert_close(r[2].Betay, rml[-1]['Betay'], rtol=1.0e-10) - assert_close(r[2].Qpx, rml[-1]['Qpx'], rtol=1.0e-5) - assert_close(r[2].Qpy, rml[-1]['Qpy'], rtol=1.0e-5) - assert_close(r[2].T1, numpy.squeeze(rml[-1]['T1']), - rtol=1.0e-8, atol=1.0e-11) - assert_close(r[2].T2, numpy.squeeze(rml[-1]['T2']), - rtol=1.0e-8, atol=1.0e-11) + assert_close(r[0].Frequency, rml[1]["Frequency"], rtol=1.0e-20) + assert_close(r[0].Voltage, rml[1]["Voltage"], rtol=1.0e-20) + assert_close(r[1].I2, rml[2]["I2"], rtol=1.0e-20) + assert_close(r[1].Length, rml[2]["Length"], rtol=1.0e-20) + assert_close(r[1].M66, rml[2]["M66"], atol=1.0e-7) + assert_close(r[1].T1, np.squeeze(rml[2]["T1"]), rtol=1.0e-8, atol=1.0e-11) + assert_close(r[1].T2, np.squeeze(rml[2]["T2"]), rtol=1.0e-8, atol=1.0e-11) + assert_close(r[2].A1, rml[-1]["A1"], rtol=0.02) + assert_close(r[2].A2, rml[-1]["A2"], rtol=0.02) + assert_close(r[2].A3, rml[-1]["A3"], rtol=0.02) + assert_close(r[2].Alphax, rml[-1]["Alphax"], rtol=1.0e-5) + assert_close(r[2].Alphay, rml[-1]["Alphay"], rtol=1.0e-6) + assert_close(r[2].Betax, rml[-1]["Betax"], rtol=1.0e-10) + assert_close(r[2].Betay, rml[-1]["Betay"], rtol=1.0e-10) + assert_close(r[2].Qpx, rml[-1]["Qpx"], rtol=1.0e-5) + assert_close(r[2].Qpy, rml[-1]["Qpy"], rtol=1.0e-5) + assert_close(r[2].T1, np.squeeze(rml[-1]["T1"]), rtol=1.0e-8, atol=1.0e-11) + assert_close(r[2].T2, np.squeeze(rml[-1]["T2"]), rtol=1.0e-8, atol=1.0e-11) if len(r) >= 4: - assert_close(r[3].Lmatp, rml[-2]['Lmatp'], - rtol=0.02, atol=1.0e-10) + assert_close(r[3].Lmatp, rml[-2]["Lmatp"], rtol=0.02, atol=1.0e-10) -@pytest.mark.parametrize('dp', (0.00, 0.01, -0.01)) -@pytest.mark.parametrize('lattices', - [pytest.lazy_fixture('hmba')]) -def test_parameters(engine, lattices, dp): - py_lattice, ml_lattice, _ = lattices +@pytest.mark.parametrize("dp", (0.00, 0.01, -0.01)) +@pytest.mark.parametrize("lattices", ["hmba"]) +def test_parameters(engine, request, lattices, dp): + py_lattice, ml_lattice, _ = request.getfixturevalue(lattices) # Test perimeter py_length = py_lattice.get_s_pos(len(py_lattice)) - ml_length = engine.findspos(ml_lattice, len(ml_lattice)+1) - assert_close(py_length, ml_length, rtol=1.E-8) + ml_length = engine.findspos(ml_lattice, len(ml_lattice) + 1) + assert_close(py_length, ml_length, rtol=1.0e-8) # test energy loss - ml_energy, ml_periods, ml_voltage, \ - ml_harms, ml_eloss = engine.pyproxy('atenergy', ml_lattice, nargout=5) - assert_close(py_lattice.rf_voltage, ml_voltage, rtol=1.E-8) - assert_close(py_lattice.energy_loss, ml_eloss, rtol=1.E-6) + ml_energy, ml_periods, ml_voltage, ml_harms, ml_eloss = engine.pyproxy( + "atenergy", ml_lattice, nargout=5 + ) + assert_close(py_lattice.rf_voltage, ml_voltage, rtol=1.0e-8) + assert_close(py_lattice.energy_loss, ml_eloss, rtol=1.0e-6) assert py_lattice.energy == ml_energy assert py_lattice.periodicity == int(ml_periods) assert py_lattice.harmonic_number == int(ml_harms) @@ -158,4 +162,4 @@ def test_parameters(engine, lattices, dp): # test momentum compaction factor py_mcf = py_lattice.get_mcf(dp) ml_mcf = engine.mcf(ml_lattice, dp) - assert_close(py_mcf, ml_mcf, rtol=1.E-8) + assert_close(py_mcf, ml_mcf, rtol=1.0e-8) diff --git a/pyat/test_matlab/test_int.py b/pyat/test_matlab/test_int.py index d73c12b9a..e9f79feeb 100644 --- a/pyat/test_matlab/test_int.py +++ b/pyat/test_matlab/test_int.py @@ -1,17 +1,17 @@ -import numpy +import numpy as np import pytest + from at.physics import get_radiation_integrals -@pytest.mark.parametrize('lattices', - [pytest.lazy_fixture('hmba')]) -def test_radiation_integrals(engine, lattices): - py_lattice, ml_lattice, _ = lattices +@pytest.mark.parametrize("lattices", ["hmba"]) +def test_radiation_integrals(engine, request, lattices): + py_lattice, ml_lattice, _ = request.getfixturevalue(lattices) # Python call py_integrals = get_radiation_integrals(py_lattice) # Matlab call - results = engine.pyproxy('ringpara', ml_lattice) - ml_integrals = numpy.squeeze(results['integrals']) + results = engine.pyproxy("ringpara", ml_lattice) + ml_integrals = np.squeeze(results["integrals"]) # Comparison - numpy.testing.assert_almost_equal(py_integrals, ml_integrals[:5]) + np.testing.assert_almost_equal(py_integrals, ml_integrals[:5]) diff --git a/pyat/test_matlab/test_orbit_matrices.py b/pyat/test_matlab/test_orbit_matrices.py index 32c6417f1..11d15dc0f 100644 --- a/pyat/test_matlab/test_orbit_matrices.py +++ b/pyat/test_matlab/test_orbit_matrices.py @@ -1,146 +1,137 @@ -from at import uint32_refpts, find_m44, find_m66 -from at import find_orbit4, find_sync_orbit, find_orbit6 -import numpy import matlab -from numpy.testing import assert_allclose as assert_close +import numpy as np import pytest +from numpy.testing import assert_allclose as assert_close + +from at import find_orbit4, find_sync_orbit, find_orbit6 +from at import uint32_refpts, find_m44, find_m66 def _py_data(ml_data): """Convert a Matlab vector to a numpy vector""" - return numpy.squeeze(numpy.asarray(ml_data)) + return np.squeeze(np.asarray(ml_data)) def _ml_refs(refpts, nelems): """Convert refpoints to Matlab""" uintrefs = uint32_refpts(refpts, nelems) - return matlab.double([float(ref+1) for ref in uintrefs]) + return matlab.double([float(ref + 1) for ref in uintrefs]) -@pytest.mark.parametrize('ct', (-0.0002, 0.0002)) -@pytest.mark.parametrize('lattices', - [pytest.lazy_fixture('dba'), - pytest.lazy_fixture('hmba')]) -def test_find_syncorbit(engine, lattices, ct): - py_lattice, ml_lattice, _ = lattices +@pytest.mark.parametrize("ct", (-0.0002, 0.0002)) +@pytest.mark.parametrize("lattices", ["dba", "hmba"]) +def test_find_syncorbit(engine, request, lattices, ct): + py_lattice, ml_lattice, _ = request.getfixturevalue(lattices) nelems = len(py_lattice) - nrefs = nelems+1 + nrefs = nelems + 1 refpts = range(nrefs) # Python call py_orb5, py_orbit5 = find_sync_orbit(py_lattice, ct, refpts) # Matlab call - ml_orbit5, ml_orb5 = engine.findsyncorbit(ml_lattice, ct, - _ml_refs(refpts, nelems), - nargout=2) - ml_orbit5 = numpy.rollaxis(numpy.asarray(ml_orbit5), -1) + ml_orbit5, ml_orb5 = engine.findsyncorbit( + ml_lattice, ct, _ml_refs(refpts, nelems), nargout=2 + ) + ml_orbit5 = np.rollaxis(np.asarray(ml_orbit5), -1) # Comparison - assert_close(py_orb5, _py_data(ml_orb5), atol=1.e-12, rtol=0) - assert_close(py_orbit5[:, :5], ml_orbit5, atol=1.e-12, rtol=0) + assert_close(py_orb5, _py_data(ml_orb5), atol=1.0e-12, rtol=0) + assert_close(py_orbit5[:, :5], ml_orbit5, atol=1.0e-12, rtol=0) -@pytest.mark.parametrize('dp', (-0.01, 0.01)) -@pytest.mark.parametrize('lattices', - [pytest.lazy_fixture('dba'), - pytest.lazy_fixture('hmba')]) -def test_find_orbit4(engine, lattices, dp): - py_lattice, ml_lattice, _ = lattices +@pytest.mark.parametrize("dp", (-0.01, 0.01)) +@pytest.mark.parametrize("lattices", ["dba", "hmba"]) +def test_find_orbit4(engine, request, lattices, dp): + py_lattice, ml_lattice, _ = request.getfixturevalue(lattices) nelems = len(py_lattice) - nrefs = nelems+1 + nrefs = nelems + 1 refpts = range(nrefs) # Python call py_orb4, py_orbit4 = find_orbit4(py_lattice, dp, refpts) # Matlab call - ml_orbit4, ml_orb4 = engine.findorbit4(ml_lattice, dp, - _ml_refs(refpts, nelems), nargout=2) - ml_orbit4 = numpy.rollaxis(numpy.asarray(ml_orbit4), -1) + ml_orbit4, ml_orb4 = engine.findorbit4( + ml_lattice, dp, _ml_refs(refpts, nelems), nargout=2 + ) + ml_orbit4 = np.rollaxis(np.asarray(ml_orbit4), -1) # Comparison assert_close(py_orb4, _py_data(ml_orb4), atol=1e-15, rtol=0) assert_close(py_orbit4[:, :4], ml_orbit4, atol=1.5e-15, rtol=0) -@pytest.mark.parametrize('lattices', - [pytest.lazy_fixture('hmba_rad')]) -def test_find_orbit6(engine, lattices): - py_lattice, ml_lattice, _ = lattices +@pytest.mark.parametrize("lattices", ["hmba_rad"]) +def test_find_orbit6(engine, request, lattices): + py_lattice, ml_lattice, _ = request.getfixturevalue(lattices) nelems = len(py_lattice) - nrefs = nelems+1 + nrefs = nelems + 1 refpts = range(nrefs) # Python call py_orb6, py_orbit6 = find_orbit6(py_lattice, refpts) # Matlab call - ml_orbit6, ml_orb6 = engine.findorbit6(ml_lattice, _ml_refs(refpts, nelems), - nargout=2) - ml_orbit6 = numpy.rollaxis(numpy.asarray(ml_orbit6), -1) + ml_orbit6, ml_orb6 = engine.findorbit6( + ml_lattice, _ml_refs(refpts, nelems), nargout=2 + ) + ml_orbit6 = np.rollaxis(np.asarray(ml_orbit6), -1) # Comparison assert_close(py_orb6, _py_data(ml_orb6), atol=5e-15, rtol=0) assert_close(py_orbit6, ml_orbit6, atol=5e-15, rtol=0) -@pytest.mark.parametrize('dp', (-0.01, 0.0, 0.01)) -@pytest.mark.parametrize('lattices', - [pytest.lazy_fixture('dba'), - pytest.lazy_fixture('hmba')]) -def test_find_m44(engine, lattices, dp): - py_lattice, ml_lattice, _ = lattices +@pytest.mark.parametrize("dp", (-0.01, 0.0, 0.01)) +@pytest.mark.parametrize("lattices", ["dba", "hmba"]) +def test_find_m44(engine, request, lattices, dp): + py_lattice, ml_lattice, _ = request.getfixturevalue(lattices) nelems = len(py_lattice) - nrefs = nelems+1 + nrefs = nelems + 1 refpts = range(nrefs) morbit = engine.zeros(6, 1) - porbit = numpy.zeros((6, 1)) + porbit = np.zeros((6, 1)) # Python call py_m44, py_mstack = find_m44(py_lattice, dp, refpts, orbit=porbit) # Matlab call - ml_m44, ml_mstack = engine.findm44(ml_lattice, dp, - _ml_refs(refpts, nelems), - 'orbit', morbit, nargout=2) - ml_mstack = numpy.rollaxis(numpy.asarray(ml_mstack).reshape((4, 4, - nrefs)), -1) + ml_m44, ml_mstack = engine.findm44( + ml_lattice, dp, _ml_refs(refpts, nelems), "orbit", morbit, nargout=2 + ) + ml_mstack = np.rollaxis(np.asarray(ml_mstack).reshape((4, 4, nrefs)), -1) # Comparison - assert_close(py_m44, numpy.asarray(ml_m44), atol=1.e-11, rtol=0) - assert_close(py_mstack, ml_mstack, atol=1.e-11, rtol=0) + assert_close(py_m44, np.asarray(ml_m44), atol=1.0e-11, rtol=0) + assert_close(py_mstack, ml_mstack, atol=1.0e-11, rtol=0) -@pytest.mark.parametrize('lattices', - [pytest.lazy_fixture('hmba_rad')]) -def test_find_m66_orb(engine, lattices): - py_lattice, ml_lattice, _ = lattices +@pytest.mark.parametrize("lattices", ["hmba_rad"]) +def test_find_m66_orb(engine, request, lattices): + py_lattice, ml_lattice, _ = request.getfixturevalue(lattices) nelems = len(py_lattice) - nrefs = nelems+1 + nrefs = nelems + 1 refpts = range(nrefs) morbit = engine.zeros(6, 1) - porbit = numpy.zeros((6, 1)) + porbit = np.zeros((6, 1)) # Python call py_m66, py_mstack = find_m66(py_lattice, refpts, orbit=porbit) # Matlab call - ml_m66, ml_mstack = engine.findm66(ml_lattice, _ml_refs(refpts, nelems), - 'orbit', morbit, nargout=2) - assert_close(py_m66, numpy.asarray(ml_m66), atol=1.e-11, rtol=0) + ml_m66, ml_mstack = engine.findm66( + ml_lattice, _ml_refs(refpts, nelems), "orbit", morbit, nargout=2 + ) + assert_close(py_m66, np.asarray(ml_m66), atol=1.0e-11, rtol=0) - ml_mstack = numpy.rollaxis(numpy.asarray(ml_mstack).reshape((6, 6, - nrefs)), -1) - assert_close(py_mstack, ml_mstack, atol=1.e-11, rtol=0) + ml_mstack = np.rollaxis(np.asarray(ml_mstack).reshape((6, 6, nrefs)), -1) + assert_close(py_mstack, ml_mstack, atol=1.0e-11, rtol=0) -@pytest.mark.parametrize('lattices', - [pytest.lazy_fixture('hmba_rad')]) -def test_find_m66(engine, lattices): - py_lattice, ml_lattice, _ = lattices +@pytest.mark.parametrize("lattices", ["hmba_rad"]) +def test_find_m66(engine, request, lattices): + py_lattice, ml_lattice, _ = request.getfixturevalue(lattices) nelems = len(py_lattice) - nrefs = nelems+1 + nrefs = nelems + 1 refpts = range(nrefs) # Python call py_m66, py_mstack = find_m66(py_lattice, refpts) # Matlab call - ml_m66, ml_mstack = engine.findm66(ml_lattice, _ml_refs(refpts, nelems), - nargout=2) - assert_close(py_m66, numpy.asarray(ml_m66), atol=1e-7, rtol=0) + ml_m66, ml_mstack = engine.findm66(ml_lattice, _ml_refs(refpts, nelems), nargout=2) + assert_close(py_m66, np.asarray(ml_m66), atol=1e-7, rtol=0) - ml_mstack = numpy.rollaxis(numpy.asarray(ml_mstack).reshape((6, 6, - nrefs)), -1) + ml_mstack = np.rollaxis(np.asarray(ml_mstack).reshape((6, 6, nrefs)), -1) assert_close(py_mstack, ml_mstack, atol=1e-7, rtol=0) diff --git a/pyproject.toml b/pyproject.toml index 6d5cec15b..6b4415f69 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,7 +40,7 @@ Home="https://atcollab.github.io/at/" [project.optional-dependencies] mpi = ["mpi4py"] plot = ["matplotlib"] -dev = ["pytest >= 2.9", "pytest-lazy-fixture", "pytest-cov", "flake8"] +dev = ["pytest >= 2.9", "pytest-cov", "flake8"] doc = ["Sphinx ~= 7.2", "myst-parser", "myst-nb", "pydata-sphinx-theme ~= 0.14", "sphinx_design", "sphinx-copybutton"]