more plotting utility code

python/energy_flow_proper/utilities.py

@@ -3,8 +3,12 @@ from types import ModuleType
 from typing import Callable, Tuple, Union
 from lmfit import minimize, Parameters
 import matplotlib.pyplot as plt
+import matplotlib.ticker as ticker
+import matplotlib
 import numpy as np
 from numpy.polynomial import Polynomial
+import functools
+from contextlib import contextmanager
 
 
 def get_n_samples(stg: ModuleType) -> int:
@@ -73,11 +77,16 @@ def fit_α(
     return w, g
 
 
+###############################################################################
+#                                  Plot Porn                                  #
+###############################################################################
+
+
 def wrap_plot(f):
-    def wrapped(*args, ax=None, **kwargs):
+    def wrapped(*args, ax=None, setup_function=plt.subplots, **kwargs):
         fig = None
         if not ax:
-            fig, ax = plt.subplots()
+            fig, ax = setup_function()
 
         ret_val = f(*args, ax=ax, **kwargs)
         return (fig, ax, ret_val) if ret_val else (fig, ax)
@@ -85,6 +94,20 @@ def wrap_plot(f):
     return wrapped
 
 
+@contextmanager
+def hiro_style(*args, **kwargs):
+    with plt.style.context("ggplot"):
+        with matplotlib.rc_context(
+            {
+                # "font.family": "serif",
+                "text.usetex": False,
+                "pgf.rcfonts": False,
+                "lines.linewidth": 1,
+            }
+        ):
+            yield True
+
+
 @wrap_plot
 def plot_complex(x, y, *args, ax=None, label="", **kwargs):
     label = label + ", " if (len(label) > 0) else ""
@@ -93,6 +116,11 @@ def plot_complex(x, y, *args, ax=None, label="", **kwargs):
     ax.legend()
 
 
+###############################################################################
+#                                 Numpy Hacks                                 #
+###############################################################################
+
+
 def e_i(i: int, size: int) -> np.ndarray:
     r"""Cartesian base vector :math:`e_i`."""