Enhance plot_var method to support direct data input and improve axis…

… labeling
msbc · Dec 1, 2024 · 030c1fb · 030c1fb
1 parent f15289e
commit 030c1fb
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 9 deletions.
diff --git a/helmeos/helm_table.py b/helmeos/helm_table.py
@@ -791,9 +791,10 @@ def full_table(self, abar=1.0, zbar=1.0, overwite=False):
         return self._full_table
 
     def plot_var(self, var, log=False, abar=1.0, zbar=1.0, fig=None, ax=None, cb=True,
-                 aspect=None, vmin=None, vmax=None, fig_opt=None, dpi=None, figsize=None,
-                 popt=None, cmap=None, cbl=None):
+                 aspect='auto', vmin=None, vmax=None, fig_opt=None, dpi=None, figsize=None,
+                 popt=None, cmap=None, cbl=None, interpolation=None):
         try:
+            import matplotlib as mpl
             import matplotlib.pyplot as plt
             from mpl_toolkits.axes_grid1 import make_axes_locatable
             from matplotlib.colors import LogNorm
@@ -803,7 +804,12 @@ def plot_var(self, var, log=False, abar=1.0, zbar=1.0, fig=None, ax=None, cb=Tru
             msg += '`pip install helmeos[plotting]`'
             raise ImportError(msg)
 
-        ft = self.full_table(abar, zbar)
+        if isinstance(var, str):
+            ft = self.full_table(abar, zbar)
+            data = ft[var]
+            cbl = cbl or var
+        else:
+            data = var
 
         if popt is None:
             popt = dict()
@@ -816,7 +822,8 @@ def plot_var(self, var, log=False, abar=1.0, zbar=1.0, fig=None, ax=None, cb=Tru
                   self.temp_log_min, self.temp_log_max]
 
         _popt = {'norm': norm, 'extent': extent, 'aspect': aspect, 'vmin': vmin,
-                 'vmax': vmax, 'cmap': cmap}
+                 'vmax': vmax, 'cmap': cmap, 'origin': 'lower',
+                 'interpolation': interpolation}
         _popt.update(popt)
 
         if fig is None and ax is None:
@@ -828,15 +835,22 @@ def plot_var(self, var, log=False, abar=1.0, zbar=1.0, fig=None, ax=None, cb=Tru
         if ax is not None:
             plt.sca(ax)
 
-        im = plt.imshow(ft[var].T, **_popt)
+        im = plt.imshow(data.T, **_popt)
         ax = plt.gca()
 
+        ax.set_xlabel('Log$_{10}$ Density (g/cm$^3$)')
+        ax.set_ylabel('Log$_{10}$ Temperature (K)')
+
+        ax.xaxis.set_ticks_position('both')
+        ax.yaxis.set_ticks_position('both')
+        ax.xaxis.set_minor_locator(mpl.ticker.AutoMinorLocator())
+        ax.yaxis.set_minor_locator(mpl.ticker.AutoMinorLocator())
+        ax.tick_params(axis='both', which='both', direction='in')
+
         if cb:
             divider = make_axes_locatable(ax)
-            cax = divider.append_axes("right", size="5%", pad=0.05)
+            cax = divider.append_axes("right", size=0.2, pad=0.05)
             cb = plt.colorbar(im, cax=cax)
-            if cbl is None:
-                cbl = var
             if cbl:
                 cb.set_label(cbl)
 

diff --git a/tests/test_plotting.py b/tests/test_plotting.py
@@ -23,7 +23,8 @@ def test_plotting(monkeypatch):
     import matplotlib.pyplot as plt
     default.plot_var('cs')
     plt.close()
-    default.plot_var('cs', log=True)
+    cs = default.full_table()['cs']
+    default.plot_var(cs, log=True)
     plt.close()
     ax = plt.subplot()
     default.plot_var('cs', ax=ax)