From 12716ae6c62bd0cab3f6fadd2ebfb3e598b22da7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mika=20M=C3=A4ki?= <agentti002x@gmail.com>
Date: Wed, 4 Dec 2024 21:59:52 +0200
Subject: [PATCH] Fix plotting the relativistic combustion figure together with
 other figures

---
 docs/fig/relativistic_combustion.py | 78 +++++++++++++++--------------
 1 file changed, 41 insertions(+), 37 deletions(-)

diff --git a/docs/fig/relativistic_combustion.py b/docs/fig/relativistic_combustion.py
index 39fdeb5d..ec1513fc 100644
--- a/docs/fig/relativistic_combustion.py
+++ b/docs/fig/relativistic_combustion.py
@@ -5,14 +5,13 @@
 
 import typing as tp
 
-import matplotlib
 import matplotlib.pyplot as plt
-import matplotlib.cm
 import matplotlib.colors
 from matplotlib.image import AxesImage
 import numpy as np
 from scipy.interpolate import interp1d
 
+from examples.utils import save
 from pttools import bubble
 
 _VIRIDIS_BIG = plt.colormaps["autumn_r"]
@@ -21,23 +20,12 @@
 COLORMAP = matplotlib.colors.ListedColormap(_NEW_COLORS)
 
 
-def setup_matplotlib(**kwargs):
-    matplotlib.rcParams.update({
-        "text.usetex": True,
-        "font.family": "serif",
-        "font.size": 28,
-        "legend.fontsize": 14,
-        "lines.linewidth": 1.75,
-        **kwargs
-    })
-
-
 def plot_bubble(ax: plt.Axes, label: str, v_wall: float, alpha: float, n_xi: int) -> AxesImage:
     v_f, enthalp, xi = bubble.sound_shell_bag(v_wall, alpha, n_xi)
     n_wall = bubble.find_v_index(xi, v_wall)
     v_fluid = interp1d(xi, v_f, fill_value=0, bounds_error=False)
 
-    xi_wall = xi[n_wall]
+    xi_wall: float = xi[n_wall]
     xvalues = np.linspace(-1.5 * xi_wall, 1.5 * xi_wall, num=4000)
     yvalues = np.linspace(-1.5 * xi_wall, 1.5 * xi_wall, num=4000)
     xxgrid, yygrid = np.meshgrid(xvalues, yvalues)
@@ -100,29 +88,45 @@ def main(
         v_walls: tp.Tuple[float, ...] = (0.44, 0.72, 0.92),
         plot_cbars: tp.Tuple[bool, ...] = (False, False, True),
         n_xi: int = 5000,
-        figsize: tp.Tuple[int, int] = (27, 9)) -> plt.Figure:
-    setup_matplotlib()
-
-    fig: plt.Figure
-    axs: np.ndarray
-    fig, axs = plt.subplots(1, len(v_walls), figsize=figsize)
-
-    labels = [
-        "subsonic deflagration" + "\n" + r"$v_\mathrm{w} \leq c_s$",
-        "supersonic deflagration" + "\n" + r"$c_s<v_\mathrm{w} < c_\mathrm{J}$",
-        "detonation" + "\n" + r"$c_s<c_\mathrm{J}\leq v_\mathrm{w}$"
-    ]
-
-    for ax, label, v_wall, plot_cbar in zip(axs, labels, v_walls, plot_cbars):
-        cs = plot_bubble(ax, label, v_wall, alpha, n_xi)
-        if plot_cbar:
-            cbar = fig.colorbar(cs, ax=axs)
-            cbar.set_label(r"$v/v_\mathrm{peak}$")
-
-    # fig.savefig("plots/all_circle.pdf", bbox_inches="tight")
-    return fig
+        figsize: tp.Tuple[int, int] = (27, 9),
+        path: str = None,
+        show: bool = False) -> plt.Figure:
+    with plt.rc_context({
+                "text.usetex": True,
+                "font.family": "serif",
+                "font.size": 28,
+                "legend.fontsize": 14,
+                "lines.linewidth": 1.75,
+            }):
+        fig: plt.Figure
+        axs: np.ndarray
+        fig, axs = plt.subplots(1, len(v_walls), figsize=figsize)
+
+        labels = [
+            "subsonic deflagration" + "\n" + r"$v_\mathrm{w} \leq c_s$",
+            "supersonic deflagration" + "\n" + r"$c_s<v_\mathrm{w} < c_\mathrm{J}$",
+            "detonation" + "\n" + r"$c_s<c_\mathrm{J}\leq v_\mathrm{w}$"
+        ]
+
+        for ax, label, v_wall, plot_cbar in zip(axs, labels, v_walls, plot_cbars):
+            cs = plot_bubble(ax, label, v_wall, alpha, n_xi)
+            if plot_cbar:
+                cbar = fig.colorbar(cs, ax=axs)
+                cbar.set_label(r"$v/v_\mathrm{peak}$")
+
+        # tight_box = fig.get_tightbbox(fig.canvas.get_renderer())
+        # fig.set(figheight=tight_box.height, figwidth=tight_box.width)
+
+        # These have to be inside the plt.rc_context environment for its settings to work.
+        if path is not None:
+            save(fig, path, bbox_inches="tight")
+        if show:
+            plt.show()
+        return fig
 
 
 if __name__ == "__main__":
-    main()
-    plt.show()
+    main(
+        show=True,
+        # path="relativistic_combustion.png"
+    )