Lint and tweak beams (#110)

* first autolint set * new round for dict, list, tuple * some more for optional, union * fix implicit optionals * refix many imports rules * fix use of literal dicts * propagate error change here * fix trailing whitespaces * miscellaneous * fix error messages not declared before raising * fix magic values use in comparisons * miscellaneous * place *args correctly * miscellaneous * use dict literal * miscellaneous * no implicit namespace * correctly propagate kwargs * clean up logic * message outside raise * message outside raise * trade map for list comprehension * cleaner checks * format and last lints * no mypy * no interrogate * rename to nemitt_[xy], accept old and let user override npart and sigma_e * formatting * remove trailing whitespaces * pytest parametrize expects tuples * dict membership test does not need .keys() * mpl * names * unecessary list comprehensions * fixtures want lists of tuples * provide 'match' argument to every pytest.raises, normalize error messages * minimal linting here * remove old helper functions * bump version for minor release * release 1.6.0 file * added to this list
fsoubelet · Aug 13, 2024 · b4ff24d · b4ff24d
1 parent fc5754e
commit b4ff24d
Show file tree

Hide file tree

Showing 100 changed files with 1,092 additions and 1,066 deletions.
diff --git a/Makefile b/Makefile
@@ -84,8 +84,8 @@ format:
 	@python -m isort examples && black -l 95 examples
 
 install: format clean
-	@echo "Installing with $(D)pip$(E) in the current environment."
-	@python -m pip install . -v
+	@echo "Installing (editable) with $(D)pip$(E) in the current environment."
+	@python -m pip install --editable . -v
 
 lines: format
 	@tokei .
@@ -113,5 +113,7 @@ alltests:
 
 # Catch-all unknow targets without returning an error. This is a POSIX-compliant syntax.
 .DEFAULT:
-	@echo "Make caught an invalid target! See help output below for available targets."
+	@echo "Make caught an invalid target."
+	@echo "See help output below for available targets."
+	@echo ""
 	@make help
diff --git a/docs/conf.py b/docs/conf.py
@@ -125,7 +125,7 @@
     "sphinx.ext.coverage",  # Collect doc coverage stats
     "sphinx.ext.doctest",  # Test snippets in the documentation
     "sphinx.ext.githubpages",  # Publish HTML docs in GitHub Pages
-    "sphinx.ext.intersphinx",  # Link to other projects’ documentation
+    "sphinx.ext.intersphinx",  # Link to other projects' documentation
     "sphinx.ext.mathjax",  # Render math via JavaScript
     "sphinx.ext.napoleon",  # Support for NumPy and Google style docstrings
     "sphinx.ext.todo",  # Support for todo items
@@ -214,7 +214,7 @@
 #
 html_static_path = ["_static"]
 
-# A dictionary of values to pass into the template engine’s context for all
+# A dictionary of values to pass into the template engine's context for all
 # pages. Single values can also be put in this dictionary using the
 # -A command-line option of sphinx-build.
 html_context = {

diff --git a/docs/release.rst b/docs/release.rst
@@ -3,6 +3,14 @@ Release Notes
 
 The full list of releases can be found in the GitHub repository's `releases page <https://github.com/fsoubelet/PyhDToolkit/releases>`_.
 
+Version 1.6.0
+-------------
+
+.. toctree::
+    :maxdepth: 2
+
+    releases/v1.6.0
+
 Version 1.5.0
 -------------
 

diff --git a/docs/releases/v1.6.0.rst b/docs/releases/v1.6.0.rst
@@ -0,0 +1,26 @@
+.. _release_1.6.0:
+
+1.6.0
+-----
+
+Release `1.6.0` brings a few changes, and a great amount of maintenance fixes to the package.
+
+Enhancements
+~~~~~~~~~~~~
+
+* The `pyhdtoolkit.cpymadtools.lhc.make_lhc_beams` function now accepts the user providing values for `npart` and `sige`. The previously hard-coded values are now used as defaults.
+
+Changes
+~~~~~~~
+
+* The `pyhdtoolkit.cpymadtools.lhc.make_lhc_beams` function now expects `nemitt_[xy]` arguments instead of the previous `emittance_[xy]`. Providing the previous name is still accepted but will raise a warning.
+
+Maintenance
+~~~~~~~~~~~
+
+* A great amount of linting has been done to ensure best practices in every line of code.
+* Switched all relevant type hints to use the latest practices.
+* Fixed the `__all__` exports in the `__init__.py` files.
+* Moved type-hint relevant imports in a `TYPE_CHECKING` block to lighten import times.
+
+See `v1.6.0 release notes on GitHub <https://github.com/fsoubelet/PyhDToolkit/releases/tag/1.6.0>`_ and the `full changes since v1.5.0 <https://github.com/fsoubelet/PyhDToolkit/compare/1.5.0...1.6.0>`_.
diff --git a/examples/demo_acd_tracking_spectra.py b/examples/demo_acd_tracking_spectra.py
@@ -6,17 +6,17 @@
 AC Dipole Driven Tracking Spectrum
 ==================================
 
-This example shows how to use the `~.lhc.install_ac_dipole_as_kicker` and 
-`~.track.track_single_particle` function to track a  particle with the 
+This example shows how to use the `~.lhc.install_ac_dipole_as_kicker` and
+`~.track.track_single_particle` function to track a  particle with the
 ``TRACK`` command of ``MAD-X``, and visualise its coordinates and spectrum.
 
-In this example we will use the LHC lattice to illustrate the ACD tracking 
+In this example we will use the LHC lattice to illustrate the ACD tracking
 workflow when using `~pyhdtoolkit.cpymadtools`.
 
 .. note::
-    This is very similar to the :ref:`free tracking example <demo-free-tracking>` 
-    with the difference that there is special care to take to install the AC Dipole 
-    element. It is recommended to read that tutorial first as this one will focus 
+    This is very similar to the :ref:`free tracking example <demo-free-tracking>`
+    with the difference that there is special care to take to install the AC Dipole
+    element. It is recommended to read that tutorial first as this one will focus
     on the specificities of the AC Dipole setup.
 
 .. important::
@@ -33,7 +33,6 @@
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
-
 from cpymad.madx import Madx
 
 from pyhdtoolkit.cpymadtools import lhc, matching, track
@@ -45,7 +44,7 @@
 
 ###############################################################################
 # Let's start by setting up the LHC in ``MAD-X``, in this case at top energy.
-# To understand the function below have a look at the :ref:`lhc setup example 
+# To understand the function below have a look at the :ref:`lhc setup example
 # <demo-lhc-setup>`.
 
 madx: Madx = lhc.prepare_lhc_run3(

diff --git a/examples/demo_aperture.py b/examples/demo_aperture.py
@@ -6,7 +6,7 @@
 Accelerator Aperture
 ====================
 
-This example shows how to use the `~.plotting.aperture.plot_aperture` and 
+This example shows how to use the `~.plotting.aperture.plot_aperture` and
 `~.plotting.aperture.plot_physical_apertures` functions to visualise the
 available aperture in your machine, with the LHC used for this example.
 
@@ -23,7 +23,6 @@
 """
 
 import matplotlib.pyplot as plt
-
 from cpymad.madx import Madx
 
 from pyhdtoolkit.cpymadtools import lhc
@@ -35,7 +34,7 @@
 plt.rcParams.update(_SPHINX_GALLERY_PARAMS)  # for readability of this tutorial
 
 ###############################################################################
-# Let's start by setting up the LHC in ``MAD-X``, in this case at injection 
+# Let's start by setting up the LHC in ``MAD-X``, in this case at injection
 # optics and energy. To understand the function below have a look at the
 # :ref:`lhc setup example <demo-lhc-setup>`.
 
@@ -86,7 +85,7 @@
 
 
 ###############################################################################
-# We can also go for a different type of aperture plot, which tries to give 
+# We can also go for a different type of aperture plot, which tries to give
 # the elements' real physical apertures, with the `~.plotting.aperture.plot_physical_apertures`
 # function:
 

diff --git a/examples/demo_beam_enveloppe.py b/examples/demo_beam_enveloppe.py
@@ -14,7 +14,6 @@
 # sphinx_gallery_thumbnail_number = 1
 import matplotlib.pyplot as plt
 import numpy as np
-
 from cpymad.madx import Madx
 
 from pyhdtoolkit.plotting.envelope import plot_beam_envelope
@@ -59,7 +58,7 @@
 madx.input(
     f"""
 circum = {circumference};
-n_cells = {n_cells}; ! Number of cells 
+n_cells = {n_cells}; ! Number of cells
 lcell = {l_cell};
 lq = {l_quad}; ! Length of a quadrupole
 ldip = {l_bend};
@@ -71,14 +70,14 @@
 mb: sbend, l=ldip, angle=2.0*pi/(4*n_cells), apertype=ellipse, aperture={{{h_gap_dipole}, {v_gap_dipole}}};
 f = {f_m};
 
-! Define quadrupoles as multipoles 
-qf: multipole, knl:={{0,1/f+qtrim_f}}; 
+! Define quadrupoles as multipoles
+qf: multipole, knl:={{0,1/f+qtrim_f}};
 qd: multipole, knl:={{0,-1/f+qtrim_d}};
-qf: quadrupole, l=lq, K1:=1/f/lq  + qtrim_f/lq, apertype=ellipse, aperture={{{r_quadrupole}, {r_quadrupole}}}; 
+qf: quadrupole, l=lq, K1:=1/f/lq  + qtrim_f/lq, apertype=ellipse, aperture={{{r_quadrupole}, {r_quadrupole}}};
 qd: quadrupole, l=lq, K1:=-1/f/lq + qtrim_d/lq, apertype=ellipse, aperture={{{r_quadrupole}, {r_quadrupole}}};
 
 ! Define the sextupoles as multipole
-! ATTENTION: must use knl:= and NOT knl= to match later! 
+! ATTENTION: must use knl:= and NOT knl= to match later!
 lsex = 0.00001; ! dummy length, only used in the sequence
 msf: multipole, knl:={{0,0,ksf}};
 msd: multipole, knl:={{0,0,ksd}};

diff --git a/examples/demo_ir_errors.py b/examples/demo_ir_errors.py
@@ -10,7 +10,7 @@
 to assign magnet errors in the Insertion Region magnets of the LHC.
 
 .. warning::
-    The implementation of this function makes it valid only for LHC IP IRs, which are 
+    The implementation of this function makes it valid only for LHC IP IRs, which are
     1, 2, 5 and 8. Other IRs have different layouts that are incompatible.
 
 .. important::
@@ -27,7 +27,6 @@
 # sphinx_gallery_thumbnail_number = 1
 import matplotlib.pyplot as plt
 import numpy as np
-
 from cpymad.madx import Madx
 
 from pyhdtoolkit.cpymadtools import lhc, matching
@@ -38,7 +37,7 @@
 plt.rcParams.update(_SPHINX_GALLERY_PARAMS)  # for readability of this tutorial
 
 ###############################################################################
-# Let's start by setting up the LHC in ``MAD-X``, in this case at injection 
+# Let's start by setting up the LHC in ``MAD-X``, in this case at injection
 # optics and energy. To understand the function below have a look at the
 # :ref:`lhc setup example <demo-lhc-setup>`.
 

diff --git a/examples/demo_lattice.py b/examples/demo_lattice.py
@@ -13,7 +13,6 @@
 of several parameters to control the plot on the example case of the LHC.
 """
 import matplotlib.pyplot as plt
-
 from cpymad.madx import Madx
 
 from pyhdtoolkit.cpymadtools import lhc, matching

diff --git a/examples/demo_lhc_crossing_schemes.py b/examples/demo_lhc_crossing_schemes.py
@@ -13,7 +13,6 @@
     This is very LHC-specific and will not work with other machines.
 """
 import matplotlib.pyplot as plt
-
 from cpymad.madx import Madx
 
 from pyhdtoolkit.cpymadtools import lhc
@@ -35,7 +34,7 @@
 
 ###############################################################################
 # Let's explicitely re-cycle both sequences to avoid having IR1 split at beginning
-# and end of lattice. Note that it is important to re-cycle both sequences from 
+# and end of lattice. Note that it is important to re-cycle both sequences from
 # the same points for the plots later on.
 
 lhc.re_cycle_sequence(madx, sequence="lhcb1", start="IP3")

diff --git a/examples/demo_lhc_rigid_waist_shift.py b/examples/demo_lhc_rigid_waist_shift.py
@@ -6,18 +6,18 @@
 LHC Rigid Waist Shift
 =====================
 
-This example shows how to use the `~.lhc.apply_lhc_rigidity_waist_shift_knob` 
-function to force a waist shift at a given IP and break the symmetry of the 
-:math:`\\beta`-functions in the Interaction Region. This is done by 
+This example shows how to use the `~.lhc.apply_lhc_rigidity_waist_shift_knob`
+function to force a waist shift at a given IP and break the symmetry of the
+:math:`\\beta`-functions in the Interaction Region. This is done by
 over-powering one triplet knob and under-powering the other, by the same
 powering delta.
 
 In :cite:t:`PRAB:Soubelet:Rigid_Waist_Shift_Method_Local_Coupling_Correction_LHC_IR`
 (2023) one can find out about studies and achievements at the LHC done with the Rigid
 Waist Shift.
 
-We will do a comparison of the interaction region situation before and after 
-applying a rigid waist shift, and look in more details at the waist shift 
+We will do a comparison of the interaction region situation before and after
+applying a rigid waist shift, and look in more details at the waist shift
 itself.
 
 .. note::
@@ -36,13 +36,12 @@
     builds, but you can use any branch you want.
 """
 # sphinx_gallery_thumbnail_number = 3
-from collections import namedtuple
 from multiprocessing import cpu_count
+from typing import NamedTuple
 
 import matplotlib.pyplot as plt
 import numpy as np
 import tfs
-
 from cpymad.madx import Madx
 from joblib import Parallel, delayed
 
@@ -325,28 +324,28 @@
 # We can use the above to determine these values for different knob settings.
 # First, let's define some structures and functions.
 
-Waist = namedtuple("Waist", ["x", "y"])
-BetasIP = namedtuple("Betas", ["x", "y"])
-Result = namedtuple("Result", ["waists", "betas"])
+Waist = NamedTuple("Waist", ["x", "y"])
+BetasIP = NamedTuple("Betas", ["x", "y"])
+Result = NamedTuple("Result", ["waists", "betas"])
 
 
 def find_waists(current_twiss: tfs.TfsDataFrame, initial_twiss: tfs.TfsDataFrame) -> Waist:
     initial = initial_twiss.copy()
-    ip_s = current_twiss.S[f"IP1"]
+    ip_s = current_twiss.S["IP1"]
     slimits = (ip_s - 10, ip_s + 10)
 
     around_ip = current_twiss[current_twiss.S.between(*slimits)]
     initial = initial[initial.S.between(*slimits)].copy()
-    hor_waist_location = around_ip.S[around_ip.BETX == around_ip.BETX.min()].iloc[0]
-    ver_waist_location = around_ip.S[around_ip.BETY == around_ip.BETY.min()].iloc[0]
+    hor_waist_location = around_ip.S[around_ip.BETX.min() == around_ip.BETX].iloc[0]
+    ver_waist_location = around_ip.S[around_ip.BETY.min() == around_ip.BETY].iloc[0]
     initial = initial_twiss.copy()
-    ip_s = current_twiss.S[f"IP1"]
+    ip_s = current_twiss.S["IP1"]
     slimits = (ip_s - 10, ip_s + 10)
 
     around_ip = current_twiss[current_twiss.S.between(*slimits)]
     initial = initial[initial.S.between(*slimits)].copy()
-    hor_waist_location = around_ip.S[around_ip.BETX == around_ip.BETX.min()].iloc[0]
-    ver_waist_location = around_ip.S[around_ip.BETY == around_ip.BETY.min()].iloc[0]
+    hor_waist_location = around_ip.S[around_ip.BETX.min() == around_ip.BETX].iloc[0]
+    ver_waist_location = around_ip.S[around_ip.BETY.min() == around_ip.BETY].iloc[0]
     return Waist(ip_s - hor_waist_location, ip_s - ver_waist_location)
 
 
@@ -363,7 +362,7 @@ def simulation(knob_value: float) -> Result:
         run=3, opticsfile="R2022a_A30cmC30cmA10mL200cm.madx", slicefactor=4, stdout=False
     ) as madx:
         lhc.add_markers_around_lhc_ip(
-            madx, sequence=f"lhcb1", ip=1, n_markers=1000, interval=0.001
+            madx, sequence="lhcb1", ip=1, n_markers=1000, interval=0.001
         )
         ref_twiss = twiss.get_twiss_tfs(madx)
         lhc.apply_lhc_rigidity_waist_shift_knob(madx, knob_value, ir=1)

diff --git a/examples/demo_lhc_setup.py b/examples/demo_lhc_setup.py
@@ -7,7 +7,7 @@
 ===============
 
 These examples show how to use the functions in `pyhdtoolkit.cpymadtools.lhc._setup`
-managers to easily and quickly set up LHC simulations. These are exported and 
+managers to easily and quickly set up LHC simulations. These are exported and
 available at the level of the `pyhdtoolkit.cpymadtools.lhc` module.
 
 .. note::
@@ -26,7 +26,6 @@
 """
 # sphinx_gallery_thumbnail_number = 4
 import matplotlib.pyplot as plt
-
 from cpymad.madx import Madx
 
 from pyhdtoolkit.cpymadtools import coupling, lhc, twiss
@@ -70,7 +69,7 @@
 # to set up the LHC simulations quickly and effortlessly:
 # `~.cpymadtools.lhc._setup.setup_lhc.prepare_lhc_run2` and
 # `~.cpymadtools.lhc._setup.setup_lhc.prepare_lhc_run3`.
-# 
+#
 # They both return a `cpyamad.Madx` instance with the desired LHC sequence and
 # optics loaded, beams defined for both ``lhcb1`` and ``lhcb2`` sequences,
 # potentially sliced lattices etc. The very minimum required at function call
@@ -126,7 +125,7 @@
 
 ###############################################################################
 # Notice we don't need to call ``madx.exit()`` as the context manager takes care
-# of that. 
+# of that.
 #
 # These quick setups, with context manager option, allow to do quick "one-shot"
 # simulations. For example, one can very quickly compare beam sizes around say
@@ -171,7 +170,7 @@
     # We'll need to call these to have aperture limitations
     madx.call("lhc/aperture.b1.madx")
     madx.call("lhc/aper_tol.b1.madx")
-    
+
     df = twiss.get_twiss_tfs(madx)
     ips = get_lhc_ips_positions(df)
     limits = (ips["IP5"] - 350, ips["IP5"] + 350)

diff --git a/examples/demo_machine_survey.py b/examples/demo_machine_survey.py
@@ -11,7 +11,6 @@
 """
 # sphinx_gallery_thumbnail_number = 2
 import matplotlib.pyplot as plt
-
 from cpymad.madx import Madx
 
 from pyhdtoolkit.cpymadtools._generators import LatticeGenerator

diff --git a/examples/demo_phase_space.py b/examples/demo_phase_space.py
@@ -17,7 +17,6 @@
 # sphinx_gallery_thumbnail_number = 2
 import matplotlib.pyplot as plt
 import numpy as np
-
 from cpymad.madx import Madx
 
 from pyhdtoolkit.cpymadtools._generators import LatticeGenerator