Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Fix parsing of BPM and CustomTransferMap from Elegant #273

Merged
merged 8 commits into from
Oct 10, 2024
Merged

Fix parsing of BPM and CustomTransferMap from Elegant #273

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
53ecd0a
Improve BPM parsing from elegant
Hespe Oct 7, 2024
5b9f2ff
Unify sbend parsing fallbacks
Hespe Oct 7, 2024
e56199a
Assign proper dtype for CustomTransferMap imported from Elegant
Hespe Oct 9, 2024
92fcd06
Update changelog
Hespe Oct 9, 2024
2ec03f1
Add name to imported TransferMap
Hespe Oct 9, 2024
0333ed0
Add test case for cavity import
Hespe Oct 9, 2024
e58fd02
Apply suggestions from code review
jank324 Oct 9, 2024
640b5bd
Split tests on Cavity and TransferMap import
Hespe Oct 10, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
51 changes: 35 additions & 16 deletions cheetah/converters/elegant.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -146,6 +146,7 @@ def convert_element(
dtype=dtype,
),
],
name=name + "_segment",
)
elif parsed["element_type"] == "rcol":
validate_understood_properties(
Expand All @@ -169,6 +170,7 @@ def convert_element(
dtype=dtype,
),
],
name=name + "_segment",
)
elif parsed["element_type"] == "quad":
validate_understood_properties(
Expand Down Expand Up @@ -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"],
Expand All @@ -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,
)
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
4 changes: 4 additions & 0 deletions tests/resources/cavity.lte
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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)
51 changes: 43 additions & 8 deletions tests/test_elegant_conversion.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import numpy as np
import pytest
import torch

Expand Down Expand Up @@ -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)
Comment on lines +40 to +44
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I feel like this should be no loop and use torch.allclose.

Copy link
Member Author

@Hespe Hespe Oct 10, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am not sure how torch.allclose would work here, but maybe there is some trick I am missing. The problem is that we have two elements with name d1 and therefore two separate length tensors, not one multidimensional tensor.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh I see ... I misunderstood the that line 🙃



def test_cavity_import():
"""Test importing an accelerating cavity defined in the Elegant file format."""
Hespe marked this conversation as resolved.
Show resolved Hide resolved
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(
Expand Down