Add test for the manzoni input

ULB-Metronu · Jan 8, 2024 · a4bc83b · a4bc83b
1 parent 1e438d0
commit a4bc83b
Show file tree

Hide file tree

Showing 11 changed files with 186 additions and 22 deletions.
diff --git a/docs/authors.rst b/docs/authors.rst
@@ -17,17 +17,21 @@ Alphabetical list of past and present authors:
 Citation
 ########
 
-Authors of papers including reference to or results obtained with Georges are encouraged to cite as follow::
-
-@article{Georges,
-	author = {Robin Tesse and C{\'e}dric Hernalsteens and Eustache Gnacadja and Nicolas Pauly and Eliott Ramoisiaux and Marion Vanwelde},
-	doi = {https://doi.org/10.1016/j.softx.2023.101579},
-	journal = {SoftwareX},
-	pages = {101579},
-	title = {Georges: A modular Python library for seamless beam dynamics simulations and optimization},
-	volume = {24},
-	year = {2023}
-   }
+Authors of papers including reference to or results obtained with Georges are encouraged to cite as follows:
+
+.. code-block:: bibtex
+
+    @article{Georges,
+        author  = {Robin Tesse and C{\'e}dric Hernalsteens and Eustache Gnacadja and Nicolas Pauly and Eliott Ramoisiaux and Marion Vanwelde},
+        doi     = {https://doi.org/10.1016/j.softx.2023.101579},
+        journal = {SoftwareX},
+        pages   = {101579},
+        title   = {Georges: A modular Python library for seamless beam dynamics simulations and optimization},
+        volume  = {24},
+        year    = {2023}
+    }
+
+
 
 List of publications
 ####################

diff --git a/docs/changelog.rst b/docs/changelog.rst
@@ -2,6 +2,14 @@
 Changelog
 *********
 
+* 2024-2
+
+    * Fix a bug when creating the Manzoni Input.
+    * Add some tests
+
+.. todolist::
+
+
 * 2024-1
 
     * Add doctest
@@ -11,8 +19,6 @@ Changelog
     * Add the paper published in Software-X
     * Update librairies
 
-.. todolist::
-
 * 2023-1
 
     * Add documentation of the PTW module

diff --git a/docs/installation.rst b/docs/installation.rst
@@ -13,7 +13,7 @@ For development purposes you can copy the source code on your computer. georges
 You can  stay on the bleeding-edge master branch or you can checkout
 a release tag::
 
-    git checkout tags/2023.1
+    git checkout tags/2024.2
 
 The installation process to get George’s running is relatively simple: the whole library is ready to
 be installed with `Poetry <https://python-poetry.org/>`_. However, depending your configuration,

diff --git a/docs/modules/manzoni/manzoni_integrators.rst b/docs/modules/manzoni/manzoni_integrators.rst
diff --git a/georges/manzoni/beam.py b/georges/manzoni/beam.py
@@ -60,6 +60,10 @@ def distribution(self, dist: np.ndarray):
 
 
 class MadXBeam(Beam):
+    """
+    Beam distribution for MADX Integrator (x, x', y, y', dpp, pt)
+    """
+
     def __init__(self, kinematics: _Kinematics, distribution: np.ndarray, first_order: bool = False):
         super().__init__(kinematics=kinematics, distribution=distribution)
         pt = Beam.compute_pt(
@@ -73,5 +77,13 @@ def __init__(self, kinematics: _Kinematics, distribution: np.ndarray, first_orde
 
 
 class TransportBeam(Beam):
+    """
+    Beam distribution for Transport Integrator (x, x', y, y', l, dpp)
+    """
+
     def __init__(self, kinematics: _Kinematics, distribution: np.ndarray):
         super().__init__(kinematics=kinematics, distribution=distribution)
+        self._distribution = np.zeros((len(distribution), 6))
+        self._distribution[:, :-1] = distribution
+        self._distribution[:, -1] = 0
+        self._distribution[:, [-2, -1]] = self._distribution[:, [-1, -2]]
diff --git a/georges/manzoni/input.py b/georges/manzoni/input.py
@@ -163,10 +163,14 @@ def set_parameters(self, element: str, parameters: Dict):
             self.sequence[self.mapper[element]].__setattr__(param, parameters[param])
         self.sequence[self.mapper[element]].freeze()
 
-    def get_parameters(self, element: str, parameters: Optional[List] = None):
+    def get_parameters(self, element: str, parameters: Optional[Union[List, str]] = None):
         if parameters is None:
             parameters = self.sequence[self.mapper[element]].attributes
-        return dict(zip(parameters, list(map(self.sequence[self.mapper[element]].__getattr__, parameters))))
+            return dict(zip(parameters, list(map(self.sequence[self.mapper[element]].__getattr__, parameters))))
+        if isinstance(parameters, List):
+            return dict(zip(parameters, list(map(self.sequence[self.mapper[element]].__getattr__, parameters))))
+        if parameters is not None and isinstance(parameters, str):
+            return self.sequence[self.mapper[element]].__getattr__(parameters)
 
     @classmethod
     def insert_thin_element(

diff --git a/georges/manzoni/kernels.py b/georges/manzoni/kernels.py
@@ -9,7 +9,7 @@
 
 
 @njit(nogil=True)  # no parallel jit found
-def batched_vector_matrix(b1: _np.ndarray, b2: _np.ndarray, matrix: _np.ndarray):  # pragma: no cover
+def batched_vector_matrix(b1: _np.ndarray, b2: _np.ndarray, matrix: _np.ndarray):
     """
 
     Args:
@@ -24,7 +24,7 @@ def batched_vector_matrix(b1: _np.ndarray, b2: _np.ndarray, matrix: _np.ndarray)
 
 
 @njit(parallel=True, nogil=True)
-def batched_vector_tensor(b1: _np.ndarray, b2: _np.ndarray, tensor: _np.ndarray):  # pragma: no cover
+def batched_vector_tensor(b1: _np.ndarray, b2: _np.ndarray, tensor: _np.ndarray):
     """
 
     Args:
@@ -51,7 +51,7 @@ def batched_vector_matrix_tensor(
     b2: _np.ndarray,
     matrix: _np.ndarray,
     tensor: _np.ndarray,
-):  # pragma: no cover
+):
     """
 
     Args:
@@ -75,5 +75,5 @@ def batched_vector_matrix_tensor(
 
 
 @njit
-def matrix_matrix(m1, m2):  # pragma: no cover
+def matrix_matrix(m1, m2):
     return _np.matmul(m1, m2)
diff --git a/sonar-project.properties b/sonar-project.properties
@@ -3,7 +3,7 @@ sonar.organization=ulb-metronu
 
 # This is the name and version displayed in the SonarCloud UI.
 sonar.projectName=georges
-sonar.projectVersion=2024.1
+sonar.projectVersion=2024.2
 
 # Path is relative to the sonar-project.properties file. Replace "\" by "/" on Windows.
 sonar.sources=georges/

diff --git a/tests/manzoni/test_input.py b/tests/manzoni/test_input.py
@@ -0,0 +1,138 @@
+import numpy.testing as npt
+
+import georges
+from georges import ureg as _ureg
+from georges.manzoni import Input
+from georges.manzoni.beam import MadXBeam, TransportBeam
+from georges.manzoni.integrators import TransportSecondOrderTaylorIntegratorExact
+from georges.manzoni.observers import LossesObserver, MeanObserver
+
+
+def test_setting_parameters():
+    B2G2 = georges.Element.SBend(
+        NAME="B2G2",
+        L=1.492 * _ureg.m,
+        ANGLE=-90 * _ureg.degrees,
+        K1=0 * _ureg.m**-2,
+        APERTYPE="RECTANGULAR",
+        APERTURE=[5 * _ureg.cm, 2.8 * _ureg.cm],
+        HGAP=0.0315 * _ureg.m,
+        E1=-0.4346 * _ureg.radians,
+        E2=-0.2889 * _ureg.radians,
+        K0=-0.97 * 1.0528125514711102 / _ureg.m,
+        FINT=0.5,
+        FINTX=0.5,
+    )
+
+    sequence = georges.PlacementSequence(name="Sequence")
+    sequence.place(B2G2, at_entry=0 * _ureg.m)
+    sequence.place_after_last(B2G2)
+
+    mi = Input.from_sequence(sequence=sequence)
+    mi.set_parameters("B2G2", parameters={"K0": 1 * _ureg.m**-1})
+
+    assert mi.get_parameters("B2G2", "K0") == 1 * _ureg.m**-1
+    assert mi.get_parameters("B2G2", ["L", "HGAP"]) == {"L": 1.492 * _ureg.m, "HGAP": 0.0315 * _ureg.m}
+    assert mi.get_parameters("B2G2") == {
+        "NAME": "B2G2",
+        "AT_ENTRY": 1.492 * _ureg.m,
+        "AT_CENTER": 2.238 * _ureg.m,
+        "AT_EXIT": 2.984 * _ureg.m,
+        "APERTYPE": "RECTANGULAR",
+        "APERTURE": [5 * _ureg.cm, 2.8 * _ureg.cm],
+        "KINEMATICS": None,
+        "ANGLE": -90 * _ureg.degrees,
+        "K0": 1 * _ureg.m**-1,
+        "K1": 0 * _ureg.m**-2,
+        "K2": 0.0 * _ureg.m**-3,
+        "L": 1.492 * _ureg.m,
+        "E1": -0.4346 * _ureg.radians,
+        "E2": -0.2889 * _ureg.radians,
+        "TILT": 0.0 * _ureg.radians,
+        "HGAP": 0.0315 * _ureg.m,
+        "FINT": 0.5,
+        "FINTX": 0.5,
+    }
+
+
+def test_integrator_setting():
+    B2G2 = georges.Element.SBend(
+        NAME="B2G2",
+        L=1.492 * _ureg.m,
+        ANGLE=-10 * _ureg.degrees,
+        K1=0 * _ureg.m**-2,
+        APERTYPE="RECTANGULAR",
+        APERTURE=[5 * _ureg.cm, 2.8 * _ureg.cm],
+        HGAP=0.0315 * _ureg.m,
+        E1=-0.4346 * _ureg.radians,
+        E2=-0.2889 * _ureg.radians,
+    )
+
+    sequence = georges.PlacementSequence(name="Sequence")
+    sequence.place(B2G2, at_entry=0 * _ureg.m)
+
+    kin = georges.Kinematics(250 * _ureg.MeV, particle=georges.particles.Proton, kinetic=True)
+    sequence.metadata.kinematics = kin
+
+    mi = Input.from_sequence(sequence=sequence)
+    mi.freeze()
+    mi.set_integrator(integrator=TransportSecondOrderTaylorIntegratorExact)
+    beam = TransportBeam(
+        kinematics=kin,
+        distribution=georges.Distribution.from_twiss_parameters(
+            n=1000,
+            x=2.5 * _ureg.mm,
+            y=2.5 * _ureg.mm,
+            emitx=7 * _ureg.mm * _ureg.mradians,
+            emity=7 * _ureg.mm * _ureg.mradians,
+        ).distribution.values,
+    )
+    beam_observer_mean = mi.track(beam=beam, observers=MeanObserver())
+    assert beam_observer_mean is not None
+
+
+def test_degrader_efficiency():
+    d1 = georges.Element.Drift(
+        NAME="D1",
+        L=0.3 * _ureg.m,
+    )
+
+    c1 = georges.Element.Degrader(
+        NAME="DEG",
+        MATERIAL=georges.fermi.materials.Beryllium,
+        L=10 * _ureg.cm,
+        WITH_LOSSES=True,
+    )
+
+    d2 = georges.Element.Drift(
+        NAME="D2",
+        L=0.3 * _ureg.m,
+    )
+
+    sequence = georges.PlacementSequence(name="Sequence")
+
+    sequence.place(d1, at_entry=0 * _ureg.m)
+    sequence.place_after_last(c1)
+    sequence.place_after_last(d2)
+
+    kin = georges.Kinematics(230 * _ureg.MeV, particle=georges.particles.Proton, kinetic=True)
+    sequence.metadata.kinematics = kin
+
+    beam = MadXBeam(
+        kinematics=kin,
+        distribution=georges.Distribution.from_5d_multigaussian_distribution(
+            n=100000,
+            xrms=0.1 * _ureg.cm,
+            yrms=0.7 * _ureg.cm,
+            pxrms=0.01,
+            pyrms=0.01,
+        ).distribution.values,
+    )
+
+    mi = Input.from_sequence(sequence=sequence)
+    efficiency = mi.compute_efficiency(input_energy=kin.ekin)
+
+    losses_observer = mi.track(beam=beam, observers=LossesObserver()).to_df()
+    efficiency_observer = losses_observer.iloc[-1]["PARTICLES_OUT"] / losses_observer.iloc[0]["PARTICLES_IN"]
+
+    npt.assert_allclose(efficiency_observer, efficiency, rtol=0.02)
diff --git a/tests/test_manzoni_tracking.py → tests/manzoni/test_manzoni_tracking.py b/tests/test_manzoni_tracking.py → tests/manzoni/test_manzoni_tracking.py
@@ -40,7 +40,7 @@ def test_manzoni_tracking():
     )
 
     B3G2 = georges.Element.SBend(
-        NAME="B2G2",
+        NAME="B3G2",
         L=1.492 * _ureg.m,
         ANGLE=-90 * _ureg.degrees,
         K1=0 * _ureg.m**-2,

diff --git a/tests/test_manzoni_twiss.py → tests/manzoni/test_manzoni_twiss.py b/tests/test_manzoni_twiss.py → tests/manzoni/test_manzoni_twiss.py