From 13d4c929298384c4015ca983d0d8f2815fe685b5 Mon Sep 17 00:00:00 2001 From: Christian Hespe Date: Tue, 27 Aug 2024 10:23:31 +0200 Subject: [PATCH 1/5] Add test for dipole broadcasting --- tests/test_dipole.py | 22 ++++++++++++++++------ 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/tests/test_dipole.py b/tests/test_dipole.py index d2e0982a..5a97daf2 100644 --- a/tests/test_dipole.py +++ b/tests/test_dipole.py @@ -1,6 +1,15 @@ +import pytest import torch -from cheetah import Dipole, Drift, ParameterBeam, ParticleBeam, Quadrupole, Segment +from cheetah import ( + Dipole, + Drift, + ParameterBeam, + ParticleBeam, + Quadrupole, + RBend, + Segment, +) def test_dipole_off(): @@ -43,22 +52,23 @@ def test_dipole_focussing(): assert not torch.allclose(outbeam_dipole_off.sigma_x, outbeam_quadrupole.sigma_x) -def test_dipole_batched_execution(): +@pytest.mark.parametrize("DipoleType", [Dipole, RBend]) +def test_dipole_batched_execution(DipoleType): """ Test that a dipole with batch dimensions behaves as expected. """ - batch_shape = torch.Size([3]) + batch_shape = torch.Size([6]) incoming = ParticleBeam.from_parameters( - num_particles=torch.tensor(1000000), + num_particles=torch.tensor(1_000_000), energy=torch.tensor([1e9]), mu_x=torch.tensor([1e-5]), ).broadcast(batch_shape) segment = Segment( [ - Dipole( + DipoleType( length=torch.tensor([0.5, 0.5, 0.5]), angle=torch.tensor([0.1, 0.2, 0.1]), - ), + ).broadcast((2,)), Drift(length=torch.tensor([0.5])).broadcast(batch_shape), ] ) From 1dd05be720d1f004098211105ebc256df1a37c00 Mon Sep 17 00:00:00 2001 From: Christian Hespe Date: Tue, 27 Aug 2024 10:27:51 +0200 Subject: [PATCH 2/5] Fix RBend broadcasting issue --- cheetah/accelerator/rbend.py | 17 +++-------------- 1 file changed, 3 insertions(+), 14 deletions(-) diff --git a/cheetah/accelerator/rbend.py b/cheetah/accelerator/rbend.py index 39192ecd..e7c21586 100644 --- a/cheetah/accelerator/rbend.py +++ b/cheetah/accelerator/rbend.py @@ -16,6 +16,7 @@ class RBend(Dipole): :param length: Length in meters. :param angle: Deflection angle in rad. + :param k1: Focussing strength in 1/m^-2. :param e1: The angle of inclination of the entrance face [rad]. :param e2: The angle of inclination of the exit face [rad]. :param tilt: Tilt of the magnet in x-y plane [rad]. @@ -30,6 +31,7 @@ def __init__( self, length: Optional[Union[torch.Tensor, nn.Parameter]], angle: Optional[Union[torch.Tensor, nn.Parameter]] = None, + k1: Optional[Union[torch.Tensor, nn.Parameter]] = None, e1: Optional[Union[torch.Tensor, nn.Parameter]] = None, e2: Optional[Union[torch.Tensor, nn.Parameter]] = None, tilt: Optional[Union[torch.Tensor, nn.Parameter]] = None, @@ -43,20 +45,7 @@ def __init__( super().__init__( length=length, angle=angle, - e1=e1, - e2=e2, - tilt=tilt, - fringe_integral=fringe_integral, - fringe_integral_exit=fringe_integral_exit, - gap=gap, - name=name, - device=device, - dtype=dtype, - ) - - super().__init__( - length=length, - angle=angle, + k1=k1, e1=e1, e2=e2, tilt=tilt, From 967f46e1dc773eeef48d5882e253743ebc5b4515 Mon Sep 17 00:00:00 2001 From: Christian Hespe Date: Tue, 27 Aug 2024 10:37:35 +0200 Subject: [PATCH 3/5] Update changelog --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index dec3153f..42d7968e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,7 +17,7 @@ - `Screen` now offers the option to use KDE for differentiable images (see #200) (@cr-xu, @roussel-ryan) - Moving `Element`s and `Beam`s to a different `device` and changing their `dtype` like with any `torch.nn.Module` is now possible (see #209) (@jank324) - `Quadrupole` now supports tracking with Cheetah's matrix-based method or with Bmad's more accurate method (see #153) (@jp-ga, @jank324) -- `Dipole` now takes a focusing moment `k1` (see #235) (@hespe) +- `Dipole` now takes a focusing moment `k1` (see #235, #247) (@hespe) ### 🐛 Bug fixes From b28a2f2a2baa65db56178e064f4435cb936696fe Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Tue, 27 Aug 2024 11:28:19 +0200 Subject: [PATCH 4/5] Apply suggestions from code review --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 42d7968e..602af756 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,7 +17,7 @@ - `Screen` now offers the option to use KDE for differentiable images (see #200) (@cr-xu, @roussel-ryan) - Moving `Element`s and `Beam`s to a different `device` and changing their `dtype` like with any `torch.nn.Module` is now possible (see #209) (@jank324) - `Quadrupole` now supports tracking with Cheetah's matrix-based method or with Bmad's more accurate method (see #153) (@jp-ga, @jank324) -- `Dipole` now takes a focusing moment `k1` (see #235, #247) (@hespe) +- `Dipole` and `Bend` now take a focusing moment `k1` (see #235, #247) (@hespe) ### 🐛 Bug fixes From ca448615ab50ec5ed7d6969df0db2e850e626641 Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Tue, 27 Aug 2024 11:28:37 +0200 Subject: [PATCH 5/5] Apply suggestions from code review --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 602af756..2b32a76c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,7 +17,7 @@ - `Screen` now offers the option to use KDE for differentiable images (see #200) (@cr-xu, @roussel-ryan) - Moving `Element`s and `Beam`s to a different `device` and changing their `dtype` like with any `torch.nn.Module` is now possible (see #209) (@jank324) - `Quadrupole` now supports tracking with Cheetah's matrix-based method or with Bmad's more accurate method (see #153) (@jp-ga, @jank324) -- `Dipole` and `Bend` now take a focusing moment `k1` (see #235, #247) (@hespe) +- `Dipole` and `RBend` now take a focusing moment `k1` (see #235, #247) (@hespe) ### 🐛 Bug fixes