Add grid feature to plt_field()

patnr · Sep 26, 2023 · 694676f · 694676f
1 parent 85a424a
commit 694676f
Showing 1 changed file with 15 additions and 1 deletion.
diff --git a/TPFA_ResSim/plotting.py b/TPFA_ResSim/plotting.py
@@ -2,6 +2,7 @@
 
 import matplotlib as mpl
 import numpy as np
+from matplotlib.ticker import MultipleLocator
 from mpl_tools import place, place_ax
 from mpl_tools.misc import axprops
 
@@ -50,7 +51,7 @@ class Plot2D:
     """Plots specialized for 2D fields."""
 
     def plt_field(self, ax, Z, style="default", wells=True,
-                  argmax=False, colorbar=True, labels=True, **kwargs):
+                  argmax=False, colorbar=True, labels=True, grid=False, **kwargs):
         """Contour-plot of the (flat) unravelled field `Z`.
 
         `kwargs` falls back to `styles[style]`, which falls back to `styles['defaults']`.
@@ -92,6 +93,19 @@ def plt_field(self, ax, Z, style="default", wells=True,
         if has_out_of_range:
             ax.set_facecolor(getattr(kwargs["cmap"], "_rgba_bad", "w"))
 
+        # Add grid (reflecting the model grid)
+        if grid:
+            n1 = 10
+            xStep = 1 + self.Nx//n1
+            yStep = 1 + self.Ny//n1
+            ax.xaxis.set_major_locator(MultipleLocator(self.hx*xStep))
+            ax.yaxis.set_major_locator(MultipleLocator(self.hy*yStep))
+            ax.xaxis.set_minor_locator(MultipleLocator(self.hx))
+            ax.yaxis.set_minor_locator(MultipleLocator(self.hy))
+            ax.xaxis.set_major_formatter('{x:.3f}')
+            ax.tick_params(which='minor', length=0, color='r')
+            ax.grid(True, which="both")
+
         # Axis lims
         ax.set_xlim((0, Lx))
         ax.set_ylim((0, Ly))