Doc tweaks (#96)

Author: fsoubelet

pyhdtoolkit/cpymadtools/coupling.py

@@ -191,13 +191,18 @@ def get_cminus_from_coupling_rdts(
         The calculated :math:`|C^{-}|` value.
 
     Examples:
+
+        To compute the :math:`|C^{-}|` taking in consideration all elements in the sequence:
+
         .. code-block:: python
 
             >>> complex_cminus = get_cminus_from_coupling_rdts(madx)
 
+
+        To simulate the calculation from a measurement, with RDTs computed at BPMs only:
+
         .. code-block:: python
 
-            >>> # To simulate the calculation from a measurement, with RDTs at BPMs only
             >>> complex_cminus = get_cminus_from_coupling_rdts(madx, patterns=["^BPM.*B[12]$"])
     """
     logger.debug(f"Getting coupling RDTs at selected elements thoughout the machine")
@@ -235,7 +240,7 @@ def match_no_coupling_through_ripkens(
     madx: Madx, sequence: str = None, location: str = None, vary_knobs: Sequence[str] = None
 ) -> None:
     """
-    ..versionadded:: 0.16.0
+    .. versionadded:: 0.16.0
 
     Matching routine to get cross-term Ripken parameters :math:`\\beta_{12}` and :math:`\\beta_{21}`
     to be 0 at a given location.

pyhdtoolkit/cpymadtools/errors.py

@@ -139,7 +139,11 @@ def misalign_lhc_ir_quadrupoles(
         .. code-block:: python
 
             >>> misalign_lhc_ir_quadrupoles(
-            ...     madx, ips=[5], quadrupoles=[7, 8, 9, 10], beam=1, sides="RL", dpsi="1E-3 + 8E-4 * TGAUSS(2.5)"
+            ...     madx, ips=[5],
+            ...     quadrupoles=[7, 8, 9, 10],
+            ...     beam=1,
+            ...     sides="RL",
+            ...     dpsi="1E-3 + 8E-4 * TGAUSS(2.5)",
             ... )
 
         For several error types on the elements, here ``DY`` and ``DPSI``:

pyhdtoolkit/cpymadtools/lhc.py

@@ -45,7 +45,7 @@ def make_lhc_beams(
         energy (float): beam energy, in [GeV]. Defaults to 6500.
         emittance_x (float): horizontal emittance in [m]. Will be used to calculate
             geometric emittance which is then fed to the ``BEAM`` command.
-        emittance_x (float): vertical emittance in [m]. Will be used to calculate
+        emittance_y (float): vertical emittance in [m]. Will be used to calculate
             geometric emittance which is then fed to the ``BEAM`` command.
         **kwargs: Any keyword argument that can be given to the ``MAD-X`` ``BEAM`` command.
 
@@ -170,7 +170,7 @@ def apply_lhc_colinearity_knob(madx: Madx, colinearity_knob_value: float = 0, ir
     """
     .. versionadded:: 0.15.0
 
-    Applies the LHC colinearity knob.
+    Applies the a trim of the LHC colinearity knob.
 
     .. note::
         If you don't know what this is, you really should not be using this function.
@@ -209,7 +209,7 @@ def apply_lhc_colinearity_knob_delta(madx: Madx, colinearity_knob_delta: float =
     .. versionadded:: 0.21.0
 
     This is essentially the same as `.apply_lhc_colinearity_knob`, but instead of a applying fixed powering
-    value, it applies to delta to the existing value.
+    value, it applies a delta to the (potentially) existing value.
 
     .. note::
         If you don't know what this is, you really should not be using this function.
@@ -249,12 +249,12 @@ def apply_lhc_rigidity_waist_shift_knob(
     """
     .. versionadded:: 0.15.0
 
-    Applies the LHC rigidity waist shift knob, moving the waist left or right of IP.
+    Applies a trim of the LHC rigidity waist shift knob, moving the waist left or right of IP.
+    The waist shift is achieved by moving all four betatron waists simltaneously: unbalancing
+    the triplet powering knobs of the left and right-hand sides of the IP.
 
     .. note::
-        If you don't know what this is, you really should not be using this function. The waist shift
-        is achieved by moving all four betatron waists simltaneously: unbalancing the triplet powering
-        knobs of the left and right-hand sides of the IP,
+        If you don't know what this is, you really should not be using this function.
 
     .. warning::
         Applying the shift will modify your tunes and is likely to flip them, making a subsequent matching
@@ -307,7 +307,7 @@ def apply_lhc_coupling_knob(
     """
     .. versionadded:: 0.15.0
 
-    Applies the LHC coupling knob to reach the desired :math:`|C^{-}|` value.
+    Applies a trim of the LHC coupling knob to reach the desired :math:`|C^{-}|` value.
 
     Args:
         madx (cpymad.madx.Madx): an instanciated `~cpymad.madx.Madx` object.
@@ -327,7 +327,7 @@ def apply_lhc_coupling_knob(
     logger.warning("You should re-match tunes & chromaticities after this coupling knob is applied")
     suffix = "_sq" if telescopic_squeeze else ""
     # NOTE: Only using this knob will give a dqmin very close to coupling_knob
-    # If one wants to also assign f"CMIS.b{beam:d}{suffix}" the dqmin be > coupling_knob
+    # If one wants to also assign f"CMIS.b{beam:d}{suffix}" the dqmin will be > coupling_knob
     knob_name = f"CMRS.b{beam:d}{suffix}"
 
     logger.trace(f"Knob '{knob_name}' is {madx.globals[knob_name]} before implementation")
@@ -346,7 +346,7 @@ def carry_colinearity_knob_over(madx: Madx, ir: int, to_left: bool = True) -> No
         madx (cpymad.madx.Madx): an instanciated `~cpymad.madx.Madx` object.
         ir (int): The Interaction Region around which to apply the change, should be
             one of [1, 2, 5, 8].
-        to_left (bool): If `True`, the magnet right of IP is powered of and its powering
+        to_left (bool): If `True`, the magnet right of IP is de-powered of and its powering
             is transferred to the magnet left of IP. If `False`, then the opposite happens.
             Defaults to `True`.
 
@@ -446,8 +446,9 @@ def vary_independent_ir_quadrupoles(
     The independent quadrupoles for which this is implemented are Q4 to Q13 included. This is useful
     to setup some specific matching involving these elements.
 
-    ..important::
-        It is necessary to have defined a ``brho`` variable when creating your beams.
+    .. important::
+        It is necessary to have defined a ``brho`` variable when creating your beams. If one has used
+        `make_lhc_beams` to create the beams, this has already been done automatically.
 
     Args:
         madx (cpymad.madx.Madx): an instanciated `~cpymad.madx.Madx` object.
@@ -827,7 +828,9 @@ def add_markers_around_lhc_ip(madx: Madx, sequence: str, ip: int, n_markers: int
     Example:
         .. code-block:: python
 
-            >>> add_markers_around_lhc_ip(madx, sequence=f"lhcb1", ip=1, n_markers=1000, interval=0.001)
+            >>> add_markers_around_lhc_ip(
+            ...     madx, sequence=f"lhcb1", ip=1, n_markers=1000, interval=0.001
+            ... )
     """
     logger.debug(f"Adding {n_markers:d} markers on each side of IP{ip:d}")
     madx.command.seqedit(sequence=sequence)
@@ -1011,15 +1014,15 @@ def get_magnets_powering(
     .. note::
         Here are below certain useful patterns for the ``LHC`` and their meaning:
 
-        * ``^mb\.`` :math:`\\rightarrow` main bends.
-        * ``^mq\.`` :math:`\\rightarrow` main quadrupoles.
-        * ``^ms\.`` :math:`\\rightarrow` main sextupoles.
-        * ``^mb[rswx]`` :math:`\\rightarrow` separation dipoles.
-        * ``^mq[mwxy]`` :math:`\\rightarrow` insertion quads.
-        * ``^mqt.1[23]`` :math:`\\rightarrow` short tuning quads (12 & 13).
-        * ``^mqtl`` :math:`\\rightarrow` long  tuning quads.
-        * ``^mcbx`` :math:`\\rightarrow` crossing scheme magnets.
-        * ``^mcb[cy]`` :math:`\\rightarrow` crossing scheme magnets.
+        * ``^mb\.`` :math:`\rightarrow` main bends.
+        * ``^mq\.`` :math:`\rightarrow` main quadrupoles.
+        * ``^ms\.`` :math:`\rightarrow` main sextupoles.
+        * ``^mb[rswx]`` :math:`\rightarrow` separation dipoles.
+        * ``^mq[mwxy]`` :math:`\rightarrow` insertion quads.
+        * ``^mqt.1[23]`` :math:`\rightarrow` short tuning quads (12 & 13).
+        * ``^mqtl`` :math:`\rightarrow` long  tuning quads.
+        * ``^mcbx`` :math:`\rightarrow` crossing scheme magnets.
+        * ``^mcb[cy]`` :math:`\rightarrow` crossing scheme magnets.
 
         To make no selection, one can give ``patterns=[""]`` and this will give back
         the results for *all* elements. One can also give a specific magnet's exact
@@ -1043,6 +1046,11 @@ def get_magnets_powering(
     Returns:
         A `~tfs.TfsDataFrame` of the ``TWISS`` table, with the relevant newly defined columns
         and including the elements matching the regex *patterns* that were provided.
+
+    Example:
+        .. code-block:: python
+
+            >>> sextupoles_powering = get_magnets_powering(madx, patterns=[r"^ms\."])
     """
     logger.debug("Computing magnets field and powering limits proportions")
     NEW_COLNAMES = ["name", "keyword", "ampere", "imax", "percent", "kn", "kmax", "integrated_field", "L"]

pyhdtoolkit/cpymadtools/matching.py

@@ -116,7 +116,7 @@ def match_tunes_and_chromaticities(
             ...     q2_target=60.32,
             ...     dq1_target=2.0,
             ...     dq2_target=2.0,
-            ...     telescopic_squeeze=True,  # influences the knobs definition
+            ...     run3=True,  # influences the knobs definition
             ... )
     """
     if accelerator and not varied_knobs:

pyhdtoolkit/cpymadtools/parameters.py

@@ -30,6 +30,11 @@ def query_beam_attributes(madx: Madx) -> MADXBeam:
 
     Returns:
         A validated `~.models.madx.MADXBeam` object.
+
+    Example:
+        .. code-block:: python
+
+            >>> beam_parameters = query_beam_attributes(madx)
     """
     logger.debug("Retrieving BEAM attributes from the MAD-X process")
     return MADXBeam(**dict(madx.beam))

pyhdtoolkit/cpymadtools/setup.py

@@ -50,7 +50,9 @@ def prepare_lhc_run2(opticsfile: str, beam: int = 1, energy: float = 6500, slice
     Example:
         .. code-block:: python
 
-            >>> madx = prepare_lhc_run2("/afs/cern.ch/eng/lhc/optics/runII/2018/PROTON/opticsfile.22", beam=2, stdout=True)
+            >>> madx = prepare_lhc_run2(
+            ...     "/afs/cern.ch/eng/lhc/optics/runII/2018/PROTON/opticsfile.22", beam=2, stdout=True
+            ... )
     """
 
     def _run2_sequence_from_opticsfile(opticsfile: Path) -> Path:
@@ -112,7 +114,9 @@ def prepare_lhc_run3(opticsfile: str, beam: int = 1, energy: float = 6800, slice
     Example:
         .. code-block:: python
 
-            >>> madx = prepare_lhc_run3("R2022a_A30cmC30cmA10mL200cm.madx", slicefactor=4, stdout=True)
+            >>> madx = prepare_lhc_run3(
+            ...     "R2022a_A30cmC30cmA10mL200cm.madx", slicefactor=4, stdout=True
+            ... )
     """
     logger.debug("Creating Run 3 setup MAD-X instance")
     echo, warn = kwargs.pop("echo", False), kwargs.pop("warn", False)

pyhdtoolkit/cpymadtools/tune.py

@@ -4,7 +4,7 @@
 Tune Utilities
 --------------
 
-Module with functions to manipulate to manipulate ``MAD-X`` functionality around the tune through
+Module with functions to manipulate ``MAD-X`` functionality around the tune through
 a `~cpymad.madx.Madx` object.
 """
 import math
@@ -55,7 +55,7 @@ def make_footprint_table(
             >>> dynap_dframe = make_footprint_table(madx)
     """
     logger.debug(f"Initiating particules up to {sigma:d} bunch sigma to create a tune footprint table")
-    small, big = 0.05, math.sqrt(1 - 0.05**2)
+    small, big = 0.05, math.sqrt(1 - 0.05 ** 2)
     sigma_multiplier, angle_multiplier = 0.1, 0.0
 
     logger.debug("Initializing particles")

pyhdtoolkit/cpymadtools/twiss.py

@@ -9,7 +9,6 @@
 """
 from typing import Sequence
 
-import numpy as np
 import tfs
 
 from cpymad.madx import Madx
@@ -70,9 +69,9 @@ def get_pattern_twiss(
             >>> triplets_df = get_pattern_twiss(
             ...     madx=madx,
             ...     patterns=[
-            ...         f"MQXA.[12345][RL]1",  # Q1 and Q3 LHC
-            ...         f"MQXB.[AB][12345][RL]1",  # Q2A and Q2B LHC
-            ...         f"MQXF[AB].[AB][12345][RL]1",  # Q1 to Q3 A and B HL-LHC
+            ...         r"MQXA.[12345][RL]1",  # Q1 and Q3 LHC
+            ...         r"MQXB.[AB][12345][RL]1",  # Q2A and Q2B LHC
+            ...         r"MQXF[AB].[AB][12345][RL]1",  # Q1 to Q3 A and B HL-LHC
             ...     ],
             ... )
     """

pyhdtoolkit/cpymadtools/utils.py

@@ -39,6 +39,12 @@ def export_madx_table(
         headers_table (str): the name of the internal table to use for headers.
             Defaults to ``SUMM``.
         **kwargs: any keyword arguments will be passed to `~tfs.writer.write_tfs`.
+
+    Example:
+        .. code-block:: python
+
+            >>> madx.command.twiss()
+            >>> export_madx_table(madx, table_name="TWISS", file_name="twiss.tfs")
     """
     file_path = Path(file_name)
     logger.debug(f"Exporting table {table_name} into '{file_path.absolute()}'")

pyhdtoolkit/maths/stats_fitting.py

@@ -34,7 +34,8 @@ def set_distributions_dict(dist_dict: Dict[st.rv_continuous, str]) -> None:
     .. versionadded:: 0.5.0
 
     Sets ``DISTRIBUTIONS`` as the provided `dict`. This allows the user to define the
-    distributions to try and fit against the data.
+    distributions to try and fit against the data. One can find an example use of this
+    function in the :ref:`gallery <demo-distributions-fitting>`.
 
     Args:
         dist_dict (Dict[st.rv_continuous, str]): dictionnary with the wanted distributions,
@@ -131,7 +132,8 @@ def make_pdf(distribution: st.rv_continuous, params: Tuple[float, ...], size: in
     .. versionadded:: 0.5.0
 
     Generates a `pandas.Series` for the distributions's Probability Distribution Function.
-    This Series will have axis values as index, and PDF values as column.
+    This Series will have axis values as index, and PDF values as column. One can find an
+    example use of this function in the :ref:`gallery <demo-distributions-fitting>`.
 
     Args:
         distribution (st.rv_continuous): a `scipy.stats` generator.

pyhdtoolkit/optics/beam.py

@@ -17,7 +17,9 @@ def compute_beam_parameters(pc_gev: float, en_x_m: float, en_y_m: float, deltap_
     """
     .. versionadded:: 0.12.0
 
-    Calculates beam parameters from provided values, for *proton* particles.
+    Calculates beam parameters from provided values, for *proton* particles. One can find
+    an example use of this function in the :ref:`beam enveloppe <demo-beam-enveloppe>`
+    example gallery.
 
     Args:
         pc_gev (float): particle momentum, in [GeV].
@@ -27,11 +29,29 @@ def compute_beam_parameters(pc_gev: float, en_x_m: float, en_y_m: float, deltap_
 
     Returns:
         A `~.optics.beam.BeamParameters` object with the calculated values.
+
+    Example:
+        .. code-block:: python
+
+            >>> params = compute_beam_parameters(1.9, 5e-6, 5e-6, 2e-3)
+            >>> print(params)
+            Beam Parameters for particle of charge 1
+            Beam momentum = 1.900 GeV/c
+            Normalized x-emittance = 5.000 mm mrad
+            Normalized y-emittance = 5.000 mm mrad
+            Momentum deviation deltap/p = 0.002
+            -> Beam total energy = 2.119 GeV
+            -> Beam kinetic energy = 1.181 GeV
+            -> Beam rigidity = 6.333 Tm
+            -> Relativistic beta = 0.89663
+            -> Relativistic gamma = 2.258
+            -> Geometrical x emittance = 2.469 mm mrad
+            -> Geometrical y emittance = 2.469 mm mrad
     """
     e0_gev = 0.9382720813
-    e_tot_gev = np.sqrt(pc_gev**2 + e0_gev**2)
+    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)
+    beta_r = pc_gev / np.sqrt(pc_gev ** 2 + e0_gev ** 2)
 
     return BeamParameters(
         pc_GeV=pc_gev,
@@ -80,7 +100,7 @@ def gamma_rel(self) -> float:
     @property
     def beta_rel(self) -> float:
         """Relativistic beta."""
-        return np.sqrt(1 + 1 / (self.gamma_rel**2))
+        return np.sqrt(1 + 1 / (self.gamma_rel ** 2))
 
     @property
     def brho(self) -> float:
@@ -123,7 +143,7 @@ def eta(self, alpha_p: float) -> float:
         Returns:
             The slip factor.
         """
-        return (1 / (self.gamma_rel**2)) - alpha_p
+        return (1 / (self.gamma_rel ** 2)) - alpha_p
 
     @staticmethod
     def gamma_transition(alpha_p: float) -> float:

pyhdtoolkit/optics/twiss.py

@@ -10,16 +10,20 @@
 
 
 def courant_snyder_transform(u_vector: np.ndarray, alpha: float, beta: float) -> np.ndarray:
-    """
+    r"""
     .. versionadded:: 0.5.0
 
     Perform the Courant-Snyder transform on regular (nonchaotic) phase-space coordinates.
     Specifically, if considering the horizontal plane and noting :math:`U = (x, px)` the
-    phase-space vector, it returns :math:`\\bar{U} = (\\bar{x}, \\bar{px})` according to
-    the transform :math:`\\bar{U} = P \\cdot U`, where::
+    phase-space vector, it returns :math:`\bar{U} = (\bar{x}, \bar{px})` according to
+    the transform :math:`\bar{U} = P \cdot U`, where
+
+    .. math::
 
-        P = [1/sqrt(beta_x)              0      ]
-            [alpha_x/sqrt(beta_x)   sqrt(beta_x)]
+        P = \begin{pmatrix}
+                \frac{1}{\sqrt{\beta_x}}          &   0              \\
+                \frac{\alpha_x}{\sqrt{\beta_x}}   &  \sqrt{\beta_x}  \\
+            \end{pmatrix}
 
     Args:
         u_vector (np.ndarray): two-dimentional array of phase-space (spatial and momenta)
@@ -29,6 +33,14 @@ def courant_snyder_transform(u_vector: np.ndarray, alpha: float, beta: float) ->
 
     Returns:
         The normalized phase-space coordinates from the Courant-Snyder transform.
+    
+    Example:
+        .. code-block:: python
+
+            >>> alfx = madx.table.twiss.alfx[0]
+            >>> betx = madx.table.twiss.betx[0]
+            >>> u = np.array([x_coords, px_coord])
+            >>> u_bar = courant_snyder_transform(u, alfx, betx)
     """
     p_matrix = np.array([[1 / np.sqrt(beta), 0], [alpha / np.sqrt(beta), np.sqrt(beta)]])
     return p_matrix @ u_vector