diff --git a/CHANGELOG.md b/CHANGELOG.md index c069ff6b..fa71be68 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,7 +23,7 @@ This is a major release with significant upgrades under the hood of Cheetah. Des - Port Bmad-X tracking methods to Cheetah for `Quadrupole`, `Drift`, and `Dipole` (see #153, #240) (@jp-ga, @jank324) - Add `TransverseDeflectingCavity` element (following the Bmad-X implementation) (see #240) (@jp-ga) - `Dipole` and `RBend` now take a focusing moment `k1` (see #235, #247) (@hespe) -- Implement a converter for lattice files imported from Elegant (see #222, #251) (@hespe) +- Implement a converter for lattice files imported from Elegant (see #222, #251, #273) (@hespe) ### 🐛 Bug fixes diff --git a/cheetah/converters/elegant.py b/cheetah/converters/elegant.py index 9f69e2f7..4927593f 100644 --- a/cheetah/converters/elegant.py +++ b/cheetah/converters/elegant.py @@ -146,6 +146,7 @@ def convert_element( dtype=dtype, ), ], + name=name + "_segment", ) elif parsed["element_type"] == "rcol": validate_understood_properties( @@ -169,6 +170,7 @@ def convert_element( dtype=dtype, ), ], + name=name + "_segment", ) elif parsed["element_type"] == "quad": validate_understood_properties( @@ -198,17 +200,27 @@ def convert_element( ) elif parsed["element_type"] == "moni": validate_understood_properties(["element_type", "group", "l"], parsed) - return cheetah.Segment( - elements=[ - cheetah.Drift( - length=torch.tensor(parsed.get("l", 0.0)), - name=name + "_drift", - device=device, - dtype=dtype, - ), - cheetah.Marker(name=name), - ] - ) + if "l" in parsed: + return cheetah.Segment( + elements=[ + cheetah.Drift( + length=torch.tensor(parsed["l"] / 2), + name=name + "_predrift", + device=device, + dtype=dtype, + ), + cheetah.BPM(name=name), + cheetah.Drift( + length=torch.tensor(parsed["l"] / 2), + name=name + "_postdrift", + device=device, + dtype=dtype, + ), + ], + name=name + "_segment", + ) + else: + return cheetah.BPM(name=name) elif parsed["element_type"] == "ematrix": validate_understood_properties( ["element_type", "l", "order", "c[1-6]", "r[1-6][1-6]", "group"], @@ -225,14 +237,21 @@ def convert_element( [ [parsed.get(f"r{i + 1}{j + 1}", 0.0) for j in range(6)] for i in range(6) - ] + ], + device=device, + dtype=dtype, ) # Add affine component (constant offset) - R[:6, 6] = torch.tensor(parsed.get(f"c{i + 1}", 0.0) for i in range(6)) + R[:6, 6] = torch.tensor( + [parsed.get(f"c{i + 1}", 0.0) for i in range(6)], + device=device, + dtype=dtype, + ) return cheetah.CustomTransferMap( length=torch.tensor(parsed["l"]), transfer_map=R, + name=name, device=device, dtype=dtype, ) @@ -343,7 +362,7 @@ def convert_element( length=torch.tensor(parsed["l"]), angle=torch.tensor(parsed.get("angle", 0.0)), k1=torch.tensor(parsed.get("k1", 0.0)), - e1=torch.tensor(parsed["e1"]), + e1=torch.tensor(parsed.get("e1", 0.0)), e2=torch.tensor(parsed.get("e2", 0.0)), tilt=torch.tensor(parsed.get("tilt", 0.0)), name=name, @@ -358,7 +377,7 @@ def convert_element( return cheetah.RBend( length=torch.tensor(parsed["l"]), angle=torch.tensor(parsed.get("angle", 0.0)), - e1=torch.tensor(parsed["e1"]), + e1=torch.tensor(parsed.get("e1", 0.0)), e2=torch.tensor(parsed.get("e2", 0.0)), tilt=torch.tensor(parsed.get("tilt", 0.0)), name=name, @@ -394,7 +413,7 @@ def convert_element( length=torch.tensor(parsed["l"]), angle=torch.tensor(parsed.get("angle", 0.0)), k1=torch.tensor(parsed.get("k1", 0.0)), - e1=torch.tensor(parsed["e1"]), + e1=torch.tensor(parsed.get("e1", 0.0)), e2=torch.tensor(parsed.get("e2", 0.0)), tilt=torch.tensor(parsed.get("tilt", 0.0)), name=name, diff --git a/tests/resources/cavity.lte b/tests/resources/cavity.lte new file mode 100644 index 00000000..0fc46540 --- /dev/null +++ b/tests/resources/cavity.lte @@ -0,0 +1,4 @@ +C1E: EMATRIX, L = 0, ORDER = 1, C2=-0.0027, C4=-0.150, R11=1, R21=0.04, R22=1, R23=0.003, R33=1, R41=0.003, R43=-0.04, R44=1, R55=1, R66=1 +C1: RFCA, L = 0.7, PHASE = 90.000000, VOLT = 16175000.000000, FREQ = 1200000000.000000 + +cavity: line=(C1E,C1) diff --git a/tests/test_elegant_conversion.py b/tests/test_elegant_conversion.py index 93b17ece..5e4ac94b 100644 --- a/tests/test_elegant_conversion.py +++ b/tests/test_elegant_conversion.py @@ -1,3 +1,4 @@ +import numpy as np import pytest import torch @@ -32,14 +33,48 @@ def test_fodo(): assert [element.name for element in converted.elements] == [ element.name for element in correct_lattice.elements ] - assert converted.q1.length == correct_lattice.q1.length - assert converted.q1.k1 == correct_lattice.q1.k1 - assert converted.q2.length == correct_lattice.q2.length - assert converted.q2.k1 == correct_lattice.q2.k1 - assert [d.length for d in converted.d1] == [d.length for d in correct_lattice.d1] - assert converted.d2.length == correct_lattice.d2.length - assert converted.s1.length == correct_lattice.s1.length - assert converted.s1.e1 == correct_lattice.s1.e1 + assert torch.isclose(converted.q1.length, correct_lattice.q1.length) + assert torch.isclose(converted.q1.k1, correct_lattice.q1.k1) + assert torch.isclose(converted.q2.length, correct_lattice.q2.length) + assert torch.isclose(converted.q2.k1, correct_lattice.q2.k1) + for i in range(2): + assert torch.isclose(converted.d1[i].length, correct_lattice.d1[i].length) + assert torch.isclose(converted.d2.length, correct_lattice.d2.length) + assert torch.isclose(converted.s1.length, correct_lattice.s1.length) + assert torch.isclose(converted.s1.e1, correct_lattice.s1.e1) + + +def test_cavity_import(): + """Test importing an accelerating cavity defined in the Elegant file format.""" + file_path = "tests/resources/cavity.lte" + converted = cheetah.Segment.from_elegant(file_path, "cavity") + + assert np.isclose(converted.c1.length, 0.7) + assert np.isclose(converted.c1.frequency, 1.2e9) + assert np.isclose(converted.c1.voltage, 16.175e6) + + # Cheetah and Elegant use different phase conventions shifted by 90 deg + assert np.isclose(converted.c1.phase, 0.0) + + +def test_custom_transfer_map_import(): + """Test importing an Elegant EMATRIX into a Cheetah CustomTransferMap.""" + file_path = "tests/resources/cavity.lte" + converted = cheetah.Segment.from_elegant(file_path, "cavity") + + correct_transfer_map = torch.tensor( + [ + [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.04, 1.0, 0.003, 0.0, 0.0, 0.0, -0.0027], + [0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0], + [0.003, 0.0, -0.04, 1.0, 0.0, 0.0, -0.15], + [0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ] + ) + + assert torch.allclose(converted.c1e._transfer_map, correct_transfer_map) @pytest.mark.parametrize(