tweaked mpl defaults andd removed numba (#41)

* update plot params, give kwargs to logger config

* force inwards ticks everywhere, and absence of grid in layout plot

* tweaks and function to create style file

* remove use of numba

* coverage for install_mpl_style

.github/workflows/tests.yml

@@ -64,7 +64,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        os: [ubuntu-16.04, ubuntu-18.04, ubuntu-20.04, macos-latest, windows-latest]
+        os: [ubuntu-18.04, ubuntu-20.04, macos-latest, windows-latest]
         python-version: [3.7, 3.8, 3.9]
       fail-fast: false
 

pyhdtoolkit/__init__.py

@@ -13,7 +13,7 @@
 __title__ = "pyhdtoolkit"
 __description__ = "An all-in-one toolkit package to easy my Python work in my PhD."
 __url__ = "https://github.com/fsoubelet/PyhDToolkit"
-__version__ = "0.9.0"
+__version__ = "0.9.1"
 __author__ = "Felix Soubelet"
 __author_email__ = "felix.soubelet@cern.ch"
 __license__ = "MIT"

pyhdtoolkit/cpymadtools/latwiss.py

@@ -21,6 +21,7 @@
 from pyhdtoolkit.utils.defaults import PLOT_PARAMS
 
 plt.rcParams.update(PLOT_PARAMS)
+plt.rcParams.update({"xtick.direction": "in", "ytick.direction": "in"})  # need to reiterate these somehow
 
 # ----- Plotters ----- #
 
@@ -341,6 +342,7 @@ def _plot_machine_layout(
     quadrupole_patches_axis.set_xlim(xlimits)
     quadrupole_patches_axis.set_title(title)
     quadrupole_patches_axis.plot(twiss_df.s, 0 * twiss_df.s, "k")  # 0-level line
+    quadrupole_patches_axis.grid(False)
 
     dipole_patches_axis = quadrupole_patches_axis.twinx()
     dipole_patches_axis.set_ylabel("$\\theta=K_{0}L$ [rad]", color="royalblue")  # dipoles in blue
@@ -419,6 +421,7 @@ def _plot_machine_layout(
             )
             plotted_elements += 1
         sextupoles_patches_axis.legend(loc=3, fontsize=16)
+        sextupoles_patches_axis.grid(False)
 
     if plot_bpms:
         logger.debug("Plotting BPM patches")
@@ -439,6 +442,7 @@ def _plot_machine_layout(
             )
             plotted_elements += 1
         bpm_patches_axis.legend(loc=4, fontsize=16)
+        bpm_patches_axis.grid(False)
 
 
 # ----- Helpers ----- #

pyhdtoolkit/optics/ripken.py

@@ -1,6 +1,5 @@
 from typing import Union
 
-import numba
 import numpy as np
 
 from loguru import logger
@@ -33,7 +32,6 @@ def lebedev_beam_size(
 # ----- JITed Calculations ----- #
 
 
-@numba.njit()
 def _beam_size(coordinates_distribution: np.ndarray, method: str = "std") -> float:
     """
     Compute beam size from particle coordinates.

pyhdtoolkit/optics/twiss.py

@@ -8,11 +8,9 @@
 This is a Python3 module implementing various functionality for optics calculations from / to
 twiss parameters.
 """
-import numba
 import numpy as np
 
 
-@numba.njit()
 def courant_snyder_transform(u_vector: np.ndarray, alpha: float, beta: float) -> np.ndarray:
     """
     Perform the Courant-Snyder transform on rergular (nonchaotic) phase-space coordinatess.

pyhdtoolkit/utils/defaults.py

@@ -10,7 +10,9 @@
 import sys
 
 from pathlib import Path
-from typing import Dict, Union
+from typing import Dict, NewType, Union
+
+import matplotlib
 
 from loguru import logger
 
@@ -30,17 +32,33 @@
     "<level>{message}</level>"
 )
 
+PlotSetting = NewType("PlotSetting", Union[float, bool, str, tuple])
+
 # Set those with matplotlib.pyplot.rcParams.update(PLOT_PARAMS).
 # Will ALWAYS be overwritten by later on definition
-PLOT_PARAMS: Dict[str, Union[float, bool, str, tuple]] = {
+PLOT_PARAMS: Dict[str, PlotSetting] = {
+    # ------ Patches ------ #
+    "patch.linewidth": 3,  # Width of patches edge lines
+    # ------ Fonts ------ #
+    "font.family": "sans-serif",  # Font family
+    "font.style": "normal",  # Style to apply to text font
+    "font.weight": "bold",  # Bold font
+    "font.size": 25,  # Default font size of elements
+    "font.sans-serif": "Helvetica",  # Sans-Serif font to use
+    # ----- Mathtext ----- #
+    "mathtext.default": "bf",  # default font for math
+    # ------ Text ------ #
+    "text.usetex": True,  # Use LaTeX for text handling (Set to False if you don't have a local installation)
+    "text.latex.preamble": r"\usepackage{amsmath}",  # \boldmath",  # Be careful with the preamble
     # ------ Axes ------ #
-    "axes.linewidth": 0.8,  # Linewidth of axes edges
-    "axes.grid": False,  # Do not display grid
-    "axes.labelsize": 30,  # Fontsize of the x and y axis labels
+    "axes.linewidth": 2,  # Linewidth of axes edges
+    "axes.grid": True,  # Do display grid
     "axes.titlesize": 30,  # Fontsize of the axes title
+    "axes.labelsize": 30,  # Fontsize of the x and y axis labels
+    "axes.labelweight": "bold",  # Bold labels
     "axes.formatter.limits": (-4, 5),  # Switch to scientific notations when order of magnitude reaches 1e3
-    # "axes.formatter.useoffset": False,  # Do not use the annoying offset on top of yticks
     "axes.formatter.use_mathtext": True,  # Format with i.e 10^{4} instead of 1e4
+    "axes.formatter.useoffset": False,  # Do not use the annoying offset on top of yticks
     # ------ Date Formats ------ #
     "date.autoformatter.year": "%Y",  # AutoDateFormatter setting for years display
     "date.autoformatter.month": "%Y-%m",  # AutoDateFormatter setting for months display
@@ -49,47 +67,71 @@
     "date.autoformatter.minute": "%d %H:%M",  # AutoDateFormatter setting for minutes display
     "date.autoformatter.second": "%H:%M:%S",  # AutoDateFormatter setting for seconds display
     "date.autoformatter.microsecond": "%M:%S.%f",  # AutoDateFormatter setting for microseconds
-    # ------ General Figure ------ #
-    "figure.autolayout": True,  # Adjust subplot params to fit the figure (tight_layout)
-    "figure.dpi": 300,  # Figure dots per inch
-    "figure.figsize": (18, 11),  # Size of the figure
-    "figure.max_open_warning": 10,  # Max number of figures to open before warning
-    "figure.titlesize": 30,  # Size of the figure title
-    # ------ Fonts ------ #
-    "font.family": "sans-serif",  # Font family
-    # "font.sans-serif": "Helvetica",  # Sans-Serif font to use
-    "font.style": "normal",  # Style to apply to text font
+    # ------ Horizontal Ticks ------ #
+    "xtick.major.size": 8,  # Size (length) of the major xtick locators
+    "xtick.minor.size": 5,  # Size (length) of the minor xtick locators
+    "xtick.major.width": 1.5,  # Width of the major xtick locators
+    "xtick.minor.width": 0.6,  # Width of the minor xtick locators
+    "xtick.labelsize": 25,  # Fontsize of the x axis tick labels
+    "xtick.direction": "in",  # Show xticks towards inside of figure
+    "xtick.minor.visible": True,  # Show minor xtick locators
+    # ------ Vertical Ticks ------ #
+    "ytick.major.size": 8,  # Size (length) of the major ytick locators
+    "ytick.minor.size": 5,  # Size (length) of the minor ytick locators
+    "ytick.major.width": 1.5,  # Width of the major ytick locators
+    "ytick.minor.width": 0.6,  # Width of the minor ytick locators
+    "ytick.labelsize": 25,  # Fontsize of the y axis tick labels
+    "ytick.direction": "in",  # Show yticks towards inside of figure
+    "ytick.minor.visible": True,  # Show minor ytick locators
+    # ----- Grid ----- #
+    "grid.linestyle": "--",  # Which linestyle for grid lines
+    "grid.linewidth": 1.3,  # Width of the grid lines
     # ------- Legend ------ #
+    "legend.loc": "best",  # Default legend location
+    "legend.frameon": True,  # Make a dedicated patch for the legend
     "legend.fancybox": True,  # Use rounded box for legend background
     "legend.fontsize": 22,  # Legend text font size
-    "legend.loc": "best",  # Default legend location
-    # ------ Lines ------ #
-    "lines.linewidth": 1,  # Line width, in points
-    "lines.markersize": 5,  # Marker size, in points
-    "lines.antialiased": True,  # Apply anti-aliasing to lines display
-    # ------ Patches ------ #
-    "patch.linewidth": 1,  # Width of patches edge lines
-    "patch.antialiased": True,  # Apply anti-aliasing to patches display
-    # ------ Paths ------ #
-    "path.simplify": True,  # Reduce file size by removing "invisible" points
+    "legend.title_fontsize": 23,  # Legend title text font size
+    # ------ Figure ------ #
+    "figure.titlesize": 35,  # Size of the figure title
+    "figure.figsize": (16, 10),  # Default size of the figures
+    "figure.dpi": 300,  # Figure dots per inch
+    "figure.subplot.left": 0.15,  # Left side of the subplots of the figure
+    "figure.subplot.right": 0.90,  # Right side of the subplots of the figure
+    "figure.subplot.bottom": 0.15,  # Bottom side of the subplots of the figure
+    "figure.subplot.top": 0.90,  # Top side of the subplots of the figure
+    "figure.autolayout": True,  # Adjust subplot params to fit the figure (tight_layout)
     # ------ Saving ------ #
     "savefig.dpi": 1000,  # Saved figure dots per inch
     "savefig.format": "pdf",  # Saved figure file format
     "savefig.bbox": "tight",  # Careful: incompatible with pipe-based animation backends
-    # ------ Text ------ #
-    "text.antialiased": True,  # Apply anti-aliasing to text elements
-    "text.color": "black",  # Default text color
-    "text.usetex": False,  # Do not use LaTeX for text handling (I don't have a local installation)
-    # ------ Ticks ------ #
-    "xtick.labelsize": 25,  # Fontsize of the x axis tick labels
-    "ytick.labelsize": 25,  # Fontsize of the y axis tick labels
 }
 
 
-def config_logger(level: str = "INFO") -> None:
+def config_logger(level: str = "INFO", **kwargs) -> None:
     """
     Resets the logger object from loguru, with `sys.stdout` as a sink and the aforedefined format.
     This comes down to personnal preference.
+    Any additional keyword argument used is transmitted to the `logger.add` call.
     """
     logger.remove()
-    logger.add(sys.stdout, format=LOGURU_FORMAT, level=level.upper())
+    logger.add(sys.stdout, format=LOGURU_FORMAT, level=level.upper(), **kwargs)
+
+
+def install_mpl_style() -> None:
+    """
+    Will create a `phd.mplstyle` file in the appropriate directory from the `PLOT_PARAMS` defined in this
+    module. This enables one to use the style without importing `PLOT_PARAMS` and updating the rcParams,
+    but instead simply using `plt.style.use("phd")`.
+    """
+    logger.info("Installing matplotlib style")
+    style_content = "\n".join(f"{option} : {setting}" for option, setting in PLOT_PARAMS.items())
+    mpl_stylelib = Path(matplotlib.get_configdir()) / "stylelib"
+
+    logger.debug("Ensuring matplotlib 'stylelib' directory exists")
+    mpl_stylelib.mkdir(parents=True, exist_ok=True)
+    style_file = mpl_stylelib / "phd.mplstyle"
+
+    logger.debug(f"Creating style file at '{style_file.absolute()}'")
+    style_file.write_text(style_content.replace("(", "").replace(")", ""))
+    logger.success("You can now use it with 'plt.style.use(\"phd\")'")

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pyhdtoolkit"
-version = "0.9.0"
+version = "0.9.1"
 description = "An all-in-one toolkit package to easy my Python work in my PhD."
 authors = ["Felix Soubelet <felix.soubelet@cern.ch>"]
 license = "MIT"
@@ -31,15 +31,13 @@ classifiers = [
 
 # Core dependencies of the package
 [tool.poetry.dependencies]
-python = "^3.7, <3.10"  # specified to help with llvm dependency
+python = "^3.7"
 numpy = "^1.19"
 pandas = "^1.0"
 matplotlib = "^3.0"
 scipy = "^1.4"
 tfs-pandas = "^2.0"
 loguru = "<1.0"
-numba = "<1.0"
-llvmlite = "^0.36"  # dependency of numba, version specified here to guarantee Python 3.9 compatibility
 cpymad = "^1.6"
 rich = "^10.0"
 pydantic = "^1.0"

tests/test_optics.py

@@ -60,17 +60,15 @@ def test_gamma_transition_raises(self):
 
 class TestRipken:
     def test_beam_size(self, _fake_coordinates):
-        # Test uses the dispatcher's 'py_func' object for coverage integrity, packaged implementation
-        # is still JIT compiled, and a JIT failure will fail the test.
-        assert np.allclose(ripken._beam_size.py_func(_fake_coordinates), _fake_coordinates.std())
+        assert np.allclose(ripken._beam_size(_fake_coordinates), _fake_coordinates.std())
         assert np.allclose(
-            ripken._beam_size.py_func(_fake_coordinates, method="rms"),
+            ripken._beam_size(_fake_coordinates, method="rms"),
             np.sqrt(np.mean(np.square(_fake_coordinates))),
         )
 
     def test_beam_size_raises(self, _fake_coordinates):
         with pytest.raises(NotImplementedError):
-            _ = ripken._beam_size.py_func(_fake_coordinates, method="not_real")
+            _ = ripken._beam_size(_fake_coordinates, method="not_real")
 
     @pytest.mark.parametrize("beta11", [0.3312])
     @pytest.mark.parametrize("beta21", [1])
@@ -82,23 +80,12 @@ def test_lebedev_size_floats(self, beta11, beta21, emit_x, emit_y):
         ) == np.sqrt(emit_x * beta11 + emit_y * beta21)
 
 
-# TODO: fix this one
-#     def test_levedev_size_arrays(self):
-# beta_11_array = np.load(INPUT_PATHS["beta11"])
-# beta_21_array = np.load(INPUT_PATHS["beta21"])
-# lebedev_sizes = np.load(INPUT_PATHS["lebedev"])
-# calculated = ripken.lebedev_beam_size(beta_11_array, beta_21_array, 3.75e-6, 3.75e-6)
-# assert np.allclose(calculated, lebedev_sizes)
-
-
 class TestTwiss:
     def test_courant_snyder_transform(self):
-        # Test uses the dispatcher's 'py_func' object for coverage integrity, packaged implementation
-        # is still JIT compiled, and a JIT failure will fail the test.
         alpha_beta = np.load(INPUT_PATHS["alpha_beta"])
         u_vector = np.load(INPUT_PATHS["u_vector"])
         u_bar_result = np.load(INPUT_PATHS["u_bar"])
-        u_transform = twiss.courant_snyder_transform.py_func(u_vector, alpha_beta[0], alpha_beta[1])
+        u_transform = twiss.courant_snyder_transform(u_vector, alpha_beta[0], alpha_beta[1])
         np.testing.assert_array_almost_equal(u_transform, u_bar_result)
 
 

tests/test_utils.py

@@ -10,7 +10,7 @@
 from typing import List
 
 import pytest
-
+import matplotlib
 from loguru import logger
 from rich.table import Table
 
@@ -94,6 +94,10 @@ def test_logger_config(self, capsys):
         logger.remove()
         logger.add(sys.stderr)
 
+    def test_mplstyle_install(self, capsys):
+        defaults.install_mpl_style()
+        assert (pathlib.Path(matplotlib.get_configdir()) / "stylelib" / "phd.mplstyle").is_file()
+
 
 class TestHTCMonitor:
     def test_read_condor_q(self, _condor_q_output, _correct_user_tasks, _correct_cluster_summary):