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