TTim plots

docs/index.rst

@@ -78,7 +78,7 @@ Quick Example
             ml.solve()
 
             # Plot head contours at t=2 days
-            ml.contour(win=[-30, 55, -30, 30], ngr=40, t=2, labels=True, decimals=1)
+            ml.plots.contour(win=[-30, 55, -30, 30], ngr=40, t=2, labels=True, decimals=1)
 
 
     .. tab-item:: Result

pyproject.toml

@@ -29,11 +29,9 @@ classifiers = [
     "Operating System :: MacOS",
     "Programming Language :: Python",
     "Programming Language :: Python :: 3",
-    "Programming Language :: Python :: 3.7",
-    "Programming Language :: Python :: 3.8",
-    "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
     "Programming Language :: Python :: 3 :: Only",
     "Topic :: Scientific/Engineering :: Hydrology",
 ]

ttim/circareasink.py

@@ -123,8 +123,10 @@ def disvecinf(self, x, y, aq=None):
             qy[:] = qr * (y - self.yc) / r
         return qx, qy
 
-    def plot(self):
-        plt.plot(
+    def plot(self, ax=None):
+        if ax is None:
+            _, ax = plt.subplots()
+        ax.plot(
             self.xc + self.R * np.cos(np.linspace(0, 2 * np.pi, 100)),
             self.yc + self.R * np.sin(np.linspace(0, 2 * np.pi, 100)),
             "k",

ttim/element.py

@@ -332,5 +332,5 @@ def run_after_solve(self):
         """
         pass
 
-    def plot(self):
+    def plot(self, ax=None):
         pass

ttim/linedoublet.py

@@ -174,8 +174,10 @@ def disvecinf(self, x, y, aq=None):
         rvy.shape = (self.nparam, aq.naq, self.model.npval)
         return rvx, rvy
 
-    def plot(self):
-        plt.plot([self.x1, self.x2], [self.y1, self.y2], "k")
+    def plot(self, ax=None):
+        if ax is None:
+            _, ax = plt.subplots()
+        ax.plot([self.x1, self.x2], [self.y1, self.y2], "k")
 
 
 class LeakyLineDoublet(LineDoubletHoBase, LeakyWallEquation):
@@ -376,5 +378,7 @@ def disvecinf(self, x, y, aq=None):
             rvy[i * ld.nparam : (i + 1) * ld.nparam, :] = qy
         return rvx, rvy
 
-    def plot(self):
-        plt.plot(self.xldlayout, self.yldlayout, "k")
+    def plot(self, ax=None):
+        if ax is None:
+            _, ax = plt.subplots()
+        ax.plot(self.xldlayout, self.yldlayout, "k")

ttim/linesink.py

@@ -168,8 +168,10 @@ def headinside(self, t):
             :, np.newaxis
         ] * self.discharge(t)
 
-    def plot(self):
-        plt.plot([self.x1, self.x2], [self.y1, self.y2], "k")
+    def plot(self, ax=None):
+        if ax is None:
+            _, ax = plt.subplots()
+        ax.plot([self.x1, self.x2], [self.y1, self.y2], "k")
 
 
 class LineSink(LineSinkBase):
@@ -442,8 +444,8 @@ def headinside(self, t, derivative=0):
             )
         return rv
 
-    def plot(self):
-        plt.plot(self.xlslayout, self.ylslayout, "k")
+    def plot(self, ax):
+        ax.plot(self.xlslayout, self.ylslayout, "k")
 
     def run_after_solve(self):
         for i in range(self.nls):
@@ -998,8 +1000,10 @@ def headinside(self, t):
             :, np.newaxis
         ] * self.discharge(t)
 
-    def plot(self):
-        plt.plot([self.x1, self.x2], [self.y1, self.y2], "k")
+    def plot(self, ax=None):
+        if ax is None:
+            _, ax = plt.subplots()
+        ax.plot([self.x1, self.x2], [self.y1, self.y2], "k")
 
 
 class HeadLineSinkHo(LineSinkHoBase, HeadEquationNores):

ttim/model.py

@@ -1,5 +1,6 @@
 # from .invlap import *
 import inspect  # Used for storing the input
+from warnings import warn
 
 import numpy as np
 
@@ -8,10 +9,10 @@
 
 # from .bessel import *
 from .invlapnumba import compute_laplace_parameters_numba, invlap, invlapcomp
-from .util import PlotTtim
+from .plots import PlotTtim
 
 
-class TimModel(PlotTtim):
+class TimModel:
     def __init__(
         self,
         kaq=[1, 1],
@@ -68,6 +69,25 @@ def __init__(
         if self.timmlmodel is not None:
             self.timmlmodel.solve()
 
+        self.plots = PlotTtim(self)
+        self.plot = self.plots.topview
+
+        # NOTE: reinstate later, after deprecation below is removed?
+        # self.xsection = self.plots.xsection
+
+    def xsection(*args, **kwargs):
+        raise DeprecationWarning(
+            "This method is deprecated. Use `ml.plots.head_along_line()` instead."
+        )
+
+    def contour(self, *args, **kwargs):
+        warn(
+            category=DeprecationWarning,
+            message="This method is deprecated. Use `ml.plots.contour()` instead.",
+            stacklevel=1,
+        )
+        self.plots.contour(*args, **kwargs)
+
     def __repr__(self):
         return "Model"
 
@@ -349,7 +369,7 @@ def velocomp(self, x, y, z, t, aq=None, layer_ltype=None):
                 )
                 qz = (h[1, 0] - h[0, 0]) / aq.c[
                     layer
-                ]  # TO DO include storage in leaky layer
+                ]  # TODO: include storage in leaky layer
             vz = qz / aq.porll[layer]
         else:  # in aquifer layer
             h = self.head(x, y, t, layers=layer, aq=aq, neglect_steady=True)
@@ -387,7 +407,7 @@ def velocomp(self, x, y, z, t, aq=None, layer_ltype=None):
                     )[:, 0]
             # this works because c[0] = 1e100 for impermeable top
             qztop = (h[1] - h[0]) / self.aq.c[layer]
-            # TO DO modify for infiltration in top aquifer
+            # TODO: modify for infiltration in top aquifer
             # if layer == 0:
             #    qztop += self.qztop(x, y)
             if layer < aq.naq - 1:

ttim/plots.py

@@ -0,0 +1,253 @@
+from typing import Optional
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+class PlotTtim:
+    def __init__(self, ml):
+        self._ml = ml
+
+    def topview(self, win=None, ax=None, figsize=None, layers=None):
+        """Plot top-view.
+
+        Parameters
+        ----------
+        win : list or tuple
+            [x1, x2, y1, y2]
+        """
+        if ax is None:
+            _, ax = plt.subplots(figsize=figsize)
+            ax.set_aspect("equal", adjustable="box")
+        if layers is not None and isinstance(layers, int):
+            layers = [layers]
+        for e in self._ml.elementlist:
+            if layers is None or np.isin(e.layers, layers):
+                e.plot(ax=ax)
+        if win is not None:
+            ax.axis(win)
+        return ax
+
+    def xsection(
+        self,
+        xy: Optional[list[tuple[float]]] = None,
+        labels=True,
+        params=False,
+        ax=None,
+    ):
+        if ax is None:
+            _, ax = plt.subplots(1, 1, figsize=(8, 4))
+
+        if xy is not None:
+            (x0, y0), (x1, y1) = xy
+            r = np.sqrt((x1 - x0) ** 2 + (y1 - y0) ** 2)
+            ax.set_xlim(0, r)
+        else:
+            r = 1.0
+
+        if labels:
+            lli = 1 if self._ml.aq.topboundary == "con" else 0
+            aqi = 0
+        else:
+            lli = None
+            aqi = None
+
+        for i in range(self._ml.aq.nlayers):
+            if self._ml.aq.ltype[i] == "l":
+                ax.axhspan(
+                    ymin=self._ml.aq.z[i + 1],
+                    ymax=self._ml.aq.z[i],
+                    color=[0.8, 0.8, 0.8],
+                )
+                if labels:
+                    ax.text(
+                        0.5 * r if not params else 0.25 * r,
+                        np.mean(self._ml.aq.z[i : i + 2]),
+                        f"leaky layer {lli}",
+                        ha="center",
+                        va="center",
+                    )
+                if params:
+                    ax.text(
+                        0.75 * r if labels else 0.5 * r,
+                        np.mean(self._ml.aq.z[i : i + 2]),
+                        (
+                            f"$c$ = {self._ml.aq.c[lli]}, "
+                            f"$S_s$ = {self._ml.aq.Sll[lli]:.2e}",
+                        ),
+                        ha="center",
+                        va="center",
+                    )
+                if labels or params:
+                    lli += 1
+
+            if labels and self._ml.aq.ltype[i] == "a":
+                ax.text(
+                    0.5 * r if not params else 0.25 * r,
+                    np.mean(self._ml.aq.z[i : i + 2]),
+                    f"aquifer {aqi}",
+                    ha="center",
+                    va="center",
+                )
+            if params and self._ml.aq.ltype[i] == "a":
+                if aqi == 0 and self._ml.aq.phreatictop:
+                    paramtxt = (
+                        f"$k_h$ = {self._ml.aq.kaq[aqi]}, "
+                        f"$S$ = {self._ml.aq.Saq[aqi]}"
+                    )
+                else:
+                    paramtxt = (
+                        f"$k_h$ = {self._ml.aq.kaq[aqi]}, "
+                        f"$S_s$ = {self._ml.aq.Saq[aqi]:.2e}"
+                    )
+                ax.text(
+                    0.75 * r if labels else 0.5 * r,
+                    np.mean(self._ml.aq.z[i : i + 2]),
+                    paramtxt,
+                    ha="center",
+                    va="center",
+                )
+            if (labels or params) and self._ml.aq.ltype[i] == "a":
+                aqi += 1
+
+        for i in range(1, self._ml.aq.nlayers):
+            if self._ml.aq.ltype[i] == "a" and self._ml.aq.ltype[i - 1] == "a":
+                ax.axhspan(
+                    ymin=self._ml.aq.z[i], ymax=self._ml.aq.z[i], color=[0.8, 0.8, 0.8]
+                )
+
+        ax.axhline(self._ml.aq.z[0], color="k", lw=0.75)
+        ax.axhline(self._ml.aq.z[-1], color="k", lw=3.0)
+        ax.set_ylabel("elevation")
+        return ax
+
+    def head_along_line(
+        self,
+        x1=0,
+        x2=1,
+        y1=0,
+        y2=0,
+        npoints=100,
+        t=1.0,
+        layers=0,
+        sstart=0,
+        color=None,
+        lw=1,
+        figsize=None,
+        ax=None,
+        legend=True,
+        grid=True,
+    ):
+        layers = np.atleast_1d(layers)
+        t = np.atleast_1d(t)
+        if ax is None:
+            _, ax = plt.subplots(figsize=figsize)
+        x = np.linspace(x1, x2, npoints)
+        y = np.linspace(y1, y2, npoints)
+        s = np.sqrt((x - x[0]) ** 2 + (y - y[0]) ** 2) + sstart
+        h = self._ml.headalongline(x, y, t, layers)
+        nlayers, ntime, npoints = h.shape
+        for i in range(nlayers):
+            for j in range(ntime):
+                lbl = f"head (t={t[j]}, layer={layers[i]})"
+                if color is None:
+                    ax.plot(s, h[i, j, :], lw=lw, label=lbl)
+                else:
+                    ax.plot(s, h[i, j, :], color, lw=lw, label=lbl)
+        if legend:
+            ax.legend(loc=(0, 1), ncol=3, frameon=False)
+        if grid:
+            ax.grid(True)
+        return ax
+
+    def contour(
+        self,
+        win,
+        ngr=20,
+        t=1,
+        layers=0,
+        levels=20,
+        layout=True,
+        labels=True,
+        decimals=1,
+        color=None,
+        ax=None,
+        figsize=None,
+        legend=True,
+    ):
+        """Contour plot.
+
+        Parameters
+        ----------
+        win : list or tuple
+            [x1, x2, y1, y2]
+        ngr : scalar, tuple or list
+            if scalar: number of grid points in x and y direction
+            if tuple or list: nx, ny, number of grid points in x and y direction
+        t : scalar
+            time
+        layers : integer, list or array
+            layers for which grid is returned
+        levels : integer or array (default 20)
+            levels that are contoured
+        layout : boolean (default True_)
+            plot layout of elements
+        labels : boolean (default True)
+            print labels along contours
+        decimals : integer (default 1)
+            number of decimals of labels along contours
+        color : str or list of strings
+            color of contour lines
+        ax : matplotlib.Axes
+            axes to plot on, default is None which creates a new figure
+        figsize : tuple of 2 values (default is mpl default)
+            size of figure
+        legend : list or boolean (default True)
+            add legend to figure
+            if list of strings: use strings as names in legend
+        """
+        x1, x2, y1, y2 = win
+        if np.isscalar(ngr):
+            nx = ny = ngr
+        else:
+            nx, ny = ngr
+        layers = np.atleast_1d(layers)
+        xg = np.linspace(x1, x2, nx)
+        yg = np.linspace(y1, y2, ny)
+        h = self._ml.headgrid(xg, yg, t, layers)
+        if ax is None:
+            _, ax = plt.subplots(figsize=figsize)
+            ax.set_aspect("equal", adjustable="box")
+        # color
+        if color is None:
+            c = plt.rcParams["axes.prop_cycle"].by_key()["color"]
+        elif isinstance(color, str):
+            c = len(layers) * [color]
+        elif isinstance(color, list):
+            c = color
+        if len(c) < len(layers):
+            n = np.ceil(self._ml.aq.naq / len(c))
+            c = n * c
+
+        # contour
+        cslist = []
+        cshandlelist = []
+        for i in range(len(layers)):
+            cs = ax.contour(
+                xg, yg, h[i, 0], levels, colors=c[i], negative_linestyles="solid"
+            )
+            cslist.append(cs)
+            handles, _ = cs.legend_elements()
+            cshandlelist.append(handles[0])
+            if labels:
+                fmt = f"%1.{decimals}f"
+                ax.clabel(cs, fmt=fmt)
+        if isinstance(legend, list):
+            ax.legend(cshandlelist, legend, loc=(0, 1), ncol=3, frameon=False)
+        elif legend:
+            legendlist = ["layer " + str(i) for i in layers]
+            ax.legend(cshandlelist, legendlist, loc=(0, 1), ncol=3, frameon=False)
+
+        if layout:
+            self.topview(win=[x1, x2, y1, y2], ax=ax)
+        return ax

ttim/well.py

@@ -160,8 +160,10 @@ def headinside(self, t, derivative=0):
             self.layers
         ] - self.resfach[:, np.newaxis] * self.discharge(t, derivative=derivative)
 
-    def plot(self):
-        plt.plot(self.xw, self.yw, "k.")
+    def plot(self, ax=None):
+        if ax is None:
+            _, ax = plt.subplots()
+        ax.plot(self.xw, self.yw, "k.")
 
     def changetrace(
         self, xyzt1, xyzt2, aq, layer, ltype, modellayer, direction, hstepmax