Merge pull request #8 from wiseodd/override-rcparams

Add defaults to `nrows` and `ncols`; add a way to override `rcParams`

Author: wiseodd

README.md

@@ -52,11 +52,12 @@ import pub_ready_plots
 with pub_ready_plots.get_context(
     width_frac=1,  # between 0 and 1
     height_frac=0.15,  # between 0 and 1
-    nrows=1,  # depending on your subplots
-    ncols=2,  # depending on your subplots
     layout="icml",  # or "iclr", "neurips", "poster-portrait", "poster-landscape"
     single_col=False,  # only works for the "icml" layout
+    nrows=1,  # depending on your subplots, default = 1
+    ncols=2,  # depending on your subplots, default = 1
     sharey=True,  # Additional keyword args for `plt.subplots`
+    override_rc_params={"lines.linewidth": 4.123}  # Overriding rcParams
 ) as (fig, axs):
     # Do whatever you want with `fig` and `axs`
     ...
@@ -74,9 +75,14 @@ Then in your LaTeX file, include the plot as follows:
 > [!IMPORTANT]
 > The argument `width=\linewidth` is **crucial**!
 
+That's it! But you should use TikZ more.
+Anyway, see the full, runnable example in [`examples/simple_plot.py`](https://github.com/wiseodd/pub-ready-plots/blob/master/examples/simple_plot.py)
+
 > [!TIP]
-> That's it! But you should use TikZ more.
-> Anyway, see the full, runnable example in [`examples/simple_plot.py`](https://github.com/wiseodd/pub-ready-plots/blob/master/examples/simple_plot.py)
+> I recommend using this library in conjunction with
+> [pypalettes]<https://github.com/JosephBARBIERDARNAL/pypalettes>
+> to avoid the generic blue-orange Matplotlib colors.
+> Distinguish your plots from others!
 
 ## Using your own styles
 

examples/advanced_usage.py

@@ -2,12 +2,37 @@
 
 import numpy as np
 import matplotlib.pyplot as plt
+from matplotlib.axes import Axes
 
+########################################################################################
+# User-specified rcParams override
+########################################################################################
+with pub_ready_plots.get_context(
+    width_frac=1,
+    height_frac=0.15,
+    layout="iclr",
+    override_rc_params={"lines.linewidth": 5},  # Pass your style overrides here!
+) as (fig, ax):
+    assert isinstance(ax, Axes)
+
+    x = np.linspace(-1, 1, 100)
+
+    ax.plot(x, np.sin(x))
+    ax.set_title("Sine")
+    ax.set_xlabel(r"$x$")
+    ax.set_ylabel(r"$\mathrm{sin}(x)$")
+
+    fig.savefig("advanced_usage_1.pdf")
+
+
+########################################################################################
+# Manual, most-flexible way to use this library
+########################################################################################
 rc_params, fig_width_in, fig_height_in = pub_ready_plots.get_mpl_rcParams(
-    width_frac=1,  # between 0 and 1
-    height_frac=0.15,  # between 0 and 1
-    layout="poster-portrait",  # or "iclr", "neurips", "poster-portrait", "poster-landscape"
-    single_col=False,  # only works for the "icml" layout
+    width_frac=1,
+    height_frac=0.15,
+    layout="poster-portrait",
+    single_col=False,
 )
 
 # You can update `rc_params` further before feeding it to `plt`, e.g.
@@ -32,4 +57,4 @@
 axs[1].set_xlabel(r"$x$")
 axs[1].set_ylabel(r"$\mathrm{cos}(x)$")
 
-fig.savefig("advanced_usage.pdf")
+fig.savefig("advanced_usage_2.pdf")

examples/simple_plot.py

@@ -1,6 +1,32 @@
 import pub_ready_plots
+
 import numpy as np
+from matplotlib.axes import Axes
 
+########################################################################################
+# Single plot (i.e. no subplots)
+########################################################################################
+with pub_ready_plots.get_context(
+    width_frac=1,  # between 0 and 1
+    height_frac=0.15,  # between 0 and 1
+    layout="iclr",  # or "iclr", "neurips", "poster-portrait", "poster-landscape"
+) as (fig, ax):
+    # Just like in `plt.subplots`, `ax` is a matplotlib Axes if
+    # nrows & ncols are not specified (both default to 1).
+    assert isinstance(ax, Axes)
+
+    x = np.linspace(-1, 1, 100)
+
+    ax.plot(x, np.sin(x))
+    ax.set_title("Sine")
+    ax.set_xlabel(r"$x$")
+    ax.set_ylabel(r"$\mathrm{sin}(x)$")
+
+    fig.savefig("simple_plot.pdf")
+
+########################################################################################
+# Multiple subplots
+########################################################################################
 with pub_ready_plots.get_context(
     width_frac=1,  # between 0 and 1
     height_frac=0.15,  # between 0 and 1
@@ -10,6 +36,9 @@
     single_col=False,  # only works for the "icml" layout
     sharey=True,  # Additional keyword args for `plt.subplots`
 ) as (fig, axs):
+    # If `nrows` or `ncols` are not 1, `axs` is a NumPy array containing Axes'
+    assert isinstance(axs, np.ndarray)
+
     x = np.linspace(-1, 1, 100)
 
     axs[0].plot(x, np.sin(x))
@@ -22,4 +51,4 @@
     axs[1].set_xlabel(r"$x$")
     axs[1].set_ylabel(r"$\mathrm{cos}(x)$")
 
-    fig.savefig("simple_plot.pdf")
+    fig.savefig("simple_subplots.pdf")

pdm.lock

@@ -13,15 +13,17 @@
 def get_context(
     width_frac: float,
     height_frac: float,
-    nrows: int,
-    ncols: int,
     layout: str = "neurips",
     single_col: bool = False,
+    nrows: int = 1,
+    ncols: int = 1,
+    override_rc_params: dict[str, Any] = dict(),
     **kwargs,
 ) -> Generator[tuple[Figure, Union[Axes, ndarray[Any, Any]]], None, None]:
     rc_params, fig_width_in, fig_height_in = get_mpl_rcParams(
         width_frac, height_frac, layout, single_col
     )
+    rc_params.update(override_rc_params)
 
     with plt.rc_context(rc_params):
         fig, axs = plt.subplots(nrows, ncols, constrained_layout=True, **kwargs)

pub_ready_plots/pub_ready_plots.py

@@ -1,6 +1,6 @@
 [project]
 name = "pub-ready-plots"
-version = "0.2.0"
+version = "0.2.1"
 description = "Easy publication-ready matplotlib plots for ML papers and posters."
 authors = [{ name = "Agustinus Kristiadi", email = "agustinus@kristia.de" }]
 classifiers = [

pyproject.toml

@@ -1,16 +1,47 @@
+from matplotlib.axes import Axes
 from pub_ready_plots.pub_ready_plots import get_context, get_mpl_rcParams
 import numpy as np
 
 
-def test_correct_func():
+def test_single_subplot():
+    with get_context(width_frac=0.5, height_frac=0.15, layout="iclr") as (fig, axs):
+        real_rc_params, fig_width_in, fig_height_in = get_mpl_rcParams(
+            width_frac=0.5, height_frac=0.15, layout="iclr"
+        )
+
+        assert np.allclose(
+            fig.get_size_inches(), (fig_width_in, fig_height_in), atol=0.001
+        )
+        assert isinstance(axs, Axes)
+
+
+def test_multi_subplots():
     nrows, ncols = 3, 2
-    with get_context(0.5, 0.15, nrows, ncols, "iclr") as (fig, axs):
+    with get_context(
+        width_frac=0.5, height_frac=0.15, nrows=nrows, ncols=ncols, layout="iclr"
+    ) as (fig, axs):
         real_rc_params, fig_width_in, fig_height_in = get_mpl_rcParams(
-            0.5, 0.15, "iclr"
+            width_frac=0.5, height_frac=0.15, layout="iclr"
         )
 
         assert np.allclose(
             fig.get_size_inches(), (fig_width_in, fig_height_in), atol=0.001
         )
         assert isinstance(axs, np.ndarray)
         assert axs.shape == (nrows, ncols)
+
+
+def test_override_rcparams():
+    LINE_WIDTH: float = 8.2329232
+
+    with get_context(
+        width_frac=0.5,
+        height_frac=0.15,
+        layout="iclr",
+        override_rc_params={"lines.linewidth": LINE_WIDTH},
+    ) as (fig, ax):
+        assert isinstance(ax, Axes)
+
+        x: np.ndarray = np.linspace(-1, 1, 100)
+        obj = ax.plot(x, np.tanh(x))
+        assert np.allclose(obj[0].get_linewidth(), LINE_WIDTH)

tests/test_context.py

@@ -8,32 +8,36 @@
 
 @pytest.mark.parametrize("layout", LAYOUTS)
 def test_correct_func(layout):
-    rc_params, width_in, height_in = get_mpl_rcParams(1, 0.15, layout)
+    rc_params, width_in, height_in = get_mpl_rcParams(
+        width_frac=1, height_frac=0.15, layout=layout
+    )
     plt.rcParams.update(rc_params)
 
 
 def test_incorrect_width():
     with pytest.raises(ValueError):
-        _, _, _ = get_mpl_rcParams(12, 0.1, "iclr")
+        _, _, _ = get_mpl_rcParams(width_frac=12, height_frac=0.1, layout="iclr")
 
     with pytest.raises(ValueError):
-        _, _, _ = get_mpl_rcParams(0, 0.1, "iclr")
+        _, _, _ = get_mpl_rcParams(width_frac=0, height_frac=0.1, layout="iclr")
 
     with pytest.raises(ValueError):
-        _, _, _ = get_mpl_rcParams(-3.2, 0.1, "iclr")
+        _, _, _ = get_mpl_rcParams(width_frac=-3.2, height_frac=0.1, layout="iclr")
 
 
 def test_incorrect_height():
     with pytest.raises(ValueError):
-        _, _, _ = get_mpl_rcParams(0.15, 12, "iclr")
+        _, _, _ = get_mpl_rcParams(width_frac=0.15, height_frac=12, layout="iclr")
 
     with pytest.raises(ValueError):
-        _, _, _ = get_mpl_rcParams(0.15, 0, "iclr")
+        _, _, _ = get_mpl_rcParams(width_frac=0.15, height_frac=0, layout="iclr")
 
     with pytest.raises(ValueError):
-        _, _, _ = get_mpl_rcParams(0.15, -1.2, "iclr")
+        _, _, _ = get_mpl_rcParams(width_frac=0.15, height_frac=-1.2, layout="iclr")
 
 
 def test_incorrect_layout():
     with pytest.raises(ValueError):
-        _, _, _ = get_mpl_rcParams(1, 1, "predatory_journal")
+        _, _, _ = get_mpl_rcParams(
+            width_frac=0.5, height_frac=0.5, layout="predatory_journal"
+        )