diff --git a/CHANGELOG.md b/CHANGELOG.md index 2b13a7511..424b1fc50 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,8 @@ ### 🐛 Bug fixes +- Fix the ocelot import method to properly import a `TDCavity` to `TransverseDeflectingCavity`. (see #675) (@cr-xu) + ### 🐆 Other ### 🌟 First Time Contributors diff --git a/cheetah/converters/ocelot.py b/cheetah/converters/ocelot.py index 7964f577b..2d8214d5f 100644 --- a/cheetah/converters/ocelot.py +++ b/cheetah/converters/ocelot.py @@ -142,12 +142,13 @@ def convert_element( sanitize_name=sanitize_name, ) elif isinstance(element, ocelot.TDCavity): - # TODO: Better replacement at some point? - return cheetah.Cavity( + return cheetah.TransverseDeflectingCavity( length=torch.tensor(element.l, **factory_kwargs), voltage=torch.tensor(element.v, **factory_kwargs) * 1e9, frequency=torch.tensor(element.freq, **factory_kwargs), - phase=torch.tensor(element.phi, **factory_kwargs), + phase=torch.tensor(element.phi, **factory_kwargs).deg2rad() + / (2 * torch.pi), + tilt=torch.tensor(element.tilt, **factory_kwargs), name=element.id, sanitize_name=sanitize_name, ) diff --git a/tests/test_ocelot_import.py b/tests/test_ocelot_import.py index 4a14a0992..248b578bc 100644 --- a/tests/test_ocelot_import.py +++ b/tests/test_ocelot_import.py @@ -104,6 +104,7 @@ def test_ocelot_lattice_import(): ocelot.Quadrupole(l=0.2), ocelot.Drift(l=1.0), ocelot.Sextupole(l=0.4), + ocelot.TDCavity(l=1.0, v=0.01, freq=1e9, phi=0.0), ] segment = cheetah.Segment.from_ocelot(cell=cell) @@ -111,6 +112,7 @@ def test_ocelot_lattice_import(): assert isinstance(segment.elements[1], cheetah.Quadrupole) assert isinstance(segment.elements[2], cheetah.Drift) assert isinstance(segment.elements[3], cheetah.Sextupole) + assert isinstance(segment.elements[4], cheetah.TransverseDeflectingCavity) assert segment.elements[0].length.device.type == "cpu" assert segment.elements[1].length.device.type == "cpu" @@ -119,3 +121,7 @@ def test_ocelot_lattice_import(): assert segment.elements[2].length.device.type == "cpu" assert segment.elements[3].length.device.type == "cpu" assert segment.elements[3].k2.device.type == "cpu" + assert segment.elements[4].length.device.type == "cpu" + assert segment.elements[4].voltage.device.type == "cpu" + assert segment.elements[4].frequency.device.type == "cpu" + assert segment.elements[4].phase.device.type == "cpu"