Multipole implementation - #399
Conversation
| outgoing_drift = drift.track(incoming) | ||
|
|
||
| assert torch.allclose( | ||
| outgoing_multipole.particles, outgoing_drift.particles, atol=1e-5 |
There was a problem hiding this comment.
Isn't this tolerance way too high?
| # Both should cause focusing/defocusing, but exact values will differ due to | ||
| # different tracking methods. Just verify they both affect the beam. | ||
| assert abs(outgoing_multipole.mu_px) > 1e-9 | ||
| assert abs(outgoing_quadrupole.mu_px) > 1e-9 | ||
|
|
||
| # The sign of px change should be the same for both | ||
| assert torch.sign(outgoing_multipole.mu_px) == torch.sign(outgoing_quadrupole.mu_px) |
There was a problem hiding this comment.
These assertions also seem odd to me. Wouldn't it be better to compare particle cloud to particle cloud?
| k1 = torch.tensor(4.2) | ||
|
|
||
| multipole = cheetah.Multipole( | ||
| length=torch.tensor(1.0), polynom_b=torch.tensor([0.0, k1]) # B1 = k1 |
There was a problem hiding this comment.
I'm not sure with the terminology, but I feel like it should be made to match Cheetah's existing terminology.
| # Verify they produce similar results | ||
| assert torch.allclose(outgoing_corrector.mu_px, outgoing_multipole.mu_px, atol=1e-9) | ||
| assert torch.allclose(outgoing_corrector.mu_py, outgoing_multipole.mu_py, atol=1e-9) | ||
| assert torch.allclose(outgoing_corrector.mu_x, outgoing_multipole.mu_x, atol=1e-9) | ||
| assert torch.allclose(outgoing_corrector.mu_y, outgoing_multipole.mu_y, atol=1e-9) |
There was a problem hiding this comment.
Same question here: This feels like comparing the particle arrays would be much better. The current version doesn't capture everything and yet requires a lot more code.
| assert torch.allclose(outgoing_corrector.mu_y, outgoing_multipole.mu_y, atol=1e-9) | ||
|
|
||
|
|
||
| def test_multipole_with_misalignment(): |
There was a problem hiding this comment.
This test seems superfluous to me, if we are already testing against a quadrupole.
| ) | ||
|
|
||
|
|
||
| def test_skew_multipole_behavior(): |
There was a problem hiding this comment.
Same here: If we already trust the multipole to behave like a quadrupole, why do this?
jank324
changed the title
3 participants
WIP: Add multipole element implementation with tests and benchmarks
Description
This PR adds a general multipole element to Cheetah's accelerator elements. It includes:
Multipoleclass incheetah/accelerator/multipole.pythat implements arbitrary-order normal and skew multipole fieldstests/test_multipole.pyverifying that the multipole implementation correctly reproduces the behavior of drift elements, dipoles, and quadrupoles when configured with appropriate coefficientsbenchmarks/benchmark_multipole_pyat.pythat compares Cheetah's multipole implementation against pyAT for validation and performance analysisMotivation and Context
Adding a general multipole element enhances Cheetah's capabilities for accurate accelerator simulations. This implementation:
The implementation includes comprehensive benchmark tests against pyAT's multipole element which all work reasonably well. I've cross-checked it with Cheetah's corrector magnet and quadrupole magnet implementations using the beam matrix tracking approach.
Note that vectorization is not yet included in this implementation.
Types of changes
Checklist
flake8(required).pytesttests pass (required).pyteston a machine with a CUDA GPU and made sure all tests pass (required).