Merge pull request #5 from wiseodd/context-manager

Context manager
wiseodd · Aug 15, 2024 · 3fa8164 · 3fa8164
2 parents 3b8d7d8 + ac19f49
commit 3fa8164
Show file tree

Hide file tree

Showing 10 changed files with 743 additions and 255 deletions.
diff --git a/README.md b/README.md
@@ -40,29 +40,25 @@ Still with me? Still want to use this library? Here's how:
 pip install pub-ready-plots
 ```
 
-## Usage
+## Quick Usage
 
 ```python
 import pub_ready_plots
 
-rc_params, fig_width_in, fig_height_in = pub_ready_plots.get_mpl_rcParams(
-    width=1,  # between 0 and 1
-    height=0.1,  # between 0 and 1
-    layout="icml",  # or "iclr", "neurips", "poster-portrait", "poster-landscape"
-    single_col=False # only works for the "icml" layout
-)
-plt.rcParams.update(rc_params)
-
-fig, axs = plt.subplots(
-    nrows,
-    ncols,
-    constrained_layout=True, # Important!
-)
-fig.set_size_inches(fig_width_in, fig_height_in)
-
-# Your plot here!
-
-plt.savefig("filename.pdf")
+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="iclr",  # or "iclr", "neurips", "poster-portrait", "poster-landscape"
+    single_col=False,  # only works for the "icml" layout
+    sharey=True,  # Additional keyword args for `plt.subplots`
+) as (fig, axs):
+    # Do whatever you want with `fig` and `axs`
+    ...
+
+    # Once your done, save it, but do NOT set `tight_layout=True`!
+    fig.savefig("filename.pdf")
 ```
 
 Then in your LaTeX file, include the plot as follows:
@@ -76,13 +72,15 @@ Then in your LaTeX file, include the plot as follows:
 
 > [!TIP]
 > That's it! But you should use TikZ more.
+> Anyway, see the full, runnable example in `examples/simple_plot.py`
 
 ## Using your own styles
 
 Two options:
 
 1. Fork this repo and modify things as you wish.
 2. Use this library and update the resulting `rc_params` dict with your styles.
+   - See `examples/advanced_usage.py`
 
 ## Other libraries
 

diff --git a/examples/advanced_usage.py b/examples/advanced_usage.py
@@ -0,0 +1,35 @@
+import pub_ready_plots
+
+import numpy as np
+import matplotlib.pyplot as plt
+
+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
+)
+
+# You can update `rc_params` further before feeding it to `plt`, e.g.
+rc_params.update({"axes.linewidth": 1})
+
+# Use the styles globally.
+# To make it local, use `with plt.rc_context(rc_params):`
+plt.rcParams.update(rc_params)
+
+fig, axs = plt.subplots(1, 2, constrained_layout=True)
+fig.set_size_inches(fig_width_in, fig_height_in)
+
+x = np.linspace(-1, 1, 100)
+
+axs[0].plot(x, np.sin(x))
+axs[0].set_title("Sine")
+axs[0].set_xlabel(r"$x$")
+axs[0].set_ylabel(r"$\mathrm{sin}(x)$")
+
+axs[1].plot(x, np.cos(x))
+axs[1].set_title("Cosine")
+axs[1].set_xlabel(r"$x$")
+axs[1].set_ylabel(r"$\mathrm{cos}(x)$")
+
+fig.savefig("advanced_usage.pdf")
diff --git a/examples/simple_plot.py b/examples/simple_plot.py
@@ -0,0 +1,25 @@
+import pub_ready_plots
+import numpy as np
+
+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="iclr",  # or "iclr", "neurips", "poster-portrait", "poster-landscape"
+    single_col=False,  # only works for the "icml" layout
+    sharey=True,  # Additional keyword args for `plt.subplots`
+) as (fig, axs):
+    x = np.linspace(-1, 1, 100)
+
+    axs[0].plot(x, np.sin(x))
+    axs[0].set_title("Sine")
+    axs[0].set_xlabel(r"$x$")
+    axs[0].set_ylabel(r"$\mathrm{sin}(x)$")
+
+    axs[1].plot(x, np.cos(x))
+    axs[1].set_title("Cosine")
+    axs[1].set_xlabel(r"$x$")
+    axs[1].set_ylabel(r"$\mathrm{cos}(x)$")
+
+    fig.savefig("simple_plot.pdf")
diff --git a/pdm.lock b/pdm.lock
diff --git a/pub_ready_plots/__init__.py b/pub_ready_plots/__init__.py
@@ -0,0 +1,3 @@
+from pub_ready_plots.pub_ready_plots import get_mpl_rcParams, get_context
+
+__all__ = ["get_mpl_rcParams", "get_context"]
diff --git a/src/pub_ready_plots/pub_ready_plots.py → pub_ready_plots/pub_ready_plots.py b/src/pub_ready_plots/pub_ready_plots.py → pub_ready_plots/pub_ready_plots.py
@@ -1,81 +1,32 @@
-from typing import Any
-
-import matplotlib as mpl
-import matplotlib.font_manager as font_manager
-
-cmfont = font_manager.FontProperties(fname=mpl.get_data_path() + "/fonts/ttf/cmr10.ttf")
-FONT_NAME_CM = cmfont.get_name()
-FONT_NAME_TNR = "Times New Roman"
-FONT_NAME_AVENIR = "Avenir Next Condensed"
-
-PAPER_FORMATS = {
-    "icml": {
-        "text_width": 6.00117,
-        "col_width": 3.25063,
-        "text_height": 8.50166,
-        "font_name": FONT_NAME_TNR,
-        "footnote_size": 8,
-        "script_size": 7,
-        "linewidth": 1.25,
-        "tick_size": 1,
-        "tick_width": 1,
-    },
-    "neurips": {
-        "text_width": 5.50107,
-        "col_width": 5.50107,
-        "text_height": 9.00177,
-        "font_name": FONT_NAME_TNR,
-        "footnote_size": 8,
-        "script_size": 7,
-        "linewidth": 1.25,
-        "tick_size": 1,
-        "tick_width": 1,
-    },
-    "iclr": {
-        "text_width": 5.50107,
-        "col_width": 5.50107,
-        "text_height": 9.00177,
-        "font_name": FONT_NAME_TNR,
-        "footnote_size": 8,
-        "script_size": 7,
-        "linewidth": 1.25,
-        "tick_size": 1,
-        "tick_width": 1,
-    },
-    "jmlr": {
-        "text_width": 6.00117,
-        "col_width": 6.00117,
-        "text_height": 8.50166,
-        "font_name": FONT_NAME_CM,
-        "footnote_size": 8,
-        "script_size": 7,
-        "linewidth": 1.25,
-        "tick_size": 1,
-        "tick_width": 1,
-    },
-    "poster-landscape": {
-        "text_width": 6.00117,
-        "col_width": 6.00117,
-        "text_height": 8.50166,
-        "font_name": FONT_NAME_AVENIR,
-        "footnote_size": 30,
-        "script_size": 23,
-        "linewidth": 3,
-        "tick_size": 4,
-        "tick_width": 2,
-    },
-    "poster-portrait": {
-        "text_width": 6.00117,
-        "col_width": 6.00117,
-        "text_height": 8.50166,
-        "font_name": FONT_NAME_AVENIR,
-        "footnote_size": 10,
-        "script_size": 8,
-        "linewidth": 1,
-        "tick_size": 1,
-        "tick_width": 1,
-    },
-}
+from contextlib import contextmanager
+from matplotlib.axes import Axes
+from matplotlib.figure import Figure
+from numpy import ndarray
+from .styles import PAPER_FORMATS
+
+import matplotlib.pyplot as plt
+
+from typing import Any, Generator, Union
+
+
+@contextmanager
+def get_context(
+    width_frac: float,
+    height_frac: float,
+    nrows: int,
+    ncols: int,
+    layout: str = "neurips",
+    single_col: bool = False,
+    **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
+    )
+
+    with plt.rc_context(rc_params):
+        fig, axs = plt.subplots(nrows, ncols, constrained_layout=True, **kwargs)
+        fig.set_size_inches(fig_width_in, fig_height_in)
+        yield (fig, axs)
 
 
 def get_mpl_rcParams(
@@ -91,7 +42,7 @@ def get_mpl_rcParams(
 
     ```python
         rc_params, fig_width_in, fig_height_in = pub_ready_plots.get_mpl_rcParams(
-            width=fig_width_frac, height=fig_height_frac, layout="icml"
+            width_frac=fig_width_frac, height_frac=fig_height_frac, layout="icml"
         )
         plt.rcParams.update(rc_params)
 
