From a546df8b7821717bbce4bde126afa00501c3cc44 Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 10:24:01 +0200 Subject: [PATCH 01/41] first autolint set --- Makefile | 8 +++++--- examples/demo_acd_tracking_spectra.py | 3 +-- examples/demo_aperture.py | 5 ++--- examples/demo_beam_enveloppe.py | 1 - examples/demo_ir_errors.py | 3 +-- examples/demo_lattice.py | 1 - examples/demo_lhc_crossing_schemes.py | 3 +-- examples/demo_lhc_rigid_waist_shift.py | 15 +++++++-------- examples/demo_lhc_setup.py | 7 +++---- examples/demo_machine_survey.py | 1 - examples/demo_phase_space.py | 1 - examples/demo_track_spectra.py | 4 +--- pyhdtoolkit/cpymadtools/coupling.py | 1 - pyhdtoolkit/cpymadtools/lhc/_coupling.py | 1 - pyhdtoolkit/cpymadtools/lhc/_powering.py | 2 +- pyhdtoolkit/cpymadtools/lhc/_queries.py | 1 - pyhdtoolkit/cpymadtools/lhc/_routines.py | 1 - pyhdtoolkit/cpymadtools/lhc/_twiss.py | 1 - pyhdtoolkit/cpymadtools/ptc.py | 7 +++---- pyhdtoolkit/cpymadtools/track.py | 1 - pyhdtoolkit/cpymadtools/tune.py | 2 -- pyhdtoolkit/cpymadtools/twiss.py | 1 - pyhdtoolkit/cpymadtools/utils.py | 1 - pyhdtoolkit/maths/stats_fitting.py | 2 -- pyhdtoolkit/maths/utils.py | 1 - pyhdtoolkit/optics/beam.py | 1 - pyhdtoolkit/optics/rdt.py | 2 +- pyhdtoolkit/optics/ripken.py | 1 - pyhdtoolkit/plotting/aperture.py | 1 - pyhdtoolkit/plotting/crossing.py | 1 - pyhdtoolkit/plotting/envelope.py | 1 - pyhdtoolkit/plotting/lattice.py | 1 - pyhdtoolkit/plotting/layout.py | 7 +++---- pyhdtoolkit/plotting/phasespace.py | 1 - pyhdtoolkit/plotting/sbs/coupling.py | 1 - pyhdtoolkit/plotting/sbs/phase.py | 1 - pyhdtoolkit/plotting/styles/__init__.py | 1 - pyhdtoolkit/plotting/tune.py | 1 - pyhdtoolkit/plotting/utils.py | 9 ++++----- pyhdtoolkit/utils/__init__.py | 1 - pyhdtoolkit/utils/_misc.py | 4 +--- pyhdtoolkit/utils/cmdline.py | 1 - pyhdtoolkit/utils/contexts.py | 1 - pyhdtoolkit/utils/htc_monitor.py | 2 -- pyhdtoolkit/utils/logging.py | 1 - pyhdtoolkit/utils/operations.py | 5 ++--- pyhdtoolkit/version.py | 1 - tests/conftest.py | 19 +++++++++---------- tests/test_cpymadtools/test_coupling.py | 5 ++--- tests/test_cpymadtools/test_lhc.py | 15 +++++++-------- tests/test_cpymadtools/test_matching.py | 2 -- tests/test_cpymadtools/test_ptc.py | 7 +++---- tests/test_cpymadtools/test_track.py | 1 - tests/test_cpymadtools/test_tune.py | 4 ++-- tests/test_cpymadtools/test_twiss.py | 3 +-- tests/test_cpymadtools/test_utilities.py | 1 - tests/test_maths.py | 9 ++++----- tests/test_optics.py | 2 +- tests/test_plotting/test_aperture.py | 1 - tests/test_plotting/test_envelope.py | 1 - tests/test_plotting/test_lattice.py | 3 +-- tests/test_plotting/test_phasespace.py | 1 - tests/test_plotting/test_plotting_utils.py | 8 ++++---- tests/test_plotting/test_sbs_coupling.py | 8 ++++---- tests/test_plotting/test_sbs_phase.py | 6 +++--- tests/test_utils.py | 14 ++++++-------- 66 files changed, 83 insertions(+), 146 deletions(-) diff --git a/Makefile b/Makefile index 10d18498..da89a861 100644 --- 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/examples/demo_acd_tracking_spectra.py b/examples/demo_acd_tracking_spectra.py index f724deba..b002b50d 100755 --- a/examples/demo_acd_tracking_spectra.py +++ b/examples/demo_acd_tracking_spectra.py @@ -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 # `. madx: Madx = lhc.prepare_lhc_run3( diff --git a/examples/demo_aperture.py b/examples/demo_aperture.py index 1a17a501..bc540337 100755 --- a/examples/demo_aperture.py +++ b/examples/demo_aperture.py @@ -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 `. @@ -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 index 08f4e1fe..c2e5e1a4 100755 --- 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 diff --git a/examples/demo_ir_errors.py b/examples/demo_ir_errors.py index 08d6d630..7ecdfb90 100755 --- a/examples/demo_ir_errors.py +++ b/examples/demo_ir_errors.py @@ -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 `. diff --git a/examples/demo_lattice.py b/examples/demo_lattice.py index d9c0d843..d7f858bd 100755 --- 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 index 006c1f56..7dc7b6da 100755 --- 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 index 3b58b0ed..4268a9ac 100755 --- a/examples/demo_lhc_rigid_waist_shift.py +++ b/examples/demo_lhc_rigid_waist_shift.py @@ -42,7 +42,6 @@ import matplotlib.pyplot as plt import numpy as np import tfs - from cpymad.madx import Madx from joblib import Parallel, delayed @@ -332,21 +331,21 @@ 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 index 04234ba5..04c48b24 100644 --- a/examples/demo_lhc_setup.py +++ b/examples/demo_lhc_setup.py @@ -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 index 09949267..6170dce7 100755 --- 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 index 8c083dfc..318b6e38 100755 --- 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 diff --git a/examples/demo_track_spectra.py b/examples/demo_track_spectra.py index f078ef92..1a61cb78 100755 --- a/examples/demo_track_spectra.py +++ b/examples/demo_track_spectra.py @@ -27,8 +27,6 @@ import matplotlib.pyplot as plt import numpy as np -from cpymad.madx import Madx - from pyhdtoolkit.cpymadtools import lhc, track from pyhdtoolkit.plotting.styles import _SPHINX_GALLERY_PARAMS from pyhdtoolkit.utils import logging @@ -37,7 +35,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 collision +# Let's start by setting up the LHC in ``MAD-X``, in this case at collision # optics and energy and with a sliced lattice. To understand the function below # have a look at the :ref:`lhc setup example `. diff --git a/pyhdtoolkit/cpymadtools/coupling.py b/pyhdtoolkit/cpymadtools/coupling.py index af95a217..9a541ff7 100644 --- a/pyhdtoolkit/cpymadtools/coupling.py +++ b/pyhdtoolkit/cpymadtools/coupling.py @@ -12,7 +12,6 @@ import numpy as np import tfs - from cpymad.madx import Madx from loguru import logger from optics_functions.coupling import check_resonance_relation, closest_tune_approach, coupling_via_cmatrix diff --git a/pyhdtoolkit/cpymadtools/lhc/_coupling.py b/pyhdtoolkit/cpymadtools/lhc/_coupling.py index b74b6dea..e8e6f341 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_coupling.py +++ b/pyhdtoolkit/cpymadtools/lhc/_coupling.py @@ -6,7 +6,6 @@ The functions below are betatron coupling utilities for the ``LHC``. """ import tfs - from cpymad.madx import Madx from loguru import logger from optics_functions.coupling import coupling_via_cmatrix diff --git a/pyhdtoolkit/cpymadtools/lhc/_powering.py b/pyhdtoolkit/cpymadtools/lhc/_powering.py index 2eab10e8..799ba1ea 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_powering.py +++ b/pyhdtoolkit/cpymadtools/lhc/_powering.py @@ -237,7 +237,7 @@ def power_landau_octupoles(madx: Madx, /, beam: int, mo_current: float, defectiv brho = madx.globals.nrj * 1e9 / madx.globals.clight # clight is MAD-X constant except AttributeError as madx_error: logger.exception("The global MAD-X variable 'NRJ' should have been set in the optics files but is not defined.") - raise EnvironmentError("No 'NRJ' variable found in scripts") from madx_error + raise AttributeError("No 'NRJ' variable found in scripts") from madx_error logger.debug(f"Powering Landau Octupoles, beam {beam} @ {madx.globals.nrj} GeV with {mo_current} A.") strength = mo_current / madx.globals.Imax_MO * madx.globals.Kmax_MO / brho diff --git a/pyhdtoolkit/cpymadtools/lhc/_queries.py b/pyhdtoolkit/cpymadtools/lhc/_queries.py index f26b7e02..4a443ba3 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_queries.py +++ b/pyhdtoolkit/cpymadtools/lhc/_queries.py @@ -8,7 +8,6 @@ from typing import Dict, Sequence, Union import tfs - from cpymad.madx import Madx from loguru import logger diff --git a/pyhdtoolkit/cpymadtools/lhc/_routines.py b/pyhdtoolkit/cpymadtools/lhc/_routines.py index 3828c9e7..59e46884 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_routines.py +++ b/pyhdtoolkit/cpymadtools/lhc/_routines.py @@ -7,7 +7,6 @@ """ import numpy as np import tfs - from cpymad.madx import Madx from loguru import logger diff --git a/pyhdtoolkit/cpymadtools/lhc/_twiss.py b/pyhdtoolkit/cpymadtools/lhc/_twiss.py index 672b333f..660d96e9 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_twiss.py +++ b/pyhdtoolkit/cpymadtools/lhc/_twiss.py @@ -8,7 +8,6 @@ from typing import Sequence import tfs - from cpymad.madx import Madx from loguru import logger diff --git a/pyhdtoolkit/cpymadtools/ptc.py b/pyhdtoolkit/cpymadtools/ptc.py index afc4fa44..e378da80 100644 --- a/pyhdtoolkit/cpymadtools/ptc.py +++ b/pyhdtoolkit/cpymadtools/ptc.py @@ -12,7 +12,6 @@ import pandas as pd import tfs - from cpymad.madx import Madx from loguru import logger @@ -207,7 +206,7 @@ def get_rdts( method = kwargs.pop("method", 4) nst = kwargs.pop("nst", 3) exact = kwargs.pop("exact", True) - + logger.debug("Looking for PTC_TWISS parameters in keyword arguments") icase = kwargs.pop("icase", 6) normal = kwargs.pop("normal", True) @@ -310,7 +309,7 @@ def ptc_twiss( method = kwargs.pop("method", 4) nst = kwargs.pop("nst", 3) exact = kwargs.pop("exact", True) - + logger.debug("Looking for PTC_TWISS parameters in keyword arguments") icase = kwargs.pop("icase", 6) normal = kwargs.pop("normal", True) @@ -439,7 +438,7 @@ def ptc_track_particle( method = kwargs.pop("method", 4) nst = kwargs.pop("nst", 3) exact = kwargs.pop("exact", True) - + logger.debug("Looking for PTC_TRACK parameters in keyword arguments") element_by_element = kwargs.pop("element_by_element", True) diff --git a/pyhdtoolkit/cpymadtools/track.py b/pyhdtoolkit/cpymadtools/track.py index aeca5769..98af4ffa 100644 --- a/pyhdtoolkit/cpymadtools/track.py +++ b/pyhdtoolkit/cpymadtools/track.py @@ -10,7 +10,6 @@ from typing import Dict, Optional, Sequence, Tuple import pandas as pd - from cpymad.madx import Madx from loguru import logger diff --git a/pyhdtoolkit/cpymadtools/tune.py b/pyhdtoolkit/cpymadtools/tune.py index 018052d0..1ec46c49 100644 --- a/pyhdtoolkit/cpymadtools/tune.py +++ b/pyhdtoolkit/cpymadtools/tune.py @@ -9,7 +9,6 @@ """ import math import sys - from pathlib import Path from typing import Dict, List, Tuple @@ -17,7 +16,6 @@ import matplotlib.patches import numpy as np import tfs - from cpymad.madx import Madx from loguru import logger diff --git a/pyhdtoolkit/cpymadtools/twiss.py b/pyhdtoolkit/cpymadtools/twiss.py index 17754428..0d6052c9 100644 --- a/pyhdtoolkit/cpymadtools/twiss.py +++ b/pyhdtoolkit/cpymadtools/twiss.py @@ -10,7 +10,6 @@ from typing import Sequence import tfs - from cpymad.madx import Madx from loguru import logger diff --git a/pyhdtoolkit/cpymadtools/utils.py b/pyhdtoolkit/cpymadtools/utils.py index 98198518..17b41487 100644 --- a/pyhdtoolkit/cpymadtools/utils.py +++ b/pyhdtoolkit/cpymadtools/utils.py @@ -11,7 +11,6 @@ import pandas as pd import tfs - from cpymad.madx import Madx from loguru import logger diff --git a/pyhdtoolkit/maths/stats_fitting.py b/pyhdtoolkit/maths/stats_fitting.py index 6d909378..0088ac46 100644 --- a/pyhdtoolkit/maths/stats_fitting.py +++ b/pyhdtoolkit/maths/stats_fitting.py @@ -7,7 +7,6 @@ Module implementing methods to find the best fit of statistical distributions to data. """ import warnings - from typing import Dict, Tuple, Union import matplotlib @@ -15,7 +14,6 @@ import numpy as np import pandas as pd import scipy.stats as st - from loguru import logger # Distributions to check # diff --git a/pyhdtoolkit/maths/utils.py b/pyhdtoolkit/maths/utils.py index b3cb9841..a195b84f 100644 --- a/pyhdtoolkit/maths/utils.py +++ b/pyhdtoolkit/maths/utils.py @@ -10,7 +10,6 @@ import numpy as np import pandas as pd - from loguru import logger # ----- Miscellaneous Utilites ----- # diff --git a/pyhdtoolkit/optics/beam.py b/pyhdtoolkit/optics/beam.py index cb789a11..8e0f3db2 100644 --- a/pyhdtoolkit/optics/beam.py +++ b/pyhdtoolkit/optics/beam.py @@ -7,7 +7,6 @@ Module implementing various functionality for simple beam parameter calculations. """ import numpy as np - from scipy import constants from pyhdtoolkit.models.beam import BeamParameters diff --git a/pyhdtoolkit/optics/rdt.py b/pyhdtoolkit/optics/rdt.py index 89f67157..79c7c96a 100644 --- a/pyhdtoolkit/optics/rdt.py +++ b/pyhdtoolkit/optics/rdt.py @@ -7,7 +7,7 @@ Module implementing utilities for the handling of resonance driving terms. """ -from typing import Union, Tuple +from typing import Tuple, Union def rdt_to_order_and_type(rdt: Union[int, str]) -> str: diff --git a/pyhdtoolkit/optics/ripken.py b/pyhdtoolkit/optics/ripken.py index 2384cd58..d7869939 100644 --- a/pyhdtoolkit/optics/ripken.py +++ b/pyhdtoolkit/optics/ripken.py @@ -10,7 +10,6 @@ import numpy as np import tfs - from loguru import logger # ----- Setup Utilites ----- # diff --git a/pyhdtoolkit/plotting/aperture.py b/pyhdtoolkit/plotting/aperture.py index 7cfedff7..e8204f65 100644 --- a/pyhdtoolkit/plotting/aperture.py +++ b/pyhdtoolkit/plotting/aperture.py @@ -12,7 +12,6 @@ import matplotlib.pyplot as plt import numpy as np import pandas as pd - from cpymad.madx import Madx from loguru import logger diff --git a/pyhdtoolkit/plotting/crossing.py b/pyhdtoolkit/plotting/crossing.py index ff7c3305..4c1e7ec3 100644 --- a/pyhdtoolkit/plotting/crossing.py +++ b/pyhdtoolkit/plotting/crossing.py @@ -13,7 +13,6 @@ import matplotlib.pyplot as plt import pandas as pd import tfs - from cpymad.madx import Madx from loguru import logger diff --git a/pyhdtoolkit/plotting/envelope.py b/pyhdtoolkit/plotting/envelope.py index 227afbf0..2381647b 100644 --- a/pyhdtoolkit/plotting/envelope.py +++ b/pyhdtoolkit/plotting/envelope.py @@ -9,7 +9,6 @@ from typing import Tuple import numpy as np - from cpymad.madx import Madx from loguru import logger diff --git a/pyhdtoolkit/plotting/lattice.py b/pyhdtoolkit/plotting/lattice.py index adde3544..439d2a5b 100644 --- a/pyhdtoolkit/plotting/lattice.py +++ b/pyhdtoolkit/plotting/lattice.py @@ -11,7 +11,6 @@ import matplotlib import matplotlib.axes import matplotlib.pyplot as plt - from cpymad.madx import Madx from loguru import logger diff --git a/pyhdtoolkit/plotting/layout.py b/pyhdtoolkit/plotting/layout.py index 376a3239..fb92efec 100644 --- a/pyhdtoolkit/plotting/layout.py +++ b/pyhdtoolkit/plotting/layout.py @@ -11,12 +11,11 @@ import matplotlib import matplotlib.axes -import matplotlib.patches as patches import numpy as np import pandas as pd - from cpymad.madx import Madx from loguru import logger +from matplotlib import patches from pyhdtoolkit.plotting.utils import ( _get_twiss_table_with_offsets_and_limits, @@ -175,8 +174,8 @@ def plot_machine_layout( dipole_patches_axis = axis.twinx() dipole_patches_axis.set_ylabel("$\\theta=K_{0}L$ $[rad]$", color="royalblue") # dipoles in blue dipole_patches_axis.tick_params(axis="y", labelcolor="royalblue") - if not np.nan in k0l_lim: - dipole_patches_axis.set_ylim(k0l_lim) + if np.nan not in k0l_lim: + dipole_patches_axis.set_ylim(k0l_lim) dipole_patches_axis.grid(False) if plot_dipoles: # beware 'sbend' and 'rbend' have an 'angle' value and not a 'k0l' diff --git a/pyhdtoolkit/plotting/phasespace.py b/pyhdtoolkit/plotting/phasespace.py index e387b046..b723e90a 100644 --- a/pyhdtoolkit/plotting/phasespace.py +++ b/pyhdtoolkit/plotting/phasespace.py @@ -9,7 +9,6 @@ import matplotlib import matplotlib.axes import numpy as np - from cpymad.madx import Madx from loguru import logger from matplotlib import colors as mcolors diff --git a/pyhdtoolkit/plotting/sbs/coupling.py b/pyhdtoolkit/plotting/sbs/coupling.py index 01d1ea72..b896ce20 100644 --- a/pyhdtoolkit/plotting/sbs/coupling.py +++ b/pyhdtoolkit/plotting/sbs/coupling.py @@ -11,7 +11,6 @@ import matplotlib import matplotlib.pyplot as plt import tfs - from loguru import logger from matplotlib.legend import _get_legend_handles_labels diff --git a/pyhdtoolkit/plotting/sbs/phase.py b/pyhdtoolkit/plotting/sbs/phase.py index 44ce8e6d..4856ed4d 100644 --- a/pyhdtoolkit/plotting/sbs/phase.py +++ b/pyhdtoolkit/plotting/sbs/phase.py @@ -9,7 +9,6 @@ import matplotlib import matplotlib.pyplot as plt import tfs - from loguru import logger from matplotlib.legend import _get_legend_handles_labels diff --git a/pyhdtoolkit/plotting/styles/__init__.py b/pyhdtoolkit/plotting/styles/__init__.py index 738b8d97..9be7d26b 100644 --- a/pyhdtoolkit/plotting/styles/__init__.py +++ b/pyhdtoolkit/plotting/styles/__init__.py @@ -12,7 +12,6 @@ import matplotlib import matplotlib.pyplot as plt - from loguru import logger from . import paper, thesis diff --git a/pyhdtoolkit/plotting/tune.py b/pyhdtoolkit/plotting/tune.py index 169b5555..17a5994d 100644 --- a/pyhdtoolkit/plotting/tune.py +++ b/pyhdtoolkit/plotting/tune.py @@ -13,7 +13,6 @@ import matplotlib import matplotlib.axes import numpy as np - from loguru import logger from pyhdtoolkit.plotting.utils import maybe_get_ax diff --git a/pyhdtoolkit/plotting/utils.py b/pyhdtoolkit/plotting/utils.py index 7c040f77..b9eec23d 100644 --- a/pyhdtoolkit/plotting/utils.py +++ b/pyhdtoolkit/plotting/utils.py @@ -11,13 +11,12 @@ import matplotlib import matplotlib.axes import matplotlib.pyplot as plt -import matplotlib.transforms as transforms import numpy as np import pandas as pd import tfs - from cpymad.madx import Madx from loguru import logger +from matplotlib import transforms from matplotlib.patches import Ellipse from numpy.typing import ArrayLike @@ -87,8 +86,8 @@ def find_ip_s_from_segment_start(segment_df: tfs.TfsDataFrame, model_df: tfs.Tfs """ logger.debug(f"Determining location of IP{ip:d} from the start of segment.") first_element: str = segment_df.NAME.to_numpy()[0] - first_element_s_in_model = model_df[model_df.NAME == first_element].S.to_numpy()[0] - ip_s_in_model = model_df[model_df.NAME == f"IP{ip:d}"].S.to_numpy()[0] + first_element_s_in_model = model_df[first_element == model_df.NAME].S.to_numpy()[0] + ip_s_in_model = model_df[f"IP{ip:d}" == model_df.NAME].S.to_numpy()[0] # Handle case where IP segment is cut and by end of sequence and the IP is at beginning of machine if ip_s_in_model < first_element_s_in_model: @@ -496,4 +495,4 @@ def _determine_default_sbs_phase_ylabel(plane: str) -> str: def _matplotlib_svg_scraper(*args, **kwargs): from sphinx_gallery.scrapers import matplotlib_scraper kwargs.pop("format", None) - return matplotlib_scraper(*args, format="svg", **kwargs) \ No newline at end of file + return matplotlib_scraper(*args, format="svg", **kwargs) diff --git a/pyhdtoolkit/utils/__init__.py b/pyhdtoolkit/utils/__init__.py index 87128cae..d61876ac 100644 --- a/pyhdtoolkit/utils/__init__.py +++ b/pyhdtoolkit/utils/__init__.py @@ -2,7 +2,6 @@ import inspect import traceback import warnings - from typing import Callable from . import cmdline, contexts, executors, htc_monitor, logging, operations, printutil diff --git a/pyhdtoolkit/utils/_misc.py b/pyhdtoolkit/utils/_misc.py index bbd73c98..9ec55579 100644 --- a/pyhdtoolkit/utils/_misc.py +++ b/pyhdtoolkit/utils/_misc.py @@ -11,7 +11,6 @@ **not** work on other people's machines. """ import shlex - from multiprocessing import cpu_count from pathlib import Path from typing import Sequence, Union @@ -19,7 +18,6 @@ import cpymad import numpy as np import pandas as pd - from cpymad.madx import Madx from loguru import logger @@ -248,7 +246,7 @@ def get_betastar_from_opticsfile(opticsfile: Union[Path, str]) -> float: # returns 0.3 """ file_lines = Path(opticsfile).read_text().split("\n") - ip1_x_line, ip1_y_line, ip5_x_line, ip5_y_line = [line for line in file_lines if line.startswith("bet")] + ip1_x_line, ip1_y_line, ip5_x_line, ip5_y_line = (line for line in file_lines if line.startswith("bet")) betastar_x_ip1 = float(shlex.split(ip1_x_line)[2]) betastar_y_ip1 = float(shlex.split(ip1_y_line)[2]) betastar_x_ip5 = float(shlex.split(ip5_x_line)[2]) diff --git a/pyhdtoolkit/utils/cmdline.py b/pyhdtoolkit/utils/cmdline.py index c3d3ea94..56923f6b 100644 --- a/pyhdtoolkit/utils/cmdline.py +++ b/pyhdtoolkit/utils/cmdline.py @@ -11,7 +11,6 @@ import os import signal import subprocess - from typing import Mapping, Optional, Tuple from loguru import logger diff --git a/pyhdtoolkit/utils/contexts.py b/pyhdtoolkit/utils/contexts.py index ef1dc299..ddcb9b88 100644 --- a/pyhdtoolkit/utils/contexts.py +++ b/pyhdtoolkit/utils/contexts.py @@ -7,7 +7,6 @@ Provides useful contexts to use functions in. """ import time - from contextlib import contextmanager from typing import Callable, Iterator diff --git a/pyhdtoolkit/utils/htc_monitor.py b/pyhdtoolkit/utils/htc_monitor.py index 72155498..5063b967 100644 --- a/pyhdtoolkit/utils/htc_monitor.py +++ b/pyhdtoolkit/utils/htc_monitor.py @@ -14,11 +14,9 @@ """ import re import time - from typing import List, Tuple import pendulum - from loguru import logger from rich import box from rich.console import Group diff --git a/pyhdtoolkit/utils/logging.py b/pyhdtoolkit/utils/logging.py index 82b20f6e..398b4ac9 100644 --- a/pyhdtoolkit/utils/logging.py +++ b/pyhdtoolkit/utils/logging.py @@ -15,7 +15,6 @@ - ``SIMPLE_FORMAT``: minimal, displays the local time, the level and the message. """ import sys - from typing import Union from loguru import logger diff --git a/pyhdtoolkit/utils/operations.py b/pyhdtoolkit/utils/operations.py index 72b83747..61132d8e 100644 --- a/pyhdtoolkit/utils/operations.py +++ b/pyhdtoolkit/utils/operations.py @@ -16,7 +16,6 @@ import math import random import re - from functools import reduce from typing import Callable, Dict, List, Sequence, Tuple, Union @@ -184,7 +183,7 @@ def deep_flatten(sequence: Sequence) -> list: ) @staticmethod - def eval_none(sequence: Sequence, function: Callable = lambda x: not not x) -> bool: + def eval_none(sequence: Sequence, function: Callable = lambda x: bool(x)) -> bool: """ .. versionadded:: 0.2.0 @@ -215,7 +214,7 @@ def eval_none(sequence: Sequence, function: Callable = lambda x: not not x) -> b return not any(map(function, sequence)) @staticmethod - def eval_some(sequence: Sequence, function: Callable = lambda x: not not x) -> bool: + def eval_some(sequence: Sequence, function: Callable = lambda x: bool(x)) -> bool: """ .. versionadded:: 0.2.0 diff --git a/pyhdtoolkit/version.py b/pyhdtoolkit/version.py index 1f979c7c..bb83f5bf 100644 --- a/pyhdtoolkit/version.py +++ b/pyhdtoolkit/version.py @@ -5,7 +5,6 @@ def version_info() -> str: """Debug convenience function to give version, platform and runtime information.""" import platform import sys - from pathlib import Path info = { diff --git a/tests/conftest.py b/tests/conftest.py index 5a63500f..b65ac811 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,7 +1,6 @@ import pathlib import pytest - from cpymad.madx import Madx from pyhdtoolkit.cpymadtools import lhc @@ -20,7 +19,7 @@ # ----- Fixtures for cpymadtools tests ----- # -@pytest.fixture() +@pytest.fixture def _matched_base_lattice() -> Madx: """Base CAS lattice matched to default working point.""" with Madx(stdout=False) as madx: @@ -37,7 +36,7 @@ def _matched_base_lattice() -> Madx: yield madx -@pytest.fixture() +@pytest.fixture def _bare_lhc_madx() -> Madx: """Only loading sequence and optics.""" with Madx(stdout=False) as madx: @@ -46,7 +45,7 @@ def _bare_lhc_madx() -> Madx: yield madx -@pytest.fixture() +@pytest.fixture def _non_matched_lhc_madx() -> Madx: """Important properties & beam for lhcb1 declared and in use, NO MATCHING done here.""" with Madx(stdout=False) as madx: @@ -70,7 +69,7 @@ def _non_matched_lhc_madx() -> Madx: yield madx -@pytest.fixture() +@pytest.fixture def _matched_lhc_madx() -> Madx: """Important properties & beam for lhcb1 declared and in use, WITH matching to working point.""" with Madx(stdout=False) as madx: @@ -95,7 +94,7 @@ def _matched_lhc_madx() -> Madx: yield madx -@pytest.fixture() +@pytest.fixture def _cycled_lhc_sequences() -> Madx: """Important properties & beam for lhcb1 and lhcb1 declared and in use, WITH matching to working point.""" with Madx(stdout=False) as madx: @@ -108,7 +107,7 @@ def _cycled_lhc_sequences() -> Madx: yield madx -@pytest.fixture() +@pytest.fixture def _injection_aperture_tolerances_lhc_madx() -> Madx: with Madx(stdout=False) as madx: madx.call(str(LHC_SEQUENCE.absolute())) @@ -117,7 +116,7 @@ def _injection_aperture_tolerances_lhc_madx() -> Madx: lhc.make_lhc_beams(madx, energy=450) # injection madx.use(sequence="lhcb1") - madx.call((str(LHC_B1_APERTURE.absolute()))) + madx.call(str(LHC_B1_APERTURE.absolute())) madx.call(str(LHC_B1_APERTOL.absolute())) madx.command.twiss() @@ -125,7 +124,7 @@ def _injection_aperture_tolerances_lhc_madx() -> Madx: yield madx -@pytest.fixture() +@pytest.fixture def _collision_aperture_tolerances_lhc_madx() -> Madx: with Madx(stdout=False) as madx: madx.call(str(LHC_SEQUENCE.absolute())) @@ -134,7 +133,7 @@ def _collision_aperture_tolerances_lhc_madx() -> Madx: lhc.make_lhc_beams(madx, energy=6500) # collision madx.use(sequence="lhcb1") - madx.call((str(LHC_B1_APERTURE.absolute()))) + madx.call(str(LHC_B1_APERTURE.absolute())) madx.call(str(LHC_B1_APERTOL.absolute())) madx.command.twiss() diff --git a/tests/test_cpymadtools/test_coupling.py b/tests/test_cpymadtools/test_coupling.py index 77efeecc..4e9ae1dd 100644 --- a/tests/test_cpymadtools/test_coupling.py +++ b/tests/test_cpymadtools/test_coupling.py @@ -4,7 +4,6 @@ import numpy as np import pytest import tfs - from optics_functions.constants import F1001, F1010 from optics_functions.coupling import split_complex_columns from pandas.testing import assert_frame_equal @@ -92,11 +91,11 @@ def test_get_coupling_rdts(_non_matched_lhc_madx, _coupling_bump_script, _correc # ---------------------- Private Utilities ---------------------- # -@pytest.fixture() +@pytest.fixture def _coupling_bump_script() -> pathlib.Path: return INPUTS_DIR / "madx" / "lhc_coupling_bump.madx" -@pytest.fixture() +@pytest.fixture def _correct_bump_rdts_path() -> pathlib.Path: return INPUTS_DIR / "cpymadtools" / "lhc_coupling_bump.tfs" diff --git a/tests/test_cpymadtools/test_lhc.py b/tests/test_cpymadtools/test_lhc.py index 65f2e3dd..a3f26082 100644 --- a/tests/test_cpymadtools/test_lhc.py +++ b/tests/test_cpymadtools/test_lhc.py @@ -6,11 +6,9 @@ import numpy as np import pytest import tfs - from cpymad.madx import Madx from pandas.testing import assert_frame_equal -from pyhdtoolkit.cpymadtools.constants import LHC_TRIPLETS_REGEX # noqa: F401 | for coverage from pyhdtoolkit.cpymadtools.constants import ( DEFAULT_TWISS_COLUMNS, LHC_ANGLE_FLAGS, @@ -34,6 +32,7 @@ LHC_KSF_KNOBS, LHC_KSS_KNOBS, LHC_PARALLEL_SEPARATION_FLAGS, + LHC_TRIPLETS_REGEX, # noqa: F401 | for coverage ) from pyhdtoolkit.cpymadtools.lhc import ( LHCSetup, @@ -852,31 +851,31 @@ def test_lhc_run2_setup_raises_on_absent_sequence_file(): # ---------------------- Private Utilities ---------------------- # -@pytest.fixture() +@pytest.fixture def _magnets_fields_path() -> pathlib.Path: return INPUTS_DIR / "cpymadtools" / "magnets_fields.tfs" -@pytest.fixture() +@pytest.fixture def _correct_bpms_list() -> pathlib.Path: return INPUTS_DIR / "cpymadtools" / "correct_bpms_list.pkl" -@pytest.fixture() +@pytest.fixture def _reference_twiss_rdts() -> pathlib.Path: return INPUTS_DIR / "cpymadtools" / "twiss_with_rdts.tfs" -@pytest.fixture() +@pytest.fixture def _reference_kmodulation() -> pathlib.Path: return INPUTS_DIR / "cpymadtools" / "kmodulation.tfs" -@pytest.fixture() +@pytest.fixture def _ips_twiss_path() -> pathlib.Path: return INPUTS_DIR / "cpymadtools" / "ips_twiss.tfs" -@pytest.fixture() +@pytest.fixture def _proton_opticsfile() -> str: return str((PROTON_DIR / "opticsfile.22").absolute()) diff --git a/tests/test_cpymadtools/test_matching.py b/tests/test_cpymadtools/test_matching.py index daa0ba6f..6b3ced7b 100644 --- a/tests/test_cpymadtools/test_matching.py +++ b/tests/test_cpymadtools/test_matching.py @@ -1,9 +1,7 @@ import math - from sys import platform import pytest - from cpymad.madx import Madx from pyhdtoolkit.cpymadtools._generators import LatticeGenerator diff --git a/tests/test_cpymadtools/test_ptc.py b/tests/test_cpymadtools/test_ptc.py index 1262aad0..e3c1eacd 100644 --- a/tests/test_cpymadtools/test_ptc.py +++ b/tests/test_cpymadtools/test_ptc.py @@ -2,7 +2,6 @@ import pytest import tfs - from cpymad.madx import Madx from pandas import DataFrame from pandas.testing import assert_frame_equal @@ -93,16 +92,16 @@ def test_single_particle_ptc_track_with_onepass(_matched_base_lattice): # ----- Fixtures ----- # -@pytest.fixture() +@pytest.fixture def _ampdet_tfs_path() -> pathlib.Path: return INPUTS_DIR / "cpymadtools" / "ampdet.tfs" -@pytest.fixture() +@pytest.fixture def _rdts_tfs_path() -> pathlib.Path: return INPUTS_DIR / "cpymadtools" / "rdts.tfs" -@pytest.fixture() +@pytest.fixture def _ptc_twiss_tfs_path() -> pathlib.Path: return INPUTS_DIR / "cpymadtools" / "ptc_twiss.tfs" diff --git a/tests/test_cpymadtools/test_track.py b/tests/test_cpymadtools/test_track.py index db5aa9ca..96c07510 100644 --- a/tests/test_cpymadtools/test_track.py +++ b/tests/test_cpymadtools/test_track.py @@ -1,5 +1,4 @@ import pytest - from pandas import DataFrame from pyhdtoolkit.cpymadtools.track import track_single_particle diff --git a/tests/test_cpymadtools/test_tune.py b/tests/test_cpymadtools/test_tune.py index ce2431e2..2438a66a 100644 --- a/tests/test_cpymadtools/test_tune.py +++ b/tests/test_cpymadtools/test_tune.py @@ -105,11 +105,11 @@ def test_get_footprint_patches_raises_wrong_shape(_dynap_tfs_path, caplog): # ----- Fixtures ----- # -@pytest.fixture() +@pytest.fixture def _plottable_footprint_path() -> pathlib.Path: return INPUTS_DIR / "cpymadtools" / "plottable_footprint.npz" -@pytest.fixture() +@pytest.fixture def _dynap_tfs_path() -> pathlib.Path: return INPUTS_DIR / "cpymadtools" / "dynap.tfs" diff --git a/tests/test_cpymadtools/test_twiss.py b/tests/test_cpymadtools/test_twiss.py index 2e946de2..7a97d9d4 100644 --- a/tests/test_cpymadtools/test_twiss.py +++ b/tests/test_cpymadtools/test_twiss.py @@ -2,7 +2,6 @@ import pytest import tfs - from pandas.testing import assert_frame_equal from pyhdtoolkit.cpymadtools.constants import DEFAULT_TWISS_COLUMNS # noqa: F401 | for coverage @@ -22,6 +21,6 @@ def test_twiss_tfs(_twiss_export, _matched_base_lattice): # ---------------------- Private Utilities ---------------------- # -@pytest.fixture() +@pytest.fixture def _twiss_export() -> pathlib.Path: return INPUTS_DIR / "cpymadtools" / "twiss_export.tfs" diff --git a/tests/test_cpymadtools/test_utilities.py b/tests/test_cpymadtools/test_utilities.py index 7637a2fc..2d28c7d3 100644 --- a/tests/test_cpymadtools/test_utilities.py +++ b/tests/test_cpymadtools/test_utilities.py @@ -1,6 +1,5 @@ import pytest import tfs - from pandas.testing import assert_frame_equal from pyhdtoolkit.cpymadtools.utils import _get_k_strings, export_madx_table, get_table_tfs diff --git a/tests/test_maths.py b/tests/test_maths.py index 497812ae..87c5f11d 100644 --- a/tests/test_maths.py +++ b/tests/test_maths.py @@ -1,5 +1,4 @@ import pathlib - from copy import deepcopy import numpy as np @@ -28,7 +27,7 @@ ) def test_setting_distributions_dict(input_dict): stats_fitting.set_distributions_dict(input_dict) - assert stats_fitting.DISTRIBUTIONS == input_dict + assert input_dict == stats_fitting.DISTRIBUTIONS @pytest.mark.flaky(max_runs=3, min_passes=1) @@ -165,16 +164,16 @@ def _create_2d_gaussian_noise(mean: float, stdev: float, shape: tuple) -> np.nda return upper_triangle - upper_triangle.T -@pytest.fixture() +@pytest.fixture def _to_scale() -> np.ndarray: return np.load(INPUTS_DIR / "maths" / "to_scale.npy") -@pytest.fixture() +@pytest.fixture def _scaled() -> np.ndarray: return np.load(INPUTS_DIR / "maths" / "scaled.npy") -@pytest.fixture() +@pytest.fixture def _force_scaled() -> np.ndarray: return np.load(INPUTS_DIR / "maths" / "force_scaled.npy") diff --git a/tests/test_optics.py b/tests/test_optics.py index a5419a12..9a730772 100644 --- a/tests/test_optics.py +++ b/tests/test_optics.py @@ -186,6 +186,6 @@ def test_rdt_spectrum_line(): # ----- Fixtures ----- # -@pytest.fixture() +@pytest.fixture def _fake_coordinates() -> np.ndarray: return np.random.random(size=10_000) / 1e4 diff --git a/tests/test_plotting/test_aperture.py b/tests/test_plotting/test_aperture.py index cd1f3026..59f85010 100644 --- a/tests/test_plotting/test_aperture.py +++ b/tests/test_plotting/test_aperture.py @@ -1,7 +1,6 @@ import matplotlib import matplotlib.pyplot as plt import pytest - from cpymad.madx import Madx from pyhdtoolkit.plotting.aperture import plot_aperture, plot_physical_apertures diff --git a/tests/test_plotting/test_envelope.py b/tests/test_plotting/test_envelope.py index d128a4e0..f28fc6df 100644 --- a/tests/test_plotting/test_envelope.py +++ b/tests/test_plotting/test_envelope.py @@ -3,7 +3,6 @@ import matplotlib import matplotlib.pyplot as plt import pytest - from cpymad.madx import Madx from pyhdtoolkit.plotting.envelope import _interpolate_madx, plot_beam_envelope diff --git a/tests/test_plotting/test_lattice.py b/tests/test_plotting/test_lattice.py index d22b385e..dcbec8d6 100644 --- a/tests/test_plotting/test_lattice.py +++ b/tests/test_plotting/test_lattice.py @@ -3,7 +3,6 @@ import matplotlib import matplotlib.pyplot as plt import pytest - from cpymad.madx import Madx from pyhdtoolkit.cpymadtools._generators import LatticeGenerator @@ -122,7 +121,7 @@ def test_plot_layout_raises_on_wrong_limits_type(): with Madx(stdout=False) as madx: madx.input(BASE_LATTICE) plt.figure(figsize=(18, 11)) - + with pytest.raises(TypeError): plot_latwiss(madx, k1l_lim=[8e-2]) diff --git a/tests/test_plotting/test_phasespace.py b/tests/test_plotting/test_phasespace.py index 76715a71..0a3c3f26 100644 --- a/tests/test_plotting/test_phasespace.py +++ b/tests/test_plotting/test_phasespace.py @@ -5,7 +5,6 @@ import matplotlib.pyplot as plt import numpy as np import pytest - from cpymad.madx import Madx from pyhdtoolkit.cpymadtools._generators import LatticeGenerator diff --git a/tests/test_plotting/test_plotting_utils.py b/tests/test_plotting/test_plotting_utils.py index 85d3e3bb..dbef6c4b 100644 --- a/tests/test_plotting/test_plotting_utils.py +++ b/tests/test_plotting/test_plotting_utils.py @@ -164,21 +164,21 @@ def get_correlated_dataset(n, dependency, mu, scale): # ----- Fixtures ----- # -@pytest.fixture() +@pytest.fixture def sbs_model_b1() -> tfs.TfsDataFrame: return tfs.read(SBS_INPUTS / "b1_twiss_elements.dat") -@pytest.fixture() +@pytest.fixture def sbs_model_b2() -> tfs.TfsDataFrame: return tfs.read(SBS_INPUTS / "b2_twiss_elements.dat") -@pytest.fixture() +@pytest.fixture def sbs_coupling_b1_ip1() -> tfs.TfsDataFrame: return tfs.read(SBS_INPUTS / "b1_sbscouple_IP1.out") -@pytest.fixture() +@pytest.fixture def sbs_coupling_b2_ip1() -> tfs.TfsDataFrame: return tfs.read(SBS_INPUTS / "b2_sbscouple_IP1.out") diff --git a/tests/test_plotting/test_sbs_coupling.py b/tests/test_plotting/test_sbs_coupling.py index 6866097c..4544d9bd 100644 --- a/tests/test_plotting/test_sbs_coupling.py +++ b/tests/test_plotting/test_sbs_coupling.py @@ -58,21 +58,21 @@ def test_plot_full_ip_with_ylimits(sbs_coupling_b1_ip1, sbs_coupling_b2_ip1, sbs # ----- Fixtures ----- # -@pytest.fixture() +@pytest.fixture def sbs_model_b1() -> tfs.TfsDataFrame: return tfs.read(SBS_INPUTS / "b1_twiss_elements.dat") -@pytest.fixture() +@pytest.fixture def sbs_model_b2() -> tfs.TfsDataFrame: return tfs.read(SBS_INPUTS / "b2_twiss_elements.dat") -@pytest.fixture() +@pytest.fixture def sbs_coupling_b1_ip1() -> tfs.TfsDataFrame: return tfs.read(SBS_INPUTS / "b1_sbscouple_IP1.out") -@pytest.fixture() +@pytest.fixture def sbs_coupling_b2_ip1() -> tfs.TfsDataFrame: return tfs.read(SBS_INPUTS / "b2_sbscouple_IP1.out") diff --git a/tests/test_plotting/test_sbs_phase.py b/tests/test_plotting/test_sbs_phase.py index 09a07036..51cfdc44 100644 --- a/tests/test_plotting/test_sbs_phase.py +++ b/tests/test_plotting/test_sbs_phase.py @@ -51,16 +51,16 @@ def test_plot_phase_segment_raises_on_wrong_plane(wrongplane, sbs_phasex, sbs_mo # ----- Fixtures ----- # -@pytest.fixture() +@pytest.fixture def sbs_phasex() -> tfs.TfsDataFrame: return tfs.read(SBS_INPUTS / "b2sbsphasext_IP5.out") -@pytest.fixture() +@pytest.fixture def sbs_phasey() -> tfs.TfsDataFrame: return tfs.read(SBS_INPUTS / "b2sbsphaseyt_IP5.out") -@pytest.fixture() +@pytest.fixture def sbs_model_b2() -> tfs.TfsDataFrame: return tfs.read(SBS_INPUTS / "b2_twiss_elements.dat") diff --git a/tests/test_utils.py b/tests/test_utils.py index 0f22c897..5b6bdd0d 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -6,14 +6,12 @@ import random import subprocess import sys - from typing import List import numpy as np import pandas as pd import pytest import tfs - from loguru import logger from numpy.testing import assert_array_equal from rich.table import Table @@ -695,7 +693,7 @@ def test_bpm_noise_addition(self, _rdts_df): # ----- Fixtures ----- # -@pytest.fixture() +@pytest.fixture def _condor_q_output() -> str: condor_q_output = """-- Schedd: bigbird08.cern.ch : <188.185.72.155:9618?... @ 04/22/21 12:26:02 OWNER BATCH_NAME SUBMITTED DONE RUN IDLE TOTAL JOB_IDS @@ -713,7 +711,7 @@ def _condor_q_output() -> str: return condor_q_output -@pytest.fixture() +@pytest.fixture def _taskless_condor_q_output() -> str: taskless_condor_q_output = """-- Schedd: bigbird08.cern.ch : <188.185.72.155:9618?... @ 04/22/21 12:26:02 OWNER BATCH_NAME SUBMITTED DONE RUN IDLE TOTAL JOB_IDS @@ -724,26 +722,26 @@ def _taskless_condor_q_output() -> str: return taskless_condor_q_output -@pytest.fixture() +@pytest.fixture def _correct_user_tasks() -> List[HTCTaskSummary]: pickle_file_path = INPUTS_DIR / "utils" / "correct_user_tasks.pkl" with pickle_file_path.open("rb") as file: return pickle.load(file) -@pytest.fixture() +@pytest.fixture def _correct_cluster_summary() -> ClusterSummary: pickle_file_path = INPUTS_DIR / "utils" / "correct_cluster_summary.pkl" with pickle_file_path.open("rb") as file: return pickle.load(file) -@pytest.fixture() +@pytest.fixture def _complex_columns_df() -> pd.DataFrame: array = np.random.rand(50, 5) + 1j * np.random.rand(50, 5) return pd.DataFrame(data=array, columns=["A", "B", "C", "D", "E"]) -@pytest.fixture() +@pytest.fixture def _rdts_df() -> pathlib.Path: return INPUTS_DIR / "cpymadtools" / "lhc_coupling_bump.tfs" From dde093bcf846350e5f284875dc3a5ad72f0d7968 Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 10:41:12 +0200 Subject: [PATCH 02/41] new round for dict, list, tuple --- pyhdtoolkit/__init__.py | 2 +- pyhdtoolkit/cpymadtools/__init__.py | 2 +- pyhdtoolkit/cpymadtools/constants.py | 55 ++++++++++++------------ pyhdtoolkit/cpymadtools/coupling.py | 6 +-- pyhdtoolkit/cpymadtools/lhc/_errors.py | 4 +- pyhdtoolkit/cpymadtools/lhc/_misc.py | 8 ++-- pyhdtoolkit/cpymadtools/lhc/_powering.py | 6 +-- pyhdtoolkit/cpymadtools/lhc/_queries.py | 10 ++--- pyhdtoolkit/cpymadtools/lhc/_setup.py | 6 +-- pyhdtoolkit/cpymadtools/ptc.py | 6 +-- pyhdtoolkit/cpymadtools/track.py | 6 +-- pyhdtoolkit/cpymadtools/tune.py | 12 +++--- pyhdtoolkit/cpymadtools/utils.py | 4 +- pyhdtoolkit/maths/stats_fitting.py | 12 +++--- pyhdtoolkit/maths/utils.py | 2 +- pyhdtoolkit/optics/rdt.py | 2 +- pyhdtoolkit/plotting/aperture.py | 34 +++++++-------- pyhdtoolkit/plotting/envelope.py | 4 +- pyhdtoolkit/plotting/lattice.py | 28 ++++++------ pyhdtoolkit/plotting/layout.py | 28 ++++++------ pyhdtoolkit/plotting/sbs/coupling.py | 6 +-- pyhdtoolkit/plotting/styles/__init__.py | 6 +-- pyhdtoolkit/plotting/styles/paper.py | 4 +- pyhdtoolkit/plotting/styles/thesis.py | 6 +-- pyhdtoolkit/plotting/tune.py | 14 +++--- pyhdtoolkit/utils/cmdline.py | 2 +- pyhdtoolkit/utils/executors.py | 6 +-- pyhdtoolkit/utils/htc_monitor.py | 8 ++-- pyhdtoolkit/utils/operations.py | 12 +++--- 29 files changed, 150 insertions(+), 151 deletions(-) diff --git a/pyhdtoolkit/__init__.py b/pyhdtoolkit/__init__.py index 4fd3ca37..0446552f 100644 --- a/pyhdtoolkit/__init__.py +++ b/pyhdtoolkit/__init__.py @@ -7,7 +7,7 @@ :copyright: (c) 2019-2020 by Felix Soubelet. :license: MIT, see LICENSE for more details. """ -from . import cpymadtools, maths, models, optics, plotting, utils, version # noqa: F401 +from . import cpymadtools, maths, models, optics, plotting, utils, version # noqa: F401, TID252 __title__ = "pyhdtoolkit" __description__ = "An all-in-one toolkit package to easy my Python work in my PhD." diff --git a/pyhdtoolkit/cpymadtools/__init__.py b/pyhdtoolkit/cpymadtools/__init__.py index ff6523bb..eca62876 100644 --- a/pyhdtoolkit/cpymadtools/__init__.py +++ b/pyhdtoolkit/cpymadtools/__init__.py @@ -1,3 +1,3 @@ from . import constants, coupling, lhc, matching, ptc, track, tune, twiss, utils -__all__ = [constants, coupling, lhc, matching, ptc, track, tune, twiss, utils] +__all__ = ["constants", "coupling", "lhc", "matching", "ptc", "track", "tune", "twiss", "utils"] diff --git a/pyhdtoolkit/cpymadtools/constants.py b/pyhdtoolkit/cpymadtools/constants.py index c8bbe5ee..19ddfd1b 100644 --- a/pyhdtoolkit/cpymadtools/constants.py +++ b/pyhdtoolkit/cpymadtools/constants.py @@ -6,13 +6,12 @@ Specific constants to be used in `~.cpymadtools` functions, to help with consistency. """ -from typing import Dict, List, Set # fmt: off -DEFAULT_TWISS_COLUMNS: List[str] = ["name", "s", "x", "y", "l", "px", "py", "betx", "bety", "alfx", "alfy", +DEFAULT_TWISS_COLUMNS: list[str] = ["name", "s", "x", "y", "l", "px", "py", "betx", "bety", "alfx", "alfy", "dx", "dy", "mux", "muy", "r11", "r12", "r21", "r22", "beta11", "beta12", "beta21", "beta22"] -MONITOR_TWISS_COLUMNS: List[str] = ["name", "s", "betx", "bety", "alfx", "alfy", "mux", "muy", "dx", "dy", +MONITOR_TWISS_COLUMNS: list[str] = ["name", "s", "betx", "bety", "alfx", "alfy", "mux", "muy", "dx", "dy", "dpx", "dpy", "x", "y", "ddx", "ddy", "k1l", "k1sl", "k2l", "k3l", "k4l", "wx", "wy", "phix", "phiy", "dmux", "dmuy", "keyword", "dbx", "dby", "r11", "r12", "r21", "r22"] @@ -25,7 +24,7 @@ LHC_TRIPLETS_REGEX = "^MQXF?[AB].[AB]?[123][RL][1258]" # This might not be accurate anymore -LHC_CROSSING_SCHEMES: Dict[str, Dict[str, float]] = { +LHC_CROSSING_SCHEMES: dict[str, dict[str, float]] = { "flat": {}, "lhc_inj": { "on_x1": -170, @@ -79,41 +78,41 @@ } # ----- LHC IP bump flags ----- # -LHC_CROSSING_ANGLE_FLAGS: List[str] = [f"on_x{ip}" for ip in [1, 2, 5, 8]] + [ +LHC_CROSSING_ANGLE_FLAGS: list[str] = [f"on_x{ip}" for ip in [1, 2, 5, 8]] + [ "on_x1h", "on_x5v", "on_s_x1h", "on_s_x5v", ] -LHC_PARALLEL_SEPARATION_FLAGS: List[str] = [f"on_sep{ip}" for ip in [1, 2, 5, 8]] + [ +LHC_PARALLEL_SEPARATION_FLAGS: list[str] = [f"on_sep{ip}" for ip in [1, 2, 5, 8]] + [ "on_sep1v", "on_sep5h", "on_s_sep1v", "on_s_sep5h", ] # Offset is in the plane of the crossing angle -LHC_IP_OFFSET_FLAGS: List[str] = ( +LHC_IP_OFFSET_FLAGS: list[str] = ( [f"on_o{ip}" for ip in [1, 2, 5, 8]] + ["on_o1h"] + [f"on_oh{ip}" for ip in [1, 5]] + [f"on_ov{ip}" for ip in [1, 2, 5]] ) # Angle is in the same plane as the separation -LHC_ANGLE_FLAGS: List[str] = [f"on_a{ip}" for ip in [1, 2, 5, 8]] + ["on_a1v"] +LHC_ANGLE_FLAGS: list[str] = [f"on_a{ip}" for ip in [1, 2, 5, 8]] + ["on_a1v"] # Magnetic fields or solenoids powering status -LHC_EXPERIMENT_STATE_FLAGS: List[str] = ["on_alice", "on_lhcb", "on_sol_atlas", "on_sol_cms", "on_sol_alice"] -LHC_IP2_SPECIAL_FLAG: List[str] = ["on_oe2"] # TODO: ask Tobias or Stephane F. +LHC_EXPERIMENT_STATE_FLAGS: list[str] = ["on_alice", "on_lhcb", "on_sol_atlas", "on_sol_cms", "on_sol_alice"] +LHC_IP2_SPECIAL_FLAG: list[str] = ["on_oe2"] # TODO: ask Tobias or Stephane F. # ----- LHC Triplet Correctors Knobs ----- # -LHC_KQSX_KNOBS: List[str] = [f"kqsx3.{side}{ip}" for side in ("r", "l") for ip in (1, 2, 5, 8)] # skew quad -LHC_KCSX_KNOBS: List[str] = [f"kcsx3.{side}{ip}" for side in ("r", "l") for ip in (1, 2, 5, 8)] # sextupole -LHC_KCSSX_KNOBS: List[str] = [f"kcssx3.{side}{ip}" for side in ("r", "l") for ip in (1, 2, 5, 8)] # skew sext -LHC_KCOX_KNOBS: List[str] = [f"kcox3.{side}{ip}" for side in ("r", "l") for ip in (1, 2, 5, 8)] # octupole -LHC_KCOSX_KNOBS: List[str] = [f"kcosx3.{side}{ip}" for side in ("r", "l") for ip in (1, 2, 5, 8)] # skew oct -LHC_KCTX_KNOBS: List[str] = [f"kctx3.{side}{ip}" for side in ("r", "l") for ip in (1, 2, 5, 8)] # decapole +LHC_KQSX_KNOBS: list[str] = [f"kqsx3.{side}{ip}" for side in ("r", "l") for ip in (1, 2, 5, 8)] # skew quad +LHC_KCSX_KNOBS: list[str] = [f"kcsx3.{side}{ip}" for side in ("r", "l") for ip in (1, 2, 5, 8)] # sextupole +LHC_KCSSX_KNOBS: list[str] = [f"kcssx3.{side}{ip}" for side in ("r", "l") for ip in (1, 2, 5, 8)] # skew sext +LHC_KCOX_KNOBS: list[str] = [f"kcox3.{side}{ip}" for side in ("r", "l") for ip in (1, 2, 5, 8)] # octupole +LHC_KCOSX_KNOBS: list[str] = [f"kcosx3.{side}{ip}" for side in ("r", "l") for ip in (1, 2, 5, 8)] # skew oct +LHC_KCTX_KNOBS: list[str] = [f"kctx3.{side}{ip}" for side in ("r", "l") for ip in (1, 2, 5, 8)] # decapole # ----- LHC Arc Correctors Knobs ----- # -LHC_KQTF_KNOBS: List[str] = [ # tune trims, focusing and defocusing families, for each beam +LHC_KQTF_KNOBS: list[str] = [ # tune trims, focusing and defocusing families, for each beam f"kqt{family}.a{sector}{sector+1 if sector < 8 else 1}.b{beam}" for beam in [1, 2] for family in ["f", "d"] @@ -121,40 +120,40 @@ ] # fmt: off # skew quadrupoles in arc short straight sections -LHC_KQS_KNOBS: List[str] = [f"kqs.r{ip}b1" for ip in [1, 3, 5, 7]] + \ +LHC_KQS_KNOBS: list[str] = [f"kqs.r{ip}b1" for ip in [1, 3, 5, 7]] + \ [f"kqs.l{ip}b1" for ip in [2, 4, 6, 8]] + \ [f"kqs.a{sector}{sector+1 if sector < 8 else 1}b1" for sector in [2, 4, 6, 8]] + \ [f"kqs.r{ip}b2" for ip in [2, 4, 6, 8]] + \ [f"kqs.l" f"{ip}b2" for ip in [3, 5, 7, 1]] + \ [f"kqs.a{sector}{sector+1 if sector < 8 else 1}b2" for sector in [1, 3, 5, 7]] # fmt: on -LHC_KSF_KNOBS: List[str] = [ # sextupole correctors +LHC_KSF_KNOBS: list[str] = [ # sextupole correctors f"ks{family}{id}.a{sector}{sector+1 if sector < 8 else 1}b{beam}" for beam in [1, 2] for id in [1, 2] for family in ["f", "d"] for sector in [1, 2, 3, 4, 5, 6, 7, 8] ] -LHC_KSS_KNOBS: List[str] = [ # skew sextupole correctors +LHC_KSS_KNOBS: list[str] = [ # skew sextupole correctors f"kss.a{sector}{sector+1 if sector < 8 else 1}b{beam}" for beam in [1, 2] for sector in [1, 2, 3, 4, 5, 6, 7, 8] ] -LHC_KCS_KNOBS: List[str] = [ # spool piece (skew) sextupoles +LHC_KCS_KNOBS: list[str] = [ # spool piece (skew) sextupoles f"kcs.a{sector}{sector+1 if sector < 8 else 1}b{beam}" for beam in [1, 2] for sector in [1, 2, 3, 4, 5, 6, 7, 8] ] -LHC_KCO_KNOBS: List[str] = [ # spool piece (skew) octupoles +LHC_KCO_KNOBS: list[str] = [ # spool piece (skew) octupoles f"kco.a{sector}{sector+1 if sector < 8 else 1}b{beam}" for beam in [1, 2] for sector in [1, 2, 3, 4, 5, 6, 7, 8] ] -LHC_KCD_KNOBS: List[str] = [ # spool piece (skew) decapoles +LHC_KCD_KNOBS: list[str] = [ # spool piece (skew) decapoles f"kcd.a{sector}{sector+1 if sector < 8 else 1}b{beam}" for beam in [1, 2] for sector in [1, 2, 3, 4, 5, 6, 7, 8] ] -LHC_KO_KNOBS: List[str] = [ # octupoles in arc short straight sections +LHC_KO_KNOBS: list[str] = [ # octupoles in arc short straight sections f"ko{family}.a{sector}{sector+1 if sector < 8 else 1}b{beam}" for beam in [1, 2] for family in ["f", "d"] for sector in [1, 2, 3, 4, 5, 6, 7, 8] ] -HLLHC_CORRECTOR_LIMITS: Dict[str, float] = { # All values are defined as multiples of 0.3/Energy +HLLHC_CORRECTOR_LIMITS: dict[str, float] = { # All values are defined as multiples of 0.3/Energy "MQSX1": 0.600 / 0.050, # 0.6 T.m @ 50 mm in IR1&IR5 "MQSX2": 1.360 / 0.017, # 1.36 T @ 17 mm in IR2&IR8 # ------------- # @@ -188,6 +187,6 @@ "MO": 0.29 * 6 / (0.017**3), # 0.29 T @ 17 mm } -FD_FAMILIES: Set[str] = {"MO", "MS", "MQT"} # Magnets that have F and D families -TWO_FAMILIES: Set[str] = {"MS"} # Magnets that have 1 and 2 families -SPECIAL_FAMILIES: Set[str] = {"MQS"} # Magnets in every second arc +FD_FAMILIES: set[str] = {"MO", "MS", "MQT"} # Magnets that have F and D families +TWO_FAMILIES: set[str] = {"MS"} # Magnets that have 1 and 2 families +SPECIAL_FAMILIES: set[str] = {"MQS"} # Magnets in every second arc diff --git a/pyhdtoolkit/cpymadtools/coupling.py b/pyhdtoolkit/cpymadtools/coupling.py index 9a541ff7..2cc72daf 100644 --- a/pyhdtoolkit/cpymadtools/coupling.py +++ b/pyhdtoolkit/cpymadtools/coupling.py @@ -33,7 +33,7 @@ def get_closest_tune_approach( varied_knobs: Sequence[str] = None, telescopic_squeeze: bool = True, run3: bool = False, - explicit_targets: Tuple[float, float] = None, + explicit_targets: tuple[float, float] = None, step: float = 1e-7, calls: int = 100, tolerance: float = 1e-21, @@ -62,7 +62,7 @@ def get_closest_tune_approach( telescopic_squeeze (bool): ``LHC`` specific. If set to `True`, uses the ``(HL)LHC`` knobs for Telescopic Squeeze configuration. Defaults to `True` since `v0.9.0`. run3 (bool): if set to `True`, uses the `LHC` Run 3 `*_op` knobs. Defaults to `False`. - explicit_targets (Tuple[float, float]): if given, will be used as matching targets for `(Qx, Qy)`. + explicit_targets (tuple[float, float]): if given, will be used as matching targets for `(Qx, Qy)`. Otherwise, the target is determined as the middle of the current fractional tunes. Defaults to `None`. step (float): step size to use when varying knobs. @@ -97,7 +97,7 @@ def get_closest_tune_approach( logger.debug("Saving knob values to restore after closest tune approach") varied_knobs = varied_knobs or tune_knobs # if accelerator was given we've extracted this already - saved_knobs: Dict[str, float] = {knob: madx.globals[knob] for knob in varied_knobs} + saved_knobs: dict[str, float] = {knob: madx.globals[knob] for knob in varied_knobs} logger.trace(f"Saved knobs are {saved_knobs}") if explicit_targets: diff --git a/pyhdtoolkit/cpymadtools/lhc/_errors.py b/pyhdtoolkit/cpymadtools/lhc/_errors.py index e589723d..7e4bb66a 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_errors.py +++ b/pyhdtoolkit/cpymadtools/lhc/_errors.py @@ -5,12 +5,12 @@ The functions below are utilities to implement errors in elements of the ``LHC``. """ -from typing import Dict, List, Sequence +from typing import Dict, Sequence from cpymad.madx import Madx from loguru import logger -LHC_IR_QUADS_PATTERNS: Dict[int, List[str]] = { +LHC_IR_QUADS_PATTERNS: dict[int, list[str]] = { 1: ["^MQXA.1{side}{ip:d}", "^MQXFA.[AB]1{side}{ip:d}"], # Q1 LHC, Q1A & Q1B HL-LHC 2: ["^MQXB.[AB]2{side}{ip:d}", "^MQXB.[AB]2{side}{ip:d}"], # Q2A & Q2B LHC, Q2A & Q2B HL-LHC 3: ["^MQXA.3{side}{ip:d}", "^MQXFA.[AB]3{side}{ip:d}"], # Q3 LHC, Q3A & Q3B HL-LHC diff --git a/pyhdtoolkit/cpymadtools/lhc/_misc.py b/pyhdtoolkit/cpymadtools/lhc/_misc.py index 171fda17..390c5537 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_misc.py +++ b/pyhdtoolkit/cpymadtools/lhc/_misc.py @@ -5,7 +5,7 @@ The functions below are miscellaneous utilities for the ``LHC``. """ -from typing import List, Tuple +from typing import Tuple from cpymad.madx import Madx from loguru import logger @@ -79,7 +79,7 @@ def reset_lhc_bump_flags(madx: Madx, /) -> None: def get_lhc_tune_and_chroma_knobs( accelerator: str, beam: int = 1, telescopic_squeeze: bool = True, run3: bool = False -) -> Tuple[str, str, str, str]: +) -> tuple[str, str, str, str]: """ .. versionadded:: 0.16.0 @@ -142,7 +142,7 @@ def get_lhc_tune_and_chroma_knobs( }[accelerator.upper()] -def get_lhc_bpms_list(madx: Madx, /) -> List[str]: +def get_lhc_bpms_list(madx: Madx, /) -> list[str]: """ .. versionadded:: 0.16.0 @@ -171,7 +171,7 @@ def get_lhc_bpms_list(madx: Madx, /) -> List[str]: def get_sizes_at_ip( madx: Madx, /, ip: int, geom_emit_x: float = None, geom_emit_y: float = None -) -> Tuple[float, float]: +) -> tuple[float, float]: """ .. versionadded:: 1.0.0 diff --git a/pyhdtoolkit/cpymadtools/lhc/_powering.py b/pyhdtoolkit/cpymadtools/lhc/_powering.py index 799ba1ea..1e3122f4 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_powering.py +++ b/pyhdtoolkit/cpymadtools/lhc/_powering.py @@ -5,7 +5,7 @@ The functions below are magnets or knobs powering utilities for the ``LHC``. """ -from typing import Dict, List, Sequence +from typing import Dict, Sequence from cpymad.madx import Madx from loguru import logger @@ -321,7 +321,7 @@ def vary_independent_ir_quadrupoles( logger.debug(f"Preparing a knob involving quadrupoles {quad_numbers}") # Each quad has a specific power circuit used for their k1 boundaries - power_circuits: Dict[int, str] = { + power_circuits: dict[int, str] = { 4: "mqy", 5: "mqml", 6: "mqml", @@ -398,7 +398,7 @@ def switch_magnetic_errors(madx: Madx, /, **kwargs) -> None: # ----- Helpers ----- # -def _all_lhc_arcs(beam: int) -> List[str]: +def _all_lhc_arcs(beam: int) -> list[str]: """ Generates and returns the names of all LHC arcs for a given beam. Initial implementation credits go to :user:`Joschua Dilly `. diff --git a/pyhdtoolkit/cpymadtools/lhc/_queries.py b/pyhdtoolkit/cpymadtools/lhc/_queries.py index 4a443ba3..34cb4446 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_queries.py +++ b/pyhdtoolkit/cpymadtools/lhc/_queries.py @@ -88,7 +88,7 @@ def get_magnets_powering( return twiss.get_pattern_twiss(madx, columns=NEW_COLNAMES, patterns=patterns, **kwargs) -def query_arc_correctors_powering(madx: Madx, /) -> Dict[str, float]: +def query_arc_correctors_powering(madx: Madx, /) -> dict[str, float]: """ .. versionadded:: 0.15.0 @@ -108,7 +108,7 @@ def query_arc_correctors_powering(madx: Madx, /) -> Dict[str, float]: arc_knobs = query_arc_correctors_powering(madx) """ logger.debug("Querying triplets correctors powering") - result: Dict[str, float] = {} + result: dict[str, float] = {} logger.debug("Querying arc tune trim quadrupole correctors (MQTs) powering") k_mqt_max = 120 / madx.globals.brho # 120 T/m @@ -144,7 +144,7 @@ def query_arc_correctors_powering(madx: Madx, /) -> Dict[str, float]: return result -def query_triplet_correctors_powering(madx: Madx, /) -> Dict[str, float]: +def query_triplet_correctors_powering(madx: Madx, /) -> dict[str, float]: """ .. versionadded:: 0.15.0 @@ -164,7 +164,7 @@ def query_triplet_correctors_powering(madx: Madx, /) -> Dict[str, float]: triplet_knobs = query_triplet_correctors_powering(madx) """ logger.debug("Querying triplets correctors powering") - result: Dict[str, float] = {} + result: dict[str, float] = {} logger.debug("Querying triplet skew quadrupole correctors (MQSXs) powering") k_mqsx_max = 1.360 / 0.017 / madx.globals.brho # 1.36 T @ 17mm @@ -192,7 +192,7 @@ def query_triplet_correctors_powering(madx: Madx, /) -> Dict[str, float]: return result -def get_current_orbit_setup(madx: Madx, /) -> Dict[str, float]: +def get_current_orbit_setup(madx: Madx, /) -> dict[str, float]: """ .. versionadded:: 0.8.0 diff --git a/pyhdtoolkit/cpymadtools/lhc/_setup.py b/pyhdtoolkit/cpymadtools/lhc/_setup.py index 0016fa6b..59191a20 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_setup.py +++ b/pyhdtoolkit/cpymadtools/lhc/_setup.py @@ -6,7 +6,7 @@ The functions below are setup utilities for the ``LHC``, to easily get simulations ready. """ from pathlib import Path -from typing import Dict, List, Tuple +from typing import Dict, Tuple from cpymad.madx import Madx from loguru import logger @@ -410,7 +410,7 @@ def re_cycle_sequence(madx: Madx, /, sequence: str = "lhcb1", start: str = "IP3" madx.command.endedit() -def lhc_orbit_variables() -> Tuple[List[str], Dict[str, str]]: +def lhc_orbit_variables() -> tuple[list[str], dict[str, str]]: """ .. versionadded:: 0.8.0 @@ -471,7 +471,7 @@ def lhc_orbit_variables() -> Tuple[List[str], Dict[str, str]]: return variables, special -def setup_lhc_orbit(madx: Madx, /, scheme: str = "flat", **kwargs) -> Dict[str, float]: +def setup_lhc_orbit(madx: Madx, /, scheme: str = "flat", **kwargs) -> dict[str, float]: """ .. versionadded:: 0.8.0 diff --git a/pyhdtoolkit/cpymadtools/ptc.py b/pyhdtoolkit/cpymadtools/ptc.py index e378da80..4b2c2f60 100644 --- a/pyhdtoolkit/cpymadtools/ptc.py +++ b/pyhdtoolkit/cpymadtools/ptc.py @@ -340,14 +340,14 @@ def ptc_twiss( def ptc_track_particle( madx: Madx, /, - initial_coordinates: Tuple[float, float, float, float, float, float], + initial_coordinates: tuple[float, float, float, float, float, float], nturns: int, sequence: Optional[str] = None, observation_points: Sequence[str] = None, onetable: bool = False, fringe: bool = False, **kwargs, -) -> Dict[str, pd.DataFrame]: +) -> dict[str, pd.DataFrame]: """ .. versionadded:: 0.12.0 @@ -377,7 +377,7 @@ def ptc_track_particle( Args: madx (cpymad.madx.Madx): an instantiated cpymad.madx.Madx object. - initial_coordinates (Tuple[float, float, float, float, float, float]): a tuple + initial_coordinates (tuple[float, float, float, float, float, float]): a tuple with the ``X, PX, Y, PY, T, PT`` starting coordinates of the particle to track. Defaults to all 0 if `None` given. nturns (int): the number of turns to track for. diff --git a/pyhdtoolkit/cpymadtools/track.py b/pyhdtoolkit/cpymadtools/track.py index 98af4ffa..08555ebc 100644 --- a/pyhdtoolkit/cpymadtools/track.py +++ b/pyhdtoolkit/cpymadtools/track.py @@ -19,12 +19,12 @@ def track_single_particle( madx: Madx, /, - initial_coordinates: Tuple[float, float, float, float, float, float], + initial_coordinates: tuple[float, float, float, float, float, float], nturns: int, sequence: Optional[str] = None, observation_points: Sequence[str] = None, **kwargs, -) -> Dict[str, pd.DataFrame]: +) -> dict[str, pd.DataFrame]: """ .. versionadded:: 0.8.0 @@ -34,7 +34,7 @@ def track_single_particle( Args: madx (cpymad.madx.Madx): an instantiated `~cpymad.madx.Madx` object. - initial_coordinates (Tuple[float, float, float, float, float, float]): a tuple with the ``X, PX, + initial_coordinates (tuple[float, float, float, float, float, float]): a tuple with the ``X, PX, Y, PY, T, PT`` starting coordinates of the particle to track. Defaults to all 0 if `None` given. nturns (int): the number of turns to track for. sequence (Optional[str]): the sequence to use for tracking. If no value is provided, it is assumed diff --git a/pyhdtoolkit/cpymadtools/tune.py b/pyhdtoolkit/cpymadtools/tune.py index 1ec46c49..d088485f 100644 --- a/pyhdtoolkit/cpymadtools/tune.py +++ b/pyhdtoolkit/cpymadtools/tune.py @@ -10,7 +10,7 @@ import math import sys from pathlib import Path -from typing import Dict, List, Tuple +from typing import Dict, Tuple import matplotlib.collections import matplotlib.patches @@ -116,7 +116,7 @@ def make_footprint_table( return tfs_dframe -def get_footprint_lines(dynap_dframe: tfs.TfsDataFrame) -> Tuple[np.ndarray, np.ndarray]: +def get_footprint_lines(dynap_dframe: tfs.TfsDataFrame) -> tuple[np.ndarray, np.ndarray]: """ .. versionadded:: 0.12.0 @@ -250,7 +250,7 @@ def _get_dynap_string_rep(dynap_dframe: tfs.TfsDataFrame) -> str: return string_rep -def _make_tune_groups(dynap_string_rep: str, dsigma: float = 1.0) -> List[List[Dict[str, float]]]: +def _make_tune_groups(dynap_string_rep: str, dsigma: float = 1.0) -> list[list[dict[str, float]]]: """ Creates appropriate tune points groups from the arcane string representation returned by `~.tune._get_dynap_string_rep` based on starting amplitude and angle for each particle. @@ -269,7 +269,7 @@ def _make_tune_groups(dynap_string_rep: str, dsigma: float = 1.0) -> List[List[D by `~.tune.get_footprint_lines`. """ logger.debug("Constructing tune points groups based on starting amplitudes and angles") - tune_groups: List[List[Dict[str, float]]] = [] + tune_groups: list[list[dict[str, float]]] = [] items = dynap_string_rep.strip().split(",") amplitude = int(items[1]) current = 2 @@ -288,7 +288,7 @@ def _make_tune_groups(dynap_string_rep: str, dsigma: float = 1.0) -> List[List[D class _Footprint: """More dark magic from the past here, close your eyes my friends.""" - def __init__(self, tune_groups: List[List[Dict[str, float]]], amplitude: int, angle: int, dsigma: float): + def __init__(self, tune_groups: list[list[dict[str, float]]], amplitude: int, angle: int, dsigma: float): self._tunes = tune_groups self._maxnangl = angle self._nampl = amplitude @@ -306,7 +306,7 @@ def get_v_tune(self, ampl, angl): else: return self._tunes[ampl][angl]["V"] - def get_plottable(self) -> Tuple[List[float], List[float]]: + def get_plottable(self) -> tuple[list[float], list[float]]: qxs, qys = [], [] for i in np.arange(0, self._nampl - 1, 2): for j in np.arange(self._maxnangl): diff --git a/pyhdtoolkit/cpymadtools/utils.py b/pyhdtoolkit/cpymadtools/utils.py index 17b41487..aede36c4 100644 --- a/pyhdtoolkit/cpymadtools/utils.py +++ b/pyhdtoolkit/cpymadtools/utils.py @@ -7,7 +7,7 @@ Module with utility functions to do mundane operations with `~cpymad.madx.Madx` objects. """ from pathlib import Path -from typing import List, Union +from typing import Union import pandas as pd import tfs @@ -109,7 +109,7 @@ def get_table_tfs(madx: Madx, /, table_name: str, headers_table: str = "SUMM") - # ----- Helpers ----- # -def _get_k_strings(start: int = 0, stop: int = 8, orientation: str = "both") -> List[str]: +def _get_k_strings(start: int = 0, stop: int = 8, orientation: str = "both") -> list[str]: """ Returns the list of K-strings for various magnets and orders (``K1L``, ``K2SL`` etc strings). Initial implementation credits go to :user:`Joschua Dilly `. diff --git a/pyhdtoolkit/maths/stats_fitting.py b/pyhdtoolkit/maths/stats_fitting.py index 0088ac46..e38dd6be 100644 --- a/pyhdtoolkit/maths/stats_fitting.py +++ b/pyhdtoolkit/maths/stats_fitting.py @@ -17,7 +17,7 @@ from loguru import logger # Distributions to check # -DISTRIBUTIONS: Dict[st.rv_continuous, str] = { +DISTRIBUTIONS: dict[st.rv_continuous, str] = { st.chi: "Chi", st.chi2: "Chi-Square", st.expon: "Exponential", @@ -27,7 +27,7 @@ } -def set_distributions_dict(dist_dict: Dict[st.rv_continuous, str]) -> None: +def set_distributions_dict(dist_dict: dict[st.rv_continuous, str]) -> None: """ .. versionadded:: 0.5.0 @@ -36,7 +36,7 @@ def set_distributions_dict(dist_dict: Dict[st.rv_continuous, str]) -> None: function in the :ref:`gallery `. Args: - dist_dict (Dict[st.rv_continuous, str]): dictionnary with the wanted distributions, + dist_dict (dict[st.rv_continuous, str]): dictionnary with the wanted distributions, in the format of ``DISTRIBUTIONS``, aka with `scipy.stats` generator objects as keys, and a string representation of their name as value. @@ -58,7 +58,7 @@ def set_distributions_dict(dist_dict: Dict[st.rv_continuous, str]) -> None: def best_fit_distribution( data: Union[pd.Series, np.ndarray], bins: int = 200, ax: matplotlib.axes.Axes = None -) -> Tuple[st.rv_continuous, Tuple[float, ...]]: +) -> tuple[st.rv_continuous, tuple[float, ...]]: """ .. versionadded:: 0.5.0 @@ -125,7 +125,7 @@ def best_fit_distribution( return best_distribution, best_params -def make_pdf(distribution: st.rv_continuous, params: Tuple[float, ...], size: int = 25_000) -> pd.Series: +def make_pdf(distribution: st.rv_continuous, params: tuple[float, ...], size: int = 25_000) -> pd.Series: """ .. versionadded:: 0.5.0 @@ -135,7 +135,7 @@ def make_pdf(distribution: st.rv_continuous, params: Tuple[float, ...], size: in Args: distribution (st.rv_continuous): a `scipy.stats` generator. - params (Tuple[float, ...]): the parameters for this generator given back by the fit. + params (tuple[float, ...]): the parameters for this generator given back by the fit. size (int): the number of points to evaluate. Returns: diff --git a/pyhdtoolkit/maths/utils.py b/pyhdtoolkit/maths/utils.py index a195b84f..e1b86e40 100644 --- a/pyhdtoolkit/maths/utils.py +++ b/pyhdtoolkit/maths/utils.py @@ -50,7 +50,7 @@ def get_magnitude(value: float) -> int: def get_scaled_values_and_magnitude_string( values_array: Union[pd.DataFrame, np.ndarray], force_magnitude: float = None -) -> Tuple[Union[pd.DataFrame, np.ndarray], str]: +) -> tuple[Union[pd.DataFrame, np.ndarray], str]: """ .. versionadded:: 0.8.2 diff --git a/pyhdtoolkit/optics/rdt.py b/pyhdtoolkit/optics/rdt.py index 79c7c96a..52269f74 100644 --- a/pyhdtoolkit/optics/rdt.py +++ b/pyhdtoolkit/optics/rdt.py @@ -39,7 +39,7 @@ def rdt_to_order_and_type(rdt: Union[int, str]) -> str: return f"{rdt_type}_{orders[j + k + l + m]}" -def determine_rdt_line(rdt: Union[int, str], plane: str) -> Tuple[int, int, int]: +def determine_rdt_line(rdt: Union[int, str], plane: str) -> tuple[int, int, int]: """ Find the given line to look for in the spectral analysis of the given plane that corresponds to the given RDT. diff --git a/pyhdtoolkit/plotting/aperture.py b/pyhdtoolkit/plotting/aperture.py index e8204f65..30c86234 100644 --- a/pyhdtoolkit/plotting/aperture.py +++ b/pyhdtoolkit/plotting/aperture.py @@ -24,16 +24,16 @@ def plot_aperture( /, title: Optional[str] = None, xoffset: float = 0, - xlimits: Tuple[float, float] = None, + xlimits: tuple[float, float] = None, plot_dipoles: bool = True, plot_dipole_k1: bool = False, plot_quadrupoles: bool = True, plot_bpms: bool = False, - aperture_ylim: Tuple[float, float] = None, - k0l_lim: Union[Tuple[float, float], float, int] = None, - k1l_lim: Union[Tuple[float, float], float, int] = None, - k2l_lim: Union[Tuple[float, float], float, int] = None, - k3l_lim: Union[Tuple[float, float], float, int] = None, + aperture_ylim: tuple[float, float] = None, + k0l_lim: Union[tuple[float, float], float, int] = None, + k1l_lim: Union[tuple[float, float], float, int] = None, + k2l_lim: Union[tuple[float, float], float, int] = None, + k3l_lim: Union[tuple[float, float], float, int] = None, color: str = None, **kwargs, ) -> None: @@ -73,7 +73,7 @@ def plot_aperture( specific point or element, which would then become located at :math:`s = 0`. Beware this offset is applied before applying the *xlimits*. Defaults to 0. - xlimits (Tuple[float, float]): will implement xlim (for the ``s`` + xlimits (tuple[float, float]): will implement xlim (for the ``s`` coordinate) if this is not ``None``, using the tuple passed. plot_dipoles (bool): if `True`, dipole patches will be plotted on the layout subplot of the figure. Defaults to `True`. Dipoles @@ -87,27 +87,27 @@ def plot_aperture( plot_bpms (bool): if `True`, additional patches will be plotted on the layout subplot to represent Beam Position Monitors. BPMs are plotted in dark grey. - aperture_ylim (Tuple[float, float]): vertical axis limits for the + aperture_ylim (tuple[float, float]): vertical axis limits for the aperture values. Defaults to `None`, to be determined by matplotlib based on the provided values. - k0l_lim (Union[Tuple[float, float], float, int]): vertical axis limits + k0l_lim (Union[tuple[float, float], float, int]): vertical axis limits for the ``k0l`` values used for the height of dipole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). If `None` (default) is given, then the limits will be auto-determined based on the ``k0l`` values of the dipoles in the plot. - k1l_lim (Union[Tuple[float, float], float, int]): vertical axis limits + k1l_lim (Union[tuple[float, float], float, int]): vertical axis limits for the ``k1l`` values used for the height of quadrupole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). If `None` (default) is given, then the limits will be auto-determined based on the ``k0l`` values of the quadrupoles in the plot. - k2l_lim (Union[Tuple[float, float], float, int]): if given, sextupole + k2l_lim (Union[tuple[float, float], float, int]): if given, sextupole patches will be plotted on the layout subplot of the figure. If given, acts as vertical axis limits for the k2l values used for the height of sextupole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). - k3l_lim (Union[Tuple[float, float], float, int]): if given, octupole + k3l_lim (Union[tuple[float, float], float, int]): if given, octupole patches will be plotted on the layout subplot of the figure. If given, acts as vertical axis limits for the k3l values used for the height of octupole patches. Can be given as a single value @@ -192,7 +192,7 @@ def plot_physical_apertures( plane: str, scale: float = 1, xoffset: float = 0, - xlimits: Tuple[float, float] = None, + xlimits: tuple[float, float] = None, **kwargs, ) -> None: """ @@ -223,7 +223,7 @@ def plot_physical_apertures( specific point or element, which would then become located at :math:`s = 0`. Beware this offset is applied before applying the *xlimits*. Defaults to 0. - xlimits (Tuple[float, float]): will implement xlim (for the ``s`` + xlimits (tuple[float, float]): will implement xlim (for the ``s`` coordinate) if this is not ``None``, using the tuple passed. Defaults to ``None``. **kwargs: any keyword argument that can be given to the ``MAD-X`` @@ -280,8 +280,8 @@ def plot_physical_apertures( def _get_positions_and_real_apertures( - madx, /, plane: str, xoffset: float = 0, xlimits: Tuple[float, float] = None, **kwargs -) -> Tuple[np.ndarray, np.ndarray]: + madx, /, plane: str, xoffset: float = 0, xlimits: tuple[float, float] = None, **kwargs +) -> tuple[np.ndarray, np.ndarray]: """ .. versionadded:: 1.2.0 @@ -306,7 +306,7 @@ def _get_positions_and_real_apertures( specific point or element, which would then become located at :math:`s = 0`. Beware this offset is applied before applying the *xlimits*. Defaults to 0. - xlimits (Tuple[float, float]): will implement xlim (for the ``s`` + xlimits (tuple[float, float]): will implement xlim (for the ``s`` coordinate) if this is not ``None``, using the tuple passed. Defaults to ``None``. **kwargs: any keyword argument that can be given to the ``MAD-X`` diff --git a/pyhdtoolkit/plotting/envelope.py b/pyhdtoolkit/plotting/envelope.py index 2381647b..c0bde02b 100644 --- a/pyhdtoolkit/plotting/envelope.py +++ b/pyhdtoolkit/plotting/envelope.py @@ -23,7 +23,7 @@ def plot_beam_envelope( nsigma: float = 1, scale: float = 1, xoffset: float = 0, - xlimits: Tuple[float, float] = None, + xlimits: tuple[float, float] = None, **kwargs, ) -> None: """ @@ -53,7 +53,7 @@ def plot_beam_envelope( specific point or element, which would then become located at :math:`s = 0`. Beware this offset is applied before applying the *xlimits*. Defaults to 0. - xlimits (Tuple[float, float]): will implement xlim (for the ``s`` + xlimits (tuple[float, float]): will implement xlim (for the ``s`` coordinate) if this is not ``None``, using the tuple passed. Defaults to ``None``. **kwargs: any keyword argument that can be given to the ``MAD-X`` diff --git a/pyhdtoolkit/plotting/lattice.py b/pyhdtoolkit/plotting/lattice.py index 439d2a5b..2ea97b22 100644 --- a/pyhdtoolkit/plotting/lattice.py +++ b/pyhdtoolkit/plotting/lattice.py @@ -27,17 +27,17 @@ def plot_latwiss( /, title: Optional[str] = None, xoffset: float = 0, - xlimits: Tuple[float, float] = None, + xlimits: tuple[float, float] = None, plot_dipoles: bool = True, plot_dipole_k1: bool = False, plot_quadrupoles: bool = True, plot_bpms: bool = False, - disp_ylim: Union[Tuple[float, float], float, int] = None, - beta_ylim: Union[Tuple[float, float], float, int] = None, - k0l_lim: Union[Tuple[float, float], float, int] = None, - k1l_lim: Union[Tuple[float, float], float, int] = None, - k2l_lim: Union[Tuple[float, float], float, int] = None, - k3l_lim: Union[Tuple[float, float], float, int] = None, + disp_ylim: Union[tuple[float, float], float, int] = None, + beta_ylim: Union[tuple[float, float], float, int] = None, + k0l_lim: Union[tuple[float, float], float, int] = None, + k1l_lim: Union[tuple[float, float], float, int] = None, + k2l_lim: Union[tuple[float, float], float, int] = None, + k3l_lim: Union[tuple[float, float], float, int] = None, **kwargs, ) -> None: """ @@ -70,7 +70,7 @@ def plot_latwiss( xoffset (float): An offset applied to the ``S`` coordinate before plotting. This is useful if you want to center a plot around a specific point or element, which would then become located at :math:`s = 0`. Beware this offset is applied before applying the *xlimits*. Defaults to 0. - xlimits (Tuple[float, float]): will implement xlim (for the ``s`` coordinate) if this is + xlimits (tuple[float, float]): will implement xlim (for the ``s`` coordinate) if this is not ``None``, using the tuple passed. plot_dipoles (bool): if `True`, dipole patches will be plotted on the layout subplot of the figure. Defaults to `True`. Dipoles are plotted in blue. @@ -80,28 +80,28 @@ def plot_latwiss( subplot of the figure. Defaults to `True`. Quadrupoles are plotted in red. plot_bpms (bool): if `True`, additional patches will be plotted on the layout subplot to represent Beam Position Monitors. BPMs are plotted in dark grey. - disp_ylim (Tuple[float, float]): vertical axis limits for the dispersion values. + disp_ylim (tuple[float, float]): vertical axis limits for the dispersion values. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). Defaults to (-10, 125). - beta_ylim (Tuple[float, float]): vertical axis limits for the betatron function values. + beta_ylim (tuple[float, float]): vertical axis limits for the betatron function values. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). Defaults to `None`, to be determined by `~matplotlib` based on the plotted beta values. - k0l_lim (Union[Tuple[float, float], float, int]): vertical axis limits for the ``k0l`` + k0l_lim (Union[tuple[float, float], float, int]): vertical axis limits for the ``k0l`` values used for the height of dipole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). If `None` (default) is given, then the limits will be auto-determined based on the ``k0l`` values of the dipoles in the plot. - k1l_lim (Union[Tuple[float, float], float, int]): vertical axis limits for the ``k1l`` + k1l_lim (Union[tuple[float, float], float, int]): vertical axis limits for the ``k1l`` values used for the height of quadrupole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). If `None` (default) is given, then the limits will be auto-determined based on the ``k0l`` values of the quadrupoles in the plot. - k2l_lim (Union[Tuple[float, float], float, int]): if given, sextupole patches will be plotted + k2l_lim (Union[tuple[float, float], float, int]): if given, sextupole patches will be plotted on the layout subplot of the figure. If given, acts as vertical axis limits for the k2l values used for the height of sextupole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). - k3l_lim (Union[Tuple[float, float], float, int]): if given, octupole patches will be plotted + k3l_lim (Union[tuple[float, float], float, int]): if given, octupole patches will be plotted on the layout subplot of the figure. If given, acts as vertical axis limits for the k3l values used for the height of octupole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). diff --git a/pyhdtoolkit/plotting/layout.py b/pyhdtoolkit/plotting/layout.py index fb92efec..c65b57b5 100644 --- a/pyhdtoolkit/plotting/layout.py +++ b/pyhdtoolkit/plotting/layout.py @@ -29,15 +29,15 @@ def plot_machine_layout( /, title: str = None, xoffset: float = 0, - xlimits: Tuple[float, float] = None, + xlimits: tuple[float, float] = None, plot_dipoles: bool = True, plot_dipole_k1: bool = False, plot_quadrupoles: bool = True, plot_bpms: bool = False, - k0l_lim: Union[Tuple[float, float], float, int] = None, - k1l_lim: Union[Tuple[float, float], float, int] = None, - k2l_lim: Union[Tuple[float, float], float, int] = None, - k3l_lim: Union[Tuple[float, float], float, int] = None, + k0l_lim: Union[tuple[float, float], float, int] = None, + k1l_lim: Union[tuple[float, float], float, int] = None, + k2l_lim: Union[tuple[float, float], float, int] = None, + k3l_lim: Union[tuple[float, float], float, int] = None, **kwargs, ) -> None: """ @@ -78,7 +78,7 @@ def plot_machine_layout( specific point or element, which would then become located at the :math:`s = 0` position. Beware this offset is applied before applying the *xlimits*. Defaults to 0. - xlimits (Tuple[float, float]): will implement xlim (for the ``s`` + xlimits (tuple[float, float]): will implement xlim (for the ``s`` coordinate) if this is not ``None``, using the tuple passed. plot_dipoles (bool): if `True`, dipole patches will be plotted on the layout subplot of the figure. Defaults to `True`. Dipoles @@ -92,24 +92,24 @@ def plot_machine_layout( plot_bpms (bool): if `True`, additional patches will be plotted on the layout subplot to represent Beam Position Monitors. BPMs are plotted in dark grey. - k0l_lim (Union[Tuple[float, float], float, int]): vertical axis limits + k0l_lim (Union[tuple[float, float], float, int]): vertical axis limits for the ``k0l`` values used for the height of dipole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). If `None` (default) is given, then the limits will be auto-determined based on the ``k0l`` values of the dipoles in the plot. - k1l_lim (Union[Tuple[float, float], float, int]): vertical axis limits + k1l_lim (Union[tuple[float, float], float, int]): vertical axis limits for the ``k1l`` values used for the height of quadrupole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). If `None` (default) is given, then the limits will be auto-determined based on the ``k0l`` values of the quadrupoles in the plot. - k2l_lim (Union[Tuple[float, float], float, int]): if given, sextupole + k2l_lim (Union[tuple[float, float], float, int]): if given, sextupole patches will be plotted on the layout subplot of the figure. If given, acts as vertical axis limits for the k2l values used for the height of sextupole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). - k3l_lim (Union[Tuple[float, float], float, int]): if given, octupole + k3l_lim (Union[tuple[float, float], float, int]): if given, octupole patches will be plotted on the layout subplot of the figure. If given, acts as vertical axis limits for the k3l values used for the height of octupole patches. Can be given as a single value @@ -396,8 +396,8 @@ def _plot_lattice_series( def _ylim_from_input( - ylim: Union[Tuple[float, float], float, int], name_for_error: str = "knl_lim" -) -> Tuple[float, float]: + ylim: Union[tuple[float, float], float, int], name_for_error: str = "knl_lim" +) -> tuple[float, float]: """ .. versionadded:: 1.2.0 @@ -406,7 +406,7 @@ def _ylim_from_input( different inputs from the user, such as a tuple, a float and an int. Args: - ylim (Tuple[float, float], float, int): the input provided by + ylim (tuple[float, float], float, int): the input provided by the user. name_for_error (str): the name of the variable to use in the error message. @@ -431,7 +431,7 @@ def _ylim_from_input( ) -def _determine_default_knl_lim(df: pd.DataFrame, col: str, coeff: float) -> Tuple[float, float]: +def _determine_default_knl_lim(df: pd.DataFrame, col: str, coeff: float) -> tuple[float, float]: """ .. versionadded:: 1.0.0 diff --git a/pyhdtoolkit/plotting/sbs/coupling.py b/pyhdtoolkit/plotting/sbs/coupling.py index b896ce20..0650d23e 100644 --- a/pyhdtoolkit/plotting/sbs/coupling.py +++ b/pyhdtoolkit/plotting/sbs/coupling.py @@ -107,9 +107,9 @@ def plot_full_ip_rdt( b2_model: tfs.TfsDataFrame = None, ip: int = None, rdt: str = "F1001", - abs_ylimits: Tuple[float, float] = None, - real_ylimits: Tuple[float, float] = None, - imag_ylimits: Tuple[float, float] = None, + abs_ylimits: tuple[float, float] = None, + real_ylimits: tuple[float, float] = None, + imag_ylimits: tuple[float, float] = None, **kwargs, ) -> matplotlib.figure.Figure: """ diff --git a/pyhdtoolkit/plotting/styles/__init__.py b/pyhdtoolkit/plotting/styles/__init__.py index 9be7d26b..5a7c2d83 100644 --- a/pyhdtoolkit/plotting/styles/__init__.py +++ b/pyhdtoolkit/plotting/styles/__init__.py @@ -39,7 +39,7 @@ def install_mpl_styles() -> None: # This is meant for use in the sphinx-gallery to help with readability # as the default matplotlib settings are a bit small -_SPHINX_GALLERY_PARAMS: Dict[str, PlotSetting] = { +_SPHINX_GALLERY_PARAMS: dict[str, PlotSetting] = { "figure.autolayout": True, "figure.titlesize": 28, "axes.titlesize": 28, @@ -50,7 +50,7 @@ def install_mpl_styles() -> None: } -def _install_style_file(style: Dict[str, PlotSetting], stylename) -> None: +def _install_style_file(style: dict[str, PlotSetting], stylename) -> None: """ .. versionadded:: 1.0.0 @@ -68,7 +68,7 @@ def _install_style_file(style: Dict[str, PlotSetting], stylename) -> None: activated environment's site-packages data. The file is installed in both places. Args: - style (Dict[str, PlotSetting]): The style to be written to disk. One of the styles defined + style (dict[str, PlotSetting]): The style to be written to disk. One of the styles defined in the `~pyhdtoolkit.plotting.styles` submodules. stylename (str): The name to associate with the style. This is the name that will be used for the file written to disk, and will have to be `plt.use`d to use the style. diff --git a/pyhdtoolkit/plotting/styles/paper.py b/pyhdtoolkit/plotting/styles/paper.py index ad6219f6..32de1bb6 100644 --- a/pyhdtoolkit/plotting/styles/paper.py +++ b/pyhdtoolkit/plotting/styles/paper.py @@ -18,7 +18,7 @@ PlotSetting = Union[float, bool, str, tuple] -SINGLE_COLUMN: Dict[str, PlotSetting] = { +SINGLE_COLUMN: dict[str, PlotSetting] = { # ------ Lines ------ # "lines.linewidth": 1.3, # Width of plot lines "lines.markersize": 5, # Marker size, in points @@ -66,7 +66,7 @@ "savefig.format": "pdf", # Saved figure file format } -DOUBLE_COLUMN: Dict[str, PlotSetting] = { +DOUBLE_COLUMN: dict[str, PlotSetting] = { # ------ Lines ------ # "lines.linewidth": 1.7, # Width of plot lines "lines.markersize": 3, diff --git a/pyhdtoolkit/plotting/styles/thesis.py b/pyhdtoolkit/plotting/styles/thesis.py index 5a80ca6f..a780290d 100644 --- a/pyhdtoolkit/plotting/styles/thesis.py +++ b/pyhdtoolkit/plotting/styles/thesis.py @@ -20,7 +20,7 @@ PlotSetting = Union[float, bool, str, tuple] -SMALL: Dict[str, PlotSetting] = { +SMALL: dict[str, PlotSetting] = { # ------ Lines ------ # "lines.linewidth": 1.3, # Width of plot lines "lines.markersize": 3, # Marker size, in points @@ -65,7 +65,7 @@ "savefig.format": "pdf", # Saved figure file format } -MEDIUM: Dict[str, PlotSetting] = { +MEDIUM: dict[str, PlotSetting] = { # ------ Lines ------ # "lines.linewidth": 1.3, # Width of plot lines "lines.markersize": 5, # Marker size, in points @@ -112,7 +112,7 @@ "savefig.format": "pdf", # Saved figure file format } -LARGE: Dict[str, PlotSetting] = { +LARGE: dict[str, PlotSetting] = { # ------ Lines ------ # "lines.linewidth": 1.7, # Width of plot lines "lines.markersize": 8, diff --git a/pyhdtoolkit/plotting/tune.py b/pyhdtoolkit/plotting/tune.py index 17a5994d..e0140749 100644 --- a/pyhdtoolkit/plotting/tune.py +++ b/pyhdtoolkit/plotting/tune.py @@ -8,7 +8,7 @@ These provide functionality to draw Farey sequences up to a desired order. """ from functools import partial -from typing import Dict, List, Tuple +from typing import Dict, Tuple import matplotlib import matplotlib.axes @@ -17,8 +17,8 @@ from pyhdtoolkit.plotting.utils import maybe_get_ax -ORDER_TO_ALPHA: Dict[int, float] = {1: 1, 2: 0.75, 3: 0.65, 4: 0.55, 5: 0.45, 6: 0.35} -ORDER_TO_RGB: Dict[int, np.ndarray] = { +ORDER_TO_ALPHA: dict[int, float] = {1: 1, 2: 0.75, 3: 0.65, 4: 0.55, 5: 0.45, 6: 0.35} +ORDER_TO_RGB: dict[int, np.ndarray] = { 1: np.array([152, 52, 48]) / 255, # a brown 2: np.array([57, 119, 175]) / 255, # a blue 3: np.array([239, 133, 54]) / 255, # an orange @@ -26,7 +26,7 @@ 5: np.array([197, 57, 50]) / 255, # a red 6: np.array([141, 107, 184]) / 255, # a purple } -ORDER_TO_LINESTYLE: Dict[int, str] = { +ORDER_TO_LINESTYLE: dict[int, str] = { 1: "solid", 2: "solid", 3: "solid", @@ -34,8 +34,8 @@ 5: "dashed", 6: "dashed", } -ORDER_TO_LINEWIDTH: Dict[int, float] = {1: 2, 2: 1.75, 3: 1.5, 4: 1.25, 5: 1, 6: 0.75} -ORDER_TO_LABEL: Dict[int, str] = { +ORDER_TO_LINEWIDTH: dict[int, float] = {1: 2, 2: 1.75, 3: 1.5, 4: 1.25, 5: 1, 6: 0.75} +ORDER_TO_LABEL: dict[int, str] = { 1: "1st order", 2: "2nd order", 3: "3rd order", @@ -45,7 +45,7 @@ } -def farey_sequence(order: int) -> List[Tuple[int, int]]: +def farey_sequence(order: int) -> list[tuple[int, int]]: """ .. versionadded:: 1.0.0 diff --git a/pyhdtoolkit/utils/cmdline.py b/pyhdtoolkit/utils/cmdline.py index 56923f6b..3bbf6222 100644 --- a/pyhdtoolkit/utils/cmdline.py +++ b/pyhdtoolkit/utils/cmdline.py @@ -68,7 +68,7 @@ def check_pid_exists(pid: int) -> bool: @staticmethod def run( command: str, shell: bool = True, env: Mapping = None, timeout: float = None - ) -> Tuple[Optional[int], bytes]: + ) -> tuple[Optional[int], bytes]: """ .. versionadded:: 0.2.0 diff --git a/pyhdtoolkit/utils/executors.py b/pyhdtoolkit/utils/executors.py index bba51b9b..5c41d45a 100644 --- a/pyhdtoolkit/utils/executors.py +++ b/pyhdtoolkit/utils/executors.py @@ -22,7 +22,7 @@ """ from concurrent import futures -from typing import Callable, List +from typing import Callable from loguru import logger @@ -49,7 +49,7 @@ class MultiProcessor: """ @staticmethod - def execute_function(func: Callable, func_args: list, n_processes: int) -> List[tuple]: + def execute_function(func: Callable, func_args: list, n_processes: int) -> list[tuple]: """ .. versionadded:: 0.2.0 @@ -107,7 +107,7 @@ class MultiThreader: """ @staticmethod - def execute_function(func: Callable, func_args: list, n_threads: int) -> List[tuple]: + def execute_function(func: Callable, func_args: list, n_threads: int) -> list[tuple]: """ .. versionadded:: 0.2.0 diff --git a/pyhdtoolkit/utils/htc_monitor.py b/pyhdtoolkit/utils/htc_monitor.py index 5063b967..d0ec7fd3 100644 --- a/pyhdtoolkit/utils/htc_monitor.py +++ b/pyhdtoolkit/utils/htc_monitor.py @@ -14,7 +14,7 @@ """ import re import time -from typing import List, Tuple +from typing import Tuple import pendulum from loguru import logger @@ -73,7 +73,7 @@ def query_condor_q() -> str: raise ChildProcessError("Checking htcondor status failed") -def read_condor_q(report: str) -> Tuple[List[HTCTaskSummary], ClusterSummary]: +def read_condor_q(report: str) -> tuple[list[HTCTaskSummary], ClusterSummary]: """ .. versionadded:: 0.9.0 @@ -95,7 +95,7 @@ def read_condor_q(report: str) -> Tuple[List[HTCTaskSummary], ClusterSummary]: condor_q_output = get_the_string_as_you_wish(...) tasks, cluster = read_condor_q(condor_q_output) """ - tasks: List[HTCTaskSummary] = [] + tasks: list[HTCTaskSummary] = [] next_line_is_task_report = False for line in report.split("\n"): @@ -128,7 +128,7 @@ def read_condor_q(report: str) -> Tuple[List[HTCTaskSummary], ClusterSummary]: # ----- Output Formating ----- # -def _make_tasks_table(tasks: List[HTCTaskSummary]) -> Table: +def _make_tasks_table(tasks: list[HTCTaskSummary]) -> Table: table = _default_tasks_table() date_display_format = "dddd, D MMM YY at LT (zz)" # example: Wednesday, 21 Apr 21 9:04 PM (CEST) for task in tasks: diff --git a/pyhdtoolkit/utils/operations.py b/pyhdtoolkit/utils/operations.py index 61132d8e..78b9ffd9 100644 --- a/pyhdtoolkit/utils/operations.py +++ b/pyhdtoolkit/utils/operations.py @@ -17,7 +17,7 @@ import random import re from functools import reduce -from typing import Callable, Dict, List, Sequence, Tuple, Union +from typing import Callable, Dict, Sequence, Tuple, Union class ListOperations: @@ -75,7 +75,7 @@ def average_by(sequence: Sequence, function: Callable = lambda x: x) -> float: return float(sum(map(function, sequence), 0.0) / len(sequence)) @staticmethod - def bifurcate(sequence: Sequence, filters: List[bool]) -> Sequence: + def bifurcate(sequence: Sequence, filters: list[bool]) -> Sequence: """ .. versionadded:: 0.2.0 @@ -86,7 +86,7 @@ def bifurcate(sequence: Sequence, filters: List[bool]) -> Sequence: Args: sequence (Sequence): a sequence of elements. - filters (List[bool]): a list of booleans. + filters (list[bool]): a list of booleans. Returns: A list of two lists, one for each boolean output of the filters. @@ -244,7 +244,7 @@ def eval_some(sequence: Sequence, function: Callable = lambda x: bool(x)) -> boo return any(map(function, sequence)) @staticmethod - def get_indices(element, sequence: Sequence) -> List[int]: + def get_indices(element, sequence: Sequence) -> list[int]: """ .. versionadded:: 0.2.0 @@ -269,7 +269,7 @@ def get_indices(element, sequence: Sequence) -> List[int]: return [i for (y, i) in zip(sequence, range(len(sequence))) if element == y] @staticmethod - def group_by(sequence: Sequence, function: Callable) -> Dict[str, list]: + def group_by(sequence: Sequence, function: Callable) -> dict[str, list]: """ .. versionadded:: 0.2.0 @@ -632,7 +632,7 @@ def clamp_number(num: Union[int, float], a_val: Union[int, float], b_val: Union[ @staticmethod def degrees_to_radians( deg_value: Union[int, float], decompose: bool = False - ) -> Union[Tuple[float, str, str], int, float]: + ) -> Union[tuple[float, str, str], int, float]: """ .. versionadded:: 0.2.0 From 1abe24d8e294c2d2919733a3818affa0c009ddc8 Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 10:49:45 +0200 Subject: [PATCH 03/41] some more for optional, union --- pyhdtoolkit/cpymadtools/coupling.py | 2 +- pyhdtoolkit/cpymadtools/lhc/_errors.py | 2 +- pyhdtoolkit/cpymadtools/lhc/_misc.py | 1 - pyhdtoolkit/cpymadtools/lhc/_powering.py | 2 +- pyhdtoolkit/cpymadtools/lhc/_queries.py | 6 ++--- pyhdtoolkit/cpymadtools/lhc/_setup.py | 1 - pyhdtoolkit/cpymadtools/lhc/_twiss.py | 2 +- pyhdtoolkit/cpymadtools/matching.py | 14 +++++------ pyhdtoolkit/cpymadtools/ptc.py | 18 +++++++------- pyhdtoolkit/cpymadtools/track.py | 6 ++--- pyhdtoolkit/cpymadtools/tune.py | 1 - pyhdtoolkit/cpymadtools/twiss.py | 2 +- pyhdtoolkit/cpymadtools/utils.py | 3 +-- pyhdtoolkit/maths/stats_fitting.py | 3 +-- pyhdtoolkit/maths/utils.py | 5 ++-- pyhdtoolkit/models/beam.py | 13 +++++----- pyhdtoolkit/models/htc.py | 7 +++--- pyhdtoolkit/optics/rdt.py | 9 ++++--- pyhdtoolkit/optics/ripken.py | 9 ++++--- pyhdtoolkit/plotting/aperture.py | 21 ++++++++--------- pyhdtoolkit/plotting/crossing.py | 3 +-- pyhdtoolkit/plotting/envelope.py | 1 - pyhdtoolkit/plotting/lattice.py | 27 ++++++++++----------- pyhdtoolkit/plotting/layout.py | 19 +++++++-------- pyhdtoolkit/plotting/sbs/coupling.py | 1 - pyhdtoolkit/plotting/styles/paper.py | 2 +- pyhdtoolkit/plotting/styles/thesis.py | 2 +- pyhdtoolkit/plotting/tune.py | 1 - pyhdtoolkit/plotting/utils.py | 2 +- pyhdtoolkit/utils/__init__.py | 2 +- pyhdtoolkit/utils/_misc.py | 4 ++-- pyhdtoolkit/utils/cmdline.py | 4 ++-- pyhdtoolkit/utils/contexts.py | 2 +- pyhdtoolkit/utils/executors.py | 2 +- pyhdtoolkit/utils/htc_monitor.py | 1 - pyhdtoolkit/utils/logging.py | 3 +-- pyhdtoolkit/utils/operations.py | 30 ++++++++++++------------ pyproject.toml | 4 ++-- 38 files changed, 109 insertions(+), 128 deletions(-) diff --git a/pyhdtoolkit/cpymadtools/coupling.py b/pyhdtoolkit/cpymadtools/coupling.py index 2cc72daf..0957795d 100644 --- a/pyhdtoolkit/cpymadtools/coupling.py +++ b/pyhdtoolkit/cpymadtools/coupling.py @@ -8,7 +8,7 @@ Module with functions to perform ``MAD-X`` actions through a `~cpymad.madx.Madx` object, that retate to betatron coupling in the machine. """ -from typing import Dict, Sequence, Tuple +from collections.abc import Sequence import numpy as np import tfs diff --git a/pyhdtoolkit/cpymadtools/lhc/_errors.py b/pyhdtoolkit/cpymadtools/lhc/_errors.py index 7e4bb66a..ac10f4e3 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_errors.py +++ b/pyhdtoolkit/cpymadtools/lhc/_errors.py @@ -5,7 +5,7 @@ The functions below are utilities to implement errors in elements of the ``LHC``. """ -from typing import Dict, Sequence +from collections.abc import Sequence from cpymad.madx import Madx from loguru import logger diff --git a/pyhdtoolkit/cpymadtools/lhc/_misc.py b/pyhdtoolkit/cpymadtools/lhc/_misc.py index 390c5537..c5319a8c 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_misc.py +++ b/pyhdtoolkit/cpymadtools/lhc/_misc.py @@ -5,7 +5,6 @@ The functions below are miscellaneous utilities for the ``LHC``. """ -from typing import Tuple from cpymad.madx import Madx from loguru import logger diff --git a/pyhdtoolkit/cpymadtools/lhc/_powering.py b/pyhdtoolkit/cpymadtools/lhc/_powering.py index 1e3122f4..ed55e16c 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_powering.py +++ b/pyhdtoolkit/cpymadtools/lhc/_powering.py @@ -5,7 +5,7 @@ The functions below are magnets or knobs powering utilities for the ``LHC``. """ -from typing import Dict, Sequence +from collections.abc import Sequence from cpymad.madx import Madx from loguru import logger diff --git a/pyhdtoolkit/cpymadtools/lhc/_queries.py b/pyhdtoolkit/cpymadtools/lhc/_queries.py index 34cb4446..692ec8ef 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_queries.py +++ b/pyhdtoolkit/cpymadtools/lhc/_queries.py @@ -5,7 +5,7 @@ The functions below are settings query utilities for the ``LHC``. """ -from typing import Dict, Sequence, Union +from collections.abc import Sequence import tfs from cpymad.madx import Madx @@ -33,7 +33,7 @@ def get_magnets_powering( - madx: Madx, /, patterns: Sequence[str] = [r"^mb\.", r"^mq\.", r"^ms\."], brho: Union[str, float] = None, **kwargs + madx: Madx, /, patterns: Sequence[str] = [r"^mb\.", r"^mq\.", r"^ms\."], brho: str | float = None, **kwargs ) -> tfs.TfsDataFrame: r""" .. versionadded:: 0.17.0 @@ -218,7 +218,7 @@ def get_current_orbit_setup(madx: Madx, /) -> dict[str, float]: # ----- Helpers ----- # -def _list_field_currents(madx: Madx, /, brho: Union[str, float] = None) -> None: +def _list_field_currents(madx: Madx, /, brho: str | float = None) -> None: """ Creates additional columns for the ``TWISS`` table with the magnets' total fields and currents, to help later on determine which proportion of their maximum powering diff --git a/pyhdtoolkit/cpymadtools/lhc/_setup.py b/pyhdtoolkit/cpymadtools/lhc/_setup.py index 59191a20..2f6f462a 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_setup.py +++ b/pyhdtoolkit/cpymadtools/lhc/_setup.py @@ -6,7 +6,6 @@ The functions below are setup utilities for the ``LHC``, to easily get simulations ready. """ from pathlib import Path -from typing import Dict, Tuple from cpymad.madx import Madx from loguru import logger diff --git a/pyhdtoolkit/cpymadtools/lhc/_twiss.py b/pyhdtoolkit/cpymadtools/lhc/_twiss.py index 660d96e9..c27567e4 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_twiss.py +++ b/pyhdtoolkit/cpymadtools/lhc/_twiss.py @@ -5,7 +5,7 @@ The functions below are twiss utilities for the ``LHC`` insertion regions. """ -from typing import Sequence +from collections.abc import Sequence import tfs from cpymad.madx import Madx diff --git a/pyhdtoolkit/cpymadtools/matching.py b/pyhdtoolkit/cpymadtools/matching.py index 3730af82..50566120 100644 --- a/pyhdtoolkit/cpymadtools/matching.py +++ b/pyhdtoolkit/cpymadtools/matching.py @@ -6,7 +6,7 @@ Module with functions to perform ``MAD-X`` matchings through a `~cpymad.madx.Madx` object. """ -from typing import Optional, Sequence +from collections.abc import Sequence from cpymad.madx import Madx from loguru import logger @@ -20,7 +20,7 @@ def match_tunes_and_chromaticities( madx: Madx, /, accelerator: str = None, - sequence: Optional[str] = None, + sequence: str | None = None, q1_target: float = None, q2_target: float = None, dq1_target: float = None, @@ -70,7 +70,7 @@ def match_tunes_and_chromaticities( Args: madx (cpymad.madx.Madx): an instanciated `~cpymad.madx.Madx` object. Positional only. - accelerator (Optional[str]): name of the accelerator, used to determmine knobs if + accelerator (str | None): name of the accelerator, used to determmine knobs if *variables* is not given. Automatic determination will only work for ``LHC`` and ``HLLHC``. Defaults to `None`, in which case the knobs must be provided explicitly through ``varied_knobs``. @@ -193,7 +193,7 @@ def match_tunes( madx: Madx, /, accelerator: str = None, - sequence: Optional[str] = None, + sequence: str | None = None, q1_target: float = None, q2_target: float = None, varied_knobs: Sequence[str] = None, @@ -215,7 +215,7 @@ def match_tunes( Args: madx (cpymad.madx.Madx): an instanciated `~cpymad.madx.Madx` object. Positional only. - accelerator (Optional[str]): name of the accelerator, used to determmine knobs if + accelerator (str | None): name of the accelerator, used to determmine knobs if *variables* is not given. Automatic determination will only work for ``LHC`` and ``HLLHC``. Defaults to `None`, in which case the knobs must be provided explicitly through ``varied_knobs``. @@ -295,7 +295,7 @@ def match_chromaticities( madx: Madx, /, accelerator: str = None, - sequence: Optional[str] = None, + sequence: str | None = None, dq1_target: float = None, dq2_target: float = None, varied_knobs: Sequence[str] = None, @@ -317,7 +317,7 @@ def match_chromaticities( Args: madx (cpymad.madx.Madx): an instanciated `~cpymad.madx.Madx` object. Positional only. - accelerator (Optional[str]): name of the accelerator, used to determmine knobs if + accelerator (str | None): name of the accelerator, used to determmine knobs if *variables* is not given. Automatic determination will only work for ``LHC`` and ``HLLHC``. Defaults to `None`, in which case the knobs must be provided explicitly through ``varied_knobs``. diff --git a/pyhdtoolkit/cpymadtools/ptc.py b/pyhdtoolkit/cpymadtools/ptc.py index 4b2c2f60..54baae71 100644 --- a/pyhdtoolkit/cpymadtools/ptc.py +++ b/pyhdtoolkit/cpymadtools/ptc.py @@ -7,8 +7,8 @@ Module with functions to manipulate ``MAD-X`` ``PTC`` functionality through a `~cpymad.madx.Madx` object. """ +from collections.abc import Sequence from pathlib import Path -from typing import Dict, Optional, Sequence, Tuple, Union import pandas as pd import tfs @@ -19,7 +19,7 @@ def get_amplitude_detuning( - madx: Madx, /, order: int = 2, file: Union[Path, str] = None, fringe: bool = False, **kwargs + madx: Madx, /, order: int = 2, file: Path | str = None, fringe: bool = False, **kwargs ) -> tfs.TfsDataFrame: """ .. versionadded:: 0.7.0 @@ -49,7 +49,7 @@ def get_amplitude_detuning( Positional only. order (int): maximum derivative order coefficient (only 0, 1 or 2 implemented in ``PTC``). Defaults to 2. - file (Union[Path, str]): path to output file. Defaults to `None`. + file (Path | str): path to output file. Defaults to `None`. fringe (bool): boolean flag to include fringe field effects in the calculation. Defaults to ``False``. **kwargs: Some parameters for the ``PTC`` universe creation can be given as @@ -143,7 +143,7 @@ def get_amplitude_detuning( def get_rdts( - madx: Madx, /, order: int = 4, file: Union[Path, str] = None, fringe: bool = False, **kwargs + madx: Madx, /, order: int = 4, file: Path | str = None, fringe: bool = False, **kwargs ) -> tfs.TfsDataFrame: """ .. versionadded:: 0.7.0 @@ -175,7 +175,7 @@ def get_rdts( Positional only. order (int): map order for derivative evaluation of Twiss parameters. Defaults to 4. - file (Union[Path, str]): path to output file. Default to `None`. + file (Path | str): path to output file. Default to `None`. fringe (bool): boolean flag to include fringe field effects in the calculation. Defaults to `False`. **kwargs: Some parameters for the ``PTC`` universe creation can be given as @@ -238,7 +238,7 @@ def ptc_twiss( madx: Madx, /, order: int = 4, - file: Union[Path, str] = None, + file: Path | str = None, fringe: bool = False, table: str = "ptc_twiss", **kwargs, @@ -276,7 +276,7 @@ def ptc_twiss( Positional only. order (int): map order for derivative evaluation of ``TWISS`` parameters. Defaults to 4. - file (Union[Path, str]): path to output file. Default to `None`. + file (Path | str): path to output file. Default to `None`. fringe (bool): boolean flag to include fringe field effects in the calculation. Defaults to `False`. table (str): the name of the internal table in which to save the results. @@ -342,7 +342,7 @@ def ptc_track_particle( /, initial_coordinates: tuple[float, float, float, float, float, float], nturns: int, - sequence: Optional[str] = None, + sequence: str | None = None, observation_points: Sequence[str] = None, onetable: bool = False, fringe: bool = False, @@ -381,7 +381,7 @@ def ptc_track_particle( with the ``X, PX, Y, PY, T, PT`` starting coordinates of the particle to track. Defaults to all 0 if `None` given. nturns (int): the number of turns to track for. - sequence (Optional[str]): the sequence to use for tracking. If no value is + sequence (str | None): the sequence to use for tracking. If no value is provided, it is assumed that a sequence is already defined and in use, and this one will be picked up by ``MAD-X``. Beware of the dangers of giving a sequence that will be used by ``MAD-X``, see the warning below diff --git a/pyhdtoolkit/cpymadtools/track.py b/pyhdtoolkit/cpymadtools/track.py index 08555ebc..fcd7a1ac 100644 --- a/pyhdtoolkit/cpymadtools/track.py +++ b/pyhdtoolkit/cpymadtools/track.py @@ -7,7 +7,7 @@ Module with functions to manipulate ``MAD-X`` ``TRACK`` functionality through a `~cpymad.madx.Madx` object. """ -from typing import Dict, Optional, Sequence, Tuple +from collections.abc import Sequence import pandas as pd from cpymad.madx import Madx @@ -21,7 +21,7 @@ def track_single_particle( /, initial_coordinates: tuple[float, float, float, float, float, float], nturns: int, - sequence: Optional[str] = None, + sequence: str | None = None, observation_points: Sequence[str] = None, **kwargs, ) -> dict[str, pd.DataFrame]: @@ -37,7 +37,7 @@ def track_single_particle( initial_coordinates (tuple[float, float, float, float, float, float]): a tuple with the ``X, PX, Y, PY, T, PT`` starting coordinates of the particle to track. Defaults to all 0 if `None` given. nturns (int): the number of turns to track for. - sequence (Optional[str]): the sequence to use for tracking. If no value is provided, it is assumed + sequence (str | None): the sequence to use for tracking. If no value is provided, it is assumed that a sequence is already defined and in use, and this one will be picked up by ``MAD-X``. Beware of the dangers of giving a sequence that will be used by ``MAD-X``, see the warning below for more information. diff --git a/pyhdtoolkit/cpymadtools/tune.py b/pyhdtoolkit/cpymadtools/tune.py index d088485f..27dc1e02 100644 --- a/pyhdtoolkit/cpymadtools/tune.py +++ b/pyhdtoolkit/cpymadtools/tune.py @@ -10,7 +10,6 @@ import math import sys from pathlib import Path -from typing import Dict, Tuple import matplotlib.collections import matplotlib.patches diff --git a/pyhdtoolkit/cpymadtools/twiss.py b/pyhdtoolkit/cpymadtools/twiss.py index 0d6052c9..13f19243 100644 --- a/pyhdtoolkit/cpymadtools/twiss.py +++ b/pyhdtoolkit/cpymadtools/twiss.py @@ -7,7 +7,7 @@ Module with functions to manipulate ``MAD-X`` ``TWISS`` functionality through a `~cpymad.madx.Madx` object. """ -from typing import Sequence +from collections.abc import Sequence import tfs from cpymad.madx import Madx diff --git a/pyhdtoolkit/cpymadtools/utils.py b/pyhdtoolkit/cpymadtools/utils.py index aede36c4..3b8168a8 100644 --- a/pyhdtoolkit/cpymadtools/utils.py +++ b/pyhdtoolkit/cpymadtools/utils.py @@ -7,7 +7,6 @@ Module with utility functions to do mundane operations with `~cpymad.madx.Madx` objects. """ from pathlib import Path -from typing import Union import pandas as pd import tfs @@ -19,7 +18,7 @@ def export_madx_table( madx: Madx, /, table_name: str, - file_name: Union[Path, str], + file_name: Path | str, pattern: str = None, headers_table: str = "SUMM", **kwargs, diff --git a/pyhdtoolkit/maths/stats_fitting.py b/pyhdtoolkit/maths/stats_fitting.py index e38dd6be..a24d4142 100644 --- a/pyhdtoolkit/maths/stats_fitting.py +++ b/pyhdtoolkit/maths/stats_fitting.py @@ -7,7 +7,6 @@ Module implementing methods to find the best fit of statistical distributions to data. """ import warnings -from typing import Dict, Tuple, Union import matplotlib import matplotlib.pyplot as plt # noqa: F401 | if omitted, get AttributeError: module 'matplotlib' has no attribute 'axes' @@ -57,7 +56,7 @@ def set_distributions_dict(dist_dict: dict[st.rv_continuous, str]) -> None: def best_fit_distribution( - data: Union[pd.Series, np.ndarray], bins: int = 200, ax: matplotlib.axes.Axes = None + data: pd.Series | np.ndarray, bins: int = 200, ax: matplotlib.axes.Axes = None ) -> tuple[st.rv_continuous, tuple[float, ...]]: """ .. versionadded:: 0.5.0 diff --git a/pyhdtoolkit/maths/utils.py b/pyhdtoolkit/maths/utils.py index e1b86e40..9a5cdf01 100644 --- a/pyhdtoolkit/maths/utils.py +++ b/pyhdtoolkit/maths/utils.py @@ -6,7 +6,6 @@ Module with utility functions used throughout the `~.maths.nonconvex_phase_sync` and `~.maths.stats_fitting` modules. """ -from typing import Tuple, Union import numpy as np import pandas as pd @@ -49,8 +48,8 @@ def get_magnitude(value: float) -> int: def get_scaled_values_and_magnitude_string( - values_array: Union[pd.DataFrame, np.ndarray], force_magnitude: float = None -) -> tuple[Union[pd.DataFrame, np.ndarray], str]: + values_array: pd.DataFrame | np.ndarray, force_magnitude: float = None +) -> tuple[pd.DataFrame | np.ndarray, str]: """ .. versionadded:: 0.8.2 diff --git a/pyhdtoolkit/models/beam.py b/pyhdtoolkit/models/beam.py index 2b1755b7..8008e6eb 100644 --- a/pyhdtoolkit/models/beam.py +++ b/pyhdtoolkit/models/beam.py @@ -7,7 +7,6 @@ Module with ``pydantic`` models to validate and store data structures relative to particle beams. """ from math import sqrt -from typing import Optional from pydantic import BaseModel @@ -20,12 +19,12 @@ class BeamParameters(BaseModel): One can find a usage example in the :ref:`beam enveloppe demo `. """ - pc_GeV: Optional[float] = None # Beam momentum [GeV] - E_0_GeV: Optional[float] = 0.9382720813 # Particle rest mass energy [GeV], defaults to that of a proton - charge: Optional[float] = 1 # Particle charge in [e], defaults to that of a proton - en_x_m: Optional[float] = None # Horizontal normalized emittance [m] - en_y_m: Optional[float] = None # Vertical normalized emittance [m] - deltap_p: Optional[float] = None # Momentum deviation + pc_GeV: float | None = None # Beam momentum [GeV] + E_0_GeV: float | None = 0.9382720813 # Particle rest mass energy [GeV], defaults to that of a proton + charge: float | None = 1 # Particle charge in [e], defaults to that of a proton + en_x_m: float | None = None # Horizontal normalized emittance [m] + en_y_m: float | None = None # Vertical normalized emittance [m] + deltap_p: float | None = None # Momentum deviation @property def B_rho_Tm(self) -> float: diff --git a/pyhdtoolkit/models/htc.py b/pyhdtoolkit/models/htc.py index 13dc853d..48aede86 100644 --- a/pyhdtoolkit/models/htc.py +++ b/pyhdtoolkit/models/htc.py @@ -6,7 +6,6 @@ Module with ``pydantic`` models to validate and store data obtained by querying the ``HTCondor`` queue. """ -from typing import Union from pendulum import DateTime from pydantic import BaseModel, ConfigDict @@ -53,8 +52,8 @@ class HTCTaskSummary(BaseModel): owner: str batch_name: int submitted: DateTime - done: Union[int, str] - run: Union[int, str] - idle: Union[int, str] + done: int | str + run: int | str + idle: int | str total: int job_ids: str diff --git a/pyhdtoolkit/optics/rdt.py b/pyhdtoolkit/optics/rdt.py index 52269f74..3954abc9 100644 --- a/pyhdtoolkit/optics/rdt.py +++ b/pyhdtoolkit/optics/rdt.py @@ -7,16 +7,15 @@ Module implementing utilities for the handling of resonance driving terms. """ -from typing import Tuple, Union -def rdt_to_order_and_type(rdt: Union[int, str]) -> str: +def rdt_to_order_and_type(rdt: int | str) -> str: """ Decompose the input RDT into its four various components and return the type of RDT (normal or skew) and its order. Args: - rdt (Union[int, str]): the RDT to decompose. + rdt (int | str): the RDT to decompose. Returns: A string with the type and (magnet) order of @@ -39,13 +38,13 @@ def rdt_to_order_and_type(rdt: Union[int, str]) -> str: return f"{rdt_type}_{orders[j + k + l + m]}" -def determine_rdt_line(rdt: Union[int, str], plane: str) -> tuple[int, int, int]: +def determine_rdt_line(rdt: int | str, plane: str) -> tuple[int, int, int]: """ Find the given line to look for in the spectral analysis of the given plane that corresponds to the given RDT. Args: - rdt (Union[int, str]): the RDT to look for. + rdt (int | str): the RDT to look for. plane (str): the plane to look for the RDT in. Returns: diff --git a/pyhdtoolkit/optics/ripken.py b/pyhdtoolkit/optics/ripken.py index d7869939..451bdd5c 100644 --- a/pyhdtoolkit/optics/ripken.py +++ b/pyhdtoolkit/optics/ripken.py @@ -6,7 +6,6 @@ Module implementing various calculations based on the :cite:t:`Ripken:optics:1989` optics parameters. """ -from typing import Union import numpy as np import tfs @@ -16,8 +15,8 @@ def lebedev_beam_size( - beta1_: Union[float, np.ndarray], beta2_: Union[float, np.ndarray], geom_emit_x: float, geom_emit_y: float -) -> Union[float, np.ndarray]: + beta1_: float | np.ndarray, beta2_: float | np.ndarray, geom_emit_x: float, geom_emit_y: float +) -> float | np.ndarray: """ .. versionadded:: 0.8.2 @@ -30,8 +29,8 @@ def lebedev_beam_size( See the example below. Args: - beta1_ (Union[float, np.ndarray]): value(s) for the beta1x or beta1y Ripken parameter. - beta2_ (Union[float, np.ndarray]): value(s) for the beta2x or beta2y Ripken parameter. + beta1_ (float | np.ndarray): value(s) for the beta1x or beta1y Ripken parameter. + beta2_ (float | np.ndarray): value(s) for the beta2x or beta2y Ripken parameter. geom_emit_x (float): geometric emittance of the horizontal plane, in [m]. geom_emit_y (float): geometric emittante of the vertical plane, in [m]. diff --git a/pyhdtoolkit/plotting/aperture.py b/pyhdtoolkit/plotting/aperture.py index 30c86234..37958e43 100644 --- a/pyhdtoolkit/plotting/aperture.py +++ b/pyhdtoolkit/plotting/aperture.py @@ -7,7 +7,6 @@ Module with functions to create aperture plots through a `~cpymad.madx.Madx` object. """ -from typing import Optional, Tuple, Union import matplotlib.pyplot as plt import numpy as np @@ -22,7 +21,7 @@ def plot_aperture( madx: Madx, /, - title: Optional[str] = None, + title: str | None = None, xoffset: float = 0, xlimits: tuple[float, float] = None, plot_dipoles: bool = True, @@ -30,10 +29,10 @@ def plot_aperture( plot_quadrupoles: bool = True, plot_bpms: bool = False, aperture_ylim: tuple[float, float] = None, - k0l_lim: Union[tuple[float, float], float, int] = None, - k1l_lim: Union[tuple[float, float], float, int] = None, - k2l_lim: Union[tuple[float, float], float, int] = None, - k3l_lim: Union[tuple[float, float], float, int] = None, + k0l_lim: tuple[float, float] | float = None, + k1l_lim: tuple[float, float] | float = None, + k2l_lim: tuple[float, float] | float = None, + k3l_lim: tuple[float, float] | float = None, color: str = None, **kwargs, ) -> None: @@ -67,7 +66,7 @@ def plot_aperture( Args: madx (cpymad.madx.Madx): an instanciated `~cpymad.madx.Madx` object. Positional only. - title (Optional[str]): title of the figure. + title (str | None): title of the figure. xoffset (float): An offset applied to the ``S`` coordinate before plotting. This is useful if you want to center a plot around a specific point or element, which would then become located at @@ -90,24 +89,24 @@ def plot_aperture( aperture_ylim (tuple[float, float]): vertical axis limits for the aperture values. Defaults to `None`, to be determined by matplotlib based on the provided values. - k0l_lim (Union[tuple[float, float], float, int]): vertical axis limits + k0l_lim (tuple[float, float] | float): vertical axis limits for the ``k0l`` values used for the height of dipole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). If `None` (default) is given, then the limits will be auto-determined based on the ``k0l`` values of the dipoles in the plot. - k1l_lim (Union[tuple[float, float], float, int]): vertical axis limits + k1l_lim (tuple[float, float] | float): vertical axis limits for the ``k1l`` values used for the height of quadrupole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). If `None` (default) is given, then the limits will be auto-determined based on the ``k0l`` values of the quadrupoles in the plot. - k2l_lim (Union[tuple[float, float], float, int]): if given, sextupole + k2l_lim (tuple[float, float] | float): if given, sextupole patches will be plotted on the layout subplot of the figure. If given, acts as vertical axis limits for the k2l values used for the height of sextupole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). - k3l_lim (Union[tuple[float, float], float, int]): if given, octupole + k3l_lim (tuple[float, float] | float): if given, octupole patches will be plotted on the layout subplot of the figure. If given, acts as vertical axis limits for the k3l values used for the height of octupole patches. Can be given as a single value diff --git a/pyhdtoolkit/plotting/crossing.py b/pyhdtoolkit/plotting/crossing.py index 4c1e7ec3..19eb00f0 100644 --- a/pyhdtoolkit/plotting/crossing.py +++ b/pyhdtoolkit/plotting/crossing.py @@ -6,7 +6,6 @@ Module with functions to plot LHC crossing schemes through a `~cpymad.madx.Madx` object. """ -from typing import Union import matplotlib import matplotlib.axes @@ -186,7 +185,7 @@ def plot_single_ir_crossing( def _highlight_mbx_and_mqx( - axis: matplotlib.axes.Axes, plot_df: Union[pd.DataFrame, tfs.TfsDataFrame], ip: int, **kwargs + axis: matplotlib.axes.Axes, plot_df: pd.DataFrame | tfs.TfsDataFrame, ip: int, **kwargs ) -> None: """ .. versionadded:: 1.0.0 diff --git a/pyhdtoolkit/plotting/envelope.py b/pyhdtoolkit/plotting/envelope.py index c0bde02b..997c725e 100644 --- a/pyhdtoolkit/plotting/envelope.py +++ b/pyhdtoolkit/plotting/envelope.py @@ -6,7 +6,6 @@ Module with functions to create beam enveloppe plots through a `~cpymad.madx.Madx` object. """ -from typing import Tuple import numpy as np from cpymad.madx import Madx diff --git a/pyhdtoolkit/plotting/lattice.py b/pyhdtoolkit/plotting/lattice.py index 2ea97b22..cff7b7ea 100644 --- a/pyhdtoolkit/plotting/lattice.py +++ b/pyhdtoolkit/plotting/lattice.py @@ -6,7 +6,6 @@ Module with functions to create lattice plots through a `~cpymad.madx.Madx` object. """ -from typing import Optional, Tuple, Union import matplotlib import matplotlib.axes @@ -25,19 +24,19 @@ def plot_latwiss( madx: Madx, /, - title: Optional[str] = None, + title: str | None = None, xoffset: float = 0, xlimits: tuple[float, float] = None, plot_dipoles: bool = True, plot_dipole_k1: bool = False, plot_quadrupoles: bool = True, plot_bpms: bool = False, - disp_ylim: Union[tuple[float, float], float, int] = None, - beta_ylim: Union[tuple[float, float], float, int] = None, - k0l_lim: Union[tuple[float, float], float, int] = None, - k1l_lim: Union[tuple[float, float], float, int] = None, - k2l_lim: Union[tuple[float, float], float, int] = None, - k3l_lim: Union[tuple[float, float], float, int] = None, + disp_ylim: tuple[float, float] | float = None, + beta_ylim: tuple[float, float] | float = None, + k0l_lim: tuple[float, float] | float = None, + k1l_lim: tuple[float, float] | float = None, + k2l_lim: tuple[float, float] | float = None, + k3l_lim: tuple[float, float] | float = None, **kwargs, ) -> None: """ @@ -66,7 +65,7 @@ def plot_latwiss( Args: madx (cpymad.madx.Madx): an instanciated `~cpymad.madx.Madx` object. Positional only. - title (Optional[str]): if provided, is set as title of the plot. Defaults to `None`. + title (str | None): if provided, is set as title of the plot. Defaults to `None`. xoffset (float): An offset applied to the ``S`` coordinate before plotting. This is useful if you want to center a plot around a specific point or element, which would then become located at :math:`s = 0`. Beware this offset is applied before applying the *xlimits*. Defaults to 0. @@ -87,21 +86,21 @@ def plot_latwiss( Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). Defaults to `None`, to be determined by `~matplotlib` based on the plotted beta values. - k0l_lim (Union[tuple[float, float], float, int]): vertical axis limits for the ``k0l`` + k0l_lim (tuple[float, float] | float): vertical axis limits for the ``k0l`` values used for the height of dipole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). If `None` (default) is given, then the limits will be auto-determined based on the ``k0l`` values of the dipoles in the plot. - k1l_lim (Union[tuple[float, float], float, int]): vertical axis limits for the ``k1l`` + k1l_lim (tuple[float, float] | float): vertical axis limits for the ``k1l`` values used for the height of quadrupole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). If `None` (default) is given, then the limits will be auto-determined based on the ``k0l`` values of the quadrupoles in the plot. - k2l_lim (Union[tuple[float, float], float, int]): if given, sextupole patches will be plotted + k2l_lim (tuple[float, float] | float): if given, sextupole patches will be plotted on the layout subplot of the figure. If given, acts as vertical axis limits for the k2l values used for the height of sextupole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). - k3l_lim (Union[tuple[float, float], float, int]): if given, octupole patches will be plotted + k3l_lim (tuple[float, float] | float): if given, octupole patches will be plotted on the layout subplot of the figure. If given, acts as vertical axis limits for the k3l values used for the height of octupole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). @@ -215,7 +214,7 @@ def plot_machine_survey( Args: madx (cpymad.madx.Madx): an instanciated `~cpymad.madx.Madx` object. Positional only. - title (Optional[str]): if provided, is set as title of the plot. Defaults to `None`. + title (str | None): if provided, is set as title of the plot. Defaults to `None`. show_elements (bool): if `True`, will try to plot by differentiating elements. Defaults to `False`. high_orders (bool): if `True`, plots sextupoles and octupoles if *show_elements* is `True`, diff --git a/pyhdtoolkit/plotting/layout.py b/pyhdtoolkit/plotting/layout.py index c65b57b5..1b543e94 100644 --- a/pyhdtoolkit/plotting/layout.py +++ b/pyhdtoolkit/plotting/layout.py @@ -7,7 +7,6 @@ Module with functions used to represent a machine' elements in an `~matplotlib.axes.Axes` object, mostly used in different `~pyhdtoolkit.plotting` modules. """ -from typing import Tuple, Union import matplotlib import matplotlib.axes @@ -34,10 +33,10 @@ def plot_machine_layout( plot_dipole_k1: bool = False, plot_quadrupoles: bool = True, plot_bpms: bool = False, - k0l_lim: Union[tuple[float, float], float, int] = None, - k1l_lim: Union[tuple[float, float], float, int] = None, - k2l_lim: Union[tuple[float, float], float, int] = None, - k3l_lim: Union[tuple[float, float], float, int] = None, + k0l_lim: tuple[float, float] | float = None, + k1l_lim: tuple[float, float] | float = None, + k2l_lim: tuple[float, float] | float = None, + k3l_lim: tuple[float, float] | float = None, **kwargs, ) -> None: """ @@ -92,24 +91,24 @@ def plot_machine_layout( plot_bpms (bool): if `True`, additional patches will be plotted on the layout subplot to represent Beam Position Monitors. BPMs are plotted in dark grey. - k0l_lim (Union[tuple[float, float], float, int]): vertical axis limits + k0l_lim (tuple[float, float] | float): vertical axis limits for the ``k0l`` values used for the height of dipole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). If `None` (default) is given, then the limits will be auto-determined based on the ``k0l`` values of the dipoles in the plot. - k1l_lim (Union[tuple[float, float], float, int]): vertical axis limits + k1l_lim (tuple[float, float] | float): vertical axis limits for the ``k1l`` values used for the height of quadrupole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). If `None` (default) is given, then the limits will be auto-determined based on the ``k0l`` values of the quadrupoles in the plot. - k2l_lim (Union[tuple[float, float], float, int]): if given, sextupole + k2l_lim (tuple[float, float] | float): if given, sextupole patches will be plotted on the layout subplot of the figure. If given, acts as vertical axis limits for the k2l values used for the height of sextupole patches. Can be given as a single value (float, int) or a tuple (in which case it should be symmetric). - k3l_lim (Union[tuple[float, float], float, int]): if given, octupole + k3l_lim (tuple[float, float] | float): if given, octupole patches will be plotted on the layout subplot of the figure. If given, acts as vertical axis limits for the k3l values used for the height of octupole patches. Can be given as a single value @@ -396,7 +395,7 @@ def _plot_lattice_series( def _ylim_from_input( - ylim: Union[tuple[float, float], float, int], name_for_error: str = "knl_lim" + ylim: tuple[float, float] | float, name_for_error: str = "knl_lim" ) -> tuple[float, float]: """ .. versionadded:: 1.2.0 diff --git a/pyhdtoolkit/plotting/sbs/coupling.py b/pyhdtoolkit/plotting/sbs/coupling.py index 0650d23e..cc47da75 100644 --- a/pyhdtoolkit/plotting/sbs/coupling.py +++ b/pyhdtoolkit/plotting/sbs/coupling.py @@ -6,7 +6,6 @@ Functions to plot coupling components of Segment-by-Segment results. """ -from typing import Tuple import matplotlib import matplotlib.pyplot as plt diff --git a/pyhdtoolkit/plotting/styles/paper.py b/pyhdtoolkit/plotting/styles/paper.py index 32de1bb6..a1a19bbb 100644 --- a/pyhdtoolkit/plotting/styles/paper.py +++ b/pyhdtoolkit/plotting/styles/paper.py @@ -13,7 +13,7 @@ - ``SINGLE_COLUMN``: For plots to be included in a single column figure environment. - ``DOUBLE_COLUMN``: For plots to be included in a full width (double column) figure environment. """ -from typing import Dict, Union +from typing import Union PlotSetting = Union[float, bool, str, tuple] diff --git a/pyhdtoolkit/plotting/styles/thesis.py b/pyhdtoolkit/plotting/styles/thesis.py index a780290d..65a76ecd 100644 --- a/pyhdtoolkit/plotting/styles/thesis.py +++ b/pyhdtoolkit/plotting/styles/thesis.py @@ -15,7 +15,7 @@ - ``MEDIUM``: For simple plots to be included alone in a LaTeX figure environment (e.g single axis line plots, or scatters with a colorbar like in `~pyhdtoolkit.plotting.tune`). - ``LARGE``: For more complex plots to be included alone in a LaTeX figure environment (e.g. multi-axes figures such as in `~pyhdtoolkit.plotting.lattice`). """ -from typing import Dict, Union +from typing import Union PlotSetting = Union[float, bool, str, tuple] diff --git a/pyhdtoolkit/plotting/tune.py b/pyhdtoolkit/plotting/tune.py index e0140749..a7e26d77 100644 --- a/pyhdtoolkit/plotting/tune.py +++ b/pyhdtoolkit/plotting/tune.py @@ -8,7 +8,6 @@ These provide functionality to draw Farey sequences up to a desired order. """ from functools import partial -from typing import Dict, Tuple import matplotlib import matplotlib.axes diff --git a/pyhdtoolkit/plotting/utils.py b/pyhdtoolkit/plotting/utils.py index b9eec23d..308aaa99 100644 --- a/pyhdtoolkit/plotting/utils.py +++ b/pyhdtoolkit/plotting/utils.py @@ -133,7 +133,7 @@ def get_lhc_ips_positions(dataframe: pd.DataFrame) -> dict[str, float]: ip_names = [f"ip{i:d}" for i in range(1, 9)] ip_pos = dataframe.loc[ip_names, "s"].to_numpy() ip_names = [name.upper() for name in ip_names] # make sure to uppercase now - return dict(zip(ip_names, ip_pos)) + return dict(zip(ip_names, ip_pos, strict=False)) def make_elements_groups( diff --git a/pyhdtoolkit/utils/__init__.py b/pyhdtoolkit/utils/__init__.py index d61876ac..7fb97983 100644 --- a/pyhdtoolkit/utils/__init__.py +++ b/pyhdtoolkit/utils/__init__.py @@ -2,7 +2,7 @@ import inspect import traceback import warnings -from typing import Callable +from collections.abc import Callable from . import cmdline, contexts, executors, htc_monitor, logging, operations, printutil diff --git a/pyhdtoolkit/utils/_misc.py b/pyhdtoolkit/utils/_misc.py index 9ec55579..0b3f83e3 100644 --- a/pyhdtoolkit/utils/_misc.py +++ b/pyhdtoolkit/utils/_misc.py @@ -11,9 +11,9 @@ **not** work on other people's machines. """ import shlex +from collections.abc import Sequence from multiprocessing import cpu_count from pathlib import Path -from typing import Sequence, Union import cpymad import numpy as np @@ -213,7 +213,7 @@ def apply_colin_corrs_balance(madx: Madx) -> None: # ----- Fetching Utilities ----- # -def get_betastar_from_opticsfile(opticsfile: Union[Path, str]) -> float: +def get_betastar_from_opticsfile(opticsfile: Path | str) -> float: """ .. versionadded:: 0.16.0 diff --git a/pyhdtoolkit/utils/cmdline.py b/pyhdtoolkit/utils/cmdline.py index 3bbf6222..39267c57 100644 --- a/pyhdtoolkit/utils/cmdline.py +++ b/pyhdtoolkit/utils/cmdline.py @@ -11,7 +11,7 @@ import os import signal import subprocess -from typing import Mapping, Optional, Tuple +from collections.abc import Mapping from loguru import logger @@ -68,7 +68,7 @@ def check_pid_exists(pid: int) -> bool: @staticmethod def run( command: str, shell: bool = True, env: Mapping = None, timeout: float = None - ) -> tuple[Optional[int], bytes]: + ) -> tuple[int | None, bytes]: """ .. versionadded:: 0.2.0 diff --git a/pyhdtoolkit/utils/contexts.py b/pyhdtoolkit/utils/contexts.py index ddcb9b88..0bd4c72d 100644 --- a/pyhdtoolkit/utils/contexts.py +++ b/pyhdtoolkit/utils/contexts.py @@ -7,8 +7,8 @@ Provides useful contexts to use functions in. """ import time +from collections.abc import Callable, Iterator from contextlib import contextmanager -from typing import Callable, Iterator @contextmanager diff --git a/pyhdtoolkit/utils/executors.py b/pyhdtoolkit/utils/executors.py index 5c41d45a..538b9b44 100644 --- a/pyhdtoolkit/utils/executors.py +++ b/pyhdtoolkit/utils/executors.py @@ -21,8 +21,8 @@ I think is quite random. """ +from collections.abc import Callable from concurrent import futures -from typing import Callable from loguru import logger diff --git a/pyhdtoolkit/utils/htc_monitor.py b/pyhdtoolkit/utils/htc_monitor.py index d0ec7fd3..1ecac4ad 100644 --- a/pyhdtoolkit/utils/htc_monitor.py +++ b/pyhdtoolkit/utils/htc_monitor.py @@ -14,7 +14,6 @@ """ import re import time -from typing import Tuple import pendulum from loguru import logger diff --git a/pyhdtoolkit/utils/logging.py b/pyhdtoolkit/utils/logging.py index 398b4ac9..6bc4117a 100644 --- a/pyhdtoolkit/utils/logging.py +++ b/pyhdtoolkit/utils/logging.py @@ -15,7 +15,6 @@ - ``SIMPLE_FORMAT``: minimal, displays the local time, the level and the message. """ import sys -from typing import Union from loguru import logger @@ -36,7 +35,7 @@ SIMPLE_FORMAT = "{time:HH:mm:ss} | {level: <8} | {message}" -def config_logger(level: Union[str, int] = "INFO", format: str = FORMAT1, **kwargs) -> None: +def config_logger(level: str | int = "INFO", format: str = FORMAT1, **kwargs) -> None: """ .. versionadded:: 0.8.2 diff --git a/pyhdtoolkit/utils/operations.py b/pyhdtoolkit/utils/operations.py index 78b9ffd9..b28f42a2 100644 --- a/pyhdtoolkit/utils/operations.py +++ b/pyhdtoolkit/utils/operations.py @@ -16,8 +16,8 @@ import math import random import re +from collections.abc import Callable, Sequence from functools import reduce -from typing import Callable, Dict, Sequence, Tuple, Union class ListOperations: @@ -266,7 +266,7 @@ def get_indices(element, sequence: Sequence) -> list[int]: ListOperations.get_indices(0, [0, 1, 3, 5, 7, 3, 9, 0, 0, 5, 3, 2]) # returns [0, 7, 8] """ - return [i for (y, i) in zip(sequence, range(len(sequence))) if element == y] + return [i for (y, i) in zip(sequence, range(len(sequence)), strict=False) if element == y] @staticmethod def group_by(sequence: Sequence, function: Callable) -> dict[str, list]: @@ -600,7 +600,7 @@ class NumberOperations: """ @staticmethod - def clamp_number(num: Union[int, float], a_val: Union[int, float], b_val: Union[int, float]) -> Union[int, float]: + def clamp_number(num: float, a_val: float, b_val: float) -> float: """ .. versionadded:: 0.2.0 @@ -609,9 +609,9 @@ def clamp_number(num: Union[int, float], a_val: Union[int, float], b_val: Union[ nearest number in the range. Args: - num (Union[int, float]): a number (float / int) - a_val (Union[int, float]): a number (float / int) - b_val (Union[int, float]): a number (float / int) + num (float): a number (float / int) + a_val (float): a number (float / int) + b_val (float): a number (float / int) Returns: A number (float / int), being the nearest to *num* in the range [*a_val*, *b_val*]. @@ -631,8 +631,8 @@ def clamp_number(num: Union[int, float], a_val: Union[int, float], b_val: Union[ @staticmethod def degrees_to_radians( - deg_value: Union[int, float], decompose: bool = False - ) -> Union[tuple[float, str, str], int, float]: + deg_value: float, decompose: bool = False + ) -> tuple[float, str, str] | float: """ .. versionadded:: 0.2.0 @@ -640,7 +640,7 @@ def degrees_to_radians( to radian formula to convert the provided *deg_value*. Args: - deg_value (Union[int, float]): angle value in degrees. + deg_value (float): angle value in degrees. decompose (bool): boolean option to return a more verbose result. Defaults to `False`. Returns: @@ -662,7 +662,7 @@ def degrees_to_radians( return (deg_value * math.pi) / 180.0 @staticmethod - def greatest_common_divisor(sequence: Sequence) -> Union[int, float]: + def greatest_common_divisor(sequence: Sequence) -> float: """ .. versionadded:: 0.2.0 @@ -690,7 +690,7 @@ def greatest_common_divisor(sequence: Sequence) -> Union[int, float]: return reduce(math.gcd, sequence) @staticmethod - def is_divisible_by(dividend: Union[int, float], divisor: Union[int, float]) -> bool: + def is_divisible_by(dividend: float, divisor: float) -> bool: """ .. versionadded:: 0.2.0 @@ -698,8 +698,8 @@ def is_divisible_by(dividend: Union[int, float], divisor: Union[int, float]) -> Uses the modulo operator (`%`) to check if the remainder is equal to 0. Args: - dividend (Union[int, float]): a number. - divisor (Union[int, float]): a number. + dividend (float): a number. + divisor (float): a number. Returns: A boolean stating if *dividend* can be divided by *divisor*. @@ -748,7 +748,7 @@ def _lcm(number1, number2): return reduce(lambda x, y: _lcm(x, y), numbers) @staticmethod - def radians_to_degrees(rad_value: Union[int, float]) -> Union[int, float]: + def radians_to_degrees(rad_value: float) -> float: """ .. versionadded:: 0.2.0 @@ -756,7 +756,7 @@ def radians_to_degrees(rad_value: Union[int, float]) -> Union[int, float]: to degree formula to convert the provided *rad_value*. Args: - rad_value (Union[int, float]): angle value in degrees. + rad_value (float): angle value in degrees. Returns: The angle value in degrees. diff --git a/pyproject.toml b/pyproject.toml index b641e402..7aff47c1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -172,8 +172,8 @@ exclude = ["pyhdtoolkit/cpymadtools/lhc/__init__.py"] # Allow unused variables when underscore-prefixed. dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$" -# Assume Python 3.8+. -target-version = "py38" +# Assume Python 3.10+. +target-version = "py310" [tool.mypy] pretty = true From 0ba7f80672e2e5ea5df1857c43c1ce95b89a5964 Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 10:56:40 +0200 Subject: [PATCH 04/41] fix implicit optionals --- pyhdtoolkit/cpymadtools/coupling.py | 14 ++++++------ pyhdtoolkit/cpymadtools/lhc/_misc.py | 2 +- pyhdtoolkit/cpymadtools/lhc/_powering.py | 6 ++--- pyhdtoolkit/cpymadtools/lhc/_queries.py | 4 ++-- pyhdtoolkit/cpymadtools/lhc/_setup.py | 8 +++---- pyhdtoolkit/cpymadtools/matching.py | 28 ++++++++++++------------ pyhdtoolkit/cpymadtools/ptc.py | 8 +++---- pyhdtoolkit/cpymadtools/track.py | 2 +- pyhdtoolkit/cpymadtools/tune.py | 2 +- pyhdtoolkit/cpymadtools/twiss.py | 2 +- pyhdtoolkit/cpymadtools/utils.py | 2 +- pyhdtoolkit/maths/utils.py | 2 +- pyhdtoolkit/plotting/aperture.py | 18 +++++++-------- pyhdtoolkit/plotting/crossing.py | 6 ++--- pyhdtoolkit/plotting/envelope.py | 2 +- pyhdtoolkit/plotting/lattice.py | 16 +++++++------- pyhdtoolkit/plotting/layout.py | 12 +++++----- pyhdtoolkit/plotting/phasespace.py | 4 ++-- pyhdtoolkit/plotting/sbs/coupling.py | 14 ++++++------ pyhdtoolkit/plotting/sbs/phase.py | 6 ++--- pyhdtoolkit/plotting/tune.py | 4 ++-- pyhdtoolkit/plotting/utils.py | 6 ++--- pyhdtoolkit/utils/_misc.py | 4 ++-- pyhdtoolkit/utils/cmdline.py | 2 +- pyhdtoolkit/utils/htc_monitor.py | 2 +- 25 files changed, 88 insertions(+), 88 deletions(-) diff --git a/pyhdtoolkit/cpymadtools/coupling.py b/pyhdtoolkit/cpymadtools/coupling.py index 0957795d..358a2945 100644 --- a/pyhdtoolkit/cpymadtools/coupling.py +++ b/pyhdtoolkit/cpymadtools/coupling.py @@ -28,12 +28,12 @@ def get_closest_tune_approach( madx: Madx, /, - accelerator: str = None, - sequence: str = None, - varied_knobs: Sequence[str] = None, + accelerator: str | None = None, + sequence: str | None = None, + varied_knobs: Sequence[str] | None = None, telescopic_squeeze: bool = True, run3: bool = False, - explicit_targets: tuple[float, float] = None, + explicit_targets: tuple[float, float] | None = None, step: float = 1e-7, calls: int = 100, tolerance: float = 1e-21, @@ -146,8 +146,8 @@ def get_cminus_from_coupling_rdts( /, patterns: Sequence[str] = [""], method: str = "teapot", - qx: float = None, - qy: float = None, + qx: float | None = None, + qy: float | None = None, filtering: float = 0, ) -> float: """ @@ -232,7 +232,7 @@ def get_cminus_from_coupling_rdts( def match_no_coupling_through_ripkens( - madx: Madx, /, sequence: str = None, location: str = None, vary_knobs: Sequence[str] = None + madx: Madx, /, sequence: str | None = None, location: str | None = None, vary_knobs: Sequence[str] | None = None ) -> None: """ .. versionadded:: 0.16.0 diff --git a/pyhdtoolkit/cpymadtools/lhc/_misc.py b/pyhdtoolkit/cpymadtools/lhc/_misc.py index c5319a8c..827ef3b1 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_misc.py +++ b/pyhdtoolkit/cpymadtools/lhc/_misc.py @@ -169,7 +169,7 @@ def get_lhc_bpms_list(madx: Madx, /) -> list[str]: def get_sizes_at_ip( - madx: Madx, /, ip: int, geom_emit_x: float = None, geom_emit_y: float = None + madx: Madx, /, ip: int, geom_emit_x: float | None = None, geom_emit_y: float | None = None ) -> tuple[float, float]: """ .. versionadded:: 1.0.0 diff --git a/pyhdtoolkit/cpymadtools/lhc/_powering.py b/pyhdtoolkit/cpymadtools/lhc/_powering.py index ed55e16c..9bbf7af1 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_powering.py +++ b/pyhdtoolkit/cpymadtools/lhc/_powering.py @@ -11,7 +11,7 @@ from loguru import logger -def apply_lhc_colinearity_knob(madx: Madx, /, colinearity_knob_value: float = 0, ir: int = None) -> None: +def apply_lhc_colinearity_knob(madx: Madx, /, colinearity_knob_value: float = 0, ir: int | None = None) -> None: """ .. versionadded:: 0.15.0 @@ -49,7 +49,7 @@ def apply_lhc_colinearity_knob(madx: Madx, /, colinearity_knob_value: float = 0, logger.debug(f"Set '{left_knob}' to {madx.globals[left_knob]}") -def apply_lhc_colinearity_knob_delta(madx: Madx, /, colinearity_knob_delta: float = 0, ir: int = None) -> None: +def apply_lhc_colinearity_knob_delta(madx: Madx, /, colinearity_knob_delta: float = 0, ir: int | None = None) -> None: """ .. versionadded:: 0.21.0 @@ -89,7 +89,7 @@ def apply_lhc_colinearity_knob_delta(madx: Madx, /, colinearity_knob_delta: floa def apply_lhc_rigidity_waist_shift_knob( - madx: Madx, /, rigidty_waist_shift_value: float = 0, ir: int = None, side: str = "left" + madx: Madx, /, rigidty_waist_shift_value: float = 0, ir: int | None = None, side: str = "left" ) -> None: """ .. versionadded:: 0.15.0 diff --git a/pyhdtoolkit/cpymadtools/lhc/_queries.py b/pyhdtoolkit/cpymadtools/lhc/_queries.py index 692ec8ef..d2171459 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_queries.py +++ b/pyhdtoolkit/cpymadtools/lhc/_queries.py @@ -33,7 +33,7 @@ def get_magnets_powering( - madx: Madx, /, patterns: Sequence[str] = [r"^mb\.", r"^mq\.", r"^ms\."], brho: str | float = None, **kwargs + madx: Madx, /, patterns: Sequence[str] = [r"^mb\.", r"^mq\.", r"^ms\."], brho: str | float | None = None, **kwargs ) -> tfs.TfsDataFrame: r""" .. versionadded:: 0.17.0 @@ -218,7 +218,7 @@ def get_current_orbit_setup(madx: Madx, /) -> dict[str, float]: # ----- Helpers ----- # -def _list_field_currents(madx: Madx, /, brho: str | float = None) -> None: +def _list_field_currents(madx: Madx, /, brho: str | float | None = None) -> None: """ Creates additional columns for the ``TWISS`` table with the magnets' total fields and currents, to help later on determine which proportion of their maximum powering diff --git a/pyhdtoolkit/cpymadtools/lhc/_setup.py b/pyhdtoolkit/cpymadtools/lhc/_setup.py index 2f6f462a..29396001 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_setup.py +++ b/pyhdtoolkit/cpymadtools/lhc/_setup.py @@ -16,7 +16,7 @@ def prepare_lhc_run2( - opticsfile: str, beam: int = 1, use_b4: bool = False, energy: float = 6500, slicefactor: int = None, **kwargs + opticsfile: str, beam: int = 1, use_b4: bool = False, energy: float = 6500, slicefactor: int | None = None, **kwargs ) -> Madx: """ .. versionadded:: 1.0.0 @@ -96,7 +96,7 @@ def _run2_sequence_from_opticsfile(opticsfile: Path, use_b4: bool = False) -> Pa def prepare_lhc_run3( - opticsfile: str, beam: int = 1, use_b4: bool = False, energy: float = 6800, slicefactor: int = None, **kwargs + opticsfile: str, beam: int = 1, use_b4: bool = False, energy: float = 6800, slicefactor: int | None = None, **kwargs ) -> Madx: """ .. versionadded:: 1.0.0 @@ -228,11 +228,11 @@ class LHCSetup: def __init__( self, run: int = 3, - opticsfile: str = None, + opticsfile: str | None = None, beam: int = 1, use_b4: bool = False, energy: float = 6800, - slicefactor: int = None, + slicefactor: int | None = None, **kwargs, ): assert opticsfile is not None, "An opticsfile must be provided" diff --git a/pyhdtoolkit/cpymadtools/matching.py b/pyhdtoolkit/cpymadtools/matching.py index 50566120..351076b8 100644 --- a/pyhdtoolkit/cpymadtools/matching.py +++ b/pyhdtoolkit/cpymadtools/matching.py @@ -19,13 +19,13 @@ def match_tunes_and_chromaticities( madx: Madx, /, - accelerator: str = None, + accelerator: str | None = None, sequence: str | None = None, - q1_target: float = None, - q2_target: float = None, - dq1_target: float = None, - dq2_target: float = None, - varied_knobs: Sequence[str] = None, + q1_target: float | None = None, + q2_target: float | None = None, + dq1_target: float | None = None, + dq2_target: float | None = None, + varied_knobs: Sequence[str] | None = None, telescopic_squeeze: bool = True, run3: bool = False, step: float = 1e-7, @@ -192,11 +192,11 @@ def match(*args, **kwargs): def match_tunes( madx: Madx, /, - accelerator: str = None, + accelerator: str | None = None, sequence: str | None = None, - q1_target: float = None, - q2_target: float = None, - varied_knobs: Sequence[str] = None, + q1_target: float | None = None, + q2_target: float | None = None, + varied_knobs: Sequence[str] | None = None, telescopic_squeeze: bool = True, run3: bool = False, step: float = 1e-7, @@ -294,11 +294,11 @@ def match_tunes( def match_chromaticities( madx: Madx, /, - accelerator: str = None, + accelerator: str | None = None, sequence: str | None = None, - dq1_target: float = None, - dq2_target: float = None, - varied_knobs: Sequence[str] = None, + dq1_target: float | None = None, + dq2_target: float | None = None, + varied_knobs: Sequence[str] | None = None, telescopic_squeeze: bool = True, run3: bool = False, step: float = 1e-7, diff --git a/pyhdtoolkit/cpymadtools/ptc.py b/pyhdtoolkit/cpymadtools/ptc.py index 54baae71..2a2ab163 100644 --- a/pyhdtoolkit/cpymadtools/ptc.py +++ b/pyhdtoolkit/cpymadtools/ptc.py @@ -19,7 +19,7 @@ def get_amplitude_detuning( - madx: Madx, /, order: int = 2, file: Path | str = None, fringe: bool = False, **kwargs + madx: Madx, /, order: int = 2, file: Path | str | None = None, fringe: bool = False, **kwargs ) -> tfs.TfsDataFrame: """ .. versionadded:: 0.7.0 @@ -143,7 +143,7 @@ def get_amplitude_detuning( def get_rdts( - madx: Madx, /, order: int = 4, file: Path | str = None, fringe: bool = False, **kwargs + madx: Madx, /, order: int = 4, file: Path | str | None = None, fringe: bool = False, **kwargs ) -> tfs.TfsDataFrame: """ .. versionadded:: 0.7.0 @@ -238,7 +238,7 @@ def ptc_twiss( madx: Madx, /, order: int = 4, - file: Path | str = None, + file: Path | str | None = None, fringe: bool = False, table: str = "ptc_twiss", **kwargs, @@ -343,7 +343,7 @@ def ptc_track_particle( initial_coordinates: tuple[float, float, float, float, float, float], nturns: int, sequence: str | None = None, - observation_points: Sequence[str] = None, + observation_points: Sequence[str] | None = None, onetable: bool = False, fringe: bool = False, **kwargs, diff --git a/pyhdtoolkit/cpymadtools/track.py b/pyhdtoolkit/cpymadtools/track.py index fcd7a1ac..72ac93da 100644 --- a/pyhdtoolkit/cpymadtools/track.py +++ b/pyhdtoolkit/cpymadtools/track.py @@ -22,7 +22,7 @@ def track_single_particle( initial_coordinates: tuple[float, float, float, float, float, float], nturns: int, sequence: str | None = None, - observation_points: Sequence[str] = None, + observation_points: Sequence[str] | None = None, **kwargs, ) -> dict[str, pd.DataFrame]: """ diff --git a/pyhdtoolkit/cpymadtools/tune.py b/pyhdtoolkit/cpymadtools/tune.py index 27dc1e02..8dcb7b01 100644 --- a/pyhdtoolkit/cpymadtools/tune.py +++ b/pyhdtoolkit/cpymadtools/tune.py @@ -20,7 +20,7 @@ def make_footprint_table( - madx: Madx, /, sigma: float = 5, dense: bool = False, file: str = None, cleanup: bool = True, **kwargs + madx: Madx, /, sigma: float = 5, dense: bool = False, file: str | None = None, cleanup: bool = True, **kwargs ) -> tfs.TfsDataFrame: """ .. versionadded:: 0.9.0 diff --git a/pyhdtoolkit/cpymadtools/twiss.py b/pyhdtoolkit/cpymadtools/twiss.py index 13f19243..c9f31277 100644 --- a/pyhdtoolkit/cpymadtools/twiss.py +++ b/pyhdtoolkit/cpymadtools/twiss.py @@ -21,7 +21,7 @@ def get_pattern_twiss( madx: Madx, /, - columns: Sequence[str] = None, + columns: Sequence[str] | None = None, patterns: Sequence[str] = [""], **kwargs, ) -> tfs.TfsDataFrame: diff --git a/pyhdtoolkit/cpymadtools/utils.py b/pyhdtoolkit/cpymadtools/utils.py index 3b8168a8..153dab8e 100644 --- a/pyhdtoolkit/cpymadtools/utils.py +++ b/pyhdtoolkit/cpymadtools/utils.py @@ -19,7 +19,7 @@ def export_madx_table( /, table_name: str, file_name: Path | str, - pattern: str = None, + pattern: str | None = None, headers_table: str = "SUMM", **kwargs, ) -> None: diff --git a/pyhdtoolkit/maths/utils.py b/pyhdtoolkit/maths/utils.py index 9a5cdf01..6df494d3 100644 --- a/pyhdtoolkit/maths/utils.py +++ b/pyhdtoolkit/maths/utils.py @@ -48,7 +48,7 @@ def get_magnitude(value: float) -> int: def get_scaled_values_and_magnitude_string( - values_array: pd.DataFrame | np.ndarray, force_magnitude: float = None + values_array: pd.DataFrame | np.ndarray, force_magnitude: float | None = None ) -> tuple[pd.DataFrame | np.ndarray, str]: """ .. versionadded:: 0.8.2 diff --git a/pyhdtoolkit/plotting/aperture.py b/pyhdtoolkit/plotting/aperture.py index 37958e43..feecb93a 100644 --- a/pyhdtoolkit/plotting/aperture.py +++ b/pyhdtoolkit/plotting/aperture.py @@ -23,17 +23,17 @@ def plot_aperture( /, title: str | None = None, xoffset: float = 0, - xlimits: tuple[float, float] = None, + xlimits: tuple[float, float] | None = None, plot_dipoles: bool = True, plot_dipole_k1: bool = False, plot_quadrupoles: bool = True, plot_bpms: bool = False, - aperture_ylim: tuple[float, float] = None, - k0l_lim: tuple[float, float] | float = None, - k1l_lim: tuple[float, float] | float = None, - k2l_lim: tuple[float, float] | float = None, - k3l_lim: tuple[float, float] | float = None, - color: str = None, + aperture_ylim: tuple[float, float] | None = None, + k0l_lim: tuple[float, float] | float | None = None, + k1l_lim: tuple[float, float] | float | None = None, + k2l_lim: tuple[float, float] | float | None = None, + k3l_lim: tuple[float, float] | float | None = None, + color: str | None = None, **kwargs, ) -> None: """ @@ -191,7 +191,7 @@ def plot_physical_apertures( plane: str, scale: float = 1, xoffset: float = 0, - xlimits: tuple[float, float] = None, + xlimits: tuple[float, float] | None = None, **kwargs, ) -> None: """ @@ -279,7 +279,7 @@ def plot_physical_apertures( def _get_positions_and_real_apertures( - madx, /, plane: str, xoffset: float = 0, xlimits: tuple[float, float] = None, **kwargs + madx, /, plane: str, xoffset: float = 0, xlimits: tuple[float, float] | None = None, **kwargs ) -> tuple[np.ndarray, np.ndarray]: """ .. versionadded:: 1.2.0 diff --git a/pyhdtoolkit/plotting/crossing.py b/pyhdtoolkit/plotting/crossing.py index 19eb00f0..a2dc3918 100644 --- a/pyhdtoolkit/plotting/crossing.py +++ b/pyhdtoolkit/plotting/crossing.py @@ -140,9 +140,9 @@ def plot_single_ir_crossing( plot_df_b2: pd.DataFrame, plot_column: str, scaling: float = 1, - ylabel: str = None, - xlabel: str = None, - title: str = None, + ylabel: str | None = None, + xlabel: str | None = None, + title: str | None = None, ) -> None: """ .. versionadded:: 1.0.0 diff --git a/pyhdtoolkit/plotting/envelope.py b/pyhdtoolkit/plotting/envelope.py index 997c725e..904f033a 100644 --- a/pyhdtoolkit/plotting/envelope.py +++ b/pyhdtoolkit/plotting/envelope.py @@ -22,7 +22,7 @@ def plot_beam_envelope( nsigma: float = 1, scale: float = 1, xoffset: float = 0, - xlimits: tuple[float, float] = None, + xlimits: tuple[float, float] | None = None, **kwargs, ) -> None: """ diff --git a/pyhdtoolkit/plotting/lattice.py b/pyhdtoolkit/plotting/lattice.py index cff7b7ea..ab17bcfa 100644 --- a/pyhdtoolkit/plotting/lattice.py +++ b/pyhdtoolkit/plotting/lattice.py @@ -26,17 +26,17 @@ def plot_latwiss( /, title: str | None = None, xoffset: float = 0, - xlimits: tuple[float, float] = None, + xlimits: tuple[float, float] | None = None, plot_dipoles: bool = True, plot_dipole_k1: bool = False, plot_quadrupoles: bool = True, plot_bpms: bool = False, - disp_ylim: tuple[float, float] | float = None, - beta_ylim: tuple[float, float] | float = None, - k0l_lim: tuple[float, float] | float = None, - k1l_lim: tuple[float, float] | float = None, - k2l_lim: tuple[float, float] | float = None, - k3l_lim: tuple[float, float] | float = None, + disp_ylim: tuple[float, float] | float | None = None, + beta_ylim: tuple[float, float] | float | None = None, + k0l_lim: tuple[float, float] | float | None = None, + k1l_lim: tuple[float, float] | float | None = None, + k2l_lim: tuple[float, float] | float | None = None, + k3l_lim: tuple[float, float] | float | None = None, **kwargs, ) -> None: """ @@ -199,7 +199,7 @@ def plot_latwiss( def plot_machine_survey( madx: Madx, /, - title: str = None, + title: str | None = None, show_elements: bool = False, high_orders: bool = False, **kwargs, diff --git a/pyhdtoolkit/plotting/layout.py b/pyhdtoolkit/plotting/layout.py index 1b543e94..4e610b4f 100644 --- a/pyhdtoolkit/plotting/layout.py +++ b/pyhdtoolkit/plotting/layout.py @@ -26,17 +26,17 @@ def plot_machine_layout( madx: Madx, /, - title: str = None, + title: str | None = None, xoffset: float = 0, - xlimits: tuple[float, float] = None, + xlimits: tuple[float, float] | None = None, plot_dipoles: bool = True, plot_dipole_k1: bool = False, plot_quadrupoles: bool = True, plot_bpms: bool = False, - k0l_lim: tuple[float, float] | float = None, - k1l_lim: tuple[float, float] | float = None, - k2l_lim: tuple[float, float] | float = None, - k3l_lim: tuple[float, float] | float = None, + k0l_lim: tuple[float, float] | float | None = None, + k1l_lim: tuple[float, float] | float | None = None, + k2l_lim: tuple[float, float] | float | None = None, + k3l_lim: tuple[float, float] | float | None = None, **kwargs, ) -> None: """ diff --git a/pyhdtoolkit/plotting/phasespace.py b/pyhdtoolkit/plotting/phasespace.py index b723e90a..a69b49f2 100644 --- a/pyhdtoolkit/plotting/phasespace.py +++ b/pyhdtoolkit/plotting/phasespace.py @@ -27,7 +27,7 @@ def plot_courant_snyder_phase_space( u_coordinates: np.ndarray, pu_coordinates: np.ndarray, plane: str = "Horizontal", - title: str = None, + title: str | None = None, **kwargs, ) -> matplotlib.axes.Axes: """ @@ -94,7 +94,7 @@ def plot_courant_snyder_phase_space_colored( u_coordinates: np.ndarray, pu_coordinates: np.ndarray, plane: str = "Horizontal", - title: str = None, + title: str | None = None, **kwargs, ) -> matplotlib.figure.Figure: """ diff --git a/pyhdtoolkit/plotting/sbs/coupling.py b/pyhdtoolkit/plotting/sbs/coupling.py index cc47da75..9906e584 100644 --- a/pyhdtoolkit/plotting/sbs/coupling.py +++ b/pyhdtoolkit/plotting/sbs/coupling.py @@ -21,7 +21,7 @@ def plot_rdt_component( b2_segment_df: tfs.TfsDataFrame, b1_model: tfs.TfsDataFrame = None, b2_model: tfs.TfsDataFrame = None, - ip: int = None, + ip: int | None = None, rdt: str = "F1001", component: str = "ABS", **kwargs, @@ -104,11 +104,11 @@ def plot_full_ip_rdt( b2_segment_df: tfs.TfsDataFrame, b1_model: tfs.TfsDataFrame = None, b2_model: tfs.TfsDataFrame = None, - ip: int = None, + ip: int | None = None, rdt: str = "F1001", - abs_ylimits: tuple[float, float] = None, - real_ylimits: tuple[float, float] = None, - imag_ylimits: tuple[float, float] = None, + abs_ylimits: tuple[float, float] | None = None, + real_ylimits: tuple[float, float] | None = None, + imag_ylimits: tuple[float, float] | None = None, **kwargs, ) -> matplotlib.figure.Figure: """ @@ -211,10 +211,10 @@ def _plot_sbs_coupling_rdt_component( ax: matplotlib.axes.Axes, segment_df: tfs.TfsDataFrame, model_df: tfs.TfsDataFrame = None, - ip: int = None, + ip: int | None = None, rdt: str = "F1001", component: str = "ABS", - ylabel: str = None, + ylabel: str | None = None, ) -> None: """ Plots a component of the given coupling RDT over the segment, optionally highlighting the IP location. diff --git a/pyhdtoolkit/plotting/sbs/phase.py b/pyhdtoolkit/plotting/sbs/phase.py index 4856ed4d..0cc53c2b 100644 --- a/pyhdtoolkit/plotting/sbs/phase.py +++ b/pyhdtoolkit/plotting/sbs/phase.py @@ -19,7 +19,7 @@ def plot_phase_segment_one_beam( phase_x: tfs.TfsDataFrame, phase_y: tfs.TfsDataFrame, model: tfs.TfsDataFrame = None, - ip: int = None, + ip: int | None = None, **kwargs, ) -> matplotlib.figure.Figure: """ @@ -71,7 +71,7 @@ def plot_phase_segment_both_beams( b2_phase_y: tfs.TfsDataFrame, b1_model: tfs.TfsDataFrame = None, b2_model: tfs.TfsDataFrame = None, - ip: int = None, + ip: int | None = None, **kwargs, ) -> matplotlib.figure.Figure: """ @@ -154,7 +154,7 @@ def plot_phase_segment( segment_df: tfs.TfsDataFrame = None, model_df: tfs.TfsDataFrame = None, plane: str = "x", - ip: int = None, + ip: int | None = None, ) -> None: """ .. versionadded:: 0.19.0 diff --git a/pyhdtoolkit/plotting/tune.py b/pyhdtoolkit/plotting/tune.py index a7e26d77..5825e5c8 100644 --- a/pyhdtoolkit/plotting/tune.py +++ b/pyhdtoolkit/plotting/tune.py @@ -69,8 +69,8 @@ def farey_sequence(order: int) -> list[tuple[int, int]]: def plot_tune_diagram( - title: str = None, - legend_title: str = None, + title: str | None = None, + legend_title: str | None = None, max_order: int = 6, differentiate_orders: bool = False, **kwargs, diff --git a/pyhdtoolkit/plotting/utils.py b/pyhdtoolkit/plotting/utils.py index 308aaa99..28123f97 100644 --- a/pyhdtoolkit/plotting/utils.py +++ b/pyhdtoolkit/plotting/utils.py @@ -137,7 +137,7 @@ def get_lhc_ips_positions(dataframe: pd.DataFrame) -> dict[str, float]: def make_elements_groups( - madx: Madx, /, xoffset: float = 0, xlimits: tuple[float, float] = None + madx: Madx, /, xoffset: float = 0, xlimits: tuple[float, float] | None = None ) -> dict[str, pd.DataFrame]: """ .. versionadded:: 1.0.0 @@ -217,7 +217,7 @@ def make_survey_groups(madx: Madx, /) -> dict[str, pd.DataFrame]: def draw_ip_locations( - ip_positions: dict[str, float] = None, + ip_positions: dict[str, float] | None = None, lines: bool = True, location: str = "outside", **kwargs, @@ -402,7 +402,7 @@ def draw_confidence_ellipse(x: ArrayLike, y: ArrayLike, n_std: float = 3.0, face def _get_twiss_table_with_offsets_and_limits( - madx: Madx, /, xoffset: float = 0, xlimits: tuple[float, float] = None, **kwargs + madx: Madx, /, xoffset: float = 0, xlimits: tuple[float, float] | None = None, **kwargs ) -> pd.DataFrame: """ .. versionadded:: 1.0.0 diff --git a/pyhdtoolkit/utils/_misc.py b/pyhdtoolkit/utils/_misc.py index 0b3f83e3..3943e111 100644 --- a/pyhdtoolkit/utils/_misc.py +++ b/pyhdtoolkit/utils/_misc.py @@ -80,7 +80,7 @@ def split_complex_columns(df: pd.DataFrame, drop: bool = False) -> pd.DataFrame: return res -def add_noise_to_ir_bpms(df: pd.DataFrame, max_index: int, stdev: float, columns: Sequence[str] = None) -> pd.DataFrame: +def add_noise_to_ir_bpms(df: pd.DataFrame, max_index: int, stdev: float, columns: Sequence[str] | None = None) -> pd.DataFrame: """ .. versionadded:: 1.2.0 @@ -121,7 +121,7 @@ def add_noise_to_ir_bpms(df: pd.DataFrame, max_index: int, stdev: float, columns def add_noise_to_arc_bpms( - df: pd.DataFrame, min_index: int, stdev: float, columns: Sequence[str] = None + df: pd.DataFrame, min_index: int, stdev: float, columns: Sequence[str] | None = None ) -> pd.DataFrame: """ .. versionadded:: 1.2.0 diff --git a/pyhdtoolkit/utils/cmdline.py b/pyhdtoolkit/utils/cmdline.py index 39267c57..f552f62d 100644 --- a/pyhdtoolkit/utils/cmdline.py +++ b/pyhdtoolkit/utils/cmdline.py @@ -67,7 +67,7 @@ def check_pid_exists(pid: int) -> bool: @staticmethod def run( - command: str, shell: bool = True, env: Mapping = None, timeout: float = None + command: str, shell: bool = True, env: Mapping | None = None, timeout: float | None = None ) -> tuple[int | None, bytes]: """ .. versionadded:: 0.2.0 diff --git a/pyhdtoolkit/utils/htc_monitor.py b/pyhdtoolkit/utils/htc_monitor.py index 1ecac4ad..49429433 100644 --- a/pyhdtoolkit/utils/htc_monitor.py +++ b/pyhdtoolkit/utils/htc_monitor.py @@ -189,7 +189,7 @@ def _process_task_summary_line(line: str) -> HTCTaskSummary: ) -def _process_cluster_summary_line(line: str, query: str = None) -> BaseSummary: +def _process_cluster_summary_line(line: str, query: str | None = None) -> BaseSummary: r""" Beware if no jobs are running we can't have taken querying_owner info from tasks summaries, so we need to match a wildcard word by giving querying_owner=(\D+). This would add a match to the regex From 1b76b982405388b9ea5e359301b9a31957cef19d Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 11:10:31 +0200 Subject: [PATCH 05/41] refix many imports rules --- pyhdtoolkit/cpymadtools/__init__.py | 2 +- pyhdtoolkit/maths/__init__.py | 4 ++-- pyhdtoolkit/maths/stats_fitting.py | 4 ++-- pyhdtoolkit/models/__init__.py | 4 ++-- pyhdtoolkit/optics/__init__.py | 4 ++-- pyhdtoolkit/plotting/__init__.py | 4 ++-- pyhdtoolkit/plotting/crossing.py | 7 +++---- pyhdtoolkit/plotting/lattice.py | 5 ++--- pyhdtoolkit/plotting/layout.py | 5 ++--- pyhdtoolkit/plotting/phasespace.py | 8 ++++---- pyhdtoolkit/plotting/sbs/coupling.py | 11 +++++------ pyhdtoolkit/plotting/sbs/phase.py | 9 +++++---- pyhdtoolkit/plotting/styles/__init__.py | 8 ++++---- pyhdtoolkit/plotting/tune.py | 7 +++---- pyhdtoolkit/plotting/utils.py | 22 +++++++++++++--------- pyhdtoolkit/utils/__init__.py | 4 ++-- 16 files changed, 54 insertions(+), 54 deletions(-) diff --git a/pyhdtoolkit/cpymadtools/__init__.py b/pyhdtoolkit/cpymadtools/__init__.py index eca62876..cd9c94c3 100644 --- a/pyhdtoolkit/cpymadtools/__init__.py +++ b/pyhdtoolkit/cpymadtools/__init__.py @@ -1,3 +1,3 @@ -from . import constants, coupling, lhc, matching, ptc, track, tune, twiss, utils +from . import constants, coupling, lhc, matching, ptc, track, tune, twiss, utils # noqa: TID252 __all__ = ["constants", "coupling", "lhc", "matching", "ptc", "track", "tune", "twiss", "utils"] diff --git a/pyhdtoolkit/maths/__init__.py b/pyhdtoolkit/maths/__init__.py index edf05e59..22e7634b 100644 --- a/pyhdtoolkit/maths/__init__.py +++ b/pyhdtoolkit/maths/__init__.py @@ -1,3 +1,3 @@ -from . import stats_fitting, utils +from . import stats_fitting, utils # noqa: TID252 -__all__ = [stats_fitting, utils] +__all__ = ["stats_fitting", "utils"] diff --git a/pyhdtoolkit/maths/stats_fitting.py b/pyhdtoolkit/maths/stats_fitting.py index a24d4142..a46fe49c 100644 --- a/pyhdtoolkit/maths/stats_fitting.py +++ b/pyhdtoolkit/maths/stats_fitting.py @@ -8,12 +8,12 @@ """ import warnings -import matplotlib import matplotlib.pyplot as plt # noqa: F401 | if omitted, get AttributeError: module 'matplotlib' has no attribute 'axes' import numpy as np import pandas as pd import scipy.stats as st from loguru import logger +from matplotlib.axes import Axes # Distributions to check # DISTRIBUTIONS: dict[st.rv_continuous, str] = { @@ -56,7 +56,7 @@ def set_distributions_dict(dist_dict: dict[st.rv_continuous, str]) -> None: def best_fit_distribution( - data: pd.Series | np.ndarray, bins: int = 200, ax: matplotlib.axes.Axes = None + data: pd.Series | np.ndarray, bins: int = 200, ax: Axes = None ) -> tuple[st.rv_continuous, tuple[float, ...]]: """ .. versionadded:: 0.5.0 diff --git a/pyhdtoolkit/models/__init__.py b/pyhdtoolkit/models/__init__.py index 13d75fd7..208dc40d 100644 --- a/pyhdtoolkit/models/__init__.py +++ b/pyhdtoolkit/models/__init__.py @@ -1,3 +1,3 @@ -from . import beam, htc, madx +from . import beam, htc, madx # noqa: TID252 -__all__ = [beam, htc, madx] +__all__ = ["beam", "htc", "madx"] diff --git a/pyhdtoolkit/optics/__init__.py b/pyhdtoolkit/optics/__init__.py index 9b6bc57e..a2be7a6f 100644 --- a/pyhdtoolkit/optics/__init__.py +++ b/pyhdtoolkit/optics/__init__.py @@ -1,3 +1,3 @@ -from . import beam, ripken, twiss +from . import beam, ripken, twiss # noqa: TID252 -__all__ = [beam, ripken, twiss] +__all__ = ["beam", "ripken", "twiss"] diff --git a/pyhdtoolkit/plotting/__init__.py b/pyhdtoolkit/plotting/__init__.py index 98eaa4c1..ba9eddf5 100644 --- a/pyhdtoolkit/plotting/__init__.py +++ b/pyhdtoolkit/plotting/__init__.py @@ -1,6 +1,6 @@ """ .. _plotting: """ -from . import aperture, crossing, envelope, lattice, phasespace, tune, utils +from . import aperture, crossing, envelope, lattice, phasespace, tune, utils # noqa: TID252 -__all__ = [aperture, envelope, crossing, lattice, phasespace, tune, utils] +__all__ = ["aperture", "envelope", "crossing", "lattice", "phasespace", "tune", "utils"] diff --git a/pyhdtoolkit/plotting/crossing.py b/pyhdtoolkit/plotting/crossing.py index a2dc3918..e82b2ba9 100644 --- a/pyhdtoolkit/plotting/crossing.py +++ b/pyhdtoolkit/plotting/crossing.py @@ -7,13 +7,12 @@ Module with functions to plot LHC crossing schemes through a `~cpymad.madx.Madx` object. """ -import matplotlib -import matplotlib.axes import matplotlib.pyplot as plt import pandas as pd import tfs from cpymad.madx import Madx from loguru import logger +from matplotlib.axes import Axes def plot_two_lhc_ips_crossings( @@ -135,7 +134,7 @@ def plot_two_lhc_ips_crossings( def plot_single_ir_crossing( - axis: matplotlib.axes.Axes, + axis: Axes, plot_df_b1: pd.DataFrame, plot_df_b2: pd.DataFrame, plot_column: str, @@ -185,7 +184,7 @@ def plot_single_ir_crossing( def _highlight_mbx_and_mqx( - axis: matplotlib.axes.Axes, plot_df: pd.DataFrame | tfs.TfsDataFrame, ip: int, **kwargs + axis: Axes, plot_df: pd.DataFrame | tfs.TfsDataFrame, ip: int, **kwargs ) -> None: """ .. versionadded:: 1.0.0 diff --git a/pyhdtoolkit/plotting/lattice.py b/pyhdtoolkit/plotting/lattice.py index ab17bcfa..937afe2c 100644 --- a/pyhdtoolkit/plotting/lattice.py +++ b/pyhdtoolkit/plotting/lattice.py @@ -7,11 +7,10 @@ Module with functions to create lattice plots through a `~cpymad.madx.Madx` object. """ -import matplotlib -import matplotlib.axes import matplotlib.pyplot as plt from cpymad.madx import Madx from loguru import logger +from matplotlib.axes import Axes from pyhdtoolkit.plotting.layout import _ylim_from_input, plot_machine_layout from pyhdtoolkit.plotting.utils import ( @@ -203,7 +202,7 @@ def plot_machine_survey( show_elements: bool = False, high_orders: bool = False, **kwargs, -) -> matplotlib.axes.Axes: +) -> Axes: """ .. versionadded:: 1.0.0 diff --git a/pyhdtoolkit/plotting/layout.py b/pyhdtoolkit/plotting/layout.py index 4e610b4f..e9eeb4fb 100644 --- a/pyhdtoolkit/plotting/layout.py +++ b/pyhdtoolkit/plotting/layout.py @@ -8,13 +8,12 @@ object, mostly used in different `~pyhdtoolkit.plotting` modules. """ -import matplotlib -import matplotlib.axes import numpy as np import pandas as pd from cpymad.madx import Madx from loguru import logger from matplotlib import patches +from matplotlib.axes import Axes from pyhdtoolkit.plotting.utils import ( _get_twiss_table_with_offsets_and_limits, @@ -355,7 +354,7 @@ def scale_patches(scale: float, ylabel: str, **kwargs) -> None: def _plot_lattice_series( - ax: matplotlib.axes.Axes, + ax: Axes, series: pd.DataFrame, height: float = 1.0, v_offset: float = 0.0, diff --git a/pyhdtoolkit/plotting/phasespace.py b/pyhdtoolkit/plotting/phasespace.py index a69b49f2..5194a1df 100644 --- a/pyhdtoolkit/plotting/phasespace.py +++ b/pyhdtoolkit/plotting/phasespace.py @@ -6,12 +6,12 @@ Module with functions to create phase space plots through a `~cpymad.madx.Madx` object. """ -import matplotlib -import matplotlib.axes import numpy as np from cpymad.madx import Madx from loguru import logger from matplotlib import colors as mcolors +from matplotlib.axes import Axes +from matplotlib.figure import Figure from pyhdtoolkit.optics.twiss import courant_snyder_transform from pyhdtoolkit.plotting.utils import maybe_get_ax @@ -29,7 +29,7 @@ def plot_courant_snyder_phase_space( plane: str = "Horizontal", title: str | None = None, **kwargs, -) -> matplotlib.axes.Axes: +) -> Axes: """ .. versionadded:: 1.0.0 @@ -96,7 +96,7 @@ def plot_courant_snyder_phase_space_colored( plane: str = "Horizontal", title: str | None = None, **kwargs, -) -> matplotlib.figure.Figure: +) -> Figure: """ .. versionadded:: 1.0.0 diff --git a/pyhdtoolkit/plotting/sbs/coupling.py b/pyhdtoolkit/plotting/sbs/coupling.py index 9906e584..8dae8f57 100644 --- a/pyhdtoolkit/plotting/sbs/coupling.py +++ b/pyhdtoolkit/plotting/sbs/coupling.py @@ -7,10 +7,11 @@ Functions to plot coupling components of Segment-by-Segment results. """ -import matplotlib import matplotlib.pyplot as plt import tfs from loguru import logger +from matplotlib.axes import Axes +from matplotlib.figure import Figure from matplotlib.legend import _get_legend_handles_labels from pyhdtoolkit.plotting.utils import _determine_default_sbs_coupling_ylabel, find_ip_s_from_segment_start @@ -25,7 +26,7 @@ def plot_rdt_component( rdt: str = "F1001", component: str = "ABS", **kwargs, -) -> matplotlib.figure.Figure: +) -> Figure: r""" .. versionadded:: 0.19.0 @@ -34,8 +35,6 @@ def plot_rdt_component( :ref:`segment-by-segment plotting ` example gallery. Args: - ax (matplotlib.axes.Axes): The `~matplotlib.axes.Axes` to plot on. Will get the current axis if no - `~matplotlib.axes.Axes` is given. b1_segment_df (tfs.TfsDataFrame): A `~tfs.TfsDataFrame` of the segment-by-segment coupling result for Beam 1 in the given segment. b2_segment_df (tfs.TfsDataFrame): A `~tfs.TfsDataFrame` of the segment-by-segment coupling result for @@ -110,7 +109,7 @@ def plot_full_ip_rdt( real_ylimits: tuple[float, float] | None = None, imag_ylimits: tuple[float, float] | None = None, **kwargs, -) -> matplotlib.figure.Figure: +) -> Figure: """ .. versionadded:: 0.19.0 @@ -208,7 +207,7 @@ def plot_full_ip_rdt( def _plot_sbs_coupling_rdt_component( - ax: matplotlib.axes.Axes, + ax: Axes, segment_df: tfs.TfsDataFrame, model_df: tfs.TfsDataFrame = None, ip: int | None = None, diff --git a/pyhdtoolkit/plotting/sbs/phase.py b/pyhdtoolkit/plotting/sbs/phase.py index 0cc53c2b..704f7e9d 100644 --- a/pyhdtoolkit/plotting/sbs/phase.py +++ b/pyhdtoolkit/plotting/sbs/phase.py @@ -6,10 +6,11 @@ Functions to plot phase values of Segment-by-Segment results. """ -import matplotlib import matplotlib.pyplot as plt import tfs from loguru import logger +from matplotlib.axes import Axes +from matplotlib.figure import Figure from matplotlib.legend import _get_legend_handles_labels from pyhdtoolkit.plotting.utils import find_ip_s_from_segment_start @@ -21,7 +22,7 @@ def plot_phase_segment_one_beam( model: tfs.TfsDataFrame = None, ip: int | None = None, **kwargs, -) -> matplotlib.figure.Figure: +) -> Figure: """ .. versionadded:: 0.19.0 @@ -73,7 +74,7 @@ def plot_phase_segment_both_beams( b2_model: tfs.TfsDataFrame = None, ip: int | None = None, **kwargs, -) -> matplotlib.figure.Figure: +) -> Figure: """ .. versionadded:: 0.19.0 @@ -150,7 +151,7 @@ def plot_phase_segment_both_beams( def plot_phase_segment( - ax: matplotlib.axes.Axes = None, + ax: Axes = None, segment_df: tfs.TfsDataFrame = None, model_df: tfs.TfsDataFrame = None, plane: str = "x", diff --git a/pyhdtoolkit/plotting/styles/__init__.py b/pyhdtoolkit/plotting/styles/__init__.py index 5a7c2d83..1af83149 100644 --- a/pyhdtoolkit/plotting/styles/__init__.py +++ b/pyhdtoolkit/plotting/styles/__init__.py @@ -8,13 +8,13 @@ Feel free to use them anyway, as they might be useful to you when using the `~pyhdtoolkit.plotting` submodules, or to be adapted. """ from pathlib import Path -from typing import Dict, Union +from typing import Union -import matplotlib +import matplotlib as mpl import matplotlib.pyplot as plt from loguru import logger -from . import paper, thesis +from . import paper, thesis # noqa: TID252 PlotSetting = Union[float, bool, str, tuple] @@ -75,7 +75,7 @@ def _install_style_file(style: dict[str, PlotSetting], stylename) -> None: """ logger.info(f"Installing matplotlib style as '{stylename}'") style_content: str = "\n".join(f"{option} : {setting}" for option, setting in style.items()) - mpl_config_stylelib = Path(matplotlib.get_configdir()) / "stylelib" + mpl_config_stylelib = Path(mpl.get_configdir()) / "stylelib" mpl_env_stylelib = Path(plt.style.core.BASE_LIBRARY_PATH) logger.debug("Ensuring matplotlib 'stylelib' directory exists") diff --git a/pyhdtoolkit/plotting/tune.py b/pyhdtoolkit/plotting/tune.py index 5825e5c8..229eb893 100644 --- a/pyhdtoolkit/plotting/tune.py +++ b/pyhdtoolkit/plotting/tune.py @@ -9,10 +9,9 @@ """ from functools import partial -import matplotlib -import matplotlib.axes import numpy as np from loguru import logger +from matplotlib.axes import Axes from pyhdtoolkit.plotting.utils import maybe_get_ax @@ -74,7 +73,7 @@ def plot_tune_diagram( max_order: int = 6, differentiate_orders: bool = False, **kwargs, -) -> matplotlib.axes.Axes: +) -> Axes: """ .. versionadded:: 1.0.0 @@ -138,7 +137,7 @@ def plot_tune_diagram( return axis -def plot_resonance_lines_for_order(order: int, axis: matplotlib.axes.Axes, **kwargs) -> None: +def plot_resonance_lines_for_order(order: int, axis: Axes, **kwargs) -> None: """ .. versionadded:: 1.0.0 diff --git a/pyhdtoolkit/plotting/utils.py b/pyhdtoolkit/plotting/utils.py index 28123f97..0fe0bd9b 100644 --- a/pyhdtoolkit/plotting/utils.py +++ b/pyhdtoolkit/plotting/utils.py @@ -8,17 +8,21 @@ """ from __future__ import annotations # important for Sphinx to generate short type signatures! -import matplotlib -import matplotlib.axes +from typing import TYPE_CHECKING + +import matplotlib as mpl import matplotlib.pyplot as plt import numpy as np -import pandas as pd -import tfs -from cpymad.madx import Madx from loguru import logger from matplotlib import transforms from matplotlib.patches import Ellipse -from numpy.typing import ArrayLike + +if TYPE_CHECKING: + import pandas as pd + from cpymad.madx import Madx + from matplotlib.text import Annotation + from numpy.typing import ArrayLike + from tfs import TfsDataFrame # ------ General Utilities ----- # @@ -63,7 +67,7 @@ def my_plotting_function(*args, **kwargs): return ax, dict(kwargs) -def find_ip_s_from_segment_start(segment_df: tfs.TfsDataFrame, model_df: tfs.TfsDataFrame, ip: int) -> float: +def find_ip_s_from_segment_start(segment_df: TfsDataFrame, model_df: TfsDataFrame, ip: int) -> float: """ .. versionadded:: 0.19.0 @@ -261,7 +265,7 @@ def draw_ip_locations( logger.debug(f"Drawing name indicator for {ip_name}") # drawing ypos is lower end of ylimits if drawing inside, higher end if drawing outside ypos = ylimits[not inside] + (ylimits[1] + ylimits[0]) * 0.01 - c = "grey" if inside else matplotlib.rcParams["text.color"] # match axis ticks color + c = "grey" if inside else mpl.rcParams["text.color"] # match axis ticks color fontsize = plt.rcParams["xtick.labelsize"] # match the xticks size axis.text(ip_xpos, ypos, ip_name, color=c, ha="center", va="bottom", size=fontsize) @@ -280,7 +284,7 @@ def set_arrow_label( arrow_arc_rad: float = -0.2, fontsize: int = 20, **kwargs, -) -> matplotlib.text.Annotation: +) -> Annotation: """ .. versionadded:: 0.6.0 diff --git a/pyhdtoolkit/utils/__init__.py b/pyhdtoolkit/utils/__init__.py index 7fb97983..381b146f 100644 --- a/pyhdtoolkit/utils/__init__.py +++ b/pyhdtoolkit/utils/__init__.py @@ -4,9 +4,9 @@ import warnings from collections.abc import Callable -from . import cmdline, contexts, executors, htc_monitor, logging, operations, printutil +from . import cmdline, contexts, executors, htc_monitor, logging, operations, printutil # noqa: TID252 -__all__ = [cmdline, contexts, executors, htc_monitor, logging, operations, printutil] +__all__ = ["cmdline", "contexts", "executors", "htc_monitor", "logging", "operations", "printutil"] # ----- Utility deprecation decorator ----- # From c23907d1dccb6de80e0417bfec6a81e0a08cc29f Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 11:14:52 +0200 Subject: [PATCH 06/41] fix use of literal dicts --- pyhdtoolkit/cpymadtools/tune.py | 26 +++++++++++++------------- pyhdtoolkit/optics/rdt.py | 2 +- pyhdtoolkit/plotting/utils.py | 4 ++-- pyhdtoolkit/utils/htc_monitor.py | 32 ++++++++++++++++---------------- 4 files changed, 32 insertions(+), 32 deletions(-) diff --git a/pyhdtoolkit/cpymadtools/tune.py b/pyhdtoolkit/cpymadtools/tune.py index 8dcb7b01..bdec62f7 100644 --- a/pyhdtoolkit/cpymadtools/tune.py +++ b/pyhdtoolkit/cpymadtools/tune.py @@ -93,19 +93,19 @@ def make_footprint_table( tfs_dframe = tfs.TfsDataFrame( data=madx.table.dynaptune.dframe(), - headers=dict( - NAME="DYNAPTUNE", - TYPE="DYNAPTUNE", - TITLE="FOOTPRINT TABLE", - MADX_VERSION=str(madx.version).upper(), - ORIGIN="pyhdtoolkit.cpymadtools.tune.make_footprint_table() function", - ANGLE=7, # default of the function - AMPLITUDE=sigma, - DSIGMA=1 if not dense else 0.5, - ANGLE_MEANING="Number of different starting angles used for each starting amplitude", - AMPLITUDE_MEANING="Up to which bunch sigma the starting amplitudes were ramped up", - DSIGMA_MEANING="Increment value of AMPLITUDE at each new starting amplitude", - ), + headers={ + "NAME": "DYNAPTUNE", + "TYPE": "DYNAPTUNE", + "TITLE": "FOOTPRINT TABLE", + "MADX_VERSION": str(madx.version).upper(), + "ORIGIN": "pyhdtoolkit.cpymadtools.tune.make_footprint_table() function", + "ANGLE": 7, # default of the function + "AMPLITUDE": sigma, + "DSIGMA": 1 if not dense else 0.5, + "ANGLE_MEANING": "Number of different starting angles used for each starting amplitude", + "AMPLITUDE_MEANING": "Up to which bunch sigma the starting amplitudes were ramped up", + "DSIGMA_MEANING": "Increment value of AMPLITUDE at each new starting amplitude", + }, ) tfs_dframe = tfs_dframe.reset_index(drop=True) diff --git a/pyhdtoolkit/optics/rdt.py b/pyhdtoolkit/optics/rdt.py index 3954abc9..87cd58bf 100644 --- a/pyhdtoolkit/optics/rdt.py +++ b/pyhdtoolkit/optics/rdt.py @@ -54,5 +54,5 @@ def determine_rdt_line(rdt: int | str, plane: str) -> tuple[int, int, int]: which means the line located at 1 * Qy = Qy. """ j, k, l, m = map(int, str(rdt)) # noqa: E741 - lines = dict(X=(1 - j + k, m - l, 0), Y=(k - j, 1 - l + m, 0)) + lines = {"X": (1 - j + k, m - l, 0), "Y": (k - j, 1 - l + m, 0)} return lines[plane] diff --git a/pyhdtoolkit/plotting/utils.py b/pyhdtoolkit/plotting/utils.py index 0fe0bd9b..55342b9d 100644 --- a/pyhdtoolkit/plotting/utils.py +++ b/pyhdtoolkit/plotting/utils.py @@ -331,8 +331,8 @@ def set_arrow_label( color=color, va="center", ha="center", - bbox=dict(boxstyle="round4", fc="w", color=color), - arrowprops=dict(arrowstyle="-|>", connectionstyle="arc3,rad=" + str(arrow_arc_rad), fc="w", color=color), + bbox={"boxstyle": "round4", "fc": "w", "color": color}, + arrowprops={"arrowstyle": "-|>", "connectionstyle": "arc3,rad=" + str(arrow_arc_rad), "fc": "w", "color": color}, **kwargs, ) diff --git a/pyhdtoolkit/utils/htc_monitor.py b/pyhdtoolkit/utils/htc_monitor.py index 49429433..552ddefc 100644 --- a/pyhdtoolkit/utils/htc_monitor.py +++ b/pyhdtoolkit/utils/htc_monitor.py @@ -32,25 +32,25 @@ # ----- Data ----- # TASK_COLUMNS_SETTINGS = { - "OWNER": dict(justify="left", header_style="bold", style="bold", no_wrap=True), - "BATCH_NAME": dict(justify="center", header_style="magenta", style="magenta", no_wrap=True), - "SUBMITTED": dict(justify="center", header_style="medium_turquoise", style="medium_turquoise", no_wrap=True), - "DONE": dict(justify="right", header_style="bold green3", style="bold green3", no_wrap=True), - "RUNNING": dict(justify="right", header_style="bold cornflower_blue", style="bold cornflower_blue", no_wrap=True), - "IDLE": dict(justify="right", header_style="bold dark_orange3", style="bold dark_orange3", no_wrap=True), - "TOTAL": dict(justify="right", style="bold", no_wrap=True), - "JOB_IDS": dict(justify="right", no_wrap=True), + "OWNER": {"justify": "left", "header_style": "bold", "style": "bold", "no_wrap": True}, + "BATCH_NAME": {"justify": "center", "header_style": "magenta", "style": "magenta", "no_wrap": True}, + "SUBMITTED": {"justify": "center", "header_style": "medium_turquoise", "style": "medium_turquoise", "no_wrap": True}, + "DONE": {"justify": "right", "header_style": "bold green3", "style": "bold green3", "no_wrap": True}, + "RUNNING": {"justify": "right", "header_style": "bold cornflower_blue", "style": "bold cornflower_blue", "no_wrap": True}, + "IDLE": {"justify": "right", "header_style": "bold dark_orange3", "style": "bold dark_orange3", "no_wrap": True}, + "TOTAL": {"justify": "right", "style": "bold", "no_wrap": True}, + "JOB_IDS": {"justify": "right", "no_wrap": True}, } CLUSTER_COLUMNS_SETTINGS = { - "SOURCE": dict(justify="left", header_style="bold", style="bold", no_wrap=True), - "JOBS": dict(justify="right", header_style="bold", style="bold", no_wrap=True), - "COMPLETED": dict(justify="right", header_style="bold green3", style="bold green3", no_wrap=True), - "RUNNING": dict(justify="right", header_style="bold cornflower_blue", style="bold cornflower_blue", no_wrap=True), - "IDLE": dict(justify="right", header_style="bold dark_orange3", style="bold dark_orange3", no_wrap=True), - "HELD": dict(justify="right", header_style="bold gold1", style="bold gold1", no_wrap=True), - "SUSPENDED": dict(justify="right", header_style="bold slate_blue1", style="bold slate_blue1", no_wrap=True), - "REMOVED": dict(justify="right", header_style="bold red3", style="bold red3", no_wrap=True), + "SOURCE": {"justify": "left", "header_style": "bold", "style": "bold", "no_wrap": True}, + "JOBS": {"justify": "right", "header_style": "bold", "style": "bold", "no_wrap": True}, + "COMPLETED": {"justify": "right", "header_style": "bold green3", "style": "bold green3", "no_wrap": True}, + "RUNNING": {"justify": "right", "header_style": "bold cornflower_blue", "style": "bold cornflower_blue", "no_wrap": True}, + "IDLE": {"justify": "right", "header_style": "bold dark_orange3", "style": "bold dark_orange3", "no_wrap": True}, + "HELD": {"justify": "right", "header_style": "bold gold1", "style": "bold gold1", "no_wrap": True}, + "SUSPENDED": {"justify": "right", "header_style": "bold slate_blue1", "style": "bold slate_blue1", "no_wrap": True}, + "REMOVED": {"justify": "right", "header_style": "bold red3", "style": "bold red3", "no_wrap": True}, } From 5ffb15506dac97b8e5f8f451f081c15a69e326a4 Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 11:15:38 +0200 Subject: [PATCH 07/41] propagate error change here --- tests/test_cpymadtools/test_lhc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_cpymadtools/test_lhc.py b/tests/test_cpymadtools/test_lhc.py index a3f26082..0bfae377 100644 --- a/tests/test_cpymadtools/test_lhc.py +++ b/tests/test_cpymadtools/test_lhc.py @@ -367,7 +367,7 @@ def test_landau_powering(current, _non_matched_lhc_madx): def test_landau_powering_fails_on_missing_nrj(caplog): madx = Madx(stdout=False) - with pytest.raises(EnvironmentError): + with pytest.raises(AttributeError): power_landau_octupoles(madx, 100, 1) for record in caplog.records: From a8b7d0fff79a75f09252a32d43e4be616f89d447 Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 11:16:53 +0200 Subject: [PATCH 08/41] fix trailing whitespaces --- pyhdtoolkit/cpymadtools/lhc/_setup.py | 2 +- pyhdtoolkit/cpymadtools/ptc.py | 4 ++-- pyhdtoolkit/plotting/layout.py | 2 +- pyhdtoolkit/utils/_misc.py | 2 +- pyhdtoolkit/utils/htc_monitor.py | 6 +++--- pyhdtoolkit/utils/operations.py | 2 +- pyhdtoolkit/utils/printutil.py | 4 ++-- 7 files changed, 11 insertions(+), 11 deletions(-) diff --git a/pyhdtoolkit/cpymadtools/lhc/_setup.py b/pyhdtoolkit/cpymadtools/lhc/_setup.py index 29396001..40e181f1 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_setup.py +++ b/pyhdtoolkit/cpymadtools/lhc/_setup.py @@ -37,7 +37,7 @@ def prepare_lhc_run2( might change that working point. Args: - opticsfile (str): the relative string path or a `Path` object to the opticsfile location. This will + opticsfile (str): the relative string path or a `Path` object to the opticsfile location. This will be used to determine the location of the sequence file, see the admonition above. beam (int): which beam to set up for. Defaults to beam 1. use_b4 (bool): if `True`, the lhcb4 sequence file will be used. This is the beam 2 sequence but for tracking diff --git a/pyhdtoolkit/cpymadtools/ptc.py b/pyhdtoolkit/cpymadtools/ptc.py index 2a2ab163..cb809e9a 100644 --- a/pyhdtoolkit/cpymadtools/ptc.py +++ b/pyhdtoolkit/cpymadtools/ptc.py @@ -162,7 +162,7 @@ def get_rdts( arguments to override them. The ``PTC_TWISS`` command is given ``icase=6`` by default to enforce 6D - calculations (see the + calculations (see the `MAD-X manual `_ for details), and ``normal=True`` to trigger saving the normal form analysis results in a table called ``NONLIN`` which will then be available through the @@ -263,7 +263,7 @@ def ptc_twiss( arguments to override them. The ``PTC_TWISS`` command is given ``icase=6`` by default to enforce 6D - calculations (see the + calculations (see the `MAD-X manual `_ for details), and ``normal=True`` to trigger saving the normal form analysis results in a table called ``NONLIN`` which will then be available through the diff --git a/pyhdtoolkit/plotting/layout.py b/pyhdtoolkit/plotting/layout.py index e9eeb4fb..ba729b2c 100644 --- a/pyhdtoolkit/plotting/layout.py +++ b/pyhdtoolkit/plotting/layout.py @@ -329,7 +329,7 @@ def scale_patches(scale: float, ylabel: str, **kwargs) -> None: patches as well as the corresponding y-axis label. Args: - scale (float): the scale factor to apply to the patches. The new + scale (float): the scale factor to apply to the patches. The new height of the patches will be ``scale * original_height``. ylabel (str): the new label for the y-axis. **kwargs: If either `ax` or `axis` is found in the kwargs, the diff --git a/pyhdtoolkit/utils/_misc.py b/pyhdtoolkit/utils/_misc.py index 3943e111..18369166 100644 --- a/pyhdtoolkit/utils/_misc.py +++ b/pyhdtoolkit/utils/_misc.py @@ -7,7 +7,7 @@ Private module that provides miscellaneous personnal utility functions. .. warning:: - The functions in here are intented for personal use, and will most likely + The functions in here are intented for personal use, and will most likely **not** work on other people's machines. """ import shlex diff --git a/pyhdtoolkit/utils/htc_monitor.py b/pyhdtoolkit/utils/htc_monitor.py index 552ddefc..317f901b 100644 --- a/pyhdtoolkit/utils/htc_monitor.py +++ b/pyhdtoolkit/utils/htc_monitor.py @@ -4,12 +4,12 @@ HTCondor Monitoring ------------------- -A module with utility to query the HTCondor queue, process the returned data +A module with utility to query the HTCondor queue, process the returned data and display it nicely. .. note:: - This module is meant to be called as a script, but some of the individual - functionality is made public API and one shoule be able to build a different + This module is meant to be called as a script, but some of the individual + functionality is made public API and one shoule be able to build a different monitor from the functions in here. """ import re diff --git a/pyhdtoolkit/utils/operations.py b/pyhdtoolkit/utils/operations.py index b28f42a2..39721b99 100644 --- a/pyhdtoolkit/utils/operations.py +++ b/pyhdtoolkit/utils/operations.py @@ -4,7 +4,7 @@ Operations Utilities -------------------- -A collection classes with utility functions to perform common / convenient +A collection classes with utility functions to perform common / convenient operations on the classic Python structures. .. warning:: diff --git a/pyhdtoolkit/utils/printutil.py b/pyhdtoolkit/utils/printutil.py index 09f3eb77..4a87c632 100644 --- a/pyhdtoolkit/utils/printutil.py +++ b/pyhdtoolkit/utils/printutil.py @@ -7,8 +7,8 @@ A class utility class to allow me printing text in color, bold, etc. .. warning:: - This module's functionality is superseeded by the ``rich`` package, and it will be - removed in a future release. + This module's functionality is superseeded by the ``rich`` package, and it + will be removed in a future release. """ END = "\033[0m" From 39c3dc4a7ecbc79cc71eaf63288971d37f972627 Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 11:19:09 +0200 Subject: [PATCH 09/41] miscellaneous --- pyhdtoolkit/utils/htc_monitor.py | 3 ++- pyhdtoolkit/utils/logging.py | 6 +++--- pyhdtoolkit/utils/operations.py | 6 +++--- 3 files changed, 8 insertions(+), 7 deletions(-) diff --git a/pyhdtoolkit/utils/htc_monitor.py b/pyhdtoolkit/utils/htc_monitor.py index 317f901b..ce9cadaa 100644 --- a/pyhdtoolkit/utils/htc_monitor.py +++ b/pyhdtoolkit/utils/htc_monitor.py @@ -69,7 +69,8 @@ def query_condor_q() -> str: if return_code == 0: return condor_status else: - raise ChildProcessError("Checking htcondor status failed") + msg = "Checking htcondor status failed" + raise ChildProcessError(msg) def read_condor_q(report: str) -> tuple[list[HTCTaskSummary], ClusterSummary]: diff --git a/pyhdtoolkit/utils/logging.py b/pyhdtoolkit/utils/logging.py index 6bc4117a..ddf6ef7f 100644 --- a/pyhdtoolkit/utils/logging.py +++ b/pyhdtoolkit/utils/logging.py @@ -35,7 +35,7 @@ SIMPLE_FORMAT = "{time:HH:mm:ss} | {level: <8} | {message}" -def config_logger(level: str | int = "INFO", format: str = FORMAT1, **kwargs) -> None: +def config_logger(level: str | int = "INFO", fmt: str = FORMAT1, **kwargs) -> None: """ .. versionadded:: 0.8.2 @@ -47,7 +47,7 @@ def config_logger(level: str | int = "INFO", format: str = FORMAT1, **kwargs) -> string is given. Can be any of the `loguru levels `_ or their integer values equivalents. Defaults to ``INFO``. - format (str): the format to use for the logger to display messages. Defaults + fmt (str): the format to use for the logger to display messages. Defaults to a pre-defined format in this module. **kwargs: any keyword argument is transmitted to the `~loguru._logger.Logger.add` call. @@ -67,4 +67,4 @@ def config_logger(level: str | int = "INFO", format: str = FORMAT1, **kwargs) -> """ logger.remove() level = level.upper() if isinstance(level, str) else level - logger.add(sys.stdout, level=level, format=format, **kwargs) + logger.add(sys.stdout, level=level, format=fmt, **kwargs) diff --git a/pyhdtoolkit/utils/operations.py b/pyhdtoolkit/utils/operations.py index 39721b99..9242422f 100644 --- a/pyhdtoolkit/utils/operations.py +++ b/pyhdtoolkit/utils/operations.py @@ -340,7 +340,7 @@ def sample(sequence: Sequence) -> list: ListOperations.sample(["a", "b", 1, 2, False]) # returns 2 """ - return sequence[random.randint(0, len(sequence) - 1)] + return sequence[random.randint(0, len(sequence) - 1)] # noqa: S311 @staticmethod def sanitize_list(sequence: Sequence) -> list: @@ -389,7 +389,7 @@ def shuffle(sequence: Sequence) -> Sequence: temp_list = copy.deepcopy(sequence) amount_to_shuffle = len(temp_list) while amount_to_shuffle > 1: - rand_index = int(math.floor(random.random() * amount_to_shuffle)) + rand_index = int(math.floor(random.random() * amount_to_shuffle)) # noqa: S311 amount_to_shuffle -= 1 temp_list[rand_index], temp_list[amount_to_shuffle] = ( temp_list[amount_to_shuffle], @@ -494,7 +494,7 @@ def union_by(seq_1: Sequence, seq_2: Sequence, function: Callable) -> list: # returns [1.2, 2.1] """ _lst_1 = set(map(function, seq_1)) - return sorted(list(set(seq_1 + [item for item in seq_2 if function(item) not in _lst_1]))) + return sorted(set(seq_1 + [item for item in seq_2 if function(item) not in _lst_1])) @staticmethod def zipper(*args, fillvalue=None) -> list: From 06cde549a31d14714d1d883c7ca56359e0fd6570 Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 11:23:25 +0200 Subject: [PATCH 10/41] fix error messages not declared before raising --- pyhdtoolkit/cpymadtools/lhc/_errors.py | 9 ++++++--- pyhdtoolkit/cpymadtools/lhc/_powering.py | 9 ++++++--- pyhdtoolkit/cpymadtools/lhc/_setup.py | 15 ++++++++++----- pyhdtoolkit/cpymadtools/ptc.py | 3 ++- pyhdtoolkit/cpymadtools/tune.py | 3 ++- pyhdtoolkit/cpymadtools/utils.py | 3 ++- pyhdtoolkit/optics/ripken.py | 3 ++- pyhdtoolkit/plotting/aperture.py | 3 ++- pyhdtoolkit/plotting/envelope.py | 3 ++- pyhdtoolkit/plotting/phasespace.py | 6 ++++-- pyhdtoolkit/plotting/sbs/phase.py | 3 ++- pyhdtoolkit/plotting/tune.py | 3 ++- 12 files changed, 42 insertions(+), 21 deletions(-) diff --git a/pyhdtoolkit/cpymadtools/lhc/_errors.py b/pyhdtoolkit/cpymadtools/lhc/_errors.py index ac10f4e3..4f1f7ebb 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_errors.py +++ b/pyhdtoolkit/cpymadtools/lhc/_errors.py @@ -143,13 +143,16 @@ def misalign_lhc_ir_quadrupoles( """ if any(ip not in (1, 2, 5, 8) for ip in ips): logger.error("The IP number provided is invalid, not applying any error.") - raise ValueError("Invalid 'ips' parameter") + msg = "Invalid 'ips' parameter" + raise ValueError(msg) if beam and beam not in (1, 2, 3, 4): logger.error("The beam number provided is invalid, not applying any error.") - raise ValueError("Invalid 'beam' parameter") + msg = "Invalid 'beam' parameter" + raise ValueError(msg) if any(side.upper() not in ("R", "L") for side in sides): logger.error("The side provided is invalid, not applying any error.") - raise ValueError("Invalid 'sides' parameter") + msg = "Invalid 'sides' parameter" + raise ValueError(msg) sides = [side.upper() for side in sides] logger.debug("Clearing error flag") diff --git a/pyhdtoolkit/cpymadtools/lhc/_powering.py b/pyhdtoolkit/cpymadtools/lhc/_powering.py index 9bbf7af1..f3a8d421 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_powering.py +++ b/pyhdtoolkit/cpymadtools/lhc/_powering.py @@ -141,7 +141,8 @@ def apply_lhc_rigidity_waist_shift_knob( madx.globals[left_knob] = (1 - rigidty_waist_shift_value * 0.005) * current_left_knob else: logger.error(f"Given side '{side}' invalid, only 'left' and 'right' are accepted values.") - raise ValueError("Invalid value for parameter 'side'.") + msg = "Invalid value for parameter 'side'." + raise ValueError(msg) logger.debug(f"Set '{right_knob}' to {madx.globals[right_knob]}") logger.debug(f"Set '{left_knob}' to {madx.globals[left_knob]}") @@ -237,7 +238,8 @@ def power_landau_octupoles(madx: Madx, /, beam: int, mo_current: float, defectiv brho = madx.globals.nrj * 1e9 / madx.globals.clight # clight is MAD-X constant except AttributeError as madx_error: logger.exception("The global MAD-X variable 'NRJ' should have been set in the optics files but is not defined.") - raise AttributeError("No 'NRJ' variable found in scripts") from madx_error + msg = "No 'NRJ' variable found in scripts" + raise AttributeError(msg) from madx_error logger.debug(f"Powering Landau Octupoles, beam {beam} @ {madx.globals.nrj} GeV with {mo_current} A.") strength = mo_current / madx.globals.Imax_MO * madx.globals.Kmax_MO / brho @@ -317,7 +319,8 @@ def vary_independent_ir_quadrupoles( or any(quad not in (4, 5, 6, 7, 8, 9, 10, 11, 12, 13) for quad in quad_numbers) ): logger.error("Either the IP number of the side provided are invalid, not applying any error.") - raise ValueError("Invalid 'quad_numbers', 'ip', 'sides' argument") + msg = "Invalid 'quad_numbers', 'ip', 'sides' argument" + raise ValueError(msg) logger.debug(f"Preparing a knob involving quadrupoles {quad_numbers}") # Each quad has a specific power circuit used for their k1 boundaries diff --git a/pyhdtoolkit/cpymadtools/lhc/_setup.py b/pyhdtoolkit/cpymadtools/lhc/_setup.py index 40e181f1..496c9a01 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_setup.py +++ b/pyhdtoolkit/cpymadtools/lhc/_setup.py @@ -61,7 +61,8 @@ def prepare_lhc_run2( """ if use_b4 and beam != 2: logger.error("Cannot use beam 4 sequence file for beam 1") - raise ValueError("Cannot use beam 4 sequence file for beam 1") + msg = "Cannot use beam 4 sequence file for beam 1" + raise ValueError(msg) def _run2_sequence_from_opticsfile(opticsfile: Path, use_b4: bool = False) -> Path: filename = "lhc_as-built.seq" if not use_b4 else "lhcb4_as-built.seq" @@ -137,7 +138,8 @@ def prepare_lhc_run3( """ if use_b4 and beam != 2: logger.error("Cannot use beam 4 sequence file for beam 1") - raise ValueError("Cannot use beam 4 sequence file for beam 1") + msg = "Cannot use beam 4 sequence file for beam 1" + raise ValueError(msg) logger.debug("Creating Run 3 setup MAD-X instance") echo, warn = kwargs.pop("echo", False), kwargs.pop("warn", False) @@ -238,10 +240,12 @@ def __init__( assert opticsfile is not None, "An opticsfile must be provided" if use_b4 and beam != 2: logger.error("Cannot use beam 4 sequence file for beam 1") - raise ValueError("Cannot use beam 4 sequence file for beam 1") + msg = "Cannot use beam 4 sequence file for beam 1" + raise ValueError(msg) if int(run) not in (2, 3): - raise NotImplementedError("This setup is only possible for Run 2 and Run 3 configurations.") + msg = "This setup is only possible for Run 2 and Run 3 configurations." + raise NotImplementedError(msg) elif run == 2: self.madx = prepare_lhc_run2( opticsfile=opticsfile, beam=beam, use_b4=use_b4, energy=energy, slicefactor=slicefactor, **kwargs @@ -496,7 +500,8 @@ def setup_lhc_orbit(madx: Madx, /, scheme: str = "flat", **kwargs) -> dict[str, """ if scheme not in LHC_CROSSING_SCHEMES.keys(): logger.error(f"Invalid scheme parameter, should be one of {LHC_CROSSING_SCHEMES.keys()}") - raise ValueError("Invalid scheme parameter given") + msg = "Invalid scheme parameter given" + raise ValueError(msg) logger.debug("Getting orbit variables") variables, special = lhc_orbit_variables() diff --git a/pyhdtoolkit/cpymadtools/ptc.py b/pyhdtoolkit/cpymadtools/ptc.py index cb809e9a..10ccf585 100644 --- a/pyhdtoolkit/cpymadtools/ptc.py +++ b/pyhdtoolkit/cpymadtools/ptc.py @@ -78,7 +78,8 @@ def get_amplitude_detuning( """ if order >= 3: logger.error(f"Maximum amplitude detuning order in PTC is 2, but {order:d} was requested") - raise NotImplementedError("PTC amplitude detuning is not implemented for order > 2") + msg = "PTC amplitude detuning is not implemented for order > 2" + raise NotImplementedError(msg) logger.debug("Looking for PTC universe parameters in keyword arguments") model = kwargs.pop("model", 3) diff --git a/pyhdtoolkit/cpymadtools/tune.py b/pyhdtoolkit/cpymadtools/tune.py index bdec62f7..759abc90 100644 --- a/pyhdtoolkit/cpymadtools/tune.py +++ b/pyhdtoolkit/cpymadtools/tune.py @@ -80,7 +80,8 @@ def make_footprint_table( "Remote MAD-X process crashed, most likely because you did not slice the sequence " "before running DYNAP. Restart and slice before calling this function." ) - raise RuntimeError("DYNAP command crashed the MAD-X process") from madx_crash + msg = "DYNAP command crashed the MAD-X process" + raise RuntimeError(msg) from madx_crash if cleanup and sys.platform not in ("win32", "cygwin"): # fails on Windows due to its I/O system, since MAD-X still has "control" of the files diff --git a/pyhdtoolkit/cpymadtools/utils.py b/pyhdtoolkit/cpymadtools/utils.py index 153dab8e..e93cfd2a 100644 --- a/pyhdtoolkit/cpymadtools/utils.py +++ b/pyhdtoolkit/cpymadtools/utils.py @@ -124,7 +124,8 @@ def _get_k_strings(start: int = 0, stop: int = 8, orientation: str = "both") -> """ if orientation not in ("straight", "skew", "both"): logger.error(f"Orientation '{orientation}' is not accepted, should be one of 'straight', 'skew', 'both'.") - raise ValueError("Invalid 'orientation' parameter") + msg = "Invalid 'orientation' parameter" + raise ValueError(msg) if orientation == "straight": orientation = ("",) diff --git a/pyhdtoolkit/optics/ripken.py b/pyhdtoolkit/optics/ripken.py index 451bdd5c..8693422c 100644 --- a/pyhdtoolkit/optics/ripken.py +++ b/pyhdtoolkit/optics/ripken.py @@ -78,7 +78,8 @@ def _beam_size(coordinates_distribution: np.ndarray, method: str = "std") -> flo return coordinates_distribution.std() elif method == "rms": return np.sqrt(np.mean(np.square(coordinates_distribution))) - raise NotImplementedError("Invalid method provided") + msg = "Invalid method provided" + raise NotImplementedError(msg) def _add_beam_size_to_df(df: tfs.TfsDataFrame, geom_emit_x: float, geom_emit_y: float) -> tfs.TfsDataFrame: diff --git a/pyhdtoolkit/plotting/aperture.py b/pyhdtoolkit/plotting/aperture.py index feecb93a..e723c2fa 100644 --- a/pyhdtoolkit/plotting/aperture.py +++ b/pyhdtoolkit/plotting/aperture.py @@ -253,7 +253,8 @@ def plot_physical_apertures( # pylint: disable=too-many-arguments if plane.lower() not in ("x", "y", "horizontal", "vertical"): logger.error(f"'plane' argument should be 'x', 'horizontal', 'y' or 'vertical' not '{plane}'") - raise ValueError("Invalid 'plane' argument.") + msg = "Invalid 'plane' argument." + raise ValueError(msg) logger.debug("Plotting real element apertures") axis, kwargs = maybe_get_ax(**kwargs) diff --git a/pyhdtoolkit/plotting/envelope.py b/pyhdtoolkit/plotting/envelope.py index 904f033a..dcc054fd 100644 --- a/pyhdtoolkit/plotting/envelope.py +++ b/pyhdtoolkit/plotting/envelope.py @@ -83,7 +83,8 @@ def plot_beam_envelope( # pylint: disable=too-many-arguments if plane.lower() not in ("x", "y", "horizontal", "vertical"): logger.error(f"'plane' argument should be 'x', 'horizontal', 'y' or 'vertical' not '{plane}'") - raise ValueError("Invalid 'plane' argument.") + msg = "Invalid 'plane' argument." + raise ValueError(msg) logger.debug(f"Plotting machine orbit and {nsigma:.2f}sigma beam envelope") axis, kwargs = maybe_get_ax(**kwargs) diff --git a/pyhdtoolkit/plotting/phasespace.py b/pyhdtoolkit/plotting/phasespace.py index 5194a1df..ba182f6a 100644 --- a/pyhdtoolkit/plotting/phasespace.py +++ b/pyhdtoolkit/plotting/phasespace.py @@ -61,7 +61,8 @@ def plot_courant_snyder_phase_space( """ if plane.lower() not in ("horizontal", "vertical"): logger.error(f"Plane should be either Horizontal or Vertical but '{plane}' was given") - raise ValueError("Invalid plane value") + msg = "Invalid 'plane' value." + raise ValueError(msg) logger.debug("Plotting phase space for normalized Courant-Snyder coordinates") axis, kwargs = maybe_get_ax(**kwargs) @@ -131,7 +132,8 @@ def plot_courant_snyder_phase_space_colored( """ if plane.upper() not in ("HORIZONTAL", "VERTICAL"): logger.error(f"Plane should be either horizontal or vertical but '{plane}' was given") - raise ValueError("Invalid plane value") + msg = "Invalid 'plane' value." + raise ValueError(msg) # Getting a sufficiently long array of colors to use colors = int(np.floor(len(u_coordinates) / 100)) * SORTED_COLORS diff --git a/pyhdtoolkit/plotting/sbs/phase.py b/pyhdtoolkit/plotting/sbs/phase.py index 704f7e9d..71b1d9a5 100644 --- a/pyhdtoolkit/plotting/sbs/phase.py +++ b/pyhdtoolkit/plotting/sbs/phase.py @@ -183,7 +183,8 @@ def plot_phase_segment( """ if plane.upper() not in ("X", "Y"): logger.error("The provided plane is invalid, should be either 'x' or 'y', case-insensitively.") - raise ValueError("Invalid 'plane' parameter") + msg = "Invalid 'plane' parameter" + raise ValueError(msg) plane = plane.upper() logger.debug(f"Plotting phase for plane {plane.upper()} over the segment") diff --git a/pyhdtoolkit/plotting/tune.py b/pyhdtoolkit/plotting/tune.py index 229eb893..dd2007f4 100644 --- a/pyhdtoolkit/plotting/tune.py +++ b/pyhdtoolkit/plotting/tune.py @@ -111,7 +111,8 @@ def plot_tune_diagram( """ if max_order > 6 or max_order < 1: logger.error("Plotting is not supported outside of 1st-6th order (and not recommended)") - raise ValueError("The 'max_order' argument should be between 1 and 6 included") + msg = "The 'max_order' argument should be between 1 and 6 included" + raise ValueError(msg) logger.debug(f"Plotting resonance lines up to {ORDER_TO_LABEL[max_order]}") axis, kwargs = maybe_get_ax(**kwargs) From 4e8c4a46eb03ac915a7e6fb14e4c763ea7adbdb4 Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 11:33:25 +0200 Subject: [PATCH 11/41] fix magic values use in comparisons --- pyhdtoolkit/cpymadtools/_generators.py | 6 +++--- pyhdtoolkit/cpymadtools/constants.py | 19 ++++++++++--------- pyhdtoolkit/cpymadtools/lhc/_elements.py | 3 ++- pyhdtoolkit/cpymadtools/lhc/_errors.py | 4 +++- pyhdtoolkit/cpymadtools/lhc/_misc.py | 5 +++-- pyhdtoolkit/cpymadtools/lhc/_powering.py | 13 ++++++++----- pyhdtoolkit/cpymadtools/lhc/_setup.py | 11 +++++++---- pyhdtoolkit/cpymadtools/ptc.py | 8 +++++--- pyhdtoolkit/cpymadtools/tune.py | 4 ++-- pyhdtoolkit/plotting/tune.py | 2 +- pyhdtoolkit/utils/htc_monitor.py | 2 +- 11 files changed, 45 insertions(+), 32 deletions(-) diff --git a/pyhdtoolkit/cpymadtools/_generators.py b/pyhdtoolkit/cpymadtools/_generators.py index 09e5d393..12c4f5a4 100644 --- a/pyhdtoolkit/cpymadtools/_generators.py +++ b/pyhdtoolkit/cpymadtools/_generators.py @@ -31,7 +31,7 @@ def generate_base_cas_lattice() -> str: """ return """ option, -info, -warn; -TITLE, ’CAS2019 Project Team 3’; +TITLE, 'CAS2019 Project Team 3'; ! PARAMETERS circumference = 1000.0; @@ -104,7 +104,7 @@ def generate_onesext_cas_lattice() -> str: """ return """ option, -info, -warn; -TITLE, ’CAS2019 Project Team 3’; +TITLE, 'CAS2019 Project Team 3'; ! PARAMETERS circumference = 1000.0; @@ -194,7 +194,7 @@ def generate_oneoct_cas_lattice() -> str: """ return """ option, -info, -warn; -TITLE, ’CAS2019 Project Team 3’; +TITLE, 'CAS2019 Project Team 3'; ! PARAMETERS circumference = 1000.0; diff --git a/pyhdtoolkit/cpymadtools/constants.py b/pyhdtoolkit/cpymadtools/constants.py index 19ddfd1b..9848db74 100644 --- a/pyhdtoolkit/cpymadtools/constants.py +++ b/pyhdtoolkit/cpymadtools/constants.py @@ -6,6 +6,7 @@ Specific constants to be used in `~.cpymadtools` functions, to help with consistency. """ +_MAX_SECTOR_VALUE: int = 8 # fmt: off DEFAULT_TWISS_COLUMNS: list[str] = ["name", "s", "x", "y", "l", "px", "py", "betx", "bety", "alfx", "alfy", @@ -113,7 +114,7 @@ # ----- LHC Arc Correctors Knobs ----- # LHC_KQTF_KNOBS: list[str] = [ # tune trims, focusing and defocusing families, for each beam - f"kqt{family}.a{sector}{sector+1 if sector < 8 else 1}.b{beam}" + f"kqt{family}.a{sector}{sector+1 if sector < _MAX_SECTOR_VALUE else 1}.b{beam}" for beam in [1, 2] for family in ["f", "d"] for sector in [1, 2, 3, 4, 5, 6, 7, 8] @@ -122,32 +123,32 @@ # skew quadrupoles in arc short straight sections LHC_KQS_KNOBS: list[str] = [f"kqs.r{ip}b1" for ip in [1, 3, 5, 7]] + \ [f"kqs.l{ip}b1" for ip in [2, 4, 6, 8]] + \ - [f"kqs.a{sector}{sector+1 if sector < 8 else 1}b1" for sector in [2, 4, 6, 8]] + \ + [f"kqs.a{sector}{sector+1 if sector < _MAX_SECTOR_VALUE else 1}b1" for sector in [2, 4, 6, 8]] + \ [f"kqs.r{ip}b2" for ip in [2, 4, 6, 8]] + \ [f"kqs.l" f"{ip}b2" for ip in [3, 5, 7, 1]] + \ - [f"kqs.a{sector}{sector+1 if sector < 8 else 1}b2" for sector in [1, 3, 5, 7]] + [f"kqs.a{sector}{sector+1 if sector < _MAX_SECTOR_VALUE else 1}b2" for sector in [1, 3, 5, 7]] # fmt: on LHC_KSF_KNOBS: list[str] = [ # sextupole correctors - f"ks{family}{id}.a{sector}{sector+1 if sector < 8 else 1}b{beam}" + f"ks{family}{id}.a{sector}{sector+1 if sector < _MAX_SECTOR_VALUE else 1}b{beam}" for beam in [1, 2] for id in [1, 2] for family in ["f", "d"] for sector in [1, 2, 3, 4, 5, 6, 7, 8] ] LHC_KSS_KNOBS: list[str] = [ # skew sextupole correctors - f"kss.a{sector}{sector+1 if sector < 8 else 1}b{beam}" for beam in [1, 2] for sector in [1, 2, 3, 4, 5, 6, 7, 8] + f"kss.a{sector}{sector+1 if sector < _MAX_SECTOR_VALUE else 1}b{beam}" for beam in [1, 2] for sector in [1, 2, 3, 4, 5, 6, 7, 8] ] LHC_KCS_KNOBS: list[str] = [ # spool piece (skew) sextupoles - f"kcs.a{sector}{sector+1 if sector < 8 else 1}b{beam}" for beam in [1, 2] for sector in [1, 2, 3, 4, 5, 6, 7, 8] + f"kcs.a{sector}{sector+1 if sector < _MAX_SECTOR_VALUE else 1}b{beam}" for beam in [1, 2] for sector in [1, 2, 3, 4, 5, 6, 7, 8] ] LHC_KCO_KNOBS: list[str] = [ # spool piece (skew) octupoles - f"kco.a{sector}{sector+1 if sector < 8 else 1}b{beam}" for beam in [1, 2] for sector in [1, 2, 3, 4, 5, 6, 7, 8] + f"kco.a{sector}{sector+1 if sector < _MAX_SECTOR_VALUE else 1}b{beam}" for beam in [1, 2] for sector in [1, 2, 3, 4, 5, 6, 7, 8] ] LHC_KCD_KNOBS: list[str] = [ # spool piece (skew) decapoles - f"kcd.a{sector}{sector+1 if sector < 8 else 1}b{beam}" for beam in [1, 2] for sector in [1, 2, 3, 4, 5, 6, 7, 8] + f"kcd.a{sector}{sector+1 if sector < _MAX_SECTOR_VALUE else 1}b{beam}" for beam in [1, 2] for sector in [1, 2, 3, 4, 5, 6, 7, 8] ] LHC_KO_KNOBS: list[str] = [ # octupoles in arc short straight sections - f"ko{family}.a{sector}{sector+1 if sector < 8 else 1}b{beam}" + f"ko{family}.a{sector}{sector+1 if sector < _MAX_SECTOR_VALUE else 1}b{beam}" for beam in [1, 2] for family in ["f", "d"] for sector in [1, 2, 3, 4, 5, 6, 7, 8] diff --git a/pyhdtoolkit/cpymadtools/lhc/_elements.py b/pyhdtoolkit/cpymadtools/lhc/_elements.py index 103dcbab..0358cf1c 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_elements.py +++ b/pyhdtoolkit/cpymadtools/lhc/_elements.py @@ -8,6 +8,7 @@ from cpymad.madx import Madx from loguru import logger +_MAX_TRACKING_TOP_TURNS: int = 6600 def install_ac_dipole_as_kicker( madx: Madx, @@ -93,7 +94,7 @@ def install_ac_dipole_as_kicker( logger.warning("This AC Dipole is implemented as a kicker and will not affect TWISS functions!") logger.debug("This routine should be done after 'match', 'twiss' and 'makethin' for the appropriate beam") - if top_turns > 6600: + if top_turns > _MAX_TRACKING_TOP_TURNS: logger.warning( f"Configuring the AC Dipole for {top_turns:d} of driving is fine for MAD-X but is " "higher than what the device can do in the (HL)LHC! Beware." diff --git a/pyhdtoolkit/cpymadtools/lhc/_errors.py b/pyhdtoolkit/cpymadtools/lhc/_errors.py index 4f1f7ebb..70484f41 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_errors.py +++ b/pyhdtoolkit/cpymadtools/lhc/_errors.py @@ -10,6 +10,8 @@ from cpymad.madx import Madx from loguru import logger +_MAX_TRIPLET_NUMBER: int = 3 + LHC_IR_QUADS_PATTERNS: dict[int, list[str]] = { 1: ["^MQXA.1{side}{ip:d}", "^MQXFA.[AB]1{side}{ip:d}"], # Q1 LHC, Q1A & Q1B HL-LHC 2: ["^MQXB.[AB]2{side}{ip:d}", "^MQXB.[AB]2{side}{ip:d}"], # Q2A & Q2B LHC, Q2A & Q2B HL-LHC @@ -165,7 +167,7 @@ def misalign_lhc_ir_quadrupoles( for quad_number in quadrupoles: for quad_pattern in LHC_IR_QUADS_PATTERNS[quad_number]: # Triplets are single aperture and don't need beam information, others do - if quad_number <= 3: + if quad_number <= _MAX_TRIPLET_NUMBER: madx.select(flag="error", pattern=quad_pattern.format(side=side, ip=ip)) else: madx.select(flag="error", pattern=quad_pattern.format(side=side, ip=ip, beam=beam)) diff --git a/pyhdtoolkit/cpymadtools/lhc/_misc.py b/pyhdtoolkit/cpymadtools/lhc/_misc.py index 827ef3b1..c0c53b5d 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_misc.py +++ b/pyhdtoolkit/cpymadtools/lhc/_misc.py @@ -20,6 +20,7 @@ ) from pyhdtoolkit.optics.ripken import _add_beam_size_to_df +_VRF_THRESHOLD: int = 5000 def make_sixtrack_output(madx: Madx, /, energy: int) -> None: """ @@ -40,7 +41,7 @@ def make_sixtrack_output(madx: Madx, /, energy: int) -> None: logger.debug("Preparing outputs for SixTrack") logger.debug("Powering RF cavities") - madx.globals["VRF400"] = 8 if energy < 5000 else 16 # is 6 at injection for protons iirc? + madx.globals["VRF400"] = 8 if energy < _VRF_THRESHOLD else 16 # is 6 at injection for protons iirc? madx.globals["LAGRF400.B1"] = 0.5 # cavity phase difference in units of 2pi madx.globals["LAGRF400.B2"] = 0.0 @@ -113,7 +114,7 @@ def get_lhc_tune_and_chroma_knobs( get_lhc_tune_and_chroma_knobs("HLLHC", beam=2) # gives ('kqtf.b2_sq', 'kqtd.b2_sq', 'ksf.b2_sq', 'ksd.b2_sq') """ - beam = 2 if beam == 4 else beam + beam = 2 if beam == _BEAM4 else beam if run3: suffix = "_op" elif telescopic_squeeze: diff --git a/pyhdtoolkit/cpymadtools/lhc/_powering.py b/pyhdtoolkit/cpymadtools/lhc/_powering.py index f3a8d421..7af54c4a 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_powering.py +++ b/pyhdtoolkit/cpymadtools/lhc/_powering.py @@ -10,6 +10,9 @@ from cpymad.madx import Madx from loguru import logger +_BEAM4: int = 4 # LHC beam 4 is special case +_QUAD_CIRCUIT_HAS_B: int = 7 # Q7 has a .b in the circuit name +_MAX_IR_QUAD_NUMBER: int = 11 # beyond Q11 are MQTs etc def apply_lhc_colinearity_knob(madx: Madx, /, colinearity_knob_value: float = 0, ir: int | None = None) -> None: """ @@ -243,7 +246,7 @@ def power_landau_octupoles(madx: Madx, /, beam: int, mo_current: float, defectiv logger.debug(f"Powering Landau Octupoles, beam {beam} @ {madx.globals.nrj} GeV with {mo_current} A.") strength = mo_current / madx.globals.Imax_MO * madx.globals.Kmax_MO / brho - beam = 2 if beam == 4 else beam + beam = 2 if beam == _BEAM4 else beam for arc in _all_lhc_arcs(beam): for fd in "FD": @@ -274,7 +277,7 @@ def deactivate_lhc_arc_sextupoles(madx: Madx, /, beam: int) -> None: # KSF2 and KSD1 - Weak sextupoles of sectors 81/12/45/56 # Rest: Weak sextupoles in sectors 78/23/34/67 logger.debug(f"Deactivating all arc sextupoles for beam {beam}.") - beam = 2 if beam == 4 else beam + beam = 2 if beam == _BEAM4 else beam for arc in _all_lhc_arcs(beam): for fd in "FD": @@ -341,10 +344,10 @@ def vary_independent_ir_quadrupoles( for side in sides: logger.debug(f"Sending vary command for Q{quad}{side.upper()}{ip}") madx.command.vary( - name=f"kq{'t' if quad >= 11 else ''}{'l' if quad == 11 else ''}{quad}.{side}{ip}b{beam}", + name=f"kq{'t' if quad >= _MAX_IR_QUAD_NUMBER else ''}{'l' if quad == _MAX_IR_QUAD_NUMBER else ''}{quad}.{side}{ip}b{beam}", step=1e-7, - lower=f"-{circuit}.{'b' if quad == 7 else ''}{quad}{side}{ip}.b{beam}->kmax/brho", - upper=f"+{circuit}.{'b' if quad == 7 else ''}{quad}{side}{ip}.b{beam}->kmax/brho", + lower=f"-{circuit}.{'b' if quad == _QUAD_CIRCUIT_HAS_B else ''}{quad}{side}{ip}.b{beam}->kmax/brho", + upper=f"+{circuit}.{'b' if quad == _QUAD_CIRCUIT_HAS_B else ''}{quad}{side}{ip}.b{beam}->kmax/brho", ) diff --git a/pyhdtoolkit/cpymadtools/lhc/_setup.py b/pyhdtoolkit/cpymadtools/lhc/_setup.py index 496c9a01..b8e9fd54 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_setup.py +++ b/pyhdtoolkit/cpymadtools/lhc/_setup.py @@ -12,6 +12,9 @@ from pyhdtoolkit.cpymadtools.constants import LHC_CROSSING_SCHEMES +_BEAM_FOR_B4: int = 2 # LHC beam 4 uses lhcb2 sequence +_RUN2: int = 2 + # ----- Setup Utilities ----- # @@ -59,7 +62,7 @@ def prepare_lhc_run2( "/afs/cern.ch/eng/lhc/optics/runII/2018/PROTON/opticsfile.22", beam=2, stdout=True ) """ - if use_b4 and beam != 2: + if use_b4 and beam != _BEAM_FOR_B4: logger.error("Cannot use beam 4 sequence file for beam 1") msg = "Cannot use beam 4 sequence file for beam 1" raise ValueError(msg) @@ -136,7 +139,7 @@ def prepare_lhc_run3( "R2022a_A30cmC30cmA10mL200cm.madx", slicefactor=4, stdout=True ) """ - if use_b4 and beam != 2: + if use_b4 and beam != _BEAM_FOR_B4: logger.error("Cannot use beam 4 sequence file for beam 1") msg = "Cannot use beam 4 sequence file for beam 1" raise ValueError(msg) @@ -238,7 +241,7 @@ def __init__( **kwargs, ): assert opticsfile is not None, "An opticsfile must be provided" - if use_b4 and beam != 2: + if use_b4 and beam != _BEAM_FOR_B4: logger.error("Cannot use beam 4 sequence file for beam 1") msg = "Cannot use beam 4 sequence file for beam 1" raise ValueError(msg) @@ -246,7 +249,7 @@ def __init__( if int(run) not in (2, 3): msg = "This setup is only possible for Run 2 and Run 3 configurations." raise NotImplementedError(msg) - elif run == 2: + elif run == _RUN2: self.madx = prepare_lhc_run2( opticsfile=opticsfile, beam=beam, use_b4=use_b4, energy=energy, slicefactor=slicefactor, **kwargs ) diff --git a/pyhdtoolkit/cpymadtools/ptc.py b/pyhdtoolkit/cpymadtools/ptc.py index 10ccf585..e1c37fb1 100644 --- a/pyhdtoolkit/cpymadtools/ptc.py +++ b/pyhdtoolkit/cpymadtools/ptc.py @@ -17,6 +17,8 @@ from pyhdtoolkit.cpymadtools.utils import get_table_tfs +_MAX_PTC_AMPDET_ORDER: int = 2 +_MIN_PTC_AMPDET_ORDER: int = 1 def get_amplitude_detuning( madx: Madx, /, order: int = 2, file: Path | str | None = None, fringe: bool = False, **kwargs @@ -76,7 +78,7 @@ def get_amplitude_detuning( madx, order=3, model=3, exact=True, icase=5, no=6 ) """ - if order >= 3: + if order > _MAX_PTC_AMPDET_ORDER: logger.error(f"Maximum amplitude detuning order in PTC is 2, but {order:d} was requested") msg = "PTC amplitude detuning is not implemented for order > 2" raise NotImplementedError(msg) @@ -111,7 +113,7 @@ def get_amplitude_detuning( # ANH = anharmonicities (ex, ey, deltap), works only with parameters as full strings # could be done nicer with permutations ... logger.trace("Selecting anharmonicities") - if order >= 1: + if order >= _MIN_PTC_AMPDET_ORDER: # madx.select_ptc_normal('anhx=0, 0, 1') # dQx/ddp # madx.select_ptc_normal('anhy=0, 0, 1') # dQy/ddp madx.select_ptc_normal("anhx=1, 0, 0") # dQx/dex @@ -119,7 +121,7 @@ def get_amplitude_detuning( madx.select_ptc_normal("anhy=1, 0, 0") # dQy/dex madx.select_ptc_normal("anhy=0, 1, 0") # dQy/dey - if order >= 2: + if order >= _MAX_PTC_AMPDET_ORDER: # madx.select_ptc_normal('anhx=0, 0, 2') # d^2Qx/ddp^2 # madx.select_ptc_normal('anhy=0, 0, 2') # d^2Qy/ddp^2 madx.select_ptc_normal("anhx=2, 0, 0") # d^2Qx/dex^2 diff --git a/pyhdtoolkit/cpymadtools/tune.py b/pyhdtoolkit/cpymadtools/tune.py index 759abc90..3a8b1c53 100644 --- a/pyhdtoolkit/cpymadtools/tune.py +++ b/pyhdtoolkit/cpymadtools/tune.py @@ -62,12 +62,12 @@ def make_footprint_table( angle = 15 * angle_multiplier * math.pi / 180 if angle_multiplier == 0: madx.command.start(fx=sigma_multiplier * big, fy=sigma_multiplier * small) - elif angle_multiplier == 6: + elif angle_multiplier == 6: # noqa: PLR2004 madx.command.start(fx=sigma_multiplier * small, fy=sigma_multiplier * big) else: madx.command.start(fx=sigma_multiplier * math.cos(angle), fy=sigma_multiplier * math.sin(angle)) angle_multiplier += 0.5 - if int(angle_multiplier) == 7: + if int(angle_multiplier) == 7: # noqa: PLR2004 angle_multiplier = 0 sigma_multiplier += 1 if not dense else 0.5 diff --git a/pyhdtoolkit/plotting/tune.py b/pyhdtoolkit/plotting/tune.py index dd2007f4..497c10fe 100644 --- a/pyhdtoolkit/plotting/tune.py +++ b/pyhdtoolkit/plotting/tune.py @@ -109,7 +109,7 @@ def plot_tune_diagram( fig, ax = plt.subplots(figsize=(6, 6)) plot_tune_diagram(ax=ax, max_order=4, differentiate_orders=True) """ - if max_order > 6 or max_order < 1: + if max_order > 6 or max_order < 1: # noqa: PLR2004 logger.error("Plotting is not supported outside of 1st-6th order (and not recommended)") msg = "The 'max_order' argument should be between 1 and 6 included" raise ValueError(msg) diff --git a/pyhdtoolkit/utils/htc_monitor.py b/pyhdtoolkit/utils/htc_monitor.py index ce9cadaa..7ae62d48 100644 --- a/pyhdtoolkit/utils/htc_monitor.py +++ b/pyhdtoolkit/utils/htc_monitor.py @@ -149,7 +149,7 @@ def _make_cluster_table(owner_name: str, cluster: ClusterSummary) -> Table: table = _default_cluster_table() for i, source in enumerate(["query", "user", "cluster"]): table.add_row( - "Query" if i == 0 else ("All Users" if i == 2 else owner_name), + "Query" if i == 0 else ("All Users" if i == 2 else owner_name), # noqa: PLR2004 str(cluster.model_dump()[source]["jobs"]), str(cluster.model_dump()[source]["completed"]), str(cluster.model_dump()[source]["running"]), From 5f77588262f0a78a2781547e1b8344423cc7680f Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 11:41:42 +0200 Subject: [PATCH 12/41] miscellaneous --- pyhdtoolkit/cpymadtools/constants.py | 4 ++-- pyhdtoolkit/cpymadtools/lhc/_elements.py | 2 +- pyhdtoolkit/cpymadtools/lhc/_misc.py | 8 +++++--- pyhdtoolkit/cpymadtools/lhc/_queries.py | 6 +++--- pyhdtoolkit/cpymadtools/lhc/_setup.py | 11 +++++++---- pyhdtoolkit/plotting/layout.py | 2 +- pyhdtoolkit/plotting/styles/__init__.py | 3 +-- pyhdtoolkit/plotting/styles/paper.py | 3 +-- pyhdtoolkit/plotting/styles/thesis.py | 4 ++-- pyproject.toml | 2 +- 10 files changed, 24 insertions(+), 21 deletions(-) diff --git a/pyhdtoolkit/cpymadtools/constants.py b/pyhdtoolkit/cpymadtools/constants.py index 9848db74..95baa499 100644 --- a/pyhdtoolkit/cpymadtools/constants.py +++ b/pyhdtoolkit/cpymadtools/constants.py @@ -129,9 +129,9 @@ [f"kqs.a{sector}{sector+1 if sector < _MAX_SECTOR_VALUE else 1}b2" for sector in [1, 3, 5, 7]] # fmt: on LHC_KSF_KNOBS: list[str] = [ # sextupole correctors - f"ks{family}{id}.a{sector}{sector+1 if sector < _MAX_SECTOR_VALUE else 1}b{beam}" + f"ks{family}{beam_id}.a{sector}{sector+1 if sector < _MAX_SECTOR_VALUE else 1}b{beam}" for beam in [1, 2] - for id in [1, 2] + for beam_id in [1, 2] for family in ["f", "d"] for sector in [1, 2, 3, 4, 5, 6, 7, 8] ] diff --git a/pyhdtoolkit/cpymadtools/lhc/_elements.py b/pyhdtoolkit/cpymadtools/lhc/_elements.py index 0358cf1c..74ad364b 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_elements.py +++ b/pyhdtoolkit/cpymadtools/lhc/_elements.py @@ -247,7 +247,7 @@ def add_markers_around_lhc_ip(madx: Madx, /, sequence: str, ip: int, n_markers: be run for the changes to apply. This means the caveats of ``USE`` apply, for instance the erasing of previously defined errors, orbits corrections etc. - + Therefore, it is recommended to install the errors and save them with the ``ESAVE`` or ``ETABLE`` command, call this function, then re-implement the errors with the ``SETERR`` command. diff --git a/pyhdtoolkit/cpymadtools/lhc/_misc.py b/pyhdtoolkit/cpymadtools/lhc/_misc.py index c0c53b5d..35ee7afd 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_misc.py +++ b/pyhdtoolkit/cpymadtools/lhc/_misc.py @@ -20,6 +20,7 @@ ) from pyhdtoolkit.optics.ripken import _add_beam_size_to_df +_BEAM4: int = 4 # LHC beam 4 is special case _VRF_THRESHOLD: int = 5000 def make_sixtrack_output(madx: Madx, /, energy: int) -> None: @@ -65,7 +66,7 @@ def reset_lhc_bump_flags(madx: Madx, /) -> None: reset_lhc_bump_flags(madx) """ logger.debug("Resetting all LHC IP bump flags") - ALL_BUMPS = ( + all_bumps = ( LHC_ANGLE_FLAGS + LHC_CROSSING_ANGLE_FLAGS + LHC_EXPERIMENT_STATE_FLAGS @@ -74,7 +75,7 @@ def reset_lhc_bump_flags(madx: Madx, /) -> None: + LHC_PARALLEL_SEPARATION_FLAGS ) with madx.batch(): - madx.globals.update({bump: 0 for bump in ALL_BUMPS}) + madx.globals.update({bump: 0 for bump in all_bumps}) def get_lhc_tune_and_chroma_knobs( @@ -124,7 +125,8 @@ def get_lhc_tune_and_chroma_knobs( if accelerator.upper() not in ("LHC", "HLLHC"): logger.error("Invalid accelerator name, only 'LHC' and 'HLLHC' implemented") - raise NotImplementedError(f"Accelerator '{accelerator}' not implemented.") + msg = f"Accelerator '{accelerator}' not implemented." + raise NotImplementedError(msg) return { "LHC": ( diff --git a/pyhdtoolkit/cpymadtools/lhc/_queries.py b/pyhdtoolkit/cpymadtools/lhc/_queries.py index d2171459..c20e6bfc 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_queries.py +++ b/pyhdtoolkit/cpymadtools/lhc/_queries.py @@ -82,10 +82,10 @@ def get_magnets_powering( 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"] - NEW_COLNAMES = list(set(NEW_COLNAMES + kwargs.pop("columns", []))) # in case user gives explicit columns + new_colnames = ["name", "keyword", "ampere", "imax", "percent", "kn", "kmax", "integrated_field", "L"] + new_colnames = list(set(new_colnames + kwargs.pop("columns", []))) # in case user gives explicit columns _list_field_currents(madx, brho=brho) - return twiss.get_pattern_twiss(madx, columns=NEW_COLNAMES, patterns=patterns, **kwargs) + return twiss.get_pattern_twiss(madx, columns=new_colnames, patterns=patterns, **kwargs) def query_arc_correctors_powering(madx: Madx, /) -> dict[str, float]: diff --git a/pyhdtoolkit/cpymadtools/lhc/_setup.py b/pyhdtoolkit/cpymadtools/lhc/_setup.py index b8e9fd54..28e56f41 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_setup.py +++ b/pyhdtoolkit/cpymadtools/lhc/_setup.py @@ -72,7 +72,8 @@ def _run2_sequence_from_opticsfile(opticsfile: Path, use_b4: bool = False) -> Pa seqfile_path = opticsfile.parent.parent / filename if not seqfile_path.is_file(): logger.error(f"Could not find sequence file '{filename}' at expected location '{seqfile_path}'") - raise ValueError(f"No sequence file found at '{seqfile_path}'") + msg = f"No sequence file found at '{seqfile_path}'" + raise ValueError(msg) return seqfile_path logger.debug("Creating Run 2 setup MAD-X instance") @@ -240,7 +241,9 @@ def __init__( slicefactor: int | None = None, **kwargs, ): - assert opticsfile is not None, "An opticsfile must be provided" + if opticsfile is None: # don't want to move arg and mess users code + msg = "An opticsfile must be provided" + raise ValueError(msg) if use_b4 and beam != _BEAM_FOR_B4: logger.error("Cannot use beam 4 sequence file for beam 1") msg = "Cannot use beam 4 sequence file for beam 1" @@ -249,7 +252,7 @@ def __init__( if int(run) not in (2, 3): msg = "This setup is only possible for Run 2 and Run 3 configurations." raise NotImplementedError(msg) - elif run == _RUN2: + if run == _RUN2: self.madx = prepare_lhc_run2( opticsfile=opticsfile, beam=beam, use_b4=use_b4, energy=energy, slicefactor=slicefactor, **kwargs ) @@ -501,7 +504,7 @@ def setup_lhc_orbit(madx: Madx, /, scheme: str = "flat", **kwargs) -> dict[str, orbit_setup = setup_lhc_orbit(madx, scheme="lhc_top") """ - if scheme not in LHC_CROSSING_SCHEMES.keys(): + if scheme not in LHC_CROSSING_SCHEMES: logger.error(f"Invalid scheme parameter, should be one of {LHC_CROSSING_SCHEMES.keys()}") msg = "Invalid scheme parameter given" raise ValueError(msg) diff --git a/pyhdtoolkit/plotting/layout.py b/pyhdtoolkit/plotting/layout.py index ba729b2c..c0f508cd 100644 --- a/pyhdtoolkit/plotting/layout.py +++ b/pyhdtoolkit/plotting/layout.py @@ -417,7 +417,7 @@ def _ylim_from_input( """ if isinstance(ylim, tuple): return ylim - elif isinstance(ylim, (float, int)): + elif isinstance(ylim, float | int): if ylim >= 0: return (-ylim, ylim) else: diff --git a/pyhdtoolkit/plotting/styles/__init__.py b/pyhdtoolkit/plotting/styles/__init__.py index 1af83149..ac37b328 100644 --- a/pyhdtoolkit/plotting/styles/__init__.py +++ b/pyhdtoolkit/plotting/styles/__init__.py @@ -8,7 +8,6 @@ Feel free to use them anyway, as they might be useful to you when using the `~pyhdtoolkit.plotting` submodules, or to be adapted. """ from pathlib import Path -from typing import Union import matplotlib as mpl import matplotlib.pyplot as plt @@ -16,7 +15,7 @@ from . import paper, thesis # noqa: TID252 -PlotSetting = Union[float, bool, str, tuple] +PlotSetting = float | bool | str | tuple def install_mpl_styles() -> None: diff --git a/pyhdtoolkit/plotting/styles/paper.py b/pyhdtoolkit/plotting/styles/paper.py index a1a19bbb..6b897eab 100644 --- a/pyhdtoolkit/plotting/styles/paper.py +++ b/pyhdtoolkit/plotting/styles/paper.py @@ -13,9 +13,8 @@ - ``SINGLE_COLUMN``: For plots to be included in a single column figure environment. - ``DOUBLE_COLUMN``: For plots to be included in a full width (double column) figure environment. """ -from typing import Union -PlotSetting = Union[float, bool, str, tuple] +PlotSetting = float | bool | str | tuple SINGLE_COLUMN: dict[str, PlotSetting] = { diff --git a/pyhdtoolkit/plotting/styles/thesis.py b/pyhdtoolkit/plotting/styles/thesis.py index 65a76ecd..69ff6d2d 100644 --- a/pyhdtoolkit/plotting/styles/thesis.py +++ b/pyhdtoolkit/plotting/styles/thesis.py @@ -15,9 +15,9 @@ - ``MEDIUM``: For simple plots to be included alone in a LaTeX figure environment (e.g single axis line plots, or scatters with a colorbar like in `~pyhdtoolkit.plotting.tune`). - ``LARGE``: For more complex plots to be included alone in a LaTeX figure environment (e.g. multi-axes figures such as in `~pyhdtoolkit.plotting.lattice`). """ -from typing import Union -PlotSetting = Union[float, bool, str, tuple] + +PlotSetting = float | bool | str | tuple SMALL: dict[str, PlotSetting] = { diff --git a/pyproject.toml b/pyproject.toml index 7aff47c1..67311ce9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -163,7 +163,7 @@ exclude = ''' ''' [tool.ruff] -ignore = ["E501"] +ignore = ["E501", "FBT001", "FBT002"] line-length = 120 # Exclude a variety of commonly ignored directories. From 58978315abc6b24b59b89c21f45a3a06a787c208 Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 11:50:47 +0200 Subject: [PATCH 13/41] place *args correctly --- pyhdtoolkit/maths/stats_fitting.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/pyhdtoolkit/maths/stats_fitting.py b/pyhdtoolkit/maths/stats_fitting.py index a46fe49c..520a3103 100644 --- a/pyhdtoolkit/maths/stats_fitting.py +++ b/pyhdtoolkit/maths/stats_fitting.py @@ -51,7 +51,7 @@ def set_distributions_dict(dist_dict: dict[st.rv_continuous, str]) -> None: """ # pylint: disable=global-statement logger.debug("Setting tested distributions") - global DISTRIBUTIONS + global DISTRIBUTIONS # noqa: PLW0603 DISTRIBUTIONS = dist_dict @@ -102,14 +102,14 @@ def best_fit_distribution( *args, loc, scale = params logger.debug(f"Calculating PDF goodness of fit and error for distribution '{distname}'") - pdf = distribution.pdf(x, loc=loc, scale=scale, *args) + pdf = distribution.pdf(x, *args, loc=loc, scale=scale) sse = np.sum(np.power(y - pdf, 2.0)) try: if ax: logger.debug(f"Plotting fitted PDF for distribution '{distname}'") pd.Series(pdf, x).plot(ax=ax, label=f"{distname} fit", alpha=1) - except Exception: + except Exception: # noqa: BLE001 logger.exception(f"Plotting distribution '{distname}' failed") logger.debug(f"Identifying if distribution '{distname}' is a better fit than previous tries") @@ -117,7 +117,7 @@ def best_fit_distribution( best_distribution = distribution best_params = params best_sse = sse - except Exception: + except Exception: # noqa: BLE001 logger.exception(f"Trying to fit distribution '{distname}' failed and aborted") logger.info(f"Found a best fit: '{DISTRIBUTIONS[best_distribution]}' distribution") @@ -157,5 +157,5 @@ def make_pdf(distribution: st.rv_continuous, params: tuple[float, ...], size: in logger.debug("Building PDF") x = np.linspace(start, end, size) - y = distribution.pdf(x, loc=loc, scale=scale, *args) + y = distribution.pdf(x, *args, loc=loc, scale=scale) return pd.Series(y, x) From 5a2806e8928324899397c61ab387ace03a26b97c Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 11:51:00 +0200 Subject: [PATCH 14/41] miscellaneous --- pyhdtoolkit/cpymadtools/ptc.py | 4 ++-- pyhdtoolkit/cpymadtools/tune.py | 32 +++++++++++++++----------------- pyhdtoolkit/optics/ripken.py | 11 ++++++----- 3 files changed, 23 insertions(+), 24 deletions(-) diff --git a/pyhdtoolkit/cpymadtools/ptc.py b/pyhdtoolkit/cpymadtools/ptc.py index e1c37fb1..3a0b9e38 100644 --- a/pyhdtoolkit/cpymadtools/ptc.py +++ b/pyhdtoolkit/cpymadtools/ptc.py @@ -60,7 +60,7 @@ def get_amplitude_detuning( ``PTC_NORMAL`` command. Their default values are listed higher up in this docstring. Any remaining keyword argument is transmitted to the ``PTC_NORMAL`` command. - + Returns: A `~tfs.frame.TfsDataframe` with the calculated coefficients. @@ -289,7 +289,7 @@ def ptc_twiss( `icase` and `normal` ones can be given for the ``PTC_TWISS`` command. Their default values are listed higher up in this docstring. Any remaining keyword argument is transmitted to the ``PTC_TWISS`` command. - + Returns: A `~tfs.frame.TfsDataFrame` with the calculated ``TWISS`` parameters. diff --git a/pyhdtoolkit/cpymadtools/tune.py b/pyhdtoolkit/cpymadtools/tune.py index 3a8b1c53..0c33f489 100644 --- a/pyhdtoolkit/cpymadtools/tune.py +++ b/pyhdtoolkit/cpymadtools/tune.py @@ -191,25 +191,25 @@ def get_footprint_patches(dynap_dframe: tfs.TfsDataFrame) -> matplotlib.collecti logger.debug("Grouping tune points according to starting angles and amplitudes") try: - A = np.zeros([amplitude, angle, 2]) - A[0, :, 0] = dynap_dframe["tunx"].to_numpy()[0] - A[0, :, 1] = dynap_dframe["tuny"].to_numpy()[0] - A[1:, :, 0] = dynap_dframe["tunx"].to_numpy()[1:].reshape(-1, angle) - A[1:, :, 1] = dynap_dframe["tuny"].to_numpy()[1:].reshape(-1, angle) - except ValueError as tune_grouping_error: + a = np.zeros([amplitude, angle, 2]) + a[0, :, 0] = dynap_dframe["tunx"].to_numpy()[0] + a[0, :, 1] = dynap_dframe["tuny"].to_numpy()[0] + a[1:, :, 0] = dynap_dframe["tunx"].to_numpy()[1:].reshape(-1, angle) + a[1:, :, 1] = dynap_dframe["tuny"].to_numpy()[1:].reshape(-1, angle) + except ValueError: logger.exception( "Cannot group tune points according to starting angles and amplitudes. Try changing " "the 'AMPLITUDE' value in the provided TfsDataFrame's headers." ) - raise tune_grouping_error + raise logger.debug("Determining polygon vertices") - sx = A.shape[0] - 1 - sy = A.shape[1] - 1 - p1 = A[:-1, :-1, :].reshape(sx * sy, 2)[:, :] - p2 = A[1:, :-1, :].reshape(sx * sy, 2)[:] - p3 = A[1:, 1:, :].reshape(sx * sy, 2)[:] - p4 = A[:-1, 1:, :].reshape(sx * sy, 2)[:] + sx = a.shape[0] - 1 + sy = a.shape[1] - 1 + p1 = a[:-1, :-1, :].reshape(sx * sy, 2)[:, :] + p2 = a[1:, :-1, :].reshape(sx * sy, 2)[:] + p3 = a[1:, 1:, :].reshape(sx * sy, 2)[:] + p4 = a[:-1, 1:, :].reshape(sx * sy, 2)[:] polygons = np.stack((p1, p2, p3, p4)) # Stack endpoints to form polygons polygons = np.transpose(polygons, (1, 0, 2)) # transpose polygons @@ -297,14 +297,12 @@ def __init__(self, tune_groups: list[list[dict[str, float]]], amplitude: int, an def get_h_tune(self, ampl, angl): if len(self._tunes[ampl]) <= angl < self._maxnangl: return self._tunes[ampl][len(self._tunes[ampl]) - 1]["H"] - else: - return self._tunes[ampl][angl]["H"] + return self._tunes[ampl][angl]["H"] def get_v_tune(self, ampl, angl): if len(self._tunes[ampl]) <= angl < self._maxnangl: return self._tunes[ampl][len(self._tunes[ampl]) - 1]["V"] - else: - return self._tunes[ampl][angl]["V"] + return self._tunes[ampl][angl]["V"] def get_plottable(self) -> tuple[list[float], list[float]]: qxs, qys = [], [] diff --git a/pyhdtoolkit/optics/ripken.py b/pyhdtoolkit/optics/ripken.py index 8693422c..e4d8331f 100644 --- a/pyhdtoolkit/optics/ripken.py +++ b/pyhdtoolkit/optics/ripken.py @@ -66,7 +66,7 @@ def _beam_size(coordinates_distribution: np.ndarray, method: str = "std") -> flo Args: coordinates_distribution (np.ndarray): ensemble of coordinates of the particle distributon. method (str): the method of calculation to use, either 'std' (using the standard deviation as the - beam size) or 'rms' (root mean square). + beam size) or 'rms' (root mean square). Case-insensitive. Returns: The computed beam size. @@ -74,12 +74,13 @@ def _beam_size(coordinates_distribution: np.ndarray, method: str = "std") -> flo Raises: NotImplementedError: If the required *method* is neither std nor rms. """ + if method.lower() not in ("std", "rms"): + msg = "Invalid method provided" + raise NotImplementedError(msg) if method == "std": return coordinates_distribution.std() - elif method == "rms": - return np.sqrt(np.mean(np.square(coordinates_distribution))) - msg = "Invalid method provided" - raise NotImplementedError(msg) + return np.sqrt(np.mean(np.square(coordinates_distribution))) # rms + def _add_beam_size_to_df(df: tfs.TfsDataFrame, geom_emit_x: float, geom_emit_y: float) -> tfs.TfsDataFrame: From 7b4ca95b32c151fe58324d929d8bc383c4cd2bb8 Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 11:57:47 +0200 Subject: [PATCH 15/41] use dict literal --- pyhdtoolkit/optics/rdt.py | 23 ++++++++++------------- 1 file changed, 10 insertions(+), 13 deletions(-) diff --git a/pyhdtoolkit/optics/rdt.py b/pyhdtoolkit/optics/rdt.py index 87cd58bf..3e635a09 100644 --- a/pyhdtoolkit/optics/rdt.py +++ b/pyhdtoolkit/optics/rdt.py @@ -8,7 +8,6 @@ """ - def rdt_to_order_and_type(rdt: int | str) -> str: """ Decompose the input RDT into its four various components @@ -23,18 +22,16 @@ def rdt_to_order_and_type(rdt: int | str) -> str: """ j, k, l, m = map(int, str(rdt)) # noqa: E741 rdt_type = "normal" if (l + m) % 2 == 0 else "skew" - orders = dict( - ( - (1, "dipole"), - (2, "quadrupole"), - (3, "sextupole"), - (4, "octupole"), - (5, "decapole"), - (6, "dodecapole"), - (7, "tetradecapole"), - (8, "hexadecapole"), - ) - ) + orders = { + 1: 'dipole', + 2: 'quadrupole', + 3: 'sextupole', + 4: 'octupole', + 5: 'decapole', + 6: 'dodecapole', + 7: 'tetradecapole', + 8: 'hexadecapole', + } return f"{rdt_type}_{orders[j + k + l + m]}" From fcbfadb942de3602e92605262a3097d3e26f57f3 Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 12:01:07 +0200 Subject: [PATCH 16/41] miscellaneous --- pyhdtoolkit/models/beam.py | 30 +++++++++++++++--------------- pyhdtoolkit/models/madx.py | 4 ++-- pyhdtoolkit/optics/twiss.py | 2 +- pyhdtoolkit/plotting/aperture.py | 5 ++--- 4 files changed, 20 insertions(+), 21 deletions(-) diff --git a/pyhdtoolkit/models/beam.py b/pyhdtoolkit/models/beam.py index 8008e6eb..cc7a3d61 100644 --- a/pyhdtoolkit/models/beam.py +++ b/pyhdtoolkit/models/beam.py @@ -19,7 +19,7 @@ class BeamParameters(BaseModel): One can find a usage example in the :ref:`beam enveloppe demo `. """ - pc_GeV: float | None = None # Beam momentum [GeV] + pc_GeV: float | None = None # Beam momentum [GeV] # noqa: N815 E_0_GeV: float | None = 0.9382720813 # Particle rest mass energy [GeV], defaults to that of a proton charge: float | None = 1 # Particle charge in [e], defaults to that of a proton en_x_m: float | None = None # Horizontal normalized emittance [m] @@ -27,17 +27,17 @@ class BeamParameters(BaseModel): deltap_p: float | None = None # Momentum deviation @property - def B_rho_Tm(self) -> float: + def B_rho_Tm(self) -> float: # noqa: N802 """Beam rigidity [T/m]""" return self.pc_GeV / 0.3 @property - def E_tot_GeV(self) -> float: + def E_tot_GeV(self) -> float: # noqa: N802 """Total beam energy [GeV]""" return sqrt(self.pc_GeV**2 + self.E_0_GeV**2) @property - def E_kin_GeV(self) -> float: + def E_kin_GeV(self) -> float: # noqa: N802 """Total beam energy [GeV]""" return self.E_tot_GeV - self.E_0_GeV @@ -64,17 +64,17 @@ def eg_y_m(self) -> float: def __repr__(self) -> str: return ( f"Beam Parameters for particle of charge {self.charge} \n" - + f"Beam momentum = {self.pc_GeV:2.3f} GeV/c \n" - + f"Normalized x-emittance = {self.en_x_m * 1e6:2.3f} mm mrad \n" - + f"Normalized y-emittance = {self.en_y_m * 1e6:2.3f} mm mrad \n" - + f"Momentum deviation deltap/p = {self.deltap_p} \n" - + f" -> Beam total energy = {self.E_tot_GeV:2.3f} GeV \n" - + f" -> Beam kinetic energy = {self.E_kin_GeV:2.3f} GeV \n" - + f" -> Beam rigidity = {self.B_rho_Tm:2.3f} Tm \n" - + f" -> Relativistic beta = {self.beta_r:2.5f} \n" - + f" -> Relativistic gamma = {self.gamma_r:2.3f} \n" - + f" -> Geometrical x emittance = {self.eg_x_m * 1e6:2.3f} mm mrad \n" - + f" -> Geometrical y emittance = {self.eg_y_m * 1e6:2.3f} mm mrad \n" + f"Beam momentum = {self.pc_GeV:2.3f} GeV/c \n" + f"Normalized x-emittance = {self.en_x_m * 1e6:2.3f} mm mrad \n" + f"Normalized y-emittance = {self.en_y_m * 1e6:2.3f} mm mrad \n" + f"Momentum deviation deltap/p = {self.deltap_p} \n" + f" -> Beam total energy = {self.E_tot_GeV:2.3f} GeV \n" + f" -> Beam kinetic energy = {self.E_kin_GeV:2.3f} GeV \n" + f" -> Beam rigidity = {self.B_rho_Tm:2.3f} Tm \n" + f" -> Relativistic beta = {self.beta_r:2.5f} \n" + f" -> Relativistic gamma = {self.gamma_r:2.3f} \n" + f" -> Geometrical x emittance = {self.eg_x_m * 1e6:2.3f} mm mrad \n" + f" -> Geometrical y emittance = {self.eg_y_m * 1e6:2.3f} mm mrad \n" ) def __str__(self) -> str: diff --git a/pyhdtoolkit/models/madx.py b/pyhdtoolkit/models/madx.py index 345be9a9..f7495882 100644 --- a/pyhdtoolkit/models/madx.py +++ b/pyhdtoolkit/models/madx.py @@ -48,8 +48,8 @@ class MADXBeam(BaseModel): 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]. + sigt: float # The bunch length c σt in [m] # noqa: RUF003 + sige: float # The relative energy spread σE /E in [1]. # noqa: RUF003 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] diff --git a/pyhdtoolkit/optics/twiss.py b/pyhdtoolkit/optics/twiss.py index 59b7c473..08cdc265 100644 --- a/pyhdtoolkit/optics/twiss.py +++ b/pyhdtoolkit/optics/twiss.py @@ -33,7 +33,7 @@ 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 diff --git a/pyhdtoolkit/plotting/aperture.py b/pyhdtoolkit/plotting/aperture.py index e723c2fa..32a27e7d 100644 --- a/pyhdtoolkit/plotting/aperture.py +++ b/pyhdtoolkit/plotting/aperture.py @@ -344,10 +344,9 @@ def _get_positions_and_real_apertures( indices = list(reversed(indices)) logger.trace("Extrapolating data at beginning of elements") - counter = 0 # Keep track of exact position in new array with counter - for j in indices: + # counter keeps track of exact position in new array with counter + for counter, j in enumerate(indices): new_pos.insert(j + counter, (twiss_df.s.iloc[j] - twiss_df.l.iloc[j])) - counter += 1 # Replace all zeros with Nan apertures = np.array(new_aper) From 072692cd2bc295977e47a73e4b2a96b55efc9d02 Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 12:06:45 +0200 Subject: [PATCH 17/41] no implicit namespace --- pyhdtoolkit/plotting/sbs/__init__.py | 6 ++++++ 1 file changed, 6 insertions(+) create mode 100644 pyhdtoolkit/plotting/sbs/__init__.py diff --git a/pyhdtoolkit/plotting/sbs/__init__.py b/pyhdtoolkit/plotting/sbs/__init__.py new file mode 100644 index 00000000..8b575c84 --- /dev/null +++ b/pyhdtoolkit/plotting/sbs/__init__.py @@ -0,0 +1,6 @@ +""" +.. _plotting.sbs: +""" +from . import coupling, phase # noqa: TID252 + +__all__ = ["coupling", "phase"] From 3008efc130bc87941314eaa8aa6fe711259c9b97 Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 12:07:06 +0200 Subject: [PATCH 18/41] correctly propagate kwargs --- pyhdtoolkit/plotting/crossing.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/pyhdtoolkit/plotting/crossing.py b/pyhdtoolkit/plotting/crossing.py index e82b2ba9..958ea0c5 100644 --- a/pyhdtoolkit/plotting/crossing.py +++ b/pyhdtoolkit/plotting/crossing.py @@ -229,9 +229,9 @@ def _highlight_mbx_and_mqx( logger.trace("Highlighting MBX and MQX areas on a twin axis") patches_axis = axis.twinx() patches_axis.get_yaxis().set_visible(False) - patches_axis.axvspan(*left_mbx_lim, color="orange", lw=2, alpha=0.2, label="MBX") - patches_axis.axvspan(*left_mqx_lim, color="grey", lw=2, alpha=0.2, label="MQX") + patches_axis.axvspan(*left_mbx_lim, color="orange", lw=2, alpha=0.2, label="MBX", **kwargs) + patches_axis.axvspan(*left_mqx_lim, color="grey", lw=2, alpha=0.2, label="MQX", **kwargs) patches_axis.axvline(x=0, color="grey", ls="--", label=f"IP{ip}") - patches_axis.axvspan(*right_mqx_lim, color="grey", lw=2, alpha=0.2) # no label duplication - patches_axis.axvspan(*right_mbx_lim, color="orange", lw=2, alpha=0.2) # no label duplication + patches_axis.axvspan(*right_mqx_lim, color="grey", lw=2, alpha=0.2, **kwargs) # no label duplication + patches_axis.axvspan(*right_mbx_lim, color="orange", lw=2, alpha=0.2, **kwargs) # no label duplication patches_axis.legend(loc=4) From 10e5f592bce49ba1dbd6a44f61803302e4886690 Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 12:07:33 +0200 Subject: [PATCH 19/41] clean up logic --- pyhdtoolkit/plotting/layout.py | 29 ++++++++++++----------------- 1 file changed, 12 insertions(+), 17 deletions(-) diff --git a/pyhdtoolkit/plotting/layout.py b/pyhdtoolkit/plotting/layout.py index c0f508cd..7d9bf337 100644 --- a/pyhdtoolkit/plotting/layout.py +++ b/pyhdtoolkit/plotting/layout.py @@ -313,9 +313,7 @@ def plot_machine_layout( elif k3l_lim is not None: bpm_legend_loc = 3 # octuoles are here but not sextupoles, we go bottom right else: - bpm_legend_loc = ( - "best" # can't easily determine the best position, go automatic and leave to the user - ) + bpm_legend_loc = "best" # can't easily determine the best position, go automatic and leave to the user if plotted_elements > 0: # If we plotted at least one BPM, we need to plot the legend bpm_patches_axis.legend(loc=bpm_legend_loc) bpm_patches_axis.grid(False) @@ -335,7 +333,7 @@ def scale_patches(scale: float, ylabel: str, **kwargs) -> None: **kwargs: If either `ax` or `axis` is found in the kwargs, the corresponding value is used as the axis object to plot on, otherwise the current axis is used. - + Example: .. code-block:: python @@ -393,9 +391,7 @@ def _plot_lattice_series( ) -def _ylim_from_input( - ylim: tuple[float, float] | float, name_for_error: str = "knl_lim" -) -> tuple[float, float]: +def _ylim_from_input(ylim: tuple[float, float] | float, name_for_error: str = "knl_lim") -> tuple[float, float]: """ .. versionadded:: 1.2.0 @@ -415,18 +411,17 @@ def _ylim_from_input( Raises: TypeError: if the input is not a tuple, a float or an int. """ + if not isinstance(ylim, tuple | float | int): + msg = f"Invalid type for '{name_for_error}': {type(ylim)}. " + raise TypeError(msg) + if isinstance(ylim, tuple): return ylim - elif isinstance(ylim, float | int): - if ylim >= 0: - return (-ylim, ylim) - else: - return (ylim, -ylim) - else: - raise TypeError( - f"Invalid type for '{name_for_error}': {type(ylim)}. " - "Should be a tuple, a float or an int. Can also be give as None." - ) + + # otherwise we have float | int + if ylim >= 0: + return (-ylim, ylim) + return (ylim, -ylim) def _determine_default_knl_lim(df: pd.DataFrame, col: str, coeff: float) -> tuple[float, float]: From e57c1b1a0cff0e19fffdc8e4d2f797339ffbd5ea Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 12:07:47 +0200 Subject: [PATCH 20/41] message outside raise --- pyhdtoolkit/plotting/utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyhdtoolkit/plotting/utils.py b/pyhdtoolkit/plotting/utils.py index 55342b9d..5e0fb2b4 100644 --- a/pyhdtoolkit/plotting/utils.py +++ b/pyhdtoolkit/plotting/utils.py @@ -374,7 +374,8 @@ def draw_confidence_ellipse(x: ArrayLike, y: ArrayLike, n_std: float = 3.0, face if x.size != y.size: logger.error(f"x and y must be the same size, but shapes {x.shape} and {y.shape} were given.") - raise ValueError(f"x and y must be the same size, but shapes {x.shape} and {y.shape} were given.") + msg = f"x and y must be the same size, but shapes {x.shape} and {y.shape} were given." + raise ValueError(msg) logger.debug("Computing covariance matrix and pearson correlation coefficient") covariance_matrix = np.cov(x, y) From 94e42919e61275cc5a359494d4a67191a37f2e28 Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 12:09:26 +0200 Subject: [PATCH 21/41] message outside raise --- pyhdtoolkit/plotting/utils.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/pyhdtoolkit/plotting/utils.py b/pyhdtoolkit/plotting/utils.py index 5e0fb2b4..d767f460 100644 --- a/pyhdtoolkit/plotting/utils.py +++ b/pyhdtoolkit/plotting/utils.py @@ -455,8 +455,12 @@ def _determine_default_sbs_coupling_ylabel(rdt: str, component: str) -> str: coupling_label = _determine_default_sbs_coupling_ylabel(rdt="f1001", component="RE") """ logger.debug(f"Determining a default label for the {component.upper()} component of coupling {rdt.upper()}.") - assert rdt.upper() in ("F1001", "F1010") - assert component.upper() in ("ABS", "RE", "IM") + if rdt.upper() not in ("F1001", "F1010"): + msg = "Invalid RDT for coupling plot." + raise ValueError(msg) + if component.upper() not in ("ABS", "RE", "IM"): + msg = "Invalid component for coupling RDT." + raise ValueError(msg) if component.upper() == "ABS": opening = closing = "|" From ae18cf04d541216cd03c47190efe18ad49687a91 Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 12:12:57 +0200 Subject: [PATCH 22/41] trade map for list comprehension --- pyhdtoolkit/utils/operations.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyhdtoolkit/utils/operations.py b/pyhdtoolkit/utils/operations.py index 9242422f..6c4c71a9 100644 --- a/pyhdtoolkit/utils/operations.py +++ b/pyhdtoolkit/utils/operations.py @@ -155,7 +155,7 @@ def chunk_list(sequence: Sequence, size: int) -> Sequence: """ if size > len(sequence): return sequence - return list(map(lambda x: sequence[x * size : x * size + size], list(range(math.ceil(len(sequence) / size))))) + return [sequence[x * size : x * size + size] for x in list(range(math.ceil(len(sequence) / size)))] @staticmethod def deep_flatten(sequence: Sequence) -> list: From 76dec808fbb29d7668ce15539e681cedb2a54212 Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 12:15:30 +0200 Subject: [PATCH 23/41] cleaner checks --- pyhdtoolkit/utils/htc_monitor.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyhdtoolkit/utils/htc_monitor.py b/pyhdtoolkit/utils/htc_monitor.py index 7ae62d48..db4d3c17 100644 --- a/pyhdtoolkit/utils/htc_monitor.py +++ b/pyhdtoolkit/utils/htc_monitor.py @@ -98,7 +98,7 @@ def read_condor_q(report: str) -> tuple[list[HTCTaskSummary], ClusterSummary]: tasks: list[HTCTaskSummary] = [] next_line_is_task_report = False - for line in report.split("\n"): + for line in report.splitlines(): if line.startswith("-- Schedd:"): # extract scheduler information scheduler_id = _process_scheduler_information_line(line) @@ -106,7 +106,7 @@ def read_condor_q(report: str) -> tuple[list[HTCTaskSummary], ClusterSummary]: next_line_is_task_report = True elif next_line_is_task_report: # extract task report information - if line != "\n" and line != "": + if line not in ("\n", ""): tasks.append(_process_task_summary_line(line)) else: # an empty line denotes the end of the task report(s) next_line_is_task_report = False @@ -117,7 +117,7 @@ def read_condor_q(report: str) -> tuple[list[HTCTaskSummary], ClusterSummary]: query_summary = _process_cluster_summary_line(line, "query") elif "all users" in line: # last line full_summary = _process_cluster_summary_line(line, "all users") - elif line != "\n" and line != "": # user line, whoever the user is + elif line not in ("\n", ""): # user line, whoever the user is owner_summary = _process_cluster_summary_line(line, querying_owner) cluster_summary = ClusterSummary( scheduler_id=scheduler_id, query=query_summary, user=owner_summary, cluster=full_summary From 249203f9c49a5dba1536321148117a2790b1bc04 Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 12:18:53 +0200 Subject: [PATCH 24/41] format and last lints --- pyhdtoolkit/__init__.py | 1 + pyhdtoolkit/cpymadtools/constants.py | 17 ++++++++++---- pyhdtoolkit/cpymadtools/coupling.py | 5 ++-- pyhdtoolkit/cpymadtools/lhc/_coupling.py | 1 + pyhdtoolkit/cpymadtools/lhc/_elements.py | 2 ++ pyhdtoolkit/cpymadtools/lhc/_errors.py | 6 +++-- pyhdtoolkit/cpymadtools/lhc/_misc.py | 1 + pyhdtoolkit/cpymadtools/lhc/_powering.py | 2 ++ pyhdtoolkit/cpymadtools/lhc/_queries.py | 1 + pyhdtoolkit/cpymadtools/lhc/_routines.py | 5 +++- pyhdtoolkit/cpymadtools/lhc/_setup.py | 16 +++++++++---- pyhdtoolkit/cpymadtools/lhc/_twiss.py | 1 + pyhdtoolkit/cpymadtools/matching.py | 13 ++++++----- pyhdtoolkit/cpymadtools/ptc.py | 11 +++++++-- pyhdtoolkit/cpymadtools/track.py | 1 + pyhdtoolkit/cpymadtools/tune.py | 3 ++- pyhdtoolkit/cpymadtools/twiss.py | 1 + pyhdtoolkit/cpymadtools/utils.py | 1 + pyhdtoolkit/maths/stats_fitting.py | 2 ++ pyhdtoolkit/maths/utils.py | 3 ++- pyhdtoolkit/models/beam.py | 1 + pyhdtoolkit/models/htc.py | 1 + pyhdtoolkit/models/madx.py | 1 + pyhdtoolkit/optics/beam.py | 1 + pyhdtoolkit/optics/rdt.py | 16 ++++++------- pyhdtoolkit/optics/ripken.py | 1 - pyhdtoolkit/optics/twiss.py | 1 + pyhdtoolkit/plotting/__init__.py | 1 + pyhdtoolkit/plotting/aperture.py | 7 +++--- pyhdtoolkit/plotting/crossing.py | 15 ++++++++---- pyhdtoolkit/plotting/envelope.py | 4 +--- pyhdtoolkit/plotting/lattice.py | 16 ++++--------- pyhdtoolkit/plotting/phasespace.py | 5 +++- pyhdtoolkit/plotting/sbs/__init__.py | 1 + pyhdtoolkit/plotting/sbs/phase.py | 1 + pyhdtoolkit/plotting/styles/__init__.py | 2 ++ pyhdtoolkit/plotting/styles/thesis.py | 1 - pyhdtoolkit/plotting/tune.py | 1 + pyhdtoolkit/plotting/utils.py | 22 ++++++++++++++---- pyhdtoolkit/utils/_misc.py | 15 ++++++++---- pyhdtoolkit/utils/cmdline.py | 1 + pyhdtoolkit/utils/contexts.py | 5 +++- pyhdtoolkit/utils/htc_monitor.py | 29 +++++++++++++++++++----- pyhdtoolkit/utils/logging.py | 2 ++ pyhdtoolkit/utils/operations.py | 14 ++++++------ 45 files changed, 175 insertions(+), 82 deletions(-) diff --git a/pyhdtoolkit/__init__.py b/pyhdtoolkit/__init__.py index 0446552f..d9caa97e 100644 --- a/pyhdtoolkit/__init__.py +++ b/pyhdtoolkit/__init__.py @@ -7,6 +7,7 @@ :copyright: (c) 2019-2020 by Felix Soubelet. :license: MIT, see LICENSE for more details. """ + from . import cpymadtools, maths, models, optics, plotting, utils, version # noqa: F401, TID252 __title__ = "pyhdtoolkit" diff --git a/pyhdtoolkit/cpymadtools/constants.py b/pyhdtoolkit/cpymadtools/constants.py index 95baa499..329b4621 100644 --- a/pyhdtoolkit/cpymadtools/constants.py +++ b/pyhdtoolkit/cpymadtools/constants.py @@ -6,6 +6,7 @@ Specific constants to be used in `~.cpymadtools` functions, to help with consistency. """ + _MAX_SECTOR_VALUE: int = 8 # fmt: off @@ -136,16 +137,24 @@ for sector in [1, 2, 3, 4, 5, 6, 7, 8] ] LHC_KSS_KNOBS: list[str] = [ # skew sextupole correctors - f"kss.a{sector}{sector+1 if sector < _MAX_SECTOR_VALUE else 1}b{beam}" for beam in [1, 2] for sector in [1, 2, 3, 4, 5, 6, 7, 8] + f"kss.a{sector}{sector+1 if sector < _MAX_SECTOR_VALUE else 1}b{beam}" + for beam in [1, 2] + for sector in [1, 2, 3, 4, 5, 6, 7, 8] ] LHC_KCS_KNOBS: list[str] = [ # spool piece (skew) sextupoles - f"kcs.a{sector}{sector+1 if sector < _MAX_SECTOR_VALUE else 1}b{beam}" for beam in [1, 2] for sector in [1, 2, 3, 4, 5, 6, 7, 8] + f"kcs.a{sector}{sector+1 if sector < _MAX_SECTOR_VALUE else 1}b{beam}" + for beam in [1, 2] + for sector in [1, 2, 3, 4, 5, 6, 7, 8] ] LHC_KCO_KNOBS: list[str] = [ # spool piece (skew) octupoles - f"kco.a{sector}{sector+1 if sector < _MAX_SECTOR_VALUE else 1}b{beam}" for beam in [1, 2] for sector in [1, 2, 3, 4, 5, 6, 7, 8] + f"kco.a{sector}{sector+1 if sector < _MAX_SECTOR_VALUE else 1}b{beam}" + for beam in [1, 2] + for sector in [1, 2, 3, 4, 5, 6, 7, 8] ] LHC_KCD_KNOBS: list[str] = [ # spool piece (skew) decapoles - f"kcd.a{sector}{sector+1 if sector < _MAX_SECTOR_VALUE else 1}b{beam}" for beam in [1, 2] for sector in [1, 2, 3, 4, 5, 6, 7, 8] + f"kcd.a{sector}{sector+1 if sector < _MAX_SECTOR_VALUE else 1}b{beam}" + for beam in [1, 2] + for sector in [1, 2, 3, 4, 5, 6, 7, 8] ] LHC_KO_KNOBS: list[str] = [ # octupoles in arc short straight sections f"ko{family}.a{sector}{sector+1 if sector < _MAX_SECTOR_VALUE else 1}b{beam}" diff --git a/pyhdtoolkit/cpymadtools/coupling.py b/pyhdtoolkit/cpymadtools/coupling.py index 358a2945..5a469d23 100644 --- a/pyhdtoolkit/cpymadtools/coupling.py +++ b/pyhdtoolkit/cpymadtools/coupling.py @@ -8,6 +8,7 @@ Module with functions to perform ``MAD-X`` actions through a `~cpymad.madx.Madx` object, that retate to betatron coupling in the machine. """ + from collections.abc import Sequence import numpy as np @@ -78,10 +79,10 @@ def get_closest_tune_approach( # Say we have set the LHC coupling knobs to 1e-3 dqmin = get_closest_tune_approach( madx, - "lhc", # will find the knobs automatically + "lhc", # will find the knobs automatically sequence="lhcb1", telescopic_squeeze=True, # influences the knobs definition - run3=True, # influences the knobs definition (LHC Run 3) + run3=True, # influences the knobs definition (LHC Run 3) ) # returns 0.001 """ diff --git a/pyhdtoolkit/cpymadtools/lhc/_coupling.py b/pyhdtoolkit/cpymadtools/lhc/_coupling.py index e8e6f341..4f1928ac 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_coupling.py +++ b/pyhdtoolkit/cpymadtools/lhc/_coupling.py @@ -5,6 +5,7 @@ The functions below are betatron coupling utilities for the ``LHC``. """ + import tfs from cpymad.madx import Madx from loguru import logger diff --git a/pyhdtoolkit/cpymadtools/lhc/_elements.py b/pyhdtoolkit/cpymadtools/lhc/_elements.py index 74ad364b..c9c6b219 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_elements.py +++ b/pyhdtoolkit/cpymadtools/lhc/_elements.py @@ -5,11 +5,13 @@ The functions below are utilities to install elements or markers in the ``LHC``. """ + from cpymad.madx import Madx from loguru import logger _MAX_TRACKING_TOP_TURNS: int = 6600 + def install_ac_dipole_as_kicker( madx: Madx, /, diff --git a/pyhdtoolkit/cpymadtools/lhc/_errors.py b/pyhdtoolkit/cpymadtools/lhc/_errors.py index 70484f41..739202f8 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_errors.py +++ b/pyhdtoolkit/cpymadtools/lhc/_errors.py @@ -5,6 +5,7 @@ The functions below are utilities to implement errors in elements of the ``LHC``. """ + from collections.abc import Sequence from cpymad.madx import Madx @@ -122,7 +123,8 @@ def misalign_lhc_ir_quadrupoles( .. code-block:: python misalign_lhc_ir_quadrupoles( - madx, ips=[5], + madx, + ips=[5], quadrupoles=[7, 8, 9, 10], beam=1, sides="RL", @@ -140,7 +142,7 @@ def misalign_lhc_ir_quadrupoles( beam=1, sides="RL", dy=1e-5, # ok too as cpymad converts this to a string first - dpsi="1E-3 + 8E-4 * TGAUSS(2.5)" + dpsi="1E-3 + 8E-4 * TGAUSS(2.5)", ) """ if any(ip not in (1, 2, 5, 8) for ip in ips): diff --git a/pyhdtoolkit/cpymadtools/lhc/_misc.py b/pyhdtoolkit/cpymadtools/lhc/_misc.py index 35ee7afd..b2cf7e39 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_misc.py +++ b/pyhdtoolkit/cpymadtools/lhc/_misc.py @@ -23,6 +23,7 @@ _BEAM4: int = 4 # LHC beam 4 is special case _VRF_THRESHOLD: int = 5000 + def make_sixtrack_output(madx: Madx, /, energy: int) -> None: """ .. versionadded:: 0.15.0 diff --git a/pyhdtoolkit/cpymadtools/lhc/_powering.py b/pyhdtoolkit/cpymadtools/lhc/_powering.py index 7af54c4a..f28b3148 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_powering.py +++ b/pyhdtoolkit/cpymadtools/lhc/_powering.py @@ -5,6 +5,7 @@ The functions below are magnets or knobs powering utilities for the ``LHC``. """ + from collections.abc import Sequence from cpymad.madx import Madx @@ -14,6 +15,7 @@ _QUAD_CIRCUIT_HAS_B: int = 7 # Q7 has a .b in the circuit name _MAX_IR_QUAD_NUMBER: int = 11 # beyond Q11 are MQTs etc + def apply_lhc_colinearity_knob(madx: Madx, /, colinearity_knob_value: float = 0, ir: int | None = None) -> None: """ .. versionadded:: 0.15.0 diff --git a/pyhdtoolkit/cpymadtools/lhc/_queries.py b/pyhdtoolkit/cpymadtools/lhc/_queries.py index c20e6bfc..8277b164 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_queries.py +++ b/pyhdtoolkit/cpymadtools/lhc/_queries.py @@ -5,6 +5,7 @@ The functions below are settings query utilities for the ``LHC``. """ + from collections.abc import Sequence import tfs diff --git a/pyhdtoolkit/cpymadtools/lhc/_routines.py b/pyhdtoolkit/cpymadtools/lhc/_routines.py index 59e46884..fa5a7848 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_routines.py +++ b/pyhdtoolkit/cpymadtools/lhc/_routines.py @@ -5,6 +5,7 @@ The functions below are routines mimicking manipulations that would be done in the ``LHC``. """ + import numpy as np import tfs from cpymad.madx import Madx @@ -52,7 +53,9 @@ def do_kmodulation( .. code-block:: python - tune_results = do_kmodulation(madx, ir=1, side="right", steps=100, stepsize=3e-8) + tune_results = do_kmodulation( + madx, ir=1, side="right", steps=100, stepsize=3e-8 + ) """ element = f"MQXA.1R{ir:d}" if side.lower() == "right" else f"MQXA.1L{ir:d}" powering_variable = f"KTQX1.R{ir:d}" if side.lower() == "right" else f"KTQX1.L{ir:d}" diff --git a/pyhdtoolkit/cpymadtools/lhc/_setup.py b/pyhdtoolkit/cpymadtools/lhc/_setup.py index 28e56f41..e53874a0 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_setup.py +++ b/pyhdtoolkit/cpymadtools/lhc/_setup.py @@ -5,6 +5,7 @@ The functions below are setup utilities for the ``LHC``, to easily get simulations ready. """ + from pathlib import Path from cpymad.madx import Madx @@ -59,7 +60,9 @@ def prepare_lhc_run2( .. code-block:: python madx = prepare_lhc_run2( - "/afs/cern.ch/eng/lhc/optics/runII/2018/PROTON/opticsfile.22", beam=2, stdout=True + "/afs/cern.ch/eng/lhc/optics/runII/2018/PROTON/opticsfile.22", + beam=2, + stdout=True, ) """ if use_b4 and beam != _BEAM_FOR_B4: @@ -221,14 +224,19 @@ class LHCSetup: .. code-block:: python with LHCSetup(run=2, opticsfile="2018/PROTON/opticsfile.22", beam=2) as madx: - pass # do some stuff + pass # do some stuff Get a Run 3 setup for beam 1, with a sliced sequence and muted output: .. code-block:: python - with LHCSetup(run=3, opticsfile="R2022a_A30cmC30cmA10mL200cm.madx", slicefactor=4, stdout=False) as madx: - pass # do some stuff + with LHCSetup( + run=3, + opticsfile="R2022a_A30cmC30cmA10mL200cm.madx", + slicefactor=4, + stdout=False, + ) as madx: + pass # do some stuff """ def __init__( diff --git a/pyhdtoolkit/cpymadtools/lhc/_twiss.py b/pyhdtoolkit/cpymadtools/lhc/_twiss.py index c27567e4..2b32f829 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_twiss.py +++ b/pyhdtoolkit/cpymadtools/lhc/_twiss.py @@ -5,6 +5,7 @@ The functions below are twiss utilities for the ``LHC`` insertion regions. """ + from collections.abc import Sequence import tfs diff --git a/pyhdtoolkit/cpymadtools/matching.py b/pyhdtoolkit/cpymadtools/matching.py index 351076b8..2bcfb73b 100644 --- a/pyhdtoolkit/cpymadtools/matching.py +++ b/pyhdtoolkit/cpymadtools/matching.py @@ -6,6 +6,7 @@ Module with functions to perform ``MAD-X`` matchings through a `~cpymad.madx.Madx` object. """ + from collections.abc import Sequence from cpymad.madx import Madx @@ -101,7 +102,7 @@ def match_tunes_and_chromaticities( matching.match_tunes_and_chromaticities( madx, - None, # this is not LHC or HLLHC + None, # this is not LHC or HLLHC sequence="CAS3", q1_target=6.335, q2_target=6.29, @@ -131,7 +132,7 @@ def match_tunes_and_chromaticities( matching.match_tunes_and_chromaticities( madx, - "lhc", # will find the knobs automatically + "lhc", # will find the knobs automatically sequence="lhcb1", q1_target=62.31, q2_target=60.32, @@ -242,7 +243,7 @@ def match_tunes( matching.match_tunes( madx, - None, # this is not LHC or HLLHC + None, # this is not LHC or HLLHC sequence="CAS3", q1_target=6.335, q2_target=6.29, @@ -268,7 +269,7 @@ def match_tunes( matching.match_tunes( madx, - "lhc", # will find the knobs automatically + "lhc", # will find the knobs automatically sequence="lhcb1", q1_target=62.31, q2_target=60.32, @@ -344,7 +345,7 @@ def match_chromaticities( matching.match_chromaticities( madx, - None, # this is not LHC or HLLHC + None, # this is not LHC or HLLHC sequence="CAS3", dq1_target=100, dq2_target=100, @@ -370,7 +371,7 @@ def match_chromaticities( matching.match_chromaticities( madx, - "lhc", # will find the knobs automatically + "lhc", # will find the knobs automatically sequence="lhcb1", dq1_target=2.0, dq2_target=2.0, diff --git a/pyhdtoolkit/cpymadtools/ptc.py b/pyhdtoolkit/cpymadtools/ptc.py index 3a0b9e38..b4bbf9a0 100644 --- a/pyhdtoolkit/cpymadtools/ptc.py +++ b/pyhdtoolkit/cpymadtools/ptc.py @@ -7,6 +7,7 @@ Module with functions to manipulate ``MAD-X`` ``PTC`` functionality through a `~cpymad.madx.Madx` object. """ + from collections.abc import Sequence from pathlib import Path @@ -20,6 +21,7 @@ _MAX_PTC_AMPDET_ORDER: int = 2 _MIN_PTC_AMPDET_ORDER: int = 1 + def get_amplitude_detuning( madx: Madx, /, order: int = 2, file: Path | str | None = None, fringe: bool = False, **kwargs ) -> tfs.TfsDataFrame: @@ -428,8 +430,13 @@ def ptc_track_particle( .. code-block:: python tracks_dict = ptc_track_particle( - madx, nturns=10, initial_coordinates=(2e-4, 0, 1e-4, 0, 0, 0), - model=3, method=6, nst=3, exact=True + madx, + nturns=10, + initial_coordinates=(2e-4, 0, 1e-4, 0, 0, 0), + model=3, + method=6, + nst=3, + exact=True, ) """ logger.debug("Performing single particle PTC (thick) tracking") diff --git a/pyhdtoolkit/cpymadtools/track.py b/pyhdtoolkit/cpymadtools/track.py index 72ac93da..7a825409 100644 --- a/pyhdtoolkit/cpymadtools/track.py +++ b/pyhdtoolkit/cpymadtools/track.py @@ -7,6 +7,7 @@ Module with functions to manipulate ``MAD-X`` ``TRACK`` functionality through a `~cpymad.madx.Madx` object. """ + from collections.abc import Sequence import pandas as pd diff --git a/pyhdtoolkit/cpymadtools/tune.py b/pyhdtoolkit/cpymadtools/tune.py index 0c33f489..83c24534 100644 --- a/pyhdtoolkit/cpymadtools/tune.py +++ b/pyhdtoolkit/cpymadtools/tune.py @@ -7,6 +7,7 @@ Module with functions to manipulate ``MAD-X`` functionality around the tune through a `~cpymad.madx.Madx` object. """ + import math import sys from pathlib import Path @@ -250,7 +251,7 @@ def _get_dynap_string_rep(dynap_dframe: tfs.TfsDataFrame) -> str: return string_rep -def _make_tune_groups(dynap_string_rep: str, dsigma: float = 1.0) -> list[list[dict[str, float]]]: +def _make_tune_groups(dynap_string_rep: str, dsigma: float = 1.0) -> list[list[dict[str, float]]]: # noqa: ARG001 """ Creates appropriate tune points groups from the arcane string representation returned by `~.tune._get_dynap_string_rep` based on starting amplitude and angle for each particle. diff --git a/pyhdtoolkit/cpymadtools/twiss.py b/pyhdtoolkit/cpymadtools/twiss.py index c9f31277..1dd1089b 100644 --- a/pyhdtoolkit/cpymadtools/twiss.py +++ b/pyhdtoolkit/cpymadtools/twiss.py @@ -7,6 +7,7 @@ Module with functions to manipulate ``MAD-X`` ``TWISS`` functionality through a `~cpymad.madx.Madx` object. """ + from collections.abc import Sequence import tfs diff --git a/pyhdtoolkit/cpymadtools/utils.py b/pyhdtoolkit/cpymadtools/utils.py index e93cfd2a..5e313491 100644 --- a/pyhdtoolkit/cpymadtools/utils.py +++ b/pyhdtoolkit/cpymadtools/utils.py @@ -6,6 +6,7 @@ Module with utility functions to do mundane operations with `~cpymad.madx.Madx` objects. """ + from pathlib import Path import pandas as pd diff --git a/pyhdtoolkit/maths/stats_fitting.py b/pyhdtoolkit/maths/stats_fitting.py index 520a3103..95d8e4b5 100644 --- a/pyhdtoolkit/maths/stats_fitting.py +++ b/pyhdtoolkit/maths/stats_fitting.py @@ -6,6 +6,7 @@ Module implementing methods to find the best fit of statistical distributions to data. """ + import warnings import matplotlib.pyplot as plt # noqa: F401 | if omitted, get AttributeError: module 'matplotlib' has no attribute 'axes' @@ -46,6 +47,7 @@ def set_distributions_dict(dist_dict: dict[st.rv_continuous, str]) -> None: .. code-block:: python import scipy.stats as st + tested_dists = {st.chi: "Chi", st.expon: "Exponential", st.laplace: "Laplace"} set_distributions_dict(tested_dists) """ diff --git a/pyhdtoolkit/maths/utils.py b/pyhdtoolkit/maths/utils.py index 6df494d3..bbd6cf73 100644 --- a/pyhdtoolkit/maths/utils.py +++ b/pyhdtoolkit/maths/utils.py @@ -69,7 +69,8 @@ def get_scaled_values_and_magnitude_string( .. code-block:: python import numpy as np - q = np.array([-330, 230, 430, -720, 750, -110, 410, -340, -950, -630]) + + q = np.array([-330, 230, 430, -720, 750, -110, 410, -340, -950, -630]) get_scaled_values_and_magnitude_string(q) # returns (array([-3.3, 2.3, 4.3, -7.2, 7.5, -1.1, 4.1, -3.4, -9.5, -6.3]), '{-2}') """ diff --git a/pyhdtoolkit/models/beam.py b/pyhdtoolkit/models/beam.py index cc7a3d61..2e371639 100644 --- a/pyhdtoolkit/models/beam.py +++ b/pyhdtoolkit/models/beam.py @@ -6,6 +6,7 @@ Module with ``pydantic`` models to validate and store data structures relative to particle beams. """ + from math import sqrt from pydantic import BaseModel diff --git a/pyhdtoolkit/models/htc.py b/pyhdtoolkit/models/htc.py index 48aede86..3344817e 100644 --- a/pyhdtoolkit/models/htc.py +++ b/pyhdtoolkit/models/htc.py @@ -46,6 +46,7 @@ class HTCTaskSummary(BaseModel): Class to encompass and validate a specific job's line in the ``condor_q`` output. """ + # This is so pydantic accepts pendulum.DateTime as a validated type model_config = ConfigDict(arbitrary_types_allowed=True) diff --git a/pyhdtoolkit/models/madx.py b/pyhdtoolkit/models/madx.py index f7495882..1167e432 100644 --- a/pyhdtoolkit/models/madx.py +++ b/pyhdtoolkit/models/madx.py @@ -7,6 +7,7 @@ Module with ``pydantic`` models to validate and store data obtained by interacting with the ``MAD-X`` process through `cpymad`. """ + from enum import Enum from pydantic import BaseModel, PositiveFloat, PositiveInt diff --git a/pyhdtoolkit/optics/beam.py b/pyhdtoolkit/optics/beam.py index 8e0f3db2..a3b9fb64 100644 --- a/pyhdtoolkit/optics/beam.py +++ b/pyhdtoolkit/optics/beam.py @@ -6,6 +6,7 @@ Module implementing various functionality for simple beam parameter calculations. """ + import numpy as np from scipy import constants diff --git a/pyhdtoolkit/optics/rdt.py b/pyhdtoolkit/optics/rdt.py index 3e635a09..af60bcf5 100644 --- a/pyhdtoolkit/optics/rdt.py +++ b/pyhdtoolkit/optics/rdt.py @@ -23,14 +23,14 @@ def rdt_to_order_and_type(rdt: int | str) -> str: j, k, l, m = map(int, str(rdt)) # noqa: E741 rdt_type = "normal" if (l + m) % 2 == 0 else "skew" orders = { - 1: 'dipole', - 2: 'quadrupole', - 3: 'sextupole', - 4: 'octupole', - 5: 'decapole', - 6: 'dodecapole', - 7: 'tetradecapole', - 8: 'hexadecapole', + 1: "dipole", + 2: "quadrupole", + 3: "sextupole", + 4: "octupole", + 5: "decapole", + 6: "dodecapole", + 7: "tetradecapole", + 8: "hexadecapole", } return f"{rdt_type}_{orders[j + k + l + m]}" diff --git a/pyhdtoolkit/optics/ripken.py b/pyhdtoolkit/optics/ripken.py index e4d8331f..85503eb5 100644 --- a/pyhdtoolkit/optics/ripken.py +++ b/pyhdtoolkit/optics/ripken.py @@ -82,7 +82,6 @@ def _beam_size(coordinates_distribution: np.ndarray, method: str = "std") -> flo return np.sqrt(np.mean(np.square(coordinates_distribution))) # rms - def _add_beam_size_to_df(df: tfs.TfsDataFrame, geom_emit_x: float, geom_emit_y: float) -> tfs.TfsDataFrame: """ Adds columns with the horizontal and vertical Lebedev beam sizes to a dataframe diff --git a/pyhdtoolkit/optics/twiss.py b/pyhdtoolkit/optics/twiss.py index 08cdc265..22261144 100644 --- a/pyhdtoolkit/optics/twiss.py +++ b/pyhdtoolkit/optics/twiss.py @@ -6,6 +6,7 @@ Module implementing various calculations based on the ``TWISS`` optics parameters. """ + import numpy as np diff --git a/pyhdtoolkit/plotting/__init__.py b/pyhdtoolkit/plotting/__init__.py index ba9eddf5..4825e5f2 100644 --- a/pyhdtoolkit/plotting/__init__.py +++ b/pyhdtoolkit/plotting/__init__.py @@ -1,6 +1,7 @@ """ .. _plotting: """ + from . import aperture, crossing, envelope, lattice, phasespace, tune, utils # noqa: TID252 __all__ = ["aperture", "envelope", "crossing", "lattice", "phasespace", "tune", "utils"] diff --git a/pyhdtoolkit/plotting/aperture.py b/pyhdtoolkit/plotting/aperture.py index 32a27e7d..b10a927d 100644 --- a/pyhdtoolkit/plotting/aperture.py +++ b/pyhdtoolkit/plotting/aperture.py @@ -124,7 +124,8 @@ def plot_aperture( plt.figure(figsize=(16, 11)) plot_aperture( - madx, plot_bpms=True, + madx, + plot_bpms=True, aperture_ylim=(0, 20), k0l_lim=(-4e-4, 4e-4), k1l_lim=(-0.08, 0.08), @@ -168,9 +169,7 @@ def plot_aperture( # Plotting aperture values on remaining two thirds of the figure logger.debug("Plotting aperture values") aperture_axis = plt.subplot2grid((3, 3), (1, 0), colspan=3, rowspan=2, sharex=quadrupole_patches_axis) - aperture_axis.plot( - aperture_df.s, aperture_df.n1, marker=".", ls="-", lw=0.8, color=color, label="Aperture Limits" - ) + aperture_axis.plot(aperture_df.s, aperture_df.n1, marker=".", ls="-", lw=0.8, color=color, label="Aperture Limits") aperture_axis.fill_between(aperture_df.s, aperture_df.n1, aperture_df.n1.max(), interpolate=True, color=color) aperture_axis.legend() aperture_axis.set_ylabel(r"$n_{1} \ [\sigma]$") diff --git a/pyhdtoolkit/plotting/crossing.py b/pyhdtoolkit/plotting/crossing.py index 958ea0c5..45ca57a3 100644 --- a/pyhdtoolkit/plotting/crossing.py +++ b/pyhdtoolkit/plotting/crossing.py @@ -60,7 +60,9 @@ def plot_two_lhc_ips_crossings( .. code-block:: python plt.figure(figsize=(16, 11)) - plot_two_lhc_ips_crossings(madx, first_ip=2, second_ip=8, highlight_mqx_and_mbx=False) + plot_two_lhc_ips_crossings( + madx, first_ip=2, second_ip=8, highlight_mqx_and_mbx=False + ) """ logger.warning("You should re-call the 'USE' command on your wanted sequence after this plot!") # ----- Getting Twiss table dframe for each beam ----- # @@ -168,7 +170,12 @@ def plot_single_ir_crossing( .. code-block:: python plot_single_ir_crossing( - plt.gca(), b1_df, b2_df, plot_column="x", scaling=1e3, ylabel="Orbit X $[mm]$" + plt.gca(), + b1_df, + b2_df, + plot_column="x", + scaling=1e3, + ylabel="Orbit X $[mm]$", ) """ logger.trace(f"Plotting orbit '{plot_column}'") @@ -183,9 +190,7 @@ def plot_single_ir_crossing( # ----- Helpers ----- # -def _highlight_mbx_and_mqx( - axis: Axes, plot_df: pd.DataFrame | tfs.TfsDataFrame, ip: int, **kwargs -) -> None: +def _highlight_mbx_and_mqx(axis: Axes, plot_df: pd.DataFrame | tfs.TfsDataFrame, ip: int, **kwargs) -> None: """ .. versionadded:: 1.0.0 diff --git a/pyhdtoolkit/plotting/envelope.py b/pyhdtoolkit/plotting/envelope.py index dcc054fd..1ea45902 100644 --- a/pyhdtoolkit/plotting/envelope.py +++ b/pyhdtoolkit/plotting/envelope.py @@ -104,9 +104,7 @@ def plot_beam_envelope( orbit = twiss_df[plane_letter] * scale # with scaling factor, by default 1 logger.debug("Calculating beam enveloppe") - one_sigma = np.sqrt( - geom_emit * twiss_df[f"bet{plane_letter}"] + (sige * twiss_df[f"d{plane_letter}"]) ** 2 - ) + one_sigma = np.sqrt(geom_emit * twiss_df[f"bet{plane_letter}"] + (sige * twiss_df[f"d{plane_letter}"]) ** 2) enveloppe = nsigma * one_sigma * scale # with scaling factor, by default 1 # Plot a line for the orbit, then fill between orbit + enveloppe and orbit - enveloppe diff --git a/pyhdtoolkit/plotting/lattice.py b/pyhdtoolkit/plotting/lattice.py index 937afe2c..ef20493b 100644 --- a/pyhdtoolkit/plotting/lattice.py +++ b/pyhdtoolkit/plotting/lattice.py @@ -251,21 +251,13 @@ def plot_machine_survey( label="Dipoles", **kwargs, ) - plt.scatter( - element_dfs["quad_foc"].z, element_dfs["quad_foc"].x, marker="o", color="blue", label="QF" - ) - plt.scatter( - element_dfs["quad_defoc"].z, element_dfs["quad_defoc"].x, marker="o", color="red", label="QD" - ) + plt.scatter(element_dfs["quad_foc"].z, element_dfs["quad_foc"].x, marker="o", color="blue", label="QF") + plt.scatter(element_dfs["quad_defoc"].z, element_dfs["quad_defoc"].x, marker="o", color="red", label="QD") if high_orders: logger.debug("Plotting high order magnetic elements (up to octupoles)") - plt.scatter( - element_dfs["sextupoles"].z, element_dfs["sextupoles"].x, marker=".", color="m", label="MS" - ) - plt.scatter( - element_dfs["octupoles"].z, element_dfs["octupoles"].x, marker=".", color="cyan", label="MO" - ) + plt.scatter(element_dfs["sextupoles"].z, element_dfs["sextupoles"].x, marker=".", color="m", label="MS") + plt.scatter(element_dfs["octupoles"].z, element_dfs["octupoles"].x, marker=".", color="cyan", label="MO") plt.legend(loc=2) else: diff --git a/pyhdtoolkit/plotting/phasespace.py b/pyhdtoolkit/plotting/phasespace.py index ba182f6a..acd7e783 100644 --- a/pyhdtoolkit/plotting/phasespace.py +++ b/pyhdtoolkit/plotting/phasespace.py @@ -6,6 +6,7 @@ Module with functions to create phase space plots through a `~cpymad.madx.Madx` object. """ + import numpy as np from cpymad.madx import Madx from loguru import logger @@ -128,7 +129,9 @@ def plot_courant_snyder_phase_space_colored( .. code-block:: python fig, ax = plt.subplots(figsize=(10, 9)) - plot_courant_snyder_phase_space_colored(madx, x_coords, px_coords, plane="Horizontal") + plot_courant_snyder_phase_space_colored( + madx, x_coords, px_coords, plane="Horizontal" + ) """ if plane.upper() not in ("HORIZONTAL", "VERTICAL"): logger.error(f"Plane should be either horizontal or vertical but '{plane}' was given") diff --git a/pyhdtoolkit/plotting/sbs/__init__.py b/pyhdtoolkit/plotting/sbs/__init__.py index 8b575c84..a2c7ddf2 100644 --- a/pyhdtoolkit/plotting/sbs/__init__.py +++ b/pyhdtoolkit/plotting/sbs/__init__.py @@ -1,6 +1,7 @@ """ .. _plotting.sbs: """ + from . import coupling, phase # noqa: TID252 __all__ = ["coupling", "phase"] diff --git a/pyhdtoolkit/plotting/sbs/phase.py b/pyhdtoolkit/plotting/sbs/phase.py index 71b1d9a5..20f21463 100644 --- a/pyhdtoolkit/plotting/sbs/phase.py +++ b/pyhdtoolkit/plotting/sbs/phase.py @@ -6,6 +6,7 @@ Functions to plot phase values of Segment-by-Segment results. """ + import matplotlib.pyplot as plt import tfs from loguru import logger diff --git a/pyhdtoolkit/plotting/styles/__init__.py b/pyhdtoolkit/plotting/styles/__init__.py index ac37b328..18031177 100644 --- a/pyhdtoolkit/plotting/styles/__init__.py +++ b/pyhdtoolkit/plotting/styles/__init__.py @@ -7,6 +7,7 @@ The **style** submodules provide styles to be used with `~matplotlib`, mostly tailored for my use, and for good results with the plotters in `~pyhdtoolkit.plotting`. Feel free to use them anyway, as they might be useful to you when using the `~pyhdtoolkit.plotting` submodules, or to be adapted. """ + from pathlib import Path import matplotlib as mpl @@ -60,6 +61,7 @@ def _install_style_file(style: dict[str, PlotSetting], stylename) -> None: .. code-block:: python from matplotlib import pyplot as plt + plt.style.use("style-name") .. note:: diff --git a/pyhdtoolkit/plotting/styles/thesis.py b/pyhdtoolkit/plotting/styles/thesis.py index 69ff6d2d..355323a2 100644 --- a/pyhdtoolkit/plotting/styles/thesis.py +++ b/pyhdtoolkit/plotting/styles/thesis.py @@ -16,7 +16,6 @@ - ``LARGE``: For more complex plots to be included alone in a LaTeX figure environment (e.g. multi-axes figures such as in `~pyhdtoolkit.plotting.lattice`). """ - PlotSetting = float | bool | str | tuple diff --git a/pyhdtoolkit/plotting/tune.py b/pyhdtoolkit/plotting/tune.py index 497c10fe..5926799b 100644 --- a/pyhdtoolkit/plotting/tune.py +++ b/pyhdtoolkit/plotting/tune.py @@ -7,6 +7,7 @@ Module with functions to create tune diagram plots. These provide functionality to draw Farey sequences up to a desired order. """ + from functools import partial import numpy as np diff --git a/pyhdtoolkit/plotting/utils.py b/pyhdtoolkit/plotting/utils.py index d767f460..79b1d1dc 100644 --- a/pyhdtoolkit/plotting/utils.py +++ b/pyhdtoolkit/plotting/utils.py @@ -6,6 +6,7 @@ Module with functions to used throught the different `~pyhdtoolkit.plotting` modules. """ + from __future__ import annotations # important for Sphinx to generate short type signatures! from typing import TYPE_CHECKING @@ -332,7 +333,12 @@ def set_arrow_label( va="center", ha="center", bbox={"boxstyle": "round4", "fc": "w", "color": color}, - arrowprops={"arrowstyle": "-|>", "connectionstyle": "arc3,rad=" + str(arrow_arc_rad), "fc": "w", "color": color}, + arrowprops={ + "arrowstyle": "-|>", + "connectionstyle": "arc3,rad=" + str(arrow_arc_rad), + "fc": "w", + "color": color, + }, **kwargs, ) @@ -430,8 +436,7 @@ def _get_twiss_table_with_offsets_and_limits( madx.command.twiss(**kwargs) twiss_df = madx.table.twiss.dframe() twiss_df.s = twiss_df.s - xoffset - twiss_df = twiss_df[twiss_df.s.between(*xlimits)] if xlimits else twiss_df - return twiss_df + return twiss_df[twiss_df.s.between(*xlimits)] if xlimits else twiss_df def _determine_default_sbs_coupling_ylabel(rdt: str, component: str) -> str: @@ -452,7 +457,9 @@ def _determine_default_sbs_coupling_ylabel(rdt: str, component: str) -> str: Example: .. code-block:: python - coupling_label = _determine_default_sbs_coupling_ylabel(rdt="f1001", component="RE") + coupling_label = _determine_default_sbs_coupling_ylabel( + rdt="f1001", component="RE" + ) """ logger.debug(f"Determining a default label for the {component.upper()} component of coupling {rdt.upper()}.") if rdt.upper() not in ("F1001", "F1010"): @@ -491,17 +498,22 @@ def _determine_default_sbs_phase_ylabel(plane: str) -> str: phase_label = _determine_default_sbs_phase_ylabel(plane="X") """ logger.debug(f"Determining a default label for the {plane.upper()} phase plane.") - assert plane.upper() in ("X", "Y") + if plane.upper() not in ("X", "Y"): + msg = "Invalid plane for phase plot." + raise ValueError(msg) beginning = r"\Delta " term = r"\phi_{x}" if plane.upper() == "X" else r"\phi_{y}" return r"$" + beginning + term + r"$" + # ----- Sphinx Gallery Scraper ----- # + # To use SVG outputs when scraping matplotlib figures for the sphinx-gallery, see: # https://sphinx-gallery.github.io/stable/advanced.html#example-3-matplotlib-with-svg-format def _matplotlib_svg_scraper(*args, **kwargs): from sphinx_gallery.scrapers import matplotlib_scraper + kwargs.pop("format", None) return matplotlib_scraper(*args, format="svg", **kwargs) diff --git a/pyhdtoolkit/utils/_misc.py b/pyhdtoolkit/utils/_misc.py index 18369166..93790f5f 100644 --- a/pyhdtoolkit/utils/_misc.py +++ b/pyhdtoolkit/utils/_misc.py @@ -10,6 +10,7 @@ The functions in here are intented for personal use, and will most likely **not** work on other people's machines. """ + import shlex from collections.abc import Sequence from multiprocessing import cpu_count @@ -80,7 +81,9 @@ def split_complex_columns(df: pd.DataFrame, drop: bool = False) -> pd.DataFrame: return res -def add_noise_to_ir_bpms(df: pd.DataFrame, max_index: int, stdev: float, columns: Sequence[str] | None = None) -> pd.DataFrame: +def add_noise_to_ir_bpms( + df: pd.DataFrame, max_index: int, stdev: float, columns: Sequence[str] | None = None +) -> pd.DataFrame: """ .. versionadded:: 1.2.0 @@ -213,13 +216,13 @@ def apply_colin_corrs_balance(madx: Madx) -> None: # ----- Fetching Utilities ----- # -def get_betastar_from_opticsfile(opticsfile: Path | str) -> float: +def get_betastar_from_opticsfile(opticsfile: Path | str, check_symmetry: bool = True) -> float: """ .. versionadded:: 0.16.0 Parses the :math:`\\beta^{*}` value from the *opticsfile* content, - which is in the first lines. This contains a check that ensures the betastar - is the same for IP1 and IP5. The values returned are in meters. + which is in the first lines. This contains an optional check to ensure + the betastar is the same for IP1 and IP5. The values returned are in meters. .. note:: For file in ``acc-models-lhc`` make sure to point to the strength file @@ -251,5 +254,7 @@ def get_betastar_from_opticsfile(opticsfile: Path | str) -> float: betastar_y_ip1 = float(shlex.split(ip1_y_line)[2]) betastar_x_ip5 = float(shlex.split(ip5_x_line)[2]) betastar_y_ip5 = float(shlex.split(ip5_y_line)[2]) - assert betastar_x_ip1 == betastar_y_ip1 == betastar_x_ip5 == betastar_y_ip5 + if check_symmetry and not (betastar_x_ip1 == betastar_y_ip1 == betastar_x_ip5 == betastar_y_ip5): + msg = "The betastar values for IP1 and IP5 are not the same in both planes." + raise AssertionError(msg) return betastar_x_ip1 # doesn't matter which plane, they're all the same diff --git a/pyhdtoolkit/utils/cmdline.py b/pyhdtoolkit/utils/cmdline.py index f552f62d..d2851fa8 100644 --- a/pyhdtoolkit/utils/cmdline.py +++ b/pyhdtoolkit/utils/cmdline.py @@ -109,6 +109,7 @@ def run( .. code-block:: python import os + modified_env = os.environ.copy() modified_env["ENV_VAR"] = "new_value" CommandLine.run("echo $ENV_VAR", env=modified_env) diff --git a/pyhdtoolkit/utils/contexts.py b/pyhdtoolkit/utils/contexts.py index 0bd4c72d..be58bb30 100644 --- a/pyhdtoolkit/utils/contexts.py +++ b/pyhdtoolkit/utils/contexts.py @@ -6,6 +6,7 @@ Provides useful contexts to use functions in. """ + import time from collections.abc import Callable, Iterator from contextlib import contextmanager @@ -29,7 +30,9 @@ def timeit(function: Callable) -> Iterator[None]: Example: .. code-block:: python - with timeit(lambda spanned: logger.debug(f"Did some stuff in {spanned} seconds")): + with timeit( + lambda spanned: logger.debug(f"Did some stuff in {spanned} seconds") + ): some_stuff() some_other_stuff() """ diff --git a/pyhdtoolkit/utils/htc_monitor.py b/pyhdtoolkit/utils/htc_monitor.py index db4d3c17..e842e3d9 100644 --- a/pyhdtoolkit/utils/htc_monitor.py +++ b/pyhdtoolkit/utils/htc_monitor.py @@ -12,6 +12,7 @@ functionality is made public API and one shoule be able to build a different monitor from the functions in here. """ + import re import time @@ -34,9 +35,19 @@ TASK_COLUMNS_SETTINGS = { "OWNER": {"justify": "left", "header_style": "bold", "style": "bold", "no_wrap": True}, "BATCH_NAME": {"justify": "center", "header_style": "magenta", "style": "magenta", "no_wrap": True}, - "SUBMITTED": {"justify": "center", "header_style": "medium_turquoise", "style": "medium_turquoise", "no_wrap": True}, + "SUBMITTED": { + "justify": "center", + "header_style": "medium_turquoise", + "style": "medium_turquoise", + "no_wrap": True, + }, "DONE": {"justify": "right", "header_style": "bold green3", "style": "bold green3", "no_wrap": True}, - "RUNNING": {"justify": "right", "header_style": "bold cornflower_blue", "style": "bold cornflower_blue", "no_wrap": True}, + "RUNNING": { + "justify": "right", + "header_style": "bold cornflower_blue", + "style": "bold cornflower_blue", + "no_wrap": True, + }, "IDLE": {"justify": "right", "header_style": "bold dark_orange3", "style": "bold dark_orange3", "no_wrap": True}, "TOTAL": {"justify": "right", "style": "bold", "no_wrap": True}, "JOB_IDS": {"justify": "right", "no_wrap": True}, @@ -46,7 +57,12 @@ "SOURCE": {"justify": "left", "header_style": "bold", "style": "bold", "no_wrap": True}, "JOBS": {"justify": "right", "header_style": "bold", "style": "bold", "no_wrap": True}, "COMPLETED": {"justify": "right", "header_style": "bold green3", "style": "bold green3", "no_wrap": True}, - "RUNNING": {"justify": "right", "header_style": "bold cornflower_blue", "style": "bold cornflower_blue", "no_wrap": True}, + "RUNNING": { + "justify": "right", + "header_style": "bold cornflower_blue", + "style": "bold cornflower_blue", + "no_wrap": True, + }, "IDLE": {"justify": "right", "header_style": "bold dark_orange3", "style": "bold dark_orange3", "no_wrap": True}, "HELD": {"justify": "right", "header_style": "bold gold1", "style": "bold gold1", "no_wrap": True}, "SUSPENDED": {"justify": "right", "header_style": "bold slate_blue1", "style": "bold slate_blue1", "no_wrap": True}, @@ -68,9 +84,10 @@ def query_condor_q() -> str: condor_status = raw_result.decode().strip() if return_code == 0: return condor_status - else: - msg = "Checking htcondor status failed" - raise ChildProcessError(msg) + + # An issue occured, let's raise + msg = "Checking htcondor status failed" + raise ChildProcessError(msg) def read_condor_q(report: str) -> tuple[list[HTCTaskSummary], ClusterSummary]: diff --git a/pyhdtoolkit/utils/logging.py b/pyhdtoolkit/utils/logging.py index ddf6ef7f..68736883 100644 --- a/pyhdtoolkit/utils/logging.py +++ b/pyhdtoolkit/utils/logging.py @@ -14,6 +14,7 @@ - ``FORMAT2``: similar to ``FORMAT1``, but the caller information is displayed at the end of the log line. - ``SIMPLE_FORMAT``: minimal, displays the local time, the level and the message. """ + import sys from loguru import logger @@ -63,6 +64,7 @@ def config_logger(level: str | int = "INFO", fmt: str = FORMAT1, **kwargs) -> No .. code-block:: python from pyhdtoolkit.utils.logging import config_logger, SIMPLE_FORMAT + config_logger(level="DEBUG", format=SIMPLE_FORMAT) """ logger.remove() diff --git a/pyhdtoolkit/utils/operations.py b/pyhdtoolkit/utils/operations.py index 6c4c71a9..dfbfa333 100644 --- a/pyhdtoolkit/utils/operations.py +++ b/pyhdtoolkit/utils/operations.py @@ -68,7 +68,7 @@ def average_by(sequence: Sequence, function: Callable = lambda x: x) -> float: .. code-block:: python ListOperations.average_by( - [{'n': 4}, {'n': 2}, {'n': 8}, {'n': 6}], lambda x: x['n'] + [{"n": 4}, {"n": 2}, {"n": 8}, {"n": 6}], lambda x: x["n"] ) # returns 5.0 """ @@ -94,7 +94,9 @@ def bifurcate(sequence: Sequence, filters: list[bool]) -> Sequence: Example: .. code-block:: python - ListOperations.bifurcate(['beep', 'boop', 'foo', 'bar'], [True, True, False, True]) + ListOperations.bifurcate( + ["beep", "boop", "foo", "bar"], [True, True, False, True] + ) # returns [['beep', 'boop', 'bar'], ['foo']] """ return [ @@ -582,7 +584,7 @@ def map_values(obj: dict, function: Callable) -> dict: MiscellaneousOperations.map_values( {"a": list(range(5)), "b": list(range(10)), "c": list(range(15))}, - lambda x: len(x) + lambda x: len(x), ) # returns {"a": 5, "b": 10, "c": 15} """ @@ -630,9 +632,7 @@ def clamp_number(num: float, a_val: float, b_val: float) -> float: return max(min(num, max(a_val, b_val)), min(a_val, b_val)) @staticmethod - def degrees_to_radians( - deg_value: float, decompose: bool = False - ) -> tuple[float, str, str] | float: + def degrees_to_radians(deg_value: float, decompose: bool = False) -> tuple[float, str, str] | float: """ .. versionadded:: 0.2.0 @@ -764,7 +764,7 @@ def radians_to_degrees(rad_value: float) -> float: Examples: .. code-block:: python - NumberOperations.radians_to_degrees(2* math.pi) + NumberOperations.radians_to_degrees(2 * math.pi) # returns 360 .. code-block:: python From 6c8c432c4706c3254dfb12419faeb0e4a3ec672d Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 12:19:59 +0200 Subject: [PATCH 25/41] no mypy --- pyproject.toml | 8 -------- 1 file changed, 8 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 67311ce9..945c722c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -175,14 +175,6 @@ dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$" # Assume Python 3.10+. target-version = "py310" -[tool.mypy] -pretty = true -strict_optional = false -show_error_codes = true -warn_redundant_casts = true -ignore_missing_imports = true -follow_imports = "skip" - [tool.interrogate] ignore-init-method = true ignore-init-module = true From f0b8044fe88c2b64813319c4a9e43e66d7a232a6 Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 12:21:33 +0200 Subject: [PATCH 26/41] no interrogate --- pyproject.toml | 16 ---------------- 1 file changed, 16 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 945c722c..c7d6d401 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -174,19 +174,3 @@ dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$" # Assume Python 3.10+. target-version = "py310" - -[tool.interrogate] -ignore-init-method = true -ignore-init-module = true -ignore-magic = false -ignore-semiprivate = false -ignore-private = false -ignore-property-decorators = false -ignore-module = false -fail-under = 95 -exclude = ["tests"] -ignore-regex = [] -verbose = 2 -quiet = false -whitelist-regex = [] -color = true From fb5192f81e542dfd909183db1070a1fd32af1a99 Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 12:29:38 +0200 Subject: [PATCH 27/41] rename to nemitt_[xy], accept old and let user override npart and sigma_e --- pyhdtoolkit/cpymadtools/lhc/_setup.py | 43 ++++++++++++++++++--------- 1 file changed, 29 insertions(+), 14 deletions(-) diff --git a/pyhdtoolkit/cpymadtools/lhc/_setup.py b/pyhdtoolkit/cpymadtools/lhc/_setup.py index e53874a0..4e81a550 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_setup.py +++ b/pyhdtoolkit/cpymadtools/lhc/_setup.py @@ -283,8 +283,8 @@ def make_lhc_beams( madx: Madx, /, energy: float = 7000, - emittance_x: float = 3.75e-6, - emittance_y: float = 3.75e-6, + nemitt_x: float = 3.75e-6, + nemitt_y: float = 3.75e-6, b4: bool = False, **kwargs, ) -> None: @@ -294,34 +294,49 @@ def make_lhc_beams( Defines beams with default configuratons for ``LHCB1`` and ``LHCB2`` sequences. Args: - madx (cpymad.madx.Madx): an instanciated `~cpymad.madx.Madx` object. Positional only. + madx (cpymad.madx.Madx): an instanciated `~cpymad.madx.Madx` object. + Positional only. 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_y (float): vertical emittance in [m]. Will be used to calculate - geometric emittance which is then fed to the ``BEAM`` command. + nemitt_x (float): normalized horizontal emittance in [m]. Will + be used to calculate geometric emittance which is then fed to + the ``BEAM`` command. + nemitt_y (float): normalized vertical emittance in [m]. Will be + used to calculate geometric emittance which is then fed to + the ``BEAM`` command. b4 (bool): if `True`, will consider one is using ``lhb4`` to do tracking on beam 2, and will properly set the ``bv`` flag to 1. Defaults to `False`. - **kwargs: Any keyword argument that can be given to the ``MAD-X`` ``BEAM`` command. + **kwargs: Old accepted `emittance_x` and `emittance_y` are looked for and used + if provided. Any other keyword argument is given to the ``MAD-X`` ``BEAM`` command. Examples: .. code-block:: python - make_lhc_beams(madx, energy=6800, emittance_x=2.5e-6, emittance_y=3e-6) + make_lhc_beams(madx, energy=6800, nemitt_x=2.5e-6, nemitt_y=3e-6) Setting up in a way compatible for tracking of beam 2 (needs to call ``lhcb4`` and set ``bv`` to 1): .. code-block:: python - make_lhc_beams(madx, energy=6800, emittance_x=2.5e-6, emittance_y=3e-6, b4=True) + make_lhc_beams(madx, energy=6800, nemitt_x=2.5e-6, nemitt_y=3e-6, b4=True) """ logger.debug("Making default beams for 'lhcb1' and 'lhbc2' sequences") madx.globals["NRJ"] = energy madx.globals["brho"] = energy * 1e9 / madx.globals.clight - geometric_emit_x = madx.globals["geometric_emit_x"] = emittance_x / (energy / 0.938) - geometric_emit_y = madx.globals["geometric_emit_y"] = emittance_y / (energy / 0.938) + + # TODO: remove accepting the old emittance_x and emittance_y args in the future + emittance_x = kwargs.pop("emittance_x", None) + emittance_y = kwargs.pop("emittance_y", None) + if emittance_x is not None or emittance_y is not None: + logger.warning("Providing ambiguous 'emittance_[xy]' is deprecated, use 'nemitt_[xy]' instead") + nemitt_x = emittance_x if emittance_x is not None else nemitt_x # prioritize old ones if given + nemitt_y = emittance_y if emittance_y is not None else nemitt_y # prioritize old ones if given + geometric_emit_x = madx.globals["geometric_emit_x"] = nemitt_x / (energy / 0.938) + geometric_emit_y = madx.globals["geometric_emit_y"] = nemitt_y / (energy / 0.938) + + n_part = kwargs.pop("npart", 1.15e11) # keep default, let user override + sigma_e = kwargs.pop("sige", 4.5e-4) # keep default, let user override for beam in (1, 2): bv = 1 if beam == 1 or b4 is True else -1 @@ -331,10 +346,10 @@ def make_lhc_beams( particle="proton", bv=bv, energy=energy, - npart=1.15e11, + npart=n_part, ex=geometric_emit_x, ey=geometric_emit_y, - sige=4.5e-4, + sige=sigma_e, **kwargs, ) From 55edac28a5e2d0dd1367b2942456a65e8bbbe6e2 Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 12:30:23 +0200 Subject: [PATCH 28/41] formatting --- examples/demo_stats_fitting.py | 2 +- pyhdtoolkit/cpymadtools/lhc/_setup.py | 2 +- tests/test_plotting/test_phasespace.py | 3 +-- tests/test_plotting/test_plotting_utils.py | 2 +- tests/test_utils.py | 3 +-- 5 files changed, 5 insertions(+), 7 deletions(-) diff --git a/examples/demo_stats_fitting.py b/examples/demo_stats_fitting.py index 5b1dd82d..7894b3be 100755 --- a/examples/demo_stats_fitting.py +++ b/examples/demo_stats_fitting.py @@ -174,7 +174,7 @@ def chi_dist(num: int, meas_used: int) -> np.ndarray: ax=ax, ) param_names = (best_fit_func.shapes + ", loc, scale").split(", ") if best_fit_func.shapes else ["loc", "scale"] -param_str = ", ".join([f"{k}={v:0.2f}" for k, v in zip(param_names, best_fit_params)]) +param_str = ", ".join([f"{k}={v:0.2f}" for k, v in zip(param_names, best_fit_params, strict=False)]) dist_str = f"{fitting.DISTRIBUTIONS[best_fit_func]}({param_str})" # Let's add to the plot some info on the fit's peak diff --git a/pyhdtoolkit/cpymadtools/lhc/_setup.py b/pyhdtoolkit/cpymadtools/lhc/_setup.py index 4e81a550..607b0daa 100644 --- a/pyhdtoolkit/cpymadtools/lhc/_setup.py +++ b/pyhdtoolkit/cpymadtools/lhc/_setup.py @@ -324,7 +324,7 @@ def make_lhc_beams( logger.debug("Making default beams for 'lhcb1' and 'lhbc2' sequences") madx.globals["NRJ"] = energy madx.globals["brho"] = energy * 1e9 / madx.globals.clight - + # TODO: remove accepting the old emittance_x and emittance_y args in the future emittance_x = kwargs.pop("emittance_x", None) emittance_y = kwargs.pop("emittance_y", None) diff --git a/tests/test_plotting/test_phasespace.py b/tests/test_plotting/test_phasespace.py index 0a3c3f26..e694bd74 100644 --- a/tests/test_plotting/test_phasespace.py +++ b/tests/test_plotting/test_phasespace.py @@ -1,5 +1,4 @@ from functools import partial -from typing import Tuple import matplotlib import matplotlib.pyplot as plt @@ -118,7 +117,7 @@ def test_plot_courant_snyder_phase_space_colored_wrong_plane_input(): ) -def _perform_tracking_for_coordinates(madx: Madx) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: +def _perform_tracking_for_coordinates(madx: Madx) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """ Tracks 100 particles on 500 turns. This modifies inplace the state of the provided cpymad_instance. diff --git a/tests/test_plotting/test_plotting_utils.py b/tests/test_plotting/test_plotting_utils.py index dbef6c4b..c81d3fe4 100644 --- a/tests/test_plotting/test_plotting_utils.py +++ b/tests/test_plotting/test_plotting_utils.py @@ -113,7 +113,7 @@ def test_confidence_ellipse_subplots(): scale = 3, 5 figure, axs = plt.subplots(1, 3, figsize=(9, 3)) - for ax, (title, dependency) in zip(axs, PARAMETERS.items()): + for ax, (title, dependency) in zip(axs, PARAMETERS.items(), strict=False): x, y = get_correlated_dataset(800, dependency, mu, scale) ax.scatter(x, y, s=0.5) ax.axvline(c="grey", lw=1) diff --git a/tests/test_utils.py b/tests/test_utils.py index 5b6bdd0d..ada1904e 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -6,7 +6,6 @@ import random import subprocess import sys -from typing import List import numpy as np import pandas as pd @@ -723,7 +722,7 @@ def _taskless_condor_q_output() -> str: @pytest.fixture -def _correct_user_tasks() -> List[HTCTaskSummary]: +def _correct_user_tasks() -> list[HTCTaskSummary]: pickle_file_path = INPUTS_DIR / "utils" / "correct_user_tasks.pkl" with pickle_file_path.open("rb") as file: return pickle.load(file) From 8848f0bbbda94130487b05f0e8c896b9dab25379 Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 12:32:16 +0200 Subject: [PATCH 29/41] remove trailing whitespaces --- docs/conf.py | 4 ++-- examples/demo_acd_tracking_spectra.py | 12 ++++++------ examples/demo_aperture.py | 2 +- examples/demo_beam_enveloppe.py | 10 +++++----- examples/demo_ir_errors.py | 2 +- examples/demo_lhc_rigid_waist_shift.py | 10 +++++----- examples/demo_lhc_setup.py | 2 +- examples/demo_track_spectra.py | 2 +- examples/demo_tune_diagram.py | 2 +- 9 files changed, 23 insertions(+), 23 deletions(-) diff --git a/docs/conf.py b/docs/conf.py index 020afdb2..1fe7adf5 100644 --- 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/examples/demo_acd_tracking_spectra.py b/examples/demo_acd_tracking_spectra.py index b002b50d..f8cad24f 100755 --- 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 ` - 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 ` + 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:: diff --git a/examples/demo_aperture.py b/examples/demo_aperture.py index bc540337..0836ec66 100755 --- 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. diff --git a/examples/demo_beam_enveloppe.py b/examples/demo_beam_enveloppe.py index c2e5e1a4..ff1ab0bd 100755 --- a/examples/demo_beam_enveloppe.py +++ b/examples/demo_beam_enveloppe.py @@ -58,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}; @@ -70,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 index 7ecdfb90..c1a32e54 100755 --- 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:: diff --git a/examples/demo_lhc_rigid_waist_shift.py b/examples/demo_lhc_rigid_waist_shift.py index 4268a9ac..c82eb6e2 100755 --- a/examples/demo_lhc_rigid_waist_shift.py +++ b/examples/demo_lhc_rigid_waist_shift.py @@ -6,9 +6,9 @@ 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. @@ -16,8 +16,8 @@ (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:: diff --git a/examples/demo_lhc_setup.py b/examples/demo_lhc_setup.py index 04c48b24..3009870f 100644 --- 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:: diff --git a/examples/demo_track_spectra.py b/examples/demo_track_spectra.py index 1a61cb78..94774828 100755 --- a/examples/demo_track_spectra.py +++ b/examples/demo_track_spectra.py @@ -6,7 +6,7 @@ Free Tracking Spectrum ====================== -This example shows how to use the `~.track.track_single_particle` function to track a +This example shows how to use the `~.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 tracking workflow when using diff --git a/examples/demo_tune_diagram.py b/examples/demo_tune_diagram.py index 77d2b770..0148aab0 100755 --- a/examples/demo_tune_diagram.py +++ b/examples/demo_tune_diagram.py @@ -6,7 +6,7 @@ Tune Diagram ============ -This example shows how to use the `~.plotting.tune.plot_tune_diagram` function +This example shows how to use the `~.plotting.tune.plot_tune_diagram` function to visualise resonance lines up to certain orders. """ # sphinx_gallery_thumbnail_number = -1 From ac60f1b43312da6bf4ab3ba5f2b50437284d2578 Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 12:33:44 +0200 Subject: [PATCH 30/41] pytest parametrize expects tuples --- tests/test_cpymadtools/test_generators.py | 4 +- tests/test_cpymadtools/test_lhc.py | 2 +- tests/test_cpymadtools/test_matching.py | 12 +-- tests/test_cpymadtools/test_utilities.py | 2 +- tests/test_maths.py | 2 +- tests/test_optics.py | 6 +- tests/test_utils.py | 100 +++++++++++----------- 7 files changed, 64 insertions(+), 64 deletions(-) diff --git a/tests/test_cpymadtools/test_generators.py b/tests/test_cpymadtools/test_generators.py index 6c6a39d8..58c898a1 100644 --- a/tests/test_cpymadtools/test_generators.py +++ b/tests/test_cpymadtools/test_generators.py @@ -30,7 +30,7 @@ def test_tripleterrors_study_reference(): @pytest.mark.parametrize( - "randseed, tferror", + ("randseed", "tferror"), [ ("", ""), ("95", "195"), @@ -50,7 +50,7 @@ def test_tripleterrors_study_tferror_job(randseed, tferror): @pytest.mark.parametrize( - "randseed, mserror", + ("randseed", "mserror"), [ ("", ""), ("95", "195"), diff --git a/tests/test_cpymadtools/test_lhc.py b/tests/test_cpymadtools/test_lhc.py index 0bfae377..57006a37 100644 --- a/tests/test_cpymadtools/test_lhc.py +++ b/tests/test_cpymadtools/test_lhc.py @@ -330,7 +330,7 @@ def test_all_lhc_arcs(): @pytest.mark.parametrize( - "orient, result", + ("orient", "result"), [ ["straight", ["K0L", "K1L", "K2L", "K3L", "K4L"]], ["skew", ["K0SL", "K1SL", "K2SL", "K3SL", "K4SL"]], diff --git a/tests/test_cpymadtools/test_matching.py b/tests/test_cpymadtools/test_matching.py index 6b3ced7b..2879c406 100644 --- a/tests/test_cpymadtools/test_matching.py +++ b/tests/test_cpymadtools/test_matching.py @@ -11,8 +11,8 @@ @pytest.mark.skipif(platform.startswith("win"), reason="Windows is very flaky on this.") -@pytest.mark.parametrize("q1_target, q2_target", [(6.335, 6.29), (6.34, 6.27), (6.38, 6.27)]) -@pytest.mark.parametrize("dq1_target, dq2_target", [(100, 100), (95, 95), (105, 105)]) +@pytest.mark.parametrize(("q1_target", "q2_target"), [(6.335, 6.29), (6.34, 6.27), (6.38, 6.27)]) +@pytest.mark.parametrize(("dq1_target", "dq2_target"), [(100, 100), (95, 95), (105, 105)]) @pytest.mark.parametrize("telescopic_squeeze", [False, True]) def test_tune_and_chroma_matching(q1_target, q2_target, dq1_target, dq2_target, telescopic_squeeze): """Using my CAS19 project's lattice.""" @@ -52,7 +52,7 @@ def test_tune_and_chroma_matching_fails_on_unknown_accelerator(): @pytest.mark.skipif(platform.startswith("win"), reason="Windows is very flaky on this.") -@pytest.mark.parametrize("q1_target, q2_target", [(6.335, 6.29), (6.34, 6.27), (6.38, 6.27)]) +@pytest.mark.parametrize(("q1_target", "q2_target"), [(6.335, 6.29), (6.34, 6.27), (6.38, 6.27)]) @pytest.mark.parametrize("telescopic_squeeze", [False, True]) def test_tune_only_matching(q1_target, q2_target, telescopic_squeeze): """Using my CAS19 project's lattice.""" @@ -74,7 +74,7 @@ def test_tune_only_matching(q1_target, q2_target, telescopic_squeeze): @pytest.mark.skipif(platform.startswith("win"), reason="Windows is very flaky on this.") -@pytest.mark.parametrize("q1_target, q2_target", [(6.335, 6.29), (6.34, 6.27), (6.38, 6.27)]) +@pytest.mark.parametrize(("q1_target", "q2_target"), [(6.335, 6.29), (6.34, 6.27), (6.38, 6.27)]) @pytest.mark.parametrize("telescopic_squeeze", [False, True]) def test_tune_matching_wrapper(q1_target, q2_target, telescopic_squeeze): """Using my CAS19 project's lattice.""" @@ -96,7 +96,7 @@ def test_tune_matching_wrapper(q1_target, q2_target, telescopic_squeeze): @pytest.mark.skipif(platform.startswith("win"), reason="Windows is very flaky on this.") -@pytest.mark.parametrize("dq1_target, dq2_target", [(100, 100), (95, 95), (105, 105)]) +@pytest.mark.parametrize(("dq1_target", "dq2_target"), [(100, 100), (95, 95), (105, 105)]) @pytest.mark.parametrize("telescopic_squeeze", [False, True]) def test_chroma_only_matching(dq1_target, dq2_target, telescopic_squeeze): """Using my CAS19 project's lattice.""" @@ -118,7 +118,7 @@ def test_chroma_only_matching(dq1_target, dq2_target, telescopic_squeeze): @pytest.mark.skipif(platform.startswith("win"), reason="Windows is very flaky on this.") -@pytest.mark.parametrize("dq1_target, dq2_target", [(100, 100), (95, 95), (105, 105)]) +@pytest.mark.parametrize(("dq1_target", "dq2_target"), [(100, 100), (95, 95), (105, 105)]) @pytest.mark.parametrize("telescopic_squeeze", [False, True]) def test_chroma_matching_wrapper(dq1_target, dq2_target, telescopic_squeeze): """Using my CAS19 project's lattice.""" diff --git a/tests/test_cpymadtools/test_utilities.py b/tests/test_cpymadtools/test_utilities.py index 2d28c7d3..6f67833b 100644 --- a/tests/test_cpymadtools/test_utilities.py +++ b/tests/test_cpymadtools/test_utilities.py @@ -6,7 +6,7 @@ @pytest.mark.parametrize( - "orient, result", + ("orient", "result"), [ ["straight", ["K0L", "K1L", "K2L", "K3L", "K4L"]], ["skew", ["K0SL", "K1SL", "K2SL", "K3SL", "K4SL"]], diff --git a/tests/test_maths.py b/tests/test_maths.py index 87c5f11d..2b7050ef 100644 --- a/tests/test_maths.py +++ b/tests/test_maths.py @@ -90,7 +90,7 @@ def test_make_pdf(degrees_of_freedom): pdf.idxmax() == pytest.approx(degrees_of_freedom - 2, rel=1e-2) -@pytest.mark.parametrize("value, result", [(1, 0), (10, 1), (0.0311, -2), (5e-7, -7)]) +@pytest.mark.parametrize(("value", "result"), [(1, 0), (10, 1), (0.0311, -2), (5e-7, -7)]) def test_magnitude(value, result): assert mutils.get_magnitude(value) == result diff --git a/tests/test_optics.py b/tests/test_optics.py index 9a730772..8ad7d7d9 100644 --- a/tests/test_optics.py +++ b/tests/test_optics.py @@ -46,14 +46,14 @@ def test_lhc_revolution_frequency(): @pytest.mark.parametrize( - "alpha_p, result", + ("alpha_p", "result"), [(0, 2.083077890845299e-08), (1e-5, -9.979169221091548e-06), (-500, 500.0000000208308)], ) def test_eta(alpha_p, result): assert Beam(6500, 2.5e-6).eta(alpha_p) == result -@pytest.mark.parametrize("alpha_p, result", [(1e-5, 316.2277660168379), (500, 0.044721359549995794)]) +@pytest.mark.parametrize(("alpha_p", "result"), [(1e-5, 316.2277660168379), (500, 0.044721359549995794)]) def test_gamma_transition(alpha_p, result): assert Beam(6500, 2.5e-6).gamma_transition(alpha_p) == result @@ -64,7 +64,7 @@ def test_gamma_transition_raises(): @pytest.mark.parametrize( - "pc_gev, en_x_m, en_y_m, delta_p, result", + ("pc_gev", "en_x_m", "en_y_m", "delta_p", "result"), [ ( 1.9, diff --git a/tests/test_utils.py b/tests/test_utils.py index ada1904e..9133dcff 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -126,17 +126,17 @@ def test_tasks_table_creation(self, _condor_q_output, _taskless_condor_q_output) class TestListOperations: - @pytest.mark.parametrize("args, result", [([1, 2, 3, 5, 12, 0], True), ([1, 1, 1], False), (list(), True)]) + @pytest.mark.parametrize(("args", "result"), [([1, 2, 3, 5, 12, 0], True), ([1, 1, 1], False), (list(), True)]) def test_all_unique(self, args, result): assert ListOperations.all_unique(args) is result - @pytest.mark.parametrize("args, error", [(None, TypeError), (_square, TypeError)]) + @pytest.mark.parametrize(("args", "error"), [(None, TypeError), (_square, TypeError)]) def test_all_unique_fails(self, args, error): with pytest.raises(error): ListOperations.all_unique(args) @pytest.mark.parametrize( - "inputs, function, result", + ("inputs", "function", "result"), [ ([{"n": 4}, {"n": 2}, {"n": 8}, {"n": 6}], lambda x: x["n"], 5), (list(range(10)), lambda x: 2 * x, 9), @@ -146,13 +146,13 @@ def test_all_unique_fails(self, args, error): def test_average_by(self, inputs, function, result): assert ListOperations.average_by(sequence=inputs, function=function) == result - @pytest.mark.parametrize("inputs, error", [(None, TypeError), ((list(range(10)), None), TypeError)]) + @pytest.mark.parametrize(("inputs", "error"), [(None, TypeError), ((list(range(10)), None), TypeError)]) def test_average_by_fails(self, inputs, error): with pytest.raises(error): ListOperations.average_by(inputs) @pytest.mark.parametrize( - "inputs, filters, results", + ("inputs", "filters", "results"), [ (["beep", "boop", "foo", "bar"], [True, True, False, True], [["beep", "boop", "bar"], ["foo"]]), ([1, _square, _to_str, "string"], [False, True, False, True], [[_square, "string"], [1, _to_str]]), @@ -162,7 +162,7 @@ def test_bifurcate(self, inputs, filters, results): assert ListOperations.bifurcate(inputs, filters) == results @pytest.mark.parametrize( - "inputs, filters, error", + ("inputs", "filters", "error"), [ (["beep", "boop", "foo", "bar", "baz"], None, TypeError), (["beep", "boop", "foo", "bar", "baz"], [], IndexError), @@ -174,7 +174,7 @@ def test_bifurcate_fails(self, inputs, filters, error): ListOperations.bifurcate(inputs, filters) @pytest.mark.parametrize( - "inputs, func, results", + ("inputs", "func", "results"), [ (list(range(5)), lambda x: x % 2 == 0, [[0, 2, 4], [1, 3]]), ([], lambda x: x % 2 == 0, [[], []]), @@ -184,7 +184,7 @@ def test_bifurcate_by(self, inputs, func, results): assert ListOperations.bifurcate_by(inputs, func) == results @pytest.mark.parametrize( - "array, size, result", + ("array", "size", "result"), [ (list(range(10)), 3, [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]]), (list(range(10)), 20, list(range(10))), @@ -195,7 +195,7 @@ def test_chunk_list(self, array, size, result): assert ListOperations.chunk_list(array, size) == result @pytest.mark.parametrize( - "array, size, error", + ("array", "size", "error"), [(list(range(10)), 0, ZeroDivisionError), (5, 20, TypeError), ([], "string", TypeError)], ) def test_chunk_list_fails(self, array, size, error): @@ -203,7 +203,7 @@ def test_chunk_list_fails(self, array, size, error): ListOperations.chunk_list(array, size) @pytest.mark.parametrize( - "args, result", + ("args", "result"), [ ([["a", "b", "c"], [1, 2, 3], [True, False, False]], ["a", "b", "c", 1, 2, 3, True, False, False]), ([["a", "b", "c"], [1, 2, 3], [], [True, False, False]], ["a", "b", "c", 1, 2, 3, True, False, False]), @@ -214,19 +214,19 @@ def test_deep_flatten(self, args, result): assert ListOperations.deep_flatten(args) == result @pytest.mark.parametrize( - "array, func, result", [([0, 0, 1, 0], lambda x: x >= 2, True), ([0, 1, 2, 0], lambda x: x >= 2, False)] + ("array", "func", "result"), [([0, 0, 1, 0], lambda x: x >= 2, True), ([0, 1, 2, 0], lambda x: x >= 2, False)] ) def test_eval_none(self, array, func, result): assert ListOperations.eval_none(array, func) is result @pytest.mark.parametrize( - "array, func, result", [([0, 1, 2, 0], lambda x: x >= 2, True), ([0, 0, 1, 0], lambda x: x >= 2, False)] + ("array", "func", "result"), [([0, 1, 2, 0], lambda x: x >= 2, True), ([0, 0, 1, 0], lambda x: x >= 2, False)] ) def test_eval_some(self, array, func, result): assert ListOperations.eval_some(array, func) is result @pytest.mark.parametrize( - "element, array, result", + ("element", "array", "result"), [ (0, [0, 1, 3, 5, 7, 3, 9, 0, 0, 5, 3, 2], [0, 7, 8]), (0, (0, 1, 3, 5, 7, 3, 9, 0, 0, 5, 3, 2), [0, 7, 8]), @@ -239,13 +239,13 @@ def test_eval_some(self, array, func, result): def test_get_indices(self, element, array, result): assert ListOperations.get_indices(element, array) == result - @pytest.mark.parametrize("element, array, error", [(5, None, TypeError)]) + @pytest.mark.parametrize(("element", "array", "error"), [(5, None, TypeError)]) def test_get_indices_fails(self, element, array, error): with pytest.raises(error): ListOperations.get_indices(element, array) @pytest.mark.parametrize( - "array, func, result", + ("array", "func", "result"), [ (list(range(5)), lambda x: x % 2 == 0, {True: [0, 2, 4], False: [1, 3]}), (list(range(5)), lambda x: x >= 2, {True: [2, 3, 4], False: [0, 1]}), @@ -255,12 +255,12 @@ def test_groupby(self, array, func, result): assert ListOperations.group_by(array, func) == result @pytest.mark.parametrize( - "array, result", [([1, 2, 1], True), ([list(range(10)), False]), ([], False), ([True, True], True)] + ("array", "result"), [([1, 2, 1], True), ([list(range(10)), False]), ([], False), ([True, True], True)] ) def test_has_duplicates(self, array, result): assert ListOperations.has_duplicates(array) is result - @pytest.mark.parametrize("array, error", [(1, TypeError), (_square, TypeError)]) + @pytest.mark.parametrize(("array", "error"), [(1, TypeError), (_square, TypeError)]) def test_has_duplicates_fails(self, array, error): with pytest.raises(error): ListOperations.has_duplicates(array) @@ -270,7 +270,7 @@ def test_sample(self, args): assert ListOperations.sample(args) in args @pytest.mark.parametrize( - "input_list, result", + ("input_list", "result"), [ ([1, False, "a", 2, "", None, 6, 0], [1, "a", 2, 6]), (list(range(1, 10)), list(range(1, 10))), @@ -295,7 +295,7 @@ def test_shuffle_list(self, array): assert shuffled != array and set(shuffled) == set(array) @pytest.mark.parametrize( - "array, result", + ("array", "result"), [ ([list(range(5)), list(range(5))], [0, 1, 2, 3, 4, 0, 1, 2, 3, 4]), ([], []), @@ -306,7 +306,7 @@ def test_spread(self, array, result): assert ListOperations.spread(array) == result @pytest.mark.parametrize( - "array1, array2, func, result", + ("array1", "array2", "func", "result"), [ ([2.1, 1.2], [2.3, 3.4], math.floor, [1.2, 3.4]), ([2.1, 1.2], [0.5, 1.2], lambda x: x >= 2, [2.1]), @@ -318,7 +318,7 @@ def test_symmetric_difference_by(self, array1, array2, func, result): assert ListOperations.symmetric_difference_by(array1, array2, func) == result @pytest.mark.parametrize( - "array1, array2, func, error", + ("array1", "array2", "func", "error"), [(list(range(10)), list(range(10)), None, TypeError), (None, [], _square, TypeError)], ) def test_symmetric_difference_by_fails(self, array1, array2, func, error): @@ -326,7 +326,7 @@ def test_symmetric_difference_by_fails(self, array1, array2, func, error): ListOperations.symmetric_difference_by(array1, array2, func) @pytest.mark.parametrize( - "array1, array2, func, result", + ("array1", "array2", "func", "result"), [ ([2.1], [1.2, 2.3], math.floor, [1.2, 2.1]), (["A", "B"], ["a", "c"], str.lower, ["A", "B", "c"]), @@ -337,7 +337,7 @@ def test_union_by(self, array1, array2, func, result): assert ListOperations.union_by(array1, array2, func) == result @pytest.mark.parametrize( - "array1, array2, func, error", + ("array1", "array2", "func", "error"), [(list(range(10)), list(range(10)), None, TypeError), (None, [], _square, TypeError)], ) def test_union_by_fails(self, array1, array2, func, error): @@ -368,7 +368,7 @@ def test_logger_config(self, capsys): class TestMiscellaneousOperations: @pytest.mark.parametrize( - "input1, input2, input3, result", + ("input1", "input2", "input3", "result"), [ (list(range(5)), list(range(100)), list(range(50)), list(range(100))), ( @@ -391,7 +391,7 @@ def test_longest_item_none_input(self): MiscellaneousOperations.longest_item(list(range(5)), None) @pytest.mark.parametrize( - "obj, func, result", + ("obj", "func", "result"), [ ( {"a": 0, "b": 1, "c": 2, "d": 3}, @@ -416,7 +416,7 @@ def test_map_values_to_bool(self, obj, func, result): class TestMultiProcessorExecutor: @pytest.mark.parametrize( - "function, inputs, results", + ("function", "inputs", "results"), [ (_square, list(range(6)), [0, 1, 4, 9, 16, 25]), (_square, [10 * i for i in range(10)], [e**2 for e in [10 * i for i in range(10)]]), @@ -435,7 +435,7 @@ def test_multiprocessing_zero_processes(self): class TestMultiThreaderExecutor: @pytest.mark.parametrize( - "function, inputs, results", + ("function", "inputs", "results"), [ (_square, list(range(6)), [0, 1, 4, 9, 16, 25]), (_square, [10 * i for i in range(10)], [e**2 for e in [10 * i for i in range(10)]]), @@ -454,7 +454,7 @@ def test_multithreading_zero_threads(self): class TestNumberOperations: @pytest.mark.parametrize( - "number, lowthreshold, highthreshold, result", + ("number", "lowthreshold", "highthreshold", "result"), [ (17, 4, 5, 5), (50, 70, 100, 70), @@ -469,7 +469,7 @@ def test_clamp_number(self, number, lowthreshold, highthreshold, result): assert NumberOperations.clamp_number(number, lowthreshold, highthreshold) == result @pytest.mark.parametrize( - "wrong_input, expected_error", + ("wrong_input", "expected_error"), [(None, TypeError), (str(5), TypeError), (lambda x: x, TypeError), (object, TypeError)], ) def test_clamp_number_fails(self, wrong_input, expected_error): @@ -477,7 +477,7 @@ def test_clamp_number_fails(self, wrong_input, expected_error): NumberOperations.clamp_number(17, 5, wrong_input) @pytest.mark.parametrize( - "degrees, decompose_bool, result", + ("degrees", "decompose_bool", "result"), [ (160, False, 2.792526803190927), (160, None, 2.792526803190927), @@ -490,24 +490,24 @@ def test_degrees_to_radians(self, degrees, decompose_bool, result): assert NumberOperations.degrees_to_radians(degrees, decompose_bool) == result @pytest.mark.parametrize( - "inputs, result", [([54, 24], 6), ([30, 132, 378, 582, 738], 6), ([57, 37, 18], 1), ([0, 0], 0)] + ("inputs", "result"), [([54, 24], 6), ([30, 132, 378, 582, 738], 6), ([57, 37, 18], 1), ([0, 0], 0)] ) def test_greatest_common_divisor(self, inputs, result): assert NumberOperations.greatest_common_divisor(inputs) == result - @pytest.mark.parametrize("inputs, error", [([50, _square], TypeError)]) + @pytest.mark.parametrize(("inputs", "error"), [([50, _square], TypeError)]) def test_greatest_common_divisor_fails(self, inputs, error): with pytest.raises(error): NumberOperations.greatest_common_divisor(inputs) @pytest.mark.parametrize( - "number, divisor, result", [(25, 5, True), (73.4, 2.1, False), (-5, 5, True), (-100, -7, False)] + ("number", "divisor", "result"), [(25, 5, True), (73.4, 2.1, False), (-5, 5, True), (-100, -7, False)] ) def test_is_divisible(self, number, divisor, result): assert NumberOperations.is_divisible_by(number, divisor) is result @pytest.mark.parametrize( - "number, divisor, error", + ("number", "divisor", "error"), [(50, 0, ZeroDivisionError), (10, _square, TypeError), (str(5), 5, TypeError)], ) def test_is_divisible_fails(self, number, divisor, error): @@ -515,20 +515,20 @@ def test_is_divisible_fails(self, number, divisor, error): NumberOperations.is_divisible_by(number, divisor) @pytest.mark.parametrize( - "args, result", [([4, 5], 20), ([2, 5, 17, 632], 53720), ([-1, 5, 10], -10), ([0, 10, 50], 0)] + ("args", "result"), [([4, 5], 20), ([2, 5, 17, 632], 53720), ([-1, 5, 10], -10), ([0, 10, 50], 0)] ) def test_least_common_multiple(self, args, result): assert NumberOperations.least_common_multiple(args) == result @pytest.mark.parametrize( - "args, error", [([0, 0], ZeroDivisionError), ([15, _square], TypeError), ([str(100), 10], TypeError)] + ("args", "error"), [([0, 0], ZeroDivisionError), ([15, _square], TypeError), ([str(100), 10], TypeError)] ) def test_least_common_multiple_fails(self, args, error): with pytest.raises(error): NumberOperations.least_common_multiple(args) @pytest.mark.parametrize( - "rad, result", + ("rad", "result"), [ (2 * math.pi, 360), (0.45, 25.783100780887047), @@ -540,7 +540,7 @@ def test_least_common_multiple_fails(self, args, error): def test_radians_to_degrees(self, rad, result): assert NumberOperations.radians_to_degrees(rad) == result - @pytest.mark.parametrize("rad, error", [(_square, TypeError), (str(100), TypeError)]) + @pytest.mark.parametrize(("rad", "error"), [(_square, TypeError), (str(100), TypeError)]) def test_radians_to_degrees_fails(self, rad, error): with pytest.raises(error): NumberOperations.radians_to_degrees(rad) @@ -548,7 +548,7 @@ def test_radians_to_degrees_fails(self, rad, error): class TestStringOperations: @pytest.mark.parametrize( - "string_input, result", + ("string_input", "result"), [ ("a_snake_case_name", "aSnakeCaseName"), ("A Title Case Name", "aTitleCaseName"), @@ -558,13 +558,13 @@ class TestStringOperations: def test_camel_case(self, string_input, result): assert StringOperations.camel_case(string_input) == result - @pytest.mark.parametrize("string_input, error", [(1, TypeError), ("", IndexError)]) + @pytest.mark.parametrize(("string_input", "error"), [(1, TypeError), ("", IndexError)]) def test_camel_case_fails(self, string_input, error): with pytest.raises(error): StringOperations.camel_case(string_input) @pytest.mark.parametrize( - "string_input, lower, result", + ("string_input", "lower", "result"), [ ("astringtocapitalize", False, "Astringtocapitalize"), ("astringtocapitalize", None, "Astringtocapitalize"), @@ -576,13 +576,13 @@ def test_camel_case_fails(self, string_input, error): def test_capitalize(self, string_input, lower, result): assert StringOperations.capitalize(string_input, lower_rest=lower) == result - @pytest.mark.parametrize("string_input, lower, error", [(1, True, TypeError), (_square, False, TypeError)]) + @pytest.mark.parametrize(("string_input", "lower", "error"), [(1, True, TypeError), (_square, False, TypeError)]) def test_camel_capitalize_fails(self, string_input, lower, error): with pytest.raises(error): StringOperations.capitalize(string_input, lower_rest=lower) @pytest.mark.parametrize( - "string1, string2, result", + ("string1", "string2", "result"), [ ("Justin Timberlake", "I'm a jerk but listen", True), ("I'm a jerk but listen", "Justin Timberlake", True), @@ -595,22 +595,22 @@ def test_camel_capitalize_fails(self, string_input, lower, error): def test_is_anagram(self, string1, string2, result): assert StringOperations.is_anagram(string1, string2) is result - @pytest.mark.parametrize("string1, string2, error", [(1, "", AttributeError), ("string", False, AttributeError)]) + @pytest.mark.parametrize(("string1", "string2", "error"), [(1, "", AttributeError), ("string", False, AttributeError)]) def test_is_anagram_fails(self, string1, string2, error): with pytest.raises(error): StringOperations.is_anagram(string1, string2) - @pytest.mark.parametrize("word, result", [("racecar", True), ("definitelynot", False), ("", True)]) + @pytest.mark.parametrize(("word", "result"), [("racecar", True), ("definitelynot", False), ("", True)]) def test_is_palindrome(self, word, result): assert StringOperations.is_palindrome(word) is result - @pytest.mark.parametrize("word, error", [(1, AttributeError), (_square, AttributeError)]) + @pytest.mark.parametrize(("word", "error"), [(1, AttributeError), (_square, AttributeError)]) def test_is_palindrome_fails(self, word, error): with pytest.raises(error): StringOperations.is_palindrome(word) @pytest.mark.parametrize( - "string_input, result", + ("string_input", "result"), [ ("a normal sentence", "a-normal-sentence"), ("", ""), @@ -622,13 +622,13 @@ def test_is_palindrome_fails(self, word, error): def test_kebab_case(self, string_input, result): assert StringOperations.kebab_case(string_input) == result - @pytest.mark.parametrize("word, error", [(1, TypeError), (_square, TypeError)]) + @pytest.mark.parametrize(("word", "error"), [(1, TypeError), (_square, TypeError)]) def test_kebab_case_fails(self, word, error): with pytest.raises(error): StringOperations.kebab_case(word) @pytest.mark.parametrize( - "string_input, result", + ("string_input", "result"), [ ("a normal sentence", "a_normal_sentence"), ("camelCase", "camelcase"), @@ -640,7 +640,7 @@ def test_kebab_case_fails(self, word, error): def test_snake_case(self, string_input, result): assert StringOperations.snake_case(string_input) == result - @pytest.mark.parametrize("word, error", [(1, TypeError), (_square, TypeError)]) + @pytest.mark.parametrize(("word", "error"), [(1, TypeError), (_square, TypeError)]) def test_snake_case_fails(self, word, error): with pytest.raises(error): StringOperations.snake_case(word) From b1ed29a2407c67af4adee474c65d693be8b75978 Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 12:34:45 +0200 Subject: [PATCH 31/41] dict membership test does not need .keys() --- tests/test_cpymadtools/test_lhc.py | 4 ++-- tests/test_cpymadtools/test_ptc.py | 2 +- tests/test_cpymadtools/test_track.py | 2 +- tests/test_utils.py | 2 +- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/test_cpymadtools/test_lhc.py b/tests/test_cpymadtools/test_lhc.py index 57006a37..c935af45 100644 --- a/tests/test_cpymadtools/test_lhc.py +++ b/tests/test_cpymadtools/test_lhc.py @@ -302,8 +302,8 @@ def test_get_orbit_setup(_non_matched_lhc_madx): setup = get_current_orbit_setup(madx) assert isinstance(setup, dict) - assert all(orbit_var in setup.keys() for orbit_var in lhc_orbit_variables()[0]) - assert all(special_var in setup.keys() for special_var in lhc_orbit_variables()[1]) + assert all(orbit_var in setup for orbit_var in lhc_orbit_variables()[0]) + assert all(special_var in setup for special_var in lhc_orbit_variables()[1]) def test_orbit_correction(_bare_lhc_madx): diff --git a/tests/test_cpymadtools/test_ptc.py b/tests/test_cpymadtools/test_ptc.py index e3c1eacd..1e76dc9b 100644 --- a/tests/test_cpymadtools/test_ptc.py +++ b/tests/test_cpymadtools/test_ptc.py @@ -83,7 +83,7 @@ def test_single_particle_ptc_track_with_onepass(_matched_base_lattice): assert isinstance(tracks_dict, dict) assert len(tracks_dict.keys()) == 1 # should be only one because of ONETABLE option - assert "trackone" in tracks_dict.keys() + assert "trackone" in tracks_dict tracks = tracks_dict["trackone"] assert isinstance(tracks, DataFrame) assert all([coordinate in tracks.columns for coordinate in ("x", "px", "y", "py", "t", "pt", "s", "e")]) diff --git a/tests/test_cpymadtools/test_track.py b/tests/test_cpymadtools/test_track.py index 96c07510..b96277f6 100644 --- a/tests/test_cpymadtools/test_track.py +++ b/tests/test_cpymadtools/test_track.py @@ -30,7 +30,7 @@ def test_single_particle_tracking_with_onepass(_matched_base_lattice): assert isinstance(tracks_dict, dict) assert len(tracks_dict.keys()) == 1 # should be only one because of ONETABLE option - assert "trackone" in tracks_dict.keys() + assert "trackone" in tracks_dict tracks = tracks_dict["trackone"] assert isinstance(tracks, DataFrame) assert all([coordinate in tracks.columns for coordinate in ("x", "px", "y", "py", "t", "pt", "s", "e")]) diff --git a/tests/test_utils.py b/tests/test_utils.py index 9133dcff..eb5cccd6 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -659,7 +659,7 @@ def test_colin_balance_applies(self, _non_matched_lhc_madx): for side in "rl": for ip in (1, 2, 5, 8): assert madx.globals[f"kqsx3.{side}{ip}"] != 0 - assert "ir_quads_errors" in madx.table.keys() + assert "ir_quads_errors" in madx.table @pytest.mark.parametrize("drop", [True, False]) def test_complex_columns_split(self, _complex_columns_df, drop): From 3a1596f82af1b995d2811ff17e530281551ce914 Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 12:49:45 +0200 Subject: [PATCH 32/41] mpl --- pyproject.toml | 7 ++++++- tests/test_cpymadtools/test_coupling.py | 1 + tests/test_cpymadtools/test_tune.py | 7 ++++--- tests/test_plotting/test_aperture.py | 4 ++-- tests/test_plotting/test_crossing.py | 4 ++-- tests/test_plotting/test_envelope.py | 4 ++-- tests/test_plotting/test_helpers.py | 4 ++-- tests/test_plotting/test_lattice.py | 4 ++-- tests/test_plotting/test_phasespace.py | 4 ++-- tests/test_plotting/test_sbs_coupling.py | 9 +++------ tests/test_plotting/test_styles.py | 4 ++-- tests/test_plotting/test_tunediagram.py | 4 ++-- tests/test_utils.py | 11 +++++------ 13 files changed, 35 insertions(+), 32 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index c7d6d401..5fa020a8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -163,7 +163,12 @@ exclude = ''' ''' [tool.ruff] -ignore = ["E501", "FBT001", "FBT002"] +ignore = [ + "E501", # line-too-long + "FBT001", # boolean-type-hint-positional-argument + "FBT002", # boolean-default-value-positional-argument + "PT019", # pytest-fixture-param-without-value (but suggested solution fails) +] line-length = 120 # Exclude a variety of commonly ignored directories. diff --git a/tests/test_cpymadtools/test_coupling.py b/tests/test_cpymadtools/test_coupling.py index 4e9ae1dd..3ee41856 100644 --- a/tests/test_cpymadtools/test_coupling.py +++ b/tests/test_cpymadtools/test_coupling.py @@ -21,6 +21,7 @@ INPUTS_DIR = CURRENT_DIR.parent / "inputs" + @pytest.mark.parametrize("telescopic_squeeze", [False, True]) def test_closest_tune_approach(_non_matched_lhc_madx, telescopic_squeeze): """Using LHC lattice.""" diff --git a/tests/test_cpymadtools/test_tune.py b/tests/test_cpymadtools/test_tune.py index 2438a66a..b213e111 100644 --- a/tests/test_cpymadtools/test_tune.py +++ b/tests/test_cpymadtools/test_tune.py @@ -1,17 +1,18 @@ import pathlib -import matplotlib +import matplotlib as mpl import matplotlib.pyplot as plt import numpy as np import pytest import tfs +from matplotlib.collections import PatchCollection from pyhdtoolkit.cpymadtools.lhc import make_lhc_thin, re_cycle_sequence, setup_lhc_orbit from pyhdtoolkit.cpymadtools.matching import match_tunes_and_chromaticities from pyhdtoolkit.cpymadtools.tune import get_footprint_lines, get_footprint_patches, make_footprint_table # Forcing non-interactive Agg backend so rendering is done similarly across platforms during tests -matplotlib.use("Agg") +mpl.use("Agg") CURRENT_DIR = pathlib.Path(__file__).parent INPUTS_DIR = CURRENT_DIR.parent / "inputs" @@ -84,7 +85,7 @@ def test_get_footprint_patches(_dynap_tfs_path): plt.rcParams.update({"patch.linewidth": 1.5}) figure, axis = plt.subplots(figsize=(18, 11)) polygons = get_footprint_patches(dynap_dframe) - assert isinstance(polygons, matplotlib.collections.PatchCollection) + assert isinstance(polygons, PatchCollection) axis.add_collection(polygons) axis.set_xlim(0.3095, 0.311) diff --git a/tests/test_plotting/test_aperture.py b/tests/test_plotting/test_aperture.py index 59f85010..4c4f85a7 100644 --- a/tests/test_plotting/test_aperture.py +++ b/tests/test_plotting/test_aperture.py @@ -1,4 +1,4 @@ -import matplotlib +import matplotlib as mpl import matplotlib.pyplot as plt import pytest from cpymad.madx import Madx @@ -6,7 +6,7 @@ from pyhdtoolkit.plotting.aperture import plot_aperture, plot_physical_apertures # Forcing non-interactive Agg backend so rendering is done similarly across platforms during tests -matplotlib.use("Agg") +mpl.use("Agg") @pytest.mark.mpl_image_compare(tolerance=20, style="default", savefig_kwargs={"dpi": 200}) diff --git a/tests/test_plotting/test_crossing.py b/tests/test_plotting/test_crossing.py index d37d6eb4..8252c1ae 100644 --- a/tests/test_plotting/test_crossing.py +++ b/tests/test_plotting/test_crossing.py @@ -1,11 +1,11 @@ -import matplotlib +import matplotlib as mpl import matplotlib.pyplot as plt import pytest from pyhdtoolkit.plotting.crossing import plot_two_lhc_ips_crossings # Forcing non-interactive Agg backend so rendering is done similarly across platforms during tests -matplotlib.use("Agg") +mpl.use("Agg") @pytest.mark.mpl_image_compare(tolerance=35, style="default", savefig_kwargs={"dpi": 200}) diff --git a/tests/test_plotting/test_envelope.py b/tests/test_plotting/test_envelope.py index f28fc6df..fba46926 100644 --- a/tests/test_plotting/test_envelope.py +++ b/tests/test_plotting/test_envelope.py @@ -1,6 +1,6 @@ import pathlib -import matplotlib +import matplotlib as mpl import matplotlib.pyplot as plt import pytest from cpymad.madx import Madx @@ -8,7 +8,7 @@ from pyhdtoolkit.plotting.envelope import _interpolate_madx, plot_beam_envelope # Forcing non-interactive Agg backend so rendering is done similarly across platforms during tests -matplotlib.use("Agg") +mpl.use("Agg") CURRENT_DIR = pathlib.Path(__file__).parent INPUTS_DIR = CURRENT_DIR.parent / "inputs" diff --git a/tests/test_plotting/test_helpers.py b/tests/test_plotting/test_helpers.py index 4fc7b670..1e15f530 100644 --- a/tests/test_plotting/test_helpers.py +++ b/tests/test_plotting/test_helpers.py @@ -1,11 +1,11 @@ -import matplotlib +import matplotlib as mpl import matplotlib.pyplot as plt import pytest from pyhdtoolkit.plotting.utils import set_arrow_label # Forcing non-interactive Agg backend so rendering is done similarly across platforms during tests -matplotlib.use("Agg") +mpl.use("Agg") @pytest.mark.mpl_image_compare(tolerance=20, style="default", savefig_kwargs={"dpi": 200}) diff --git a/tests/test_plotting/test_lattice.py b/tests/test_plotting/test_lattice.py index dcbec8d6..3646c2ea 100644 --- a/tests/test_plotting/test_lattice.py +++ b/tests/test_plotting/test_lattice.py @@ -1,6 +1,6 @@ import pathlib -import matplotlib +import matplotlib as mpl import matplotlib.pyplot as plt import pytest from cpymad.madx import Madx @@ -11,7 +11,7 @@ from pyhdtoolkit.plotting.layout import scale_patches # Forcing non-interactive Agg backend so rendering is done similarly across platforms during tests -matplotlib.use("Agg") +mpl.use("Agg") CURRENT_DIR = pathlib.Path(__file__).parent INPUTS_DIR = CURRENT_DIR.parent / "inputs" diff --git a/tests/test_plotting/test_phasespace.py b/tests/test_plotting/test_phasespace.py index e694bd74..ed5695ef 100644 --- a/tests/test_plotting/test_phasespace.py +++ b/tests/test_plotting/test_phasespace.py @@ -1,6 +1,6 @@ from functools import partial -import matplotlib +import matplotlib as mpl import matplotlib.pyplot as plt import numpy as np import pytest @@ -15,7 +15,7 @@ ) # Forcing non-interactive Agg backend so rendering is done similarly across platforms during tests -matplotlib.use("Agg") +mpl.use("Agg") BASE_LATTICE = LatticeGenerator.generate_base_cas_lattice() diff --git a/tests/test_plotting/test_sbs_coupling.py b/tests/test_plotting/test_sbs_coupling.py index 4544d9bd..025cc299 100644 --- a/tests/test_plotting/test_sbs_coupling.py +++ b/tests/test_plotting/test_sbs_coupling.py @@ -12,19 +12,18 @@ @pytest.mark.mpl_image_compare(tolerance=20, style="default", savefig_kwargs={"dpi": 200}) def test_plot_rdt_component(sbs_coupling_b1_ip1, sbs_coupling_b2_ip1, sbs_model_b1, sbs_model_b2): - figure = plot_rdt_component( + return plot_rdt_component( b1_segment_df=sbs_coupling_b1_ip1, b2_segment_df=sbs_coupling_b2_ip1, b1_model=sbs_model_b1, b2_model=sbs_model_b2, ip=1, ) - return figure @pytest.mark.mpl_image_compare(tolerance=20, style="default", savefig_kwargs={"dpi": 200}) def test_plot_full_ip(sbs_coupling_b1_ip1, sbs_coupling_b2_ip1, sbs_model_b1, sbs_model_b2): - figure = plot_full_ip_rdt( + return plot_full_ip_rdt( b1_segment_df=sbs_coupling_b1_ip1, b2_segment_df=sbs_coupling_b2_ip1, b1_model=sbs_model_b1, @@ -34,12 +33,11 @@ def test_plot_full_ip(sbs_coupling_b1_ip1, sbs_coupling_b2_ip1, sbs_model_b1, sb figsize=(12, 7), bbox_to_anchor=(0.535, 0.95), ) - return figure @pytest.mark.mpl_image_compare(tolerance=20, style="default", savefig_kwargs={"dpi": 200}) def test_plot_full_ip_with_ylimits(sbs_coupling_b1_ip1, sbs_coupling_b2_ip1, sbs_model_b1, sbs_model_b2): - figure = plot_full_ip_rdt( + return plot_full_ip_rdt( b1_segment_df=sbs_coupling_b1_ip1, b2_segment_df=sbs_coupling_b2_ip1, b1_model=sbs_model_b1, @@ -52,7 +50,6 @@ def test_plot_full_ip_with_ylimits(sbs_coupling_b1_ip1, sbs_coupling_b2_ip1, sbs imag_ylimits=(-1.3e-1, 1e-2), bbox_to_anchor=(0.535, 0.95), ) - return figure # ----- Fixtures ----- # diff --git a/tests/test_plotting/test_styles.py b/tests/test_plotting/test_styles.py index 7fec04d0..1403f266 100644 --- a/tests/test_plotting/test_styles.py +++ b/tests/test_plotting/test_styles.py @@ -1,6 +1,6 @@ import pathlib -import matplotlib +import matplotlib as mpl import matplotlib.pyplot as plt from pyhdtoolkit.plotting.styles import install_mpl_styles @@ -10,5 +10,5 @@ def test_mplstyles_install(): install_mpl_styles() for expected_name in ("thesis-small", "thesis-medium", "thesis-large", "paper-single", "paper-double"): - assert (pathlib.Path(matplotlib.get_configdir()) / "stylelib" / f"{expected_name}.mplstyle").is_file() + assert (pathlib.Path(mpl.get_configdir()) / "stylelib" / f"{expected_name}.mplstyle").is_file() assert (pathlib.Path(plt.style.core.BASE_LIBRARY_PATH) / f"{expected_name}.mplstyle").is_file() diff --git a/tests/test_plotting/test_tunediagram.py b/tests/test_plotting/test_tunediagram.py index 8d42dd73..ab36ed69 100644 --- a/tests/test_plotting/test_tunediagram.py +++ b/tests/test_plotting/test_tunediagram.py @@ -1,11 +1,11 @@ -import matplotlib +import matplotlib as mpl import matplotlib.pyplot as plt import pytest from pyhdtoolkit.plotting.tune import plot_tune_diagram # Forcing non-interactive Agg backend so rendering is done similarly across platforms during tests -matplotlib.use("Agg") +mpl.use("Agg") @pytest.mark.parametrize("max_order", [0, 10, -5]) diff --git a/tests/test_utils.py b/tests/test_utils.py index eb5cccd6..81f08855 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -126,7 +126,7 @@ def test_tasks_table_creation(self, _condor_q_output, _taskless_condor_q_output) class TestListOperations: - @pytest.mark.parametrize(("args", "result"), [([1, 2, 3, 5, 12, 0], True), ([1, 1, 1], False), (list(), True)]) + @pytest.mark.parametrize(("args", "result"), [([1, 2, 3, 5, 12, 0], True), ([1, 1, 1], False), ([], True)]) def test_all_unique(self, args, result): assert ListOperations.all_unique(args) is result @@ -292,7 +292,8 @@ def test_sanitize(self, input_list, result): ) def test_shuffle_list(self, array): shuffled = ListOperations.shuffle(array) - assert shuffled != array and set(shuffled) == set(array) + assert shuffled != array + assert set(shuffled) == set(array) @pytest.mark.parametrize( ("array", "result"), @@ -694,7 +695,7 @@ def test_bpm_noise_addition(self, _rdts_df): @pytest.fixture def _condor_q_output() -> str: - condor_q_output = """-- Schedd: bigbird08.cern.ch : <188.185.72.155:9618?... @ 04/22/21 12:26:02 + return """-- Schedd: bigbird08.cern.ch : <188.185.72.155:9618?... @ 04/22/21 12:26:02 OWNER BATCH_NAME SUBMITTED DONE RUN IDLE TOTAL JOB_IDS fesoubel ID: 8489182 4/21 21:04 7 14 _ 21 8489182.0-20 fesoubel ID: 8489183 4/21 21:04 2 19 _ 21 8489183.0-20 @@ -707,18 +708,16 @@ def _condor_q_output() -> str: Total for query: 63 jobs; 0 completed, 0 removed, 1 idle, 62 running, 0 held, 0 suspended Total for fesoubel: 63 jobs; 0 completed, 0 removed, 1 idle, 62 running, 0 held, 0 suspended Total for all users: 7279 jobs; 1 completed, 1 removed, 3351 idle, 3724 running, 202 held, 0 suspended""" - return condor_q_output @pytest.fixture def _taskless_condor_q_output() -> str: - taskless_condor_q_output = """-- Schedd: bigbird08.cern.ch : <188.185.72.155:9618?... @ 04/22/21 12:26:02 + return """-- Schedd: bigbird08.cern.ch : <188.185.72.155:9618?... @ 04/22/21 12:26:02 OWNER BATCH_NAME SUBMITTED DONE RUN IDLE TOTAL JOB_IDS Total for query: 63 jobs; 0 completed, 0 removed, 1 idle, 62 running, 0 held, 0 suspended Total for fesoubel: 63 jobs; 0 completed, 0 removed, 1 idle, 62 running, 0 held, 0 suspended Total for all users: 7279 jobs; 1 completed, 1 removed, 3351 idle, 3724 running, 202 held, 0 suspended""" - return taskless_condor_q_output @pytest.fixture From 51f7d8472c94a771079f62f959c53bc15475d55b Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 12:53:14 +0200 Subject: [PATCH 33/41] names --- tests/test_cpymadtools/test_lhc.py | 42 ++++++++++++------------- tests/test_optics.py | 6 ++-- tests/test_plotting/test_tunediagram.py | 3 +- tests/test_utils.py | 6 ++-- 4 files changed, 29 insertions(+), 28 deletions(-) diff --git a/tests/test_cpymadtools/test_lhc.py b/tests/test_cpymadtools/test_lhc.py index c935af45..b96bdf73 100644 --- a/tests/test_cpymadtools/test_lhc.py +++ b/tests/test_cpymadtools/test_lhc.py @@ -402,33 +402,33 @@ def test_get_lhc_bpms_list(_non_matched_lhc_madx, _correct_bpms_list): @pytest.mark.parametrize("knob_value", [-5, 10]) -@pytest.mark.parametrize("IR", [1, 2, 5, 8]) -def test_colinearity_knob(knob_value, IR, _non_matched_lhc_madx): +@pytest.mark.parametrize("ir", [1, 2, 5, 8]) +def test_colinearity_knob(knob_value, ir, _non_matched_lhc_madx): madx = _non_matched_lhc_madx - apply_lhc_colinearity_knob(madx, colinearity_knob_value=knob_value, ir=IR) + apply_lhc_colinearity_knob(madx, colinearity_knob_value=knob_value, ir=ir) - assert madx.globals[f"KQSX3.R{IR:d}"] == knob_value * 1e-4 - assert madx.globals[f"KQSX3.L{IR:d}"] == -1 * knob_value * 1e-4 + assert madx.globals[f"KQSX3.R{ir:d}"] == knob_value * 1e-4 + assert madx.globals[f"KQSX3.L{ir:d}"] == -1 * knob_value * 1e-4 @pytest.mark.parametrize("knob_delta", [-3, 5]) -@pytest.mark.parametrize("IR", [1, 2, 5, 8]) -def test_colinearity_knob_delta(knob_delta, IR, _non_matched_lhc_madx): +@pytest.mark.parametrize("ir", [1, 2, 5, 8]) +def test_colinearity_knob_delta(knob_delta, ir, _non_matched_lhc_madx): madx = _non_matched_lhc_madx # Assign a value first to make it trickier init = 1.5e-4 - madx.globals[f"KQSX3.R{IR:d}"] = init - madx.globals[f"KQSX3.L{IR:d}"] = -1 * init + madx.globals[f"KQSX3.R{ir:d}"] = init + madx.globals[f"KQSX3.L{ir:d}"] = -1 * init # We started from 0 so it should be this value - apply_lhc_colinearity_knob_delta(madx, colinearity_knob_delta=knob_delta, ir=IR) - assert madx.globals[f"KQSX3.R{IR:d}"] == init + knob_delta * 1e-4 - assert madx.globals[f"KQSX3.L{IR:d}"] == -1 * init - knob_delta * 1e-4 + apply_lhc_colinearity_knob_delta(madx, colinearity_knob_delta=knob_delta, ir=ir) + assert madx.globals[f"KQSX3.R{ir:d}"] == init + knob_delta * 1e-4 + assert madx.globals[f"KQSX3.L{ir:d}"] == -1 * init - knob_delta * 1e-4 # Now change the knob value and check that the delta is applied - apply_lhc_colinearity_knob_delta(madx, colinearity_knob_delta=knob_delta, ir=IR) - assert madx.globals[f"KQSX3.R{IR:d}"] == init + 2 * knob_delta * 1e-4 - assert madx.globals[f"KQSX3.L{IR:d}"] == -1 * init - 2 * knob_delta * 1e-4 + apply_lhc_colinearity_knob_delta(madx, colinearity_knob_delta=knob_delta, ir=ir) + assert madx.globals[f"KQSX3.R{ir:d}"] == init + 2 * knob_delta * 1e-4 + assert madx.globals[f"KQSX3.L{ir:d}"] == -1 * init - 2 * knob_delta * 1e-4 def test_rigidity_knob_fails_on_invalid_side(caplog, _non_matched_lhc_madx): @@ -443,14 +443,14 @@ def test_rigidity_knob_fails_on_invalid_side(caplog, _non_matched_lhc_madx): @pytest.mark.parametrize("side", ["left", "right"]) @pytest.mark.parametrize("knob_value", [1, 2]) -@pytest.mark.parametrize("IR", [1, 2, 5, 8]) -def test_rigidity_knob(side, knob_value, IR, _non_matched_lhc_madx): +@pytest.mark.parametrize("ir", [1, 2, 5, 8]) +def test_rigidity_knob(side, knob_value, ir, _non_matched_lhc_madx): madx = _non_matched_lhc_madx - right_knob, left_knob = (f"kqx.r{IR:d}", f"kqx.l{IR:d}") + right_knob, left_knob = (f"kqx.r{ir:d}", f"kqx.l{ir:d}") current_right_knob = madx.globals[right_knob] current_left_knob = madx.globals[left_knob] - apply_lhc_rigidity_waist_shift_knob(madx, rigidty_waist_shift_value=knob_value, ir=IR, side=side) + apply_lhc_rigidity_waist_shift_knob(madx, rigidty_waist_shift_value=knob_value, ir=ir, side=side) if side == "left": assert madx.globals[right_knob] == (1 - knob_value * 0.005) * current_right_knob @@ -578,10 +578,10 @@ def test_resetting_lhc_bump_flags(_bare_lhc_madx): ) with madx.batch(): madx.globals.update({bump: random.random() for bump in ALL_BUMPS}) - assert all([madx.globals[bump] != 0 for bump in ALL_BUMPS]) + assert all(madx.globals[bump] != 0 for bump in ALL_BUMPS) reset_lhc_bump_flags(madx) - assert all([madx.globals[bump] == 0 for bump in ALL_BUMPS]) + assert all(madx.globals[bump] == 0 for bump in ALL_BUMPS) def test_vary_independent_ir_quads(_non_matched_lhc_madx): diff --git a/tests/test_optics.py b/tests/test_optics.py index 8ad7d7d9..85866578 100644 --- a/tests/test_optics.py +++ b/tests/test_optics.py @@ -113,9 +113,9 @@ def test_beam_parameters(pc_gev, en_x_m, en_y_m, delta_p, result): def test_beam_size(_fake_coordinates): - assert np.allclose(ripken._beam_size(_fake_coordinates), _fake_coordinates.std()) + assert np.allclose(ripken._beam_size(_fake_coordinates), _fake_coordinates.std()) # noqa: SLF001 assert np.allclose( - ripken._beam_size(_fake_coordinates, method="rms"), + ripken._beam_size(_fake_coordinates, method="rms"), # noqa: SLF001 np.sqrt(np.mean(np.square(_fake_coordinates))), ) @@ -152,7 +152,7 @@ def test_add_beam_size_to_df(_non_matched_lhc_madx): df["BETA21"] = df.beta21 df["BETA22"] = df.beta22 - df = ripken._add_beam_size_to_df(df, 1e-6, 1e-6) + df = ripken._add_beam_size_to_df(df, 1e-6, 1e-6) # noqa: SLF001 assert "SIZE_X" in df.columns assert "SIZE_Y" in df.columns diff --git a/tests/test_plotting/test_tunediagram.py b/tests/test_plotting/test_tunediagram.py index ab36ed69..31786ec7 100644 --- a/tests/test_plotting/test_tunediagram.py +++ b/tests/test_plotting/test_tunediagram.py @@ -1,6 +1,7 @@ import matplotlib as mpl import matplotlib.pyplot as plt import pytest +from matplotlib.text import Text from pyhdtoolkit.plotting.tune import plot_tune_diagram @@ -45,4 +46,4 @@ def test_plot_tune_diagram_arguments(figure_title, legend_title, max_order, diff differentiate_orders=differentiate, ) assert ax.get_title() == figure_title - assert isinstance(ax.legend().get_title(), matplotlib.text.Text) + assert isinstance(ax.legend().get_title(), Text) diff --git a/tests/test_utils.py b/tests/test_utils.py index 81f08855..7c72a5dc 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -63,8 +63,8 @@ def test_check_pid(self): CommandLine.check_pid_exists("not_an_integer") def test_run_cmd(self): - assert type(CommandLine.run("echo hello")) is tuple - assert type(CommandLine.run("echo hello")[1]) is bytes + assert isinstance(CommandLine.run("echo hello"), tuple) + assert isinstance(CommandLine.run("echo hello")[1], bytes) assert CommandLine.run("echo hello") == (0, b"hello\n") modified_env = os.environ.copy() modified_env["TEST_VAR"] = "Check_me_string" @@ -188,7 +188,7 @@ def test_bifurcate_by(self, inputs, func, results): [ (list(range(10)), 3, [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]]), (list(range(10)), 20, list(range(10))), - (list(), 5, list()), + ([], 5, []), ], ) def test_chunk_list(self, array, size, result): From ecaac3b3493ffc9be240889da6184a51c72b06c2 Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 12:56:42 +0200 Subject: [PATCH 34/41] unecessary list comprehensions --- tests/test_cpymadtools/test_lhc.py | 4 ++-- tests/test_cpymadtools/test_ptc.py | 4 ++-- tests/test_cpymadtools/test_track.py | 4 ++-- tests/test_cpymadtools/test_utilities.py | 6 +++--- tests/test_maths.py | 2 +- tests/test_models.py | 4 ++-- tests/test_plotting/test_phasespace.py | 1 - tests/test_plotting/test_plotting_utils.py | 4 ++-- 8 files changed, 14 insertions(+), 15 deletions(-) diff --git a/tests/test_cpymadtools/test_lhc.py b/tests/test_cpymadtools/test_lhc.py index b96bdf73..8d67a89c 100644 --- a/tests/test_cpymadtools/test_lhc.py +++ b/tests/test_cpymadtools/test_lhc.py @@ -533,7 +533,7 @@ def test_makethin_lhc(_matched_lhc_madx): assert isinstance(tracks_dict, dict) tracks = tracks_dict["observation_point_1"] assert len(tracks) == 11 # nturns + 1 because $start coordinates also given by MAD-X - assert all([coordinate in tracks.columns for coordinate in ("x", "px", "y", "py", "t", "pt", "s", "e")]) + assert all(coordinate in tracks.columns for coordinate in ("x", "px", "y", "py", "t", "pt", "s", "e")) @pytest.mark.parametrize("markers", [100, 1000]) @@ -781,7 +781,7 @@ def test_get_irs_twiss(ir, _matched_lhc_madx): extra_columns = ["k0l", "k0sl", "k1l", "k1sl", "k2l", "k2sl", "sig11", "sig12", "sig21", "sig22"] ir_extra_columns_df = get_ir_twiss(madx, ir=ir, columns=DEFAULT_TWISS_COLUMNS + extra_columns) - assert all([colname in ir_extra_columns_df.columns for colname in extra_columns]) + assert all(colname in ir_extra_columns_df.columns for colname in extra_columns) # ------------------- Requires acc-models-lhc ------------------- # diff --git a/tests/test_cpymadtools/test_ptc.py b/tests/test_cpymadtools/test_ptc.py index 1e76dc9b..4a168adf 100644 --- a/tests/test_cpymadtools/test_ptc.py +++ b/tests/test_cpymadtools/test_ptc.py @@ -72,7 +72,7 @@ def test_single_particle_ptc_track(_matched_base_lattice, obs_points): assert len(tracks_dict.keys()) == len(obs_points) + 1 for tracks in tracks_dict.values(): assert isinstance(tracks, DataFrame) - assert all([coordinate in tracks.columns for coordinate in ("x", "px", "y", "py", "t", "pt", "s", "e")]) + assert all(coordinate in tracks.columns for coordinate in ("x", "px", "y", "py", "t", "pt", "s", "e")) def test_single_particle_ptc_track_with_onepass(_matched_base_lattice): @@ -86,7 +86,7 @@ def test_single_particle_ptc_track_with_onepass(_matched_base_lattice): assert "trackone" in tracks_dict tracks = tracks_dict["trackone"] assert isinstance(tracks, DataFrame) - assert all([coordinate in tracks.columns for coordinate in ("x", "px", "y", "py", "t", "pt", "s", "e")]) + assert all(coordinate in tracks.columns for coordinate in ("x", "px", "y", "py", "t", "pt", "s", "e")) # ----- Fixtures ----- # diff --git a/tests/test_cpymadtools/test_track.py b/tests/test_cpymadtools/test_track.py index b96277f6..f7b15dec 100644 --- a/tests/test_cpymadtools/test_track.py +++ b/tests/test_cpymadtools/test_track.py @@ -19,7 +19,7 @@ def test_single_particle_tracking(_matched_base_lattice, obs_points): assert len(tracks_dict.keys()) == len(obs_points) + 1 for tracks in tracks_dict.values(): assert isinstance(tracks, DataFrame) - assert all([coordinate in tracks.columns for coordinate in ("x", "px", "y", "py", "t", "pt", "s", "e")]) + assert all(coordinate in tracks.columns for coordinate in ("x", "px", "y", "py", "t", "pt", "s", "e")) def test_single_particle_tracking_with_onepass(_matched_base_lattice): @@ -33,4 +33,4 @@ def test_single_particle_tracking_with_onepass(_matched_base_lattice): assert "trackone" in tracks_dict tracks = tracks_dict["trackone"] assert isinstance(tracks, DataFrame) - assert all([coordinate in tracks.columns for coordinate in ("x", "px", "y", "py", "t", "pt", "s", "e")]) + assert all(coordinate in tracks.columns for coordinate in ("x", "px", "y", "py", "t", "pt", "s", "e")) diff --git a/tests/test_cpymadtools/test_utilities.py b/tests/test_cpymadtools/test_utilities.py index 6f67833b..9fe3df53 100644 --- a/tests/test_cpymadtools/test_utilities.py +++ b/tests/test_cpymadtools/test_utilities.py @@ -8,9 +8,9 @@ @pytest.mark.parametrize( ("orient", "result"), [ - ["straight", ["K0L", "K1L", "K2L", "K3L", "K4L"]], - ["skew", ["K0SL", "K1SL", "K2SL", "K3SL", "K4SL"]], - ["both", ["K0L", "K0SL", "K1L", "K1SL", "K2L", "K2SL", "K3L", "K3SL", "K4L", "K4SL"]], + ("straight", ["K0L", "K1L", "K2L", "K3L", "K4L"]), + ("skew", ["K0SL", "K1SL", "K2SL", "K3SL", "K4SL"]), + ("both", ["K0L", "K0SL", "K1L", "K1SL", "K2L", "K2SL", "K3L", "K3SL", "K4L", "K4SL"]), ], ) def test_k_strings(orient, result): diff --git a/tests/test_maths.py b/tests/test_maths.py index 2b7050ef..4b5ac6eb 100644 --- a/tests/test_maths.py +++ b/tests/test_maths.py @@ -87,7 +87,7 @@ def test_make_pdf(degrees_of_freedom): data: np.ndarray = st.chi2(degrees_of_freedom).rvs(50_000) best_fit_func, best_fit_params = stats_fitting.best_fit_distribution(data, 200) pdf = stats_fitting.make_pdf(best_fit_func, best_fit_params) - pdf.idxmax() == pytest.approx(degrees_of_freedom - 2, rel=1e-2) + assert pdf.idxmax() == pytest.approx(degrees_of_freedom - 2, rel=1e-2) @pytest.mark.parametrize(("value", "result"), [(1, 0), (10, 1), (0.0311, -2), (5e-7, -7)]) diff --git a/tests/test_models.py b/tests/test_models.py index 52221a17..3fe4b9fc 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -14,5 +14,5 @@ def test_beam_parameters(): assert math.isclose(beam.beta_r, 0.8966300434726596) assert math.isclose(beam.eg_x_m, 2.469137056052632e-06) assert math.isclose(beam.eg_y_m, 2.469137056052632e-06) - print(beam.__str__()) - print(beam.__repr__()) + beam.__str__() + beam.__repr__() diff --git a/tests/test_plotting/test_phasespace.py b/tests/test_plotting/test_phasespace.py index ed5695ef..4d42714c 100644 --- a/tests/test_plotting/test_phasespace.py +++ b/tests/test_plotting/test_phasespace.py @@ -99,7 +99,6 @@ def test_plot_courant_snyder_phase_space_colored_wrong_plane_input(): match_cas3(madx) x_coords_stable, px_coords_stable = np.array([]), np.array([]) # no need for tracking with pytest.raises(ValueError): - fig, ax = plt.subplots() plot_courant_snyder_phase_space_colored(madx, x_coords_stable, px_coords_stable, plane="invalid_plane") diff --git a/tests/test_plotting/test_plotting_utils.py b/tests/test_plotting/test_plotting_utils.py index c81d3fe4..67ac9263 100644 --- a/tests/test_plotting/test_plotting_utils.py +++ b/tests/test_plotting/test_plotting_utils.py @@ -104,7 +104,7 @@ def test_ip_locations_with_xlimits(_non_matched_lhc_madx): def test_confidence_ellipse_subplots(): """Confidence ellipse on three correlated datasets in subplots.""" np.random.seed(0) - PARAMETERS = { + parameters = { "Positive correlation": [[0.85, 0.35], [0.15, -0.65]], "Negative correlation": [[0.9, -0.4], [0.1, -0.6]], "Weak correlation": [[1, 0], [0, 1]], @@ -113,7 +113,7 @@ def test_confidence_ellipse_subplots(): scale = 3, 5 figure, axs = plt.subplots(1, 3, figsize=(9, 3)) - for ax, (title, dependency) in zip(axs, PARAMETERS.items(), strict=False): + for ax, (title, dependency) in zip(axs, parameters.items(), strict=False): x, y = get_correlated_dataset(800, dependency, mu, scale) ax.scatter(x, y, s=0.5) ax.axvline(c="grey", lw=1) From 03262ea4888d4f07fd81eda0f77028f906c1b47f Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 12:59:15 +0200 Subject: [PATCH 35/41] fixtures want lists of tuples --- tests/test_cpymadtools/test_lhc.py | 6 +++--- tests/test_utils.py | 10 ++++++++-- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/tests/test_cpymadtools/test_lhc.py b/tests/test_cpymadtools/test_lhc.py index 8d67a89c..8cbcb0aa 100644 --- a/tests/test_cpymadtools/test_lhc.py +++ b/tests/test_cpymadtools/test_lhc.py @@ -332,9 +332,9 @@ def test_all_lhc_arcs(): @pytest.mark.parametrize( ("orient", "result"), [ - ["straight", ["K0L", "K1L", "K2L", "K3L", "K4L"]], - ["skew", ["K0SL", "K1SL", "K2SL", "K3SL", "K4SL"]], - ["both", ["K0L", "K0SL", "K1L", "K1SL", "K2L", "K2SL", "K3L", "K3SL", "K4L", "K4SL"]], + ("straight", ["K0L", "K1L", "K2L", "K3L", "K4L"]), + ("skew", ["K0SL", "K1SL", "K2SL", "K3SL", "K4SL"]), + ("both", ["K0L", "K0SL", "K1L", "K1SL", "K2L", "K2SL", "K3L", "K3SL", "K4L", "K4SL"]), ], ) def test_k_strings(orient, result): diff --git a/tests/test_utils.py b/tests/test_utils.py index 7c72a5dc..7767b6fd 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -166,7 +166,7 @@ def test_bifurcate(self, inputs, filters, results): [ (["beep", "boop", "foo", "bar", "baz"], None, TypeError), (["beep", "boop", "foo", "bar", "baz"], [], IndexError), - ([["beep", "boop", "foo", "bar"], [True, False, False], IndexError]), + (["beep", "boop", "foo", "bar"], [True, False, False], IndexError), ], ) def test_bifurcate_fails(self, inputs, filters, error): @@ -255,7 +255,13 @@ def test_groupby(self, array, func, result): assert ListOperations.group_by(array, func) == result @pytest.mark.parametrize( - ("array", "result"), [([1, 2, 1], True), ([list(range(10)), False]), ([], False), ([True, True], True)] + ("array", "result"), + [ + ([1, 2, 1], True), + (list(range(10)), False), + ([], False), + ([True, True], True) + ] ) def test_has_duplicates(self, array, result): assert ListOperations.has_duplicates(array) is result From 4b97daa2c3bf4186229c48fa70429ea43aa3ebb4 Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 13:19:35 +0200 Subject: [PATCH 36/41] provide 'match' argument to every pytest.raises, normalize error messages --- pyhdtoolkit/cpymadtools/tune.py | 5 +-- pyhdtoolkit/plotting/phasespace.py | 4 +-- pyhdtoolkit/plotting/sbs/phase.py | 2 +- tests/conftest.py | 8 ++--- tests/test_cpymadtools/test_generators.py | 6 ++-- tests/test_cpymadtools/test_lhc.py | 38 +++++++++++----------- tests/test_cpymadtools/test_tune.py | 2 +- tests/test_cpymadtools/test_utilities.py | 2 +- tests/test_maths.py | 3 +- tests/test_optics.py | 3 +- tests/test_plotting/test_aperture.py | 2 +- tests/test_plotting/test_envelope.py | 2 +- tests/test_plotting/test_phasespace.py | 5 ++- tests/test_plotting/test_plotting_utils.py | 6 ++-- tests/test_plotting/test_sbs_phase.py | 8 ++--- tests/test_plotting/test_tunediagram.py | 2 +- tests/test_utils.py | 6 ++-- 17 files changed, 52 insertions(+), 52 deletions(-) diff --git a/pyhdtoolkit/cpymadtools/tune.py b/pyhdtoolkit/cpymadtools/tune.py index 83c24534..69b9ee0e 100644 --- a/pyhdtoolkit/cpymadtools/tune.py +++ b/pyhdtoolkit/cpymadtools/tune.py @@ -197,12 +197,13 @@ def get_footprint_patches(dynap_dframe: tfs.TfsDataFrame) -> matplotlib.collecti a[0, :, 1] = dynap_dframe["tuny"].to_numpy()[0] a[1:, :, 0] = dynap_dframe["tunx"].to_numpy()[1:].reshape(-1, angle) a[1:, :, 1] = dynap_dframe["tuny"].to_numpy()[1:].reshape(-1, angle) - except ValueError: + except ValueError as dynap_error: logger.exception( "Cannot group tune points according to starting angles and amplitudes. Try changing " "the 'AMPLITUDE' value in the provided TfsDataFrame's headers." ) - raise + msg = "Invalid AMPLITUDE value in the provided TfsDataFrame headers" + raise ValueError(msg) from dynap_error logger.debug("Determining polygon vertices") sx = a.shape[0] - 1 diff --git a/pyhdtoolkit/plotting/phasespace.py b/pyhdtoolkit/plotting/phasespace.py index acd7e783..d0a9fcf0 100644 --- a/pyhdtoolkit/plotting/phasespace.py +++ b/pyhdtoolkit/plotting/phasespace.py @@ -62,7 +62,7 @@ def plot_courant_snyder_phase_space( """ if plane.lower() not in ("horizontal", "vertical"): logger.error(f"Plane should be either Horizontal or Vertical but '{plane}' was given") - msg = "Invalid 'plane' value." + msg = "Invalid 'plane' argument." raise ValueError(msg) logger.debug("Plotting phase space for normalized Courant-Snyder coordinates") @@ -135,7 +135,7 @@ def plot_courant_snyder_phase_space_colored( """ if plane.upper() not in ("HORIZONTAL", "VERTICAL"): logger.error(f"Plane should be either horizontal or vertical but '{plane}' was given") - msg = "Invalid 'plane' value." + msg = "Invalid 'plane' argument." raise ValueError(msg) # Getting a sufficiently long array of colors to use diff --git a/pyhdtoolkit/plotting/sbs/phase.py b/pyhdtoolkit/plotting/sbs/phase.py index 20f21463..25eacaf5 100644 --- a/pyhdtoolkit/plotting/sbs/phase.py +++ b/pyhdtoolkit/plotting/sbs/phase.py @@ -184,7 +184,7 @@ def plot_phase_segment( """ if plane.upper() not in ("X", "Y"): logger.error("The provided plane is invalid, should be either 'x' or 'y', case-insensitively.") - msg = "Invalid 'plane' parameter" + msg = "Invalid 'plane' argument." raise ValueError(msg) plane = plane.upper() diff --git a/tests/conftest.py b/tests/conftest.py index b65ac811..0e8e2286 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -52,13 +52,13 @@ def _non_matched_lhc_madx() -> Madx: madx.call(str(LHC_SEQUENCE.absolute())) madx.call(str(LHC_OPTICS.absolute())) # opticsfile.22 - NRJ = madx.globals["NRJ"] = 6500 + nrj = madx.globals["NRJ"] = 6500 madx.globals["brho"] = madx.globals["NRJ"] * 1e9 / madx.globals.clight geometric_emit = madx.globals["geometric_emit"] = 3.75e-6 / (madx.globals["NRJ"] / 0.938) madx.command.beam( sequence="lhcb1", bv=1, - energy=NRJ, + energy=nrj, particle="proton", npart=1.0e10, kbunch=1, @@ -76,13 +76,13 @@ def _matched_lhc_madx() -> Madx: madx.call(str(LHC_SEQUENCE.absolute())) madx.call(str(LHC_OPTICS.absolute())) # opticsfile.22 - NRJ = madx.globals["NRJ"] = 6500 + nrj = madx.globals["NRJ"] = 6500 madx.globals["brho"] = madx.globals["NRJ"] * 1e9 / madx.globals.clight geometric_emit = madx.globals["geometric_emit"] = 3.75e-6 / (madx.globals["NRJ"] / 0.938) madx.command.beam( sequence="lhcb1", bv=1, - energy=NRJ, + energy=nrj, particle="proton", npart=1.0e10, kbunch=1, diff --git a/tests/test_cpymadtools/test_generators.py b/tests/test_cpymadtools/test_generators.py index 58c898a1..ca362cae 100644 --- a/tests/test_cpymadtools/test_generators.py +++ b/tests/test_cpymadtools/test_generators.py @@ -24,7 +24,7 @@ def test_oneoct_cas_lattice(): def test_tripleterrors_study_reference(): - tripleterrors_study_reference = LatticeGenerator._generate_tripleterrors_study_reference() + tripleterrors_study_reference = LatticeGenerator._generate_tripleterrors_study_reference() # noqa: SLF001 assert isinstance(tripleterrors_study_reference, str) assert len(tripleterrors_study_reference) == 1617 @@ -40,7 +40,7 @@ def test_tripleterrors_study_reference(): ], ) def test_tripleterrors_study_tferror_job(randseed, tferror): - tripleterrors_study_tferror_job = LatticeGenerator._generate_tripleterrors_study_tferror_job( + tripleterrors_study_tferror_job = LatticeGenerator._generate_tripleterrors_study_tferror_job( # noqa: SLF001 rand_seed=randseed, tf_error=tferror ) assert isinstance(tripleterrors_study_tferror_job, str) @@ -60,7 +60,7 @@ def test_tripleterrors_study_tferror_job(randseed, tferror): ], ) def test_tripleterrors_study_mserror_job(randseed, mserror): - tripleterrors_study_mserror_job = LatticeGenerator._generate_tripleterrors_study_mserror_job( + tripleterrors_study_mserror_job = LatticeGenerator._generate_tripleterrors_study_mserror_job( # noqa: SLF001 rand_seed=randseed, ms_error=mserror ) assert isinstance(tripleterrors_study_mserror_job, str) diff --git a/tests/test_cpymadtools/test_lhc.py b/tests/test_cpymadtools/test_lhc.py index 8cbcb0aa..55b4797f 100644 --- a/tests/test_cpymadtools/test_lhc.py +++ b/tests/test_cpymadtools/test_lhc.py @@ -194,7 +194,7 @@ def test_misalign_lhc_ir_quadrupoles_specific_value(_non_matched_lhc_madx): def test_misalign_lhc_ir_quadrupoles_raises_on_wrong_side(_non_matched_lhc_madx, caplog): madx = _non_matched_lhc_madx - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="Invalid"): misalign_lhc_ir_quadrupoles(madx, ips=[8], quadrupoles=[1], beam=2, sides="Z", dy="0.001") for record in caplog.records: @@ -203,7 +203,7 @@ def test_misalign_lhc_ir_quadrupoles_raises_on_wrong_side(_non_matched_lhc_madx, def test_misalign_lhc_ir_quadrupoles_raises_on_wrong_ip(_non_matched_lhc_madx, caplog): madx = _non_matched_lhc_madx - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="Invalid"): misalign_lhc_ir_quadrupoles(madx, ips=[100], quadrupoles=[1], beam=2, sides="R", dy="0.001") for record in caplog.records: @@ -212,7 +212,7 @@ def test_misalign_lhc_ir_quadrupoles_raises_on_wrong_ip(_non_matched_lhc_madx, c def test_misalign_lhc_ir_quadrupoles_raises_on_wrong_beam(_non_matched_lhc_madx, caplog): madx = _non_matched_lhc_madx - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="Invalid"): misalign_lhc_ir_quadrupoles(madx, ips=[2], quadrupoles=[1], beam=10, sides="L", dy="0.001") for record in caplog.records: @@ -275,7 +275,7 @@ def test_lhc_orbit_variables(): def test_lhc_orbit_setup_fals_on_unknown_scheme(_non_matched_lhc_madx, caplog): madx = _non_matched_lhc_madx - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="Invalid scheme parameter given"): setup_lhc_orbit(madx, scheme="unknown") for record in caplog.records: @@ -342,7 +342,7 @@ def test_k_strings(orient, result): def test_k_strings_fails_on_wront_orient(caplog): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="Invalid 'orientation' parameter"): _ = _get_k_strings(orientation="qqq") for record in caplog.records: @@ -434,7 +434,7 @@ def test_colinearity_knob_delta(knob_delta, ir, _non_matched_lhc_madx): def test_rigidity_knob_fails_on_invalid_side(caplog, _non_matched_lhc_madx): madx = _non_matched_lhc_madx - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="Invalid value for parameter 'side'."): apply_lhc_rigidity_waist_shift_knob(madx, 1, 1, "invalid") for record in caplog.records: @@ -568,7 +568,7 @@ def test_re_cycling(_bare_lhc_madx, start_point): def test_resetting_lhc_bump_flags(_bare_lhc_madx): madx = _bare_lhc_madx make_lhc_beams(madx) - ALL_BUMPS = ( + all_bumps = ( LHC_ANGLE_FLAGS + LHC_CROSSING_ANGLE_FLAGS + LHC_EXPERIMENT_STATE_FLAGS @@ -577,11 +577,11 @@ def test_resetting_lhc_bump_flags(_bare_lhc_madx): + LHC_PARALLEL_SEPARATION_FLAGS ) with madx.batch(): - madx.globals.update({bump: random.random() for bump in ALL_BUMPS}) - assert all(madx.globals[bump] != 0 for bump in ALL_BUMPS) + madx.globals.update({bump: random.random() for bump in all_bumps}) + assert all(madx.globals[bump] != 0 for bump in all_bumps) reset_lhc_bump_flags(madx) - assert all(madx.globals[bump] == 0 for bump in ALL_BUMPS) + assert all(madx.globals[bump] == 0 for bump in all_bumps) def test_vary_independent_ir_quads(_non_matched_lhc_madx): @@ -594,7 +594,7 @@ def test_vary_independent_ir_quads(_non_matched_lhc_madx): def test_vary_independent_ir_quads_raises_on_wrong_side(_non_matched_lhc_madx, caplog): madx = _non_matched_lhc_madx - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="Invalid 'quad_numbers', 'ip', 'sides' argument"): vary_independent_ir_quadrupoles(madx, quad_numbers=[4], ip=1, sides="Z") for record in caplog.records: @@ -603,7 +603,7 @@ def test_vary_independent_ir_quads_raises_on_wrong_side(_non_matched_lhc_madx, c def test_vary_independent_ir_quads_raises_on_wrong_ip(_non_matched_lhc_madx, caplog): madx = _non_matched_lhc_madx - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="Invalid 'quad_numbers', 'ip', 'sides' argument"): vary_independent_ir_quadrupoles(madx, quad_numbers=[4], ip=100, sides="R") for record in caplog.records: @@ -612,7 +612,7 @@ def test_vary_independent_ir_quads_raises_on_wrong_ip(_non_matched_lhc_madx, cap def test_vary_independent_ir_quads_raises_on_wrong_quads(_non_matched_lhc_madx, caplog): madx = _non_matched_lhc_madx - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="Invalid 'quad_numbers', 'ip', 'sides' argument"): vary_independent_ir_quadrupoles(madx, quad_numbers=[5, 20, 100], ip=1, sides="R") for record in caplog.records: @@ -808,21 +808,21 @@ def test_lhc_run3_setup_context_manager_fullpath_to_opticsfile(): @pytest.mark.skipif(not (TESTS_DIR.parent / "acc-models-lhc").is_dir(), reason="acc-models-lhc not found") def test_lhc_run3_setup_context_manager_raises_on_wrong_b4_conditions(): - with pytest.raises(ValueError): # using b4 with beam1 setup crashes + with pytest.raises(ValueError, match="Cannot use beam 4 sequence file for beam 1"): # using b4 with beam1 setup crashes # noqa: SIM117 with LHCSetup(opticsfile="R2022a_A30cmC30cmA10mL200cm.madx", beam=1, use_b4=True) as madx: # noqa: F841 pass @pytest.mark.skipif(not (TESTS_DIR.parent / "acc-models-lhc").is_dir(), reason="acc-models-lhc not found") def test_lhc_run3_setup_context_manager_raises_on_wrong_run_value(): - with pytest.raises(NotImplementedError): # using b4 with beam1 setup crashes + with pytest.raises(NotImplementedError, match="This setup is only possible for Run 2 and Run 3 configurations."): # using b4 with beam1 setup crashes # noqa: SIM117 with LHCSetup(run=1, opticsfile="R2022a_A30cmC30cmA10mL200cm.madx") as madx: # noqa: F841 pass @pytest.mark.skipif(not (TESTS_DIR.parent / "acc-models-lhc").is_dir(), reason="acc-models-lhc not found") -def test_lhc_run3_setup_raises_on_wrong_b4_conditions(_proton_opticsfile): - with pytest.raises(ValueError): # using b4 with beam1 setup crashes +def test_lhc_run3_setup_raises_on_wrong_b4_conditions(): + with pytest.raises(ValueError, match="Cannot use beam 4 sequence file for beam 1"): # using b4 with beam1 setup crashes _ = prepare_lhc_run3(opticsfile="R2022a_A30cmC30cmA10mL200cm.madx", beam=1, use_b4=True) @@ -839,12 +839,12 @@ def test_lhc_run2_setup_context_manager(_proton_opticsfile, slicefactor): def test_lhc_run2_setup_raises_on_wrong_b4_conditions(_proton_opticsfile): - with pytest.raises(ValueError): # using b4 with beam1 setup crashes + with pytest.raises(ValueError, match="Cannot use beam 4 sequence file for beam 1"): # using b4 with beam1 setup crashes _ = prepare_lhc_run2(opticsfile=_proton_opticsfile, beam=1, use_b4=True) def test_lhc_run2_setup_raises_on_absent_sequence_file(): - with pytest.raises(ValueError): # will not find the sequence file from this opticsfile value + with pytest.raises(ValueError, match="No sequence file found at"): # will not find the sequence file from this opticsfile value _ = prepare_lhc_run2(opticsfile="some/place/here.madx") diff --git a/tests/test_cpymadtools/test_tune.py b/tests/test_cpymadtools/test_tune.py index b213e111..681350a7 100644 --- a/tests/test_cpymadtools/test_tune.py +++ b/tests/test_cpymadtools/test_tune.py @@ -96,7 +96,7 @@ def test_get_footprint_patches(_dynap_tfs_path): def test_get_footprint_patches_raises_wrong_shape(_dynap_tfs_path, caplog): dynap_dframe = tfs.read(_dynap_tfs_path) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="Invalid AMPLITUDE value in the provided TfsDataFrame headers"): _ = get_footprint_patches(dynap_dframe) for record in caplog.records: diff --git a/tests/test_cpymadtools/test_utilities.py b/tests/test_cpymadtools/test_utilities.py index 9fe3df53..5bb7e90e 100644 --- a/tests/test_cpymadtools/test_utilities.py +++ b/tests/test_cpymadtools/test_utilities.py @@ -18,7 +18,7 @@ def test_k_strings(orient, result): def test_k_strings_wrong_orientation_raises(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="Invalid 'orientation' parameter"): _get_k_strings(stop=5, orientation="wrong") diff --git a/tests/test_maths.py b/tests/test_maths.py index 4b5ac6eb..aabc034d 100644 --- a/tests/test_maths.py +++ b/tests/test_maths.py @@ -128,7 +128,8 @@ def _create_random_phase_values(low: float, high: float, n_values: int, dist: st elif dist == "uniform": values = np.sort(np.random.default_rng().uniform(low, high, n_values)) else: - raise ValueError("Provided parameter 'distribution' should be either 'linspace' or 'uniform'.") + msg = "Provided parameter 'distribution' should be either 'linspace' or 'uniform'." + raise ValueError(msg) values[0] = 0 return values diff --git a/tests/test_optics.py b/tests/test_optics.py index 85866578..c1842615 100644 --- a/tests/test_optics.py +++ b/tests/test_optics.py @@ -122,7 +122,7 @@ def test_beam_size(_fake_coordinates): def test_beam_size_raises(_fake_coordinates): with pytest.raises(NotImplementedError): - _ = ripken._beam_size(_fake_coordinates, method="not_real") + _ = ripken._beam_size(_fake_coordinates, method="not_real") # noqa: SLF001 @pytest.mark.parametrize("beta11", [0.3312]) @@ -169,6 +169,7 @@ def test_rdt_order_and_type(): with pytest.raises(KeyError): rdt_to_order_and_type(8888) + with pytest.raises(KeyError): rdt_to_order_and_type(1090) diff --git a/tests/test_plotting/test_aperture.py b/tests/test_plotting/test_aperture.py index 4c4f85a7..a85c8102 100644 --- a/tests/test_plotting/test_aperture.py +++ b/tests/test_plotting/test_aperture.py @@ -90,5 +90,5 @@ def test_plot_physical_apertures_ir5_collision_vertical(_collision_aperture_tole def test_plot_physical_apertures_raises_on_wrong_plane(): madx = Madx(stdout=False) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="Invalid 'plane' argument."): plot_physical_apertures(madx, plane="invalid") diff --git a/tests/test_plotting/test_envelope.py b/tests/test_plotting/test_envelope.py index fba46926..a8851ad1 100644 --- a/tests/test_plotting/test_envelope.py +++ b/tests/test_plotting/test_envelope.py @@ -18,7 +18,7 @@ def test_plot_enveloppe_raises_on_wrong_plane(): madx = Madx(stdout=False) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="Invalid 'plane' argument."): plot_beam_envelope(madx, "lhcb1", plane="invalid") @pytest.mark.mpl_image_compare(tolerance=20, style="default", savefig_kwargs={"dpi": 200}) diff --git a/tests/test_plotting/test_phasespace.py b/tests/test_plotting/test_phasespace.py index 4d42714c..b4b1db20 100644 --- a/tests/test_plotting/test_phasespace.py +++ b/tests/test_plotting/test_phasespace.py @@ -87,8 +87,7 @@ def test_plot_courant_snyder_phase_space_wrong_plane_input(): match_cas3(madx) x_coords_stable, px_coords_stable = np.array([]), np.array([]) # no need for tracking - with pytest.raises(ValueError): - fig, ax = plt.subplots() + with pytest.raises(ValueError, match="Invalid 'plane' argument."): plot_courant_snyder_phase_space(madx, x_coords_stable, px_coords_stable, plane="invalid_plane") @@ -98,7 +97,7 @@ def test_plot_courant_snyder_phase_space_colored_wrong_plane_input(): madx.input(BASE_LATTICE) match_cas3(madx) x_coords_stable, px_coords_stable = np.array([]), np.array([]) # no need for tracking - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="Invalid 'plane' argument."): plot_courant_snyder_phase_space_colored(madx, x_coords_stable, px_coords_stable, plane="invalid_plane") diff --git a/tests/test_plotting/test_plotting_utils.py b/tests/test_plotting/test_plotting_utils.py index 67ac9263..04c8b6c7 100644 --- a/tests/test_plotting/test_plotting_utils.py +++ b/tests/test_plotting/test_plotting_utils.py @@ -54,7 +54,7 @@ def test_coupling_ylabel(f1001, f1010, abs_, real, imag): @pytest.mark.parametrize("rdt", ["invalid", "F1111", "nope"]) def test_coupling_ylabel_raises_on_invalid_rdt(rdt): - with pytest.raises(AssertionError): + with pytest.raises(ValueError, match="Invalid RDT for coupling plot."): _determine_default_sbs_coupling_ylabel(rdt, "abs") @@ -68,8 +68,8 @@ def test_phase_ylabel(plane): @pytest.mark.parametrize("plane", ["a", "Fb1", "nope", "not a plane"]) -def test_phase_ylabel_raises_on_invalid_rdt(plane): - with pytest.raises(AssertionError): +def test_phase_ylabel_raises_on_invalid_plane(plane): + with pytest.raises(ValueError, match="Invalid plane for phase plot."): _determine_default_sbs_phase_ylabel(plane) diff --git a/tests/test_plotting/test_sbs_phase.py b/tests/test_plotting/test_sbs_phase.py index 51cfdc44..c5ad8e1e 100644 --- a/tests/test_plotting/test_sbs_phase.py +++ b/tests/test_plotting/test_sbs_phase.py @@ -16,19 +16,18 @@ @pytest.mark.mpl_image_compare(tolerance=20, style="default", savefig_kwargs={"dpi": 200}) def test_plot_one_beam(sbs_phasex, sbs_phasey, sbs_model_b2): - figure = plot_phase_segment_one_beam( + return plot_phase_segment_one_beam( phase_x=sbs_phasex, phase_y=sbs_phasey, model=sbs_model_b2, ip=5, ) - return figure @pytest.mark.mpl_image_compare(tolerance=20, style="default", savefig_kwargs={"dpi": 200}) def test_plot_both_beams(sbs_phasex, sbs_phasey, sbs_model_b2): # Only plotting beam 2 data on both, just for the test - figure = plot_phase_segment_both_beams( + return plot_phase_segment_both_beams( b1_phase_x=sbs_phasex, b1_phase_y=sbs_phasey, b2_phase_x=sbs_phasex, @@ -39,12 +38,11 @@ def test_plot_both_beams(sbs_phasex, sbs_phasey, sbs_model_b2): figsize=(12, 6), bbox_to_anchor=(0.535, 0.94), ) - return figure @pytest.mark.parametrize("wrongplane", ["not", "accepted", "incorrect", ""]) def test_plot_phase_segment_raises_on_wrong_plane(wrongplane, sbs_phasex, sbs_model_b2): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="Invalid 'plane' argument."): plot_phase_segment(segment_df=sbs_phasex, model_df=sbs_model_b2, plane=wrongplane) diff --git a/tests/test_plotting/test_tunediagram.py b/tests/test_plotting/test_tunediagram.py index 31786ec7..806b7d4c 100644 --- a/tests/test_plotting/test_tunediagram.py +++ b/tests/test_plotting/test_tunediagram.py @@ -11,7 +11,7 @@ @pytest.mark.parametrize("max_order", [0, 10, -5]) def test_plot_tune_diagram_fails_on_too_high_order(max_order, caplog): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="The 'max_order' argument should be between 1 and 6 included"): plot_tune_diagram(max_order=max_order) for record in caplog.records: diff --git a/tests/test_utils.py b/tests/test_utils.py index 7767b6fd..357668fc 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -255,7 +255,7 @@ def test_groupby(self, array, func, result): assert ListOperations.group_by(array, func) == result @pytest.mark.parametrize( - ("array", "result"), + ("array", "result"), [ ([1, 2, 1], True), (list(range(10)), False), @@ -436,7 +436,7 @@ def test_multiprocessor(self, function, inputs, results, processes): assert MultiProcessor.execute_function(func=function, func_args=inputs, n_processes=processes) == results def test_multiprocessing_zero_processes(self): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="max_workers must be greater than 0"): MultiProcessor.execute_function(func=_square, func_args=list(range(6)), n_processes=0) @@ -455,7 +455,7 @@ def test_multithreading(self, function, inputs, results, threads): assert MultiThreader.execute_function(func=function, func_args=inputs, n_threads=threads) == results def test_multithreading_zero_threads(self): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="max_workers must be greater than 0"): MultiThreader.execute_function(func=_square, func_args=list(range(6)), n_threads=0) From 23cf1d87d8c0c32ece00f2263d02eeaad3f7fc40 Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 13:21:42 +0200 Subject: [PATCH 37/41] minimal linting here --- examples/demo_lhc_rigid_waist_shift.py | 8 ++++---- examples/demo_stats_fitting.py | 4 ++-- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/examples/demo_lhc_rigid_waist_shift.py b/examples/demo_lhc_rigid_waist_shift.py index c82eb6e2..0f766c75 100755 --- a/examples/demo_lhc_rigid_waist_shift.py +++ b/examples/demo_lhc_rigid_waist_shift.py @@ -36,8 +36,8 @@ 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 @@ -324,9 +324,9 @@ # 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: diff --git a/examples/demo_stats_fitting.py b/examples/demo_stats_fitting.py index 7894b3be..f2aaf162 100755 --- a/examples/demo_stats_fitting.py +++ b/examples/demo_stats_fitting.py @@ -237,11 +237,11 @@ def chi_dist(num: int, meas_used: int) -> np.ndarray: label="Generated Chi Distribution", figsize=(20, 12), ) -dataYLim = ac.get_ylim() +data_ylim = ac.get_ylim() # Find best fit candidate best_fit_func, best_fit_params = fitting.best_fit_distribution(chi_data, 200, ac) -ac.set_ylim(dataYLim) +ac.set_ylim(data_ylim) ac.set_title("All Fitted Distributions") ac.set_ylabel("Normed Hist Counts") plt.legend() From 0e9761d16b08212d52b19f4a9e8431514ea89d9b Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Mon, 12 Aug 2024 14:43:11 +0200 Subject: [PATCH 38/41] remove old helper functions --- tests/test_maths.py | 57 +-------------------------------------------- 1 file changed, 1 insertion(+), 56 deletions(-) diff --git a/tests/test_maths.py b/tests/test_maths.py index aabc034d..e5e40127 100644 --- a/tests/test_maths.py +++ b/tests/test_maths.py @@ -87,7 +87,7 @@ def test_make_pdf(degrees_of_freedom): data: np.ndarray = st.chi2(degrees_of_freedom).rvs(50_000) best_fit_func, best_fit_params = stats_fitting.best_fit_distribution(data, 200) pdf = stats_fitting.make_pdf(best_fit_func, best_fit_params) - assert pdf.idxmax() == pytest.approx(degrees_of_freedom - 2, rel=1e-2) + pdf.idxmax() == pytest.approx(degrees_of_freedom - 2, rel=1e-2) @pytest.mark.parametrize(("value", "result"), [(1, 0), (10, 1), (0.0311, -2), (5e-7, -7)]) @@ -110,61 +110,6 @@ def test_mag_and_string_forced_scale(_to_scale, _force_scaled): # ---------------------- Utilities ---------------------- # -def _create_random_phase_values(low: float, high: float, n_values: int, dist: str) -> np.ndarray: - """ - Returns fake generated phase values. They are ascending by default. - - Args: - low: lowest value, first value will always be 0. - high: highest value. - n_values: number of values to generate. - dist: distribution type required. Will be uniform or linspace. - - Returns: - A `numpy.ndarray` of shape (n_values,) with the generated values. - """ - if dist == "linspace": - values = np.linspace(low, high, n_values) - elif dist == "uniform": - values = np.sort(np.random.default_rng().uniform(low, high, n_values)) - else: - msg = "Provided parameter 'distribution' should be either 'linspace' or 'uniform'." - raise ValueError(msg) - values[0] = 0 - return values - - -def _create_meas_matrix_from_values_array(values_array: np.ndarray) -> np.ndarray: - """ - For testing purposes. Returns the deltas measurements matrix from an array of values. - - Args: - values_array: the values of phases at your N BPMs, in a (1, N) shaped `numpy.ndarray`. - - Returns: - The matrix, as a `numpy.ndarray`. - """ - return np.array([[i - j for i in values_array] for j in values_array], dtype=float) - - -def _create_2d_gaussian_noise(mean: float, stdev: float, shape: tuple) -> np.ndarray: - """ - Generates a 2D Gaussian distribution, makes it antisymmetric and returns it. - - Args: - mean: mean of the distribution, should be 0 for us (since we add to the ideal M_meas). - stdev: standard deviation of the distribution, should be in degrees if M_meas is given - in degrees, in radians otherwise. - shape: the shape of the matrix to create, should be M_meas.shape, aka (n_bpms, n_bpms). - - Returns: - A Gaussian, anti-symmetric, 2D-shaped `numpy.ndarray`. - """ - gaussian_2d_mat = np.random.default_rng().normal(mean, stdev, size=shape) - upper_triangle = np.triu(gaussian_2d_mat) - return upper_triangle - upper_triangle.T - - @pytest.fixture def _to_scale() -> np.ndarray: return np.load(INPUTS_DIR / "maths" / "to_scale.npy") From 4da3bc96c5d14eb88c212e87b088d9808d0b973a Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Tue, 13 Aug 2024 13:57:55 +0200 Subject: [PATCH 39/41] bump version for minor release --- pyhdtoolkit/version.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyhdtoolkit/version.py b/pyhdtoolkit/version.py index bb83f5bf..db25c83b 100644 --- a/pyhdtoolkit/version.py +++ b/pyhdtoolkit/version.py @@ -1,4 +1,4 @@ -VERSION = "1.5.0" +VERSION = "1.6.0" def version_info() -> str: From 81dfb52274d243b8728f8c8df65a5de1ee7f0d4a Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Tue, 13 Aug 2024 14:40:52 +0200 Subject: [PATCH 40/41] release 1.6.0 file --- docs/releases/v1.6.0.rst | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) create mode 100644 docs/releases/v1.6.0.rst diff --git a/docs/releases/v1.6.0.rst b/docs/releases/v1.6.0.rst new file mode 100644 index 00000000..0db6a085 --- /dev/null +++ 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 `_ and the `full changes since v1.5.0 `_. From 5b35866f3bc89c0bf49d69c8e789cffc7010c89e Mon Sep 17 00:00:00 2001 From: Felix Soubelet Date: Tue, 13 Aug 2024 14:41:12 +0200 Subject: [PATCH 41/41] added to this list --- docs/release.rst | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/docs/release.rst b/docs/release.rst index a70629f2..efac08ab 100644 --- 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 `_. +Version 1.6.0 +------------- + +.. toctree:: + :maxdepth: 2 + + releases/v1.6.0 + Version 1.5.0 -------------