Beam functionalit and models (#49)

.codeclimate.yml

@@ -0,0 +1,52 @@
+version: "2"
+checks:
+  argument-count:
+    enabled: true
+    config:
+      threshold: 6
+  complex-logic:
+    enabled: true
+    config:
+      threshold: 10
+  file-lines:
+    enabled: true
+    config:
+      threshold: 1000
+  method-complexity:
+    enabled: true
+    config:
+      threshold: 10
+  method-count:
+    enabled: true
+    config:
+      threshold: 20
+  method-lines:
+    enabled: true
+    config:
+      threshold: 50
+  nested-control-flow:
+    enabled: true
+    config:
+      threshold: 5
+  return-statements:
+    enabled: true
+    config:
+      threshold: 5
+  similar-code:
+    enabled: true
+  identical-code:
+    enabled: true
+plugins:
+  pep8:
+    enabled: false
+  sonar-python:
+    enabled: true
+    config:
+      tests_patterns:
+        - tests/**
+exclude_patterns:
+- "config/"
+- "dist/"
+- "sdist"
+- "**/test/"
+- "**/tests/"

pyhdtoolkit/cpymadtools/__init__.py

@@ -12,7 +12,7 @@
 from .latwiss import plot_latwiss, plot_machine_survey
 from .matching import get_closest_tune_approach, get_lhc_tune_and_chroma_knobs, match_tunes_and_chromaticities
 from .orbit import correct_lhc_orbit, get_current_orbit_setup, lhc_orbit_variables, setup_lhc_orbit
-from .parameters import beam_parameters
+from .parameters import query_beam_attributes
 from .plotters import AperturePlotter, DynamicAperturePlotter, PhaseSpacePlotter, TuneDiagramPlotter
 from .ptc import get_amplitude_detuning, get_rdts
 from .special import (

pyhdtoolkit/cpymadtools/matching.py

@@ -112,7 +112,7 @@ def match(*args, **kwargs):
         madx.command.lmdif(calls=calls, tolerance=tolerance)
         madx.command.endmatch()
         logger.trace("Performing routine TWISS")
-        madx.twiss()  # prevents errors if the user forget to do so before querying tables
+        madx.twiss()  # prevents errors if the user forgets to TWISS before querying tables
 
     if q1_target is not None and q2_target is not None and dq1_target is not None and dq2_target is not None:
         logger.info(

pyhdtoolkit/cpymadtools/parameters.py

@@ -7,65 +7,26 @@
 
 A module with functions to compute different beam and machine parameters.
 """
-from typing import Dict
-
-import numpy as np
 
+from cpymad.madx import Madx
 from loguru import logger
 
+from pyhdtoolkit.models.madx import MADXBeam
+
 # ----- Utilities ----- #
 
 
-def beam_parameters(
-    pc_gev: float, en_x_m: float, en_y_m: float, deltap_p: float, verbose: bool = False,
-) -> Dict[str, float]:
-    """Calculate beam parameters from provided values.
+def query_beam_attributes(madx: Madx) -> MADXBeam:
+    """
+    Returns all `BEAM` attributes from the `MAD-X` process based on the currently defined beam. If no beam
+    has been defined at function call, then `MAD-X` will return all the default values. See the `MAD-X`
+    manual for details.
 
     Args:
-        pc_gev (float): particle momentum.
-        en_x_m (float): horizontal emittance, in meters.
-        en_y_m (float): vertical emittance, in meters.
-        deltap_p (float): momentum deviation.
-        verbose (bool): bolean, whether to print out a summary of parameters or not.
+        madx (cpymad.madx.Madx): an instanciated cpymad Madx object.
 
     Returns:
-        A dictionnary with the calculated values.
+        A validated `MADXBeam` object.
     """
-    e0_gev = 0.9382720813
-    e_tot_gev = np.sqrt(pc_gev ** 2 + e0_gev ** 2)
-    gamma_r = e_tot_gev / e0_gev
-    beta_r = pc_gev / np.sqrt(pc_gev ** 2 + e0_gev ** 2)
-
-    parameters: Dict[str, float] = {
-        "pc_GeV": pc_gev,  # Particle momentum [GeV]
-        "B_rho_Tm": 3.3356 * pc_gev,  # Beam rigidity [T/m]
-        "E_0_GeV": e0_gev,  # Rest mass energy [GeV]
-        "E_tot_GeV": e_tot_gev,  # Total beam energy [GeV]
-        "E_kin_GeV": e_tot_gev - e0_gev,  # Kinectic beam energy [GeV]
-        "gamma_r": gamma_r,  # Relativistic gamma
-        "beta_r": beta_r,  # Relativistic beta
-        "en_x_m": en_x_m,  # Horizontal emittance [m]
-        "en_y_m": en_y_m,  # Vertical emittance [m]
-        "eg_x_m": en_x_m / gamma_r / beta_r,  # Horizontal geometrical emittance
-        "eg_y_m": en_y_m / gamma_r / beta_r,  # Vertical geometrical emittance
-        "deltap_p": deltap_p,  # Momentum deviation
-    }
-
-    if verbose:
-        logger.trace("Outputing computed parameter values")
-        print(
-            f"""Particle type: proton
-        Beam momentum = {parameters["pc_GeV"]:2.3f} GeV/c
-        Normalized x-emittance = {parameters["en_x_m"] * 1e6:2.3f} mm mrad
-        Normalized y-emittance = {parameters["en_y_m"] * 1e6:2.3f} mm mrad
-        Momentum deviation deltap/p = {parameters["deltap_p"]}
-        -> Beam total energy = {parameters["E_tot_GeV"]:2.3f} GeV
-        -> Beam kinetic energy = {parameters["E_kin_GeV"]:2.3f} GeV
-        -> Beam rigidity = {parameters["B_rho_Tm"]:2.3f} Tm
-        -> Relativistic beta = {parameters["beta_r"]:2.5f}
-        -> Relativistic gamma = {parameters["gamma_r"]:2.3f}
-        -> Geometrical x emittance = {parameters["eg_x_m"] * 1e6:2.3f} mm mrad
-        -> Geometrical y emittance = {parameters["eg_y_m"] * 1e6:2.3f} mm mrad
-        """
-        )
-    return parameters
+    logger.info("Retrieving BEAM attributes from the MAD-X process")
+    return MADXBeam(**dict(madx.beam))

pyhdtoolkit/cpymadtools/plotters.py

@@ -9,7 +9,7 @@
 simulation results.
 """
 from pathlib import Path
-from typing import Dict, Tuple
+from typing import Tuple
 
 import matplotlib
 import matplotlib.pyplot as plt
@@ -20,6 +20,7 @@
 from loguru import logger
 from matplotlib import colors as mcolors
 
+from pyhdtoolkit.models.beam import BeamParameters
 from pyhdtoolkit.optics.twiss import courant_snyder_transform
 from pyhdtoolkit.utils.defaults import PLOT_PARAMS
 
@@ -40,7 +41,7 @@ class AperturePlotter:
     @staticmethod
     def plot_aperture(
         madx: Madx,
-        beam_params: Dict[str, float],
+        beam_params: BeamParameters,
         figsize: Tuple[int, int] = (13, 20),
         xlimits: Tuple[float, float] = None,
         hplane_ylim: Tuple[float, float] = (-0.12, 0.12),
@@ -53,8 +54,8 @@ def plot_aperture(
 
         Args:
             madx (cpymad.madx.Madx): an instanciated cpymad Madx object.
-            beam_params (Dict[str, float]): a beam_parameters dictionary obtained through
-                cpymadtools.helpers.beam_parameters.
+            beam_params (BeamParameters): a validated BeamParameters object from
+                `pyhdtoolkit.optics.beam.compute_beam_parameters`.
             figsize (str): size of the figure, defaults to (15, 15).
             xlimits (Tuple[float, float]): will implement xlim (for the s coordinate) if this is
                 not None, using the tuple passed.
@@ -83,15 +84,15 @@ def plot_aperture(
 
         logger.debug("Getting Twiss dframe from cpymad")
         twiss_hr: pd.DataFrame = madx.table.twiss.dframe()
-        twiss_hr["betatronic_envelope_x"] = np.sqrt(twiss_hr.betx.values * beam_params["eg_y_m"])
-        twiss_hr["betatronic_envelope_y"] = np.sqrt(twiss_hr.bety.values * beam_params["eg_y_m"])
-        twiss_hr["dispersive_envelope_x"] = twiss_hr.dx.values * beam_params["deltap_p"]
-        twiss_hr["dispersive_envelope_y"] = twiss_hr.dy.values * beam_params["deltap_p"]
+        twiss_hr["betatronic_envelope_x"] = np.sqrt(twiss_hr.betx.values * beam_params.eg_y_m)
+        twiss_hr["betatronic_envelope_y"] = np.sqrt(twiss_hr.bety.values * beam_params.eg_y_m)
+        twiss_hr["dispersive_envelope_x"] = twiss_hr.dx.values * beam_params.deltap_p
+        twiss_hr["dispersive_envelope_y"] = twiss_hr.dy.values * beam_params.deltap_p
         twiss_hr["envelope_x"] = np.sqrt(
-            twiss_hr.betatronic_envelope_x.values ** 2 + (twiss_hr.dx.values * beam_params["deltap_p"]) ** 2
+            twiss_hr.betatronic_envelope_x.values ** 2 + (twiss_hr.dx.values * beam_params.deltap_p) ** 2
         )
         twiss_hr["envelope_y"] = np.sqrt(
-            twiss_hr.betatronic_envelope_y.values ** 2 + (twiss_hr.dy.values * beam_params["deltap_p"]) ** 2
+            twiss_hr.betatronic_envelope_y.values ** 2 + (twiss_hr.dy.values * beam_params.deltap_p) ** 2
         )
         machine = twiss_hr[twiss_hr.apertype == "ellipse"]
 
@@ -113,7 +114,7 @@ def plot_aperture(
         axis1.set_ylim(hplane_ylim)
         axis1.set_ylabel("x [m]")
         axis1.set_xlabel("s [m]")
-        axis1.set_title(f"Horizontal aperture at {beam_params['pc_GeV']} GeV/c")
+        axis1.set_title(f"Horizontal aperture at {beam_params.pc_GeV} GeV/c")
 
         logger.debug("Plotting the vertical aperture")
         axis2 = plt.subplot2grid((3, 3), (1, 0), colspan=3, rowspan=1, sharex=axis1)
@@ -134,7 +135,7 @@ def plot_aperture(
         axis2.set_ylim(vplane_ylim)
         axis2.set_ylabel("y [m]")
         axis2.set_xlabel("s [m]")
-        axis2.set_title(f"Vertical aperture at {beam_params['pc_GeV']} GeV/c")
+        axis2.set_title(f"Vertical aperture at {beam_params.pc_GeV} GeV/c")
 
         logger.debug("Plotting the stay-clear envelope")
         axis3 = plt.subplot2grid((3, 3), (2, 0), colspan=3, rowspan=1, sharex=axis1)
@@ -144,7 +145,7 @@ def plot_aperture(
         axis3.set_ylabel("n1")
         axis3.set_xlabel("s [m]")
         axis3.legend(loc="best")
-        axis3.set_title(f"Stay-clear envelope at {beam_params['pc_GeV']} GeV/c")
+        axis3.set_title(f"Stay-clear envelope at {beam_params.pc_GeV} GeV/c")
 
         if savefig:
             logger.info(f"Saving aperture plot at '{Path(savefig).absolute()}'")

pyhdtoolkit/cpymadtools/track.py

@@ -67,9 +67,7 @@ def track_single_particle(
         logger.trace(f"Setting observation point for tracking with OBSERVE at element '{element}'")
         madx.command.observe(place=element)
 
-    madx.command.start(
-        X=start[0], PX=start[1], Y=start[2], PY=start[3], T=start[4], PT=start[5],
-    )
+    madx.command.start(X=start[0], PX=start[1], Y=start[2], PY=start[3], T=start[4], PT=start[5])
     madx.command.run(turns=nturns)
     madx.command.endtrack()
     if onetable:  # user asked for ONETABLE, there will only be one table 'trackone' given back by MAD-X

pyhdtoolkit/models/__init__.py

@@ -0,0 +1,11 @@
+"""
+models package
+~~~~~~~~~~~~~~
+These are various `pydantic` models useful for data parsing and validation in `pyhdtoolkit`.
+
+:copyright: (c) 2021 by Felix Soubelet.
+:license: MIT, see LICENSE for more details.
+"""
+from .beam import BeamParameters
+from .htc import BaseSummary, ClusterSummary, HTCTaskSummary
+from .madx import MADXBeam

pyhdtoolkit/models/beam.py

@@ -0,0 +1,32 @@
+"""
+Module models.beam
+------------------
+
+Created on 2021.08.03
+:author: Felix Soubelet (felix.soubelet@cern.ch)
+
+A module with `pydantic` models to validate and store data structures used in the `beam` module.
+"""
+from typing import Optional
+
+from pydantic import BaseModel
+
+
+class BeamParameters(BaseModel):
+    """
+    Class to encompass, validate and manipulate properties of a particle beam.
+    """
+
+    pc_GeV: Optional[float]  # Beam momentum [GeV]
+    B_rho_Tm: Optional[float]  # Beam rigidity [T/m]
+    E_0_GeV: Optional[float]  # Particle rest mass energy [GeV]
+    charge: Optional[float]  # Particle charge in [e]
+    E_tot_GeV: Optional[float]  # Total beam energy [GeV]
+    E_kin_GeV: Optional[float]  # Kinectic beam energy [GeV]
+    gamma_r: Optional[float]  # Relativistic gamma
+    beta_r: Optional[float]  # Relativistic beta
+    en_x_m: Optional[float]  # Horizontal normalized emittance [m]
+    en_y_m: Optional[float]  # Vertical normalized emittance [m]
+    eg_x_m: Optional[float]  # Horizontal geometrical emittance
+    eg_y_m: Optional[float]  # Vertical geometrical emittance
+    deltap_p: Optional[float]  # Momentum deviation

pyhdtoolkit/models/madx.py

@@ -8,4 +8,44 @@
 A module with `pydantic` models to validate and store data obtained by interacting with the `MAD-X` process
 through `cpymad`.
 """
-from pydantic import BaseModel
+from enum import Enum
+
+from pydantic import BaseModel, PositiveFloat, PositiveInt
+
+
+class ParticleEnum(str, Enum):
+    """Validator Enum defining the accepted particle names in `MAD-X` beams."""
+
+    positron = "positron"
+    electron = "electron"
+    proton = "proton"
+    antiproton = "antiproton"
+    posmuon = "posmuon"  # positive muons
+    negmuon = "negmuon"  # negative muons
+    ions = "ions"
+
+
+class MADXBeam(BaseModel):
+    """This is a class to encompass and validate `BEAM` attributes from the `MAD-X` process."""
+
+    particle: ParticleEnum  # The name of particles in the beam
+    mass: PositiveFloat  # The rest mass of the particles in the beam in [GeV]
+    charge: float  # The electrical charge of the particles in the beam in units of qp, the proton charge
+    energy: PositiveFloat  # Total energy per particle in [GeV]
+    pc: float  # Particle momentum times the speed of light, in [GeV]
+    gamma: float  # Relativistic factor in [1]
+    beta: float  # Relativistic beta in [1]
+    brho: float  # Magnetic rigidity of the particles in [T.m]
+    ex: float  # Horizontal emittance in [m]
+    ey: float  # Vertical emittance in [m]
+    et: float  # Longitudinal emittance in [m]
+    exn: float  # Normalized horizontal emittance in [m] (beta * gamma * ex)
+    eyn: float  # Normalized vertical emittance in [m]  *beta * gamma * ey)
+    sigt: float  # The bunch length c σt in [m]
+    sige: float  # The relative energy spread σE /E in [1].
+    kbunch: PositiveInt  # The number of particle bunches in the machine in [1]
+    npart: PositiveInt  # The number of particles per bunch in [1]
+    bcurrent: float  # The bunch current, in [A]
+    bunched: bool  # A logical flag. If set, the beam is treated as bunched whenever this makes sense
+    radiate: bool  # A logical flag. If set, synchrotron radiation is considered in all dipole magnets
+    bv: int  # integer specifying the direction of the particle movement in the beam line; either +1 or -1

pyhdtoolkit/optics/__init__.py

@@ -7,4 +7,6 @@
 :license: MIT, see LICENSE for more details.
 """
 
-from .beam import Beam
+from .beam import Beam, compute_beam_parameters
+from .ripken import lebedev_beam_size
+from .twiss import courant_snyder_transform

pyhdtoolkit/optics/beam.py

@@ -11,6 +11,42 @@
 
 from scipy import constants
 
+from pyhdtoolkit.models.beam import BeamParameters
+
+
+def compute_beam_parameters(pc_gev: float, en_x_m: float, en_y_m: float, deltap_p: float) -> BeamParameters:
+    """
+    Calculate beam parameters from provided values, for proton particles.
+
+    Args:
+        pc_gev (float): particle momentum.
+        en_x_m (float): horizontal emittance, in meters.
+        en_y_m (float): vertical emittance, in meters.
+        deltap_p (float): momentum deviation.
+
+    Returns:
+        A `BeamParameters` object with the calculated values.
+    """
+    e0_gev = 0.9382720813
+    e_tot_gev = np.sqrt(pc_gev ** 2 + e0_gev ** 2)
+    gamma_r = e_tot_gev / e0_gev
+    beta_r = pc_gev / np.sqrt(pc_gev ** 2 + e0_gev ** 2)
+
+    return BeamParameters(
+        pc_GeV=pc_gev,
+        B_rho_Tm=3.3356 * pc_gev,
+        E_0_GeV=e0_gev,
+        E_tot_GeV=e_tot_gev,
+        E_kin_GeV=e_tot_gev - e0_gev,
+        gamma_r=gamma_r,
+        beta_r=beta_r,
+        en_x_m=en_x_m,
+        en_y_m=en_y_m,
+        eg_x_m=en_x_m / gamma_r / beta_r,
+        eg_y_m=en_y_m / gamma_r / beta_r,
+        deltap_p=deltap_p,
+    )
+
 
 class Beam:
     """

pyhdtoolkit/optics/ripken.py

@@ -29,9 +29,6 @@ def lebedev_beam_size(
     return np.sqrt(geom_emit_x * beta1_ + geom_emit_y * beta2_)
 
 
-# ----- JITed Calculations ----- #
-
-
 def _beam_size(coordinates_distribution: np.ndarray, method: str = "std") -> float:
     """
     Compute beam size from particle coordinates.