@@ -149,7 +100,7 @@ def get_mpl_rcParams(
         "font.serif": format["font_name"],
         "mathtext.fontset": "stixsans" if is_poster else "cm",
         "lines.linewidth": format["linewidth"],
-        "axes.linewidth": 1,
+        "axes.linewidth": 0.5,
         "axes.titlesize": format["footnote_size"],
         "axes.labelsize": format["script_size"],
         "axes.unicode_minus": False,

diff --git a/pub_ready_plots/styles.py b/pub_ready_plots/styles.py
@@ -0,0 +1,76 @@
+import matplotlib as mpl
+import matplotlib.font_manager as font_manager
+
+cmfont = font_manager.FontProperties(fname=mpl.get_data_path() + "/fonts/ttf/cmr10.ttf")
+FONT_NAME_CM = cmfont.get_name()
+FONT_NAME_TNR = "Times New Roman"
+FONT_NAME_AVENIR = "Avenir Next Condensed"
+
+PAPER_FORMATS = {
+    "icml": {
+        "text_width": 6.00117,
+        "col_width": 3.25063,
+        "text_height": 8.50166,
+        "font_name": FONT_NAME_TNR,
+        "footnote_size": 8,
+        "script_size": 7,
+        "linewidth": 1.25,
+        "tick_size": 1.5,
+        "tick_width": 0.5,
+    },
+    "neurips": {
+        "text_width": 5.50107,
+        "col_width": 5.50107,
+        "text_height": 9.00177,
+        "font_name": FONT_NAME_TNR,
+        "footnote_size": 8,
+        "script_size": 7,
+        "linewidth": 1.25,
+        "tick_size": 1.5,
+        "tick_width": 0.5,
+    },
+    "iclr": {
+        "text_width": 5.50107,
+        "col_width": 5.50107,
+        "text_height": 9.00177,
+        "font_name": FONT_NAME_TNR,
+        "footnote_size": 8,
+        "script_size": 7,
+        "linewidth": 1.25,
+        "tick_size": 1.5,
+        "tick_width": 0.5,
+    },
+    "jmlr": {
+        "text_width": 6.00117,
+        "col_width": 6.00117,
+        "text_height": 8.50166,
+        "font_name": FONT_NAME_CM,
+        "footnote_size": 8,
+        "script_size": 7,
+        "linewidth": 1.25,
+        "tick_size": 1.5,
+        "tick_width": 0.5,
+    },
+    "poster-landscape": {
+        "text_width": 6.00117,
+        "col_width": 6.00117,
+        "text_height": 8.50166,
+        "font_name": FONT_NAME_AVENIR,
+        "footnote_size": 30,
+        "script_size": 23,
+        "linewidth": 3,
+        "tick_size": 4,
+        "tick_width": 2,
+    },
+    "poster-portrait": {
+        "text_width": 6.00117,
+        "col_width": 6.00117,
+        "text_height": 8.50166,
+        "font_name": FONT_NAME_AVENIR,
+        "footnote_size": 10,
+        "script_size": 8,
+        "linewidth": 1,
+        "tick_size": 1.5,
+        "tick_width": 0.5,
+    },
+}
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,21 +1,29 @@
 [project]
 name = "pub-ready-plots"
-version = "0.1.0"
+version = "0.2.0"
 description = "Easy publication-ready matplotlib plots for ML papers and posters."
 authors = [{ name = "Agustinus Kristiadi", email = "[email protected]" }]
+classifiers = [
+  "Development Status :: 3 - Alpha",
+  "License :: OSI Approved :: MIT License",
+  "Operating System :: OS Independent",
+  "Programming Language :: Python :: 3.9",
+]
+requires-python = ">=3.9"
 dependencies = ["matplotlib>=3.9.0"]
-requires-python = ">=3.10"
 readme = "README.md"
 license = { text = "MIT" }
 
 [build-system]
 requires = ["pdm-backend"]
 build-backend = "pdm.backend"
 
-
 [tool.pdm]
 distribution = true
 
+[tool.pdm.dev-dependencies]
+test = ["pytest>=8.2.2", "pytest-cov>=5.0.0"]
+lint = ["ruff>=0.5.0"]
+
 [tool.pytest.ini_options]
-pythonpath = ["src"]
-testpaths = ["tests"]
+testpaths = "tests"
diff --git a/src/pub_ready_plots/__init__.py b/src/pub_ready_plots/__init__.py
diff --git a/tests/test_context.py b/tests/test_context.py
@@ -0,0 +1,16 @@
+from pub_ready_plots.pub_ready_plots import get_context, get_mpl_rcParams
+import numpy as np
+
+
+def test_correct_func():
+    nrows, ncols = 3, 2
+    with get_context(0.5, 0.15, nrows, ncols, "iclr") as (fig, axs):
+        real_rc_params, fig_width_in, fig_height_in = get_mpl_rcParams(
+            0.5, 0.15, "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)