First attempt at embarrasingly parallel execution to compute head grids.

timflow/bessel/besselnumba.py

@@ -824,7 +824,7 @@ def Fp(x, y, z1, z2, biga, order, d1, d2, a, b, nt):
 
 
 @numba.njit(nogil=True, cache=True)
-def bessells_int_ho(x, y, z1, z2, lab, order, d1, d2, nt=20):
+def bessells_int_ho(x, y, z1, z2, lab, order, d1, d2, nt=20, out=None):
     """Docs.
 
     To come here
@@ -836,13 +836,17 @@ def bessells_int_ho(x, y, z1, z2, lab, order, d1, d2, nt=20):
     exprange = np.exp(-complex(0, 2) * ang * nrange)
     ahat = a * exprange
     bhat = (b - a * complex(0, 2) * ang) * exprange
-
-    omega = Fp(x, y, z1, z2, biga, order, d1, d2, ahat, bhat, nt)
-    return -L / (4 * np.pi) * omega
+    if out is None:
+        omega = np.zeros(order + 1, dtype=np.complex128)
+    else:
+        omega = out
+    omega[:] = Fp(x, y, z1, z2, biga, order, d1, d2, ahat, bhat, nt)
+    omega *= -L / (4 * np.pi)
+    return omega
 
 
 @numba.njit(nogil=True, cache=True)
-def bessells_gauss_ho(x, y, z1, z2, lab, order):
+def bessells_gauss_ho(x, y, z1, z2, lab, order, out=None):
     """bessells_gauss_ho.
 
     implicit none
@@ -864,8 +868,10 @@ def bessells_gauss_ho(x, y, z1, z2, lab, order):
     for n in range(8):
         x0 = bigz.real - xg[n]
         k0[n] = besselk0(x0, bigz.imag, biglab)
-
-    omega = np.zeros(order + 1, dtype=np.complex128)
+    if out is None:
+        omega = np.zeros(order + 1, dtype=np.complex128)
+    else:
+        omega = out
     for p in range(order + 1):
         omega[p] = complex(0, 0)
         for n in range(8):
@@ -876,7 +882,7 @@ def bessells_gauss_ho(x, y, z1, z2, lab, order):
 
 
 @numba.njit(nogil=True, cache=True)
-def bessells_gauss_ho_d1d2(x, y, z1, z2, lab, order, d1, d2):
+def bessells_gauss_ho_d1d2(x, y, z1, z2, lab, order, d1, d2, out=None, work=None):
     """Returns integral from d1 to d2 along real axis.
 
     While strength is still Delta^order from -1 to +1.
@@ -890,12 +896,20 @@ def bessells_gauss_ho_d1d2(x, y, z1, z2, lab, order, d1, d2):
     real(kind=8) :: xp, yp, dc, fac
     complex(kind=8) :: z1p,z2p,bigz1,bigz2
     """
-    omega = np.zeros(order + 1, dtype=np.complex128)
+    if out is None:
+        omega = np.zeros(order + 1, dtype=np.complex128)
+    else:
+        omega = out
+        omega[:] = 0.0
+    if work is None:
+        omegac = np.zeros(order + 1, dtype=np.complex128)
+    else:
+        omegac = work
     bigz1 = complex(d1, 0)
     bigz2 = complex(d2, 0)
     z1p = 0.5 * (z2 - z1) * bigz1 + 0.5 * (z1 + z2)
     z2p = 0.5 * (z2 - z1) * bigz2 + 0.5 * (z1 + z2)
-    omegac = bessells_gauss_ho(x, y, z1p, z2p, lab, order)
+    bessells_gauss_ho(x, y, z1p, z2p, lab, order, out=omegac)
     dc = (d1 + d2) / (d2 - d1)
     for n in range(order + 1):
         for m in range(n + 1):
@@ -905,7 +919,7 @@ def bessells_gauss_ho_d1d2(x, y, z1, z2, lab, order, d1, d2):
 
 
 @numba.njit(nogil=True, cache=True)
-def bessellsv2(x, y, z1, z2, lab, order, R):
+def bessellsv2(x, y, z1, z2, lab, order, R, out=None):
     """bessellsv2.
 
     implicit none
@@ -920,16 +934,19 @@ def bessellsv2(x, y, z1, z2, lab, order, R):
     """
     nlab = len(lab)
     nterms = order + 1
-    omega = np.zeros((order + 1, nlab), dtype=np.complex128)
+    if out is None:
+        omega = np.zeros((nterms, nlab), dtype=np.complex128)
+    else:
+        omega = out
     # Check if endpoints need to be adjusted using the largest lambda (the first one)
     d1, d2 = find_d1d2(z1, z2, complex(x, y), R * np.abs(lab[0]))
     for n in range(nlab):
-        omega[: nterms + 1, n] = bessells(x, y, z1, z2, lab[n], order, d1, d2)
+        bessells(x, y, z1, z2, lab[n], order, d1, d2, out=omega[: nterms + 1, n])
     return omega
 
 
 @numba.njit(nogil=True, cache=True)
-def bessells(x, y, z1, z2, lab, order, d1in, d2in):
+def bessells(x, y, z1, z2, lab, order, d1in, d2in, out=None):
     """Bessells.
 
     implicit none
@@ -943,20 +960,27 @@ def bessells(x, y, z1, z2, lab, order, d1in, d2in):
     real(kind=8) :: Lnear, L, d1, d2, delta
     complex(kind=8) :: z, delz, za, zb
     """
-    omega = np.zeros(order + 1, dtype=np.complex128)
+    if out is None:
+        omega = np.zeros(order + 1, dtype=np.complex128)
+    else:
+        omega = out
+        omega[:] = 0.0
     Lnear = 3
     z = complex(x, y)
     L = np.abs(z2 - z1)
     if L < Lnear * np.abs(lab):  # No need to break integral up
         if np.abs(z - 0.5 * (z1 + z2)) < 0.5 * Lnear * L:  # Do integration
-            omega = bessells_int_ho(x, y, z1, z2, lab, order, d1in, d2in)
+            bessells_int_ho(x, y, z1, z2, lab, order, d1in, d2in, out=omega)
         else:
-            omega = bessells_gauss_ho_d1d2(x, y, z1, z2, lab, order, d1in, d2in)
+            bessells_gauss_ho_d1d2(x, y, z1, z2, lab, order, d1in, d2in, out=omega)
     else:  # Break integral up in parts
         Nls = int(np.ceil(L / (Lnear * np.abs(lab))))
         delta = 2 / Nls
         delz = (z2 - z1) / Nls
         L = np.abs(delz)
+        # temporary buffers for integration
+        tmp = np.empty(order + 1, dtype=np.complex128)
+        tmp2 = np.empty(order + 1, dtype=np.complex128)
         for n in range(1, Nls + 1):
             d1 = -1 + (n - 1) * delta
             d2 = -1 + n * delta
@@ -967,9 +991,12 @@ def bessells(x, y, z1, z2, lab, order, d1in, d2in):
             za = z1 + (n - 1) * delz
             zb = z1 + n * delz
             if np.abs(z - 0.5 * (za + zb)) < 0.5 * Lnear * L:  # Do integration
-                omega = omega + bessells_int_ho(x, y, z1, z2, lab, order, d1, d2)
+                bessells_int_ho(x, y, z1, z2, lab, order, d1, d2, out=tmp)
             else:
-                omega = omega + bessells_gauss_ho_d1d2(x, y, z1, z2, lab, order, d1, d2)
+                bessells_gauss_ho_d1d2(
+                    x, y, z1, z2, lab, order, d1, d2, out=tmp, work=tmp2
+                )
+            omega += tmp
     return omega
 
 

timflow/steady/aquifer.py

@@ -9,6 +9,8 @@
 
 """
 
+import inspect  # Used for storing the input
+
 import numpy as np
 import pandas as pd
 
@@ -130,6 +132,9 @@ def add_element(self, e):
     def isinside(self, x, y):
         raise Exception("Must overload AquiferData.isinside()")
 
+    def storeinput(self, frame):
+        self.inputargs, _, _, self.inputvalues = inspect.getargvalues(frame)
+
     def findlayer(self, z):
         """Returns layer-number, layer-type and model-layer-number."""
         if z > self.z[0]:

timflow/steady/circareasink.py

@@ -7,6 +7,8 @@
     CircAreaSink(ml, xc=0, yc=0, R=50, N=0.001, layer=0)
 """
 
+import inspect  # Used for storing the input
+
 import numpy as np
 from scipy.special import i0, i1, k0, k1
 
@@ -41,6 +43,7 @@ class CircAreaSink(Element):
     def __init__(
         self, model, xc=0, yc=0, R=1, N=0.001, layer=0, name="CircAreasink", label=None
     ):
+        self.storeinput(inspect.currentframe())
         Element.__init__(
             self, model, nparam=1, nunknowns=0, layers=layer, name=name, label=label
         )

timflow/steady/constant.py

@@ -8,6 +8,8 @@
     Constant(ml, xr=0, yr=0, hr=10.0, layer=0)
 """
 
+import inspect  # Used for storing the input
+
 import numpy as np
 
 from timflow.steady.element import Element
@@ -94,6 +96,7 @@ class Constant(ConstantBase, PotentialEquation):
     """
 
     def __init__(self, model, xr=0, yr=0, hr=0.0, layer=0, label=None):
+        self.storeinput(inspect.currentframe())
         ConstantBase.__init__(
             self, model, xr=xr, yr=yr, hr=hr, layer=layer, name="Constant", label=label
         )

timflow/steady/element.py

@@ -9,6 +9,8 @@ def initialize(self):
             pass
 """
 
+import inspect  # Used for storing the input
+
 import numpy as np
 
 __all__ = ["Element"]
@@ -138,6 +140,9 @@ def headinside(self):
     def setparams(self, sol):
         raise Exception("Must overload Element.setparams()")
 
+    def storeinput(self, frame):
+        self.inputargs, _, _, self.inputvalues = inspect.getargvalues(frame)
+
     # def stoptrace(self, xyz, layer, ltype, step, direction):
     #    return False, 0
 
@@ -156,3 +161,20 @@ def qztop(self, x, y, aq):
 
     def plot(self, **kwargs):
         pass
+
+    def write(self):
+        rv = self.name + "(" + self.model.modelname + ",\n"
+        for key in self.inputargs[2:]:  # The first two are ignored
+            if isinstance(self.inputvalues[key], np.ndarray):
+                rv += (
+                    key
+                    + " = "
+                    + np.array2string(self.inputvalues[key], separator=",")
+                    + ",\n"
+                )
+            elif isinstance(self.inputvalues[key], str):
+                rv += key + " = '" + self.inputvalues[key] + "',\n"
+            else:
+                rv += key + " = " + str(self.inputvalues[key]) + ",\n"
+        rv += ")\n"
+        return rv

timflow/steady/inhomogeneity.py

@@ -13,6 +13,7 @@
     )
 """
 
+import inspect  # Used for storing the input
 from warnings import warn
 
 import numpy as np
@@ -217,6 +218,7 @@ def __init__(
             assert topboundary[:4] == "conf", (
                 "Error: infiltration can only be added if topboundary='conf'"
             )
+        self.storeinput(inspect.currentframe())
         (
             kaq,
             c,
@@ -313,6 +315,7 @@ def __init__(
             assert topboundary[:4] == "conf", (
                 "Error: infiltration can only be added if topboundary='conf'"
             )
+        self.storeinput(inspect.currentframe())
         kaq, kzoverkh, c, npor, ltype = param_3d(
             kaq, z, kzoverkh, npor, topboundary, topres
         )

timflow/steady/inhomogeneity1d.py

@@ -10,6 +10,7 @@
 
 """
 
+import inspect  # user for storing the input
 import warnings
 
 import matplotlib.pyplot as plt
@@ -353,6 +354,7 @@ def __init__(
             c = []
         if z is None:
             z = [1, 0]
+        self.storeinput(inspect.currentframe())
         (
             kaq,
             c,
@@ -431,6 +433,7 @@ def __init__(
     ):
         if z is None:
             z = [1, 0]
+        self.storeinput(inspect.currentframe())
         (
             kaq,
             kzoverkh,

timflow/steady/linedoublet.py

@@ -7,6 +7,7 @@
     ImpLineDoublet(ml, x1=-10, y1=0, x2=10, y2=0, layers=0)
 """
 
+import inspect  # Used for storing the input
 import warnings
 
 import matplotlib.pyplot as plt
@@ -227,6 +228,7 @@ def __init__(
         label=None,
         addtomodel=True,
     ):
+        self.storeinput(inspect.currentframe())
         LineDoubletHoBase.__init__(
             self,
             model,
@@ -300,6 +302,7 @@ def __init__(
         label=None,
         addtomodel=True,
     ):
+        self.storeinput(inspect.currentframe())
         LineDoubletHoBase.__init__(
             self,
             model,
@@ -451,6 +454,7 @@ class ImpermeableWallString(LineDoubletStringBase, DisvecEquation):
     def __init__(self, model, xy=None, layers=0, order=0, label=None):
         if xy is None:
             xy = [(-1, 0), (1, 0)]
+        self.storeinput(inspect.currentframe())
         LineDoubletStringBase.__init__(
             self,
             model,
@@ -503,6 +507,7 @@ class LeakyWallString(LineDoubletStringBase, LeakyWallEquation):
     def __init__(self, model, xy=None, res=np.inf, layers=0, order=0, label=None):
         if xy is None:
             xy = [(-1, 0), (1, 0)]
+        self.storeinput(inspect.currentframe())
         LineDoubletStringBase.__init__(
             self,
             model,

timflow/steady/linedoublet1d.py

@@ -8,6 +8,7 @@
 
 """
 
+import inspect  # Used for storing the input
 import warnings
 
 import matplotlib.pyplot as plt
@@ -126,6 +127,7 @@ class ImpermeableWall1D(LineDoublet1D, DisvecEquation):
     """Create 1D impermeable wall."""
 
     def __init__(self, model, xld=0, layers=0, label=None):
+        self.storeinput(inspect.currentframe())
         LineDoublet1D.__init__(
             self,
             model,
@@ -171,6 +173,7 @@ class LeakyWall1D(LineDoublet1D, LeakyWallEquation):
     tiny = 1e-6
 
     def __init__(self, model, xld=0, res=np.inf, layers=0, label=None):
+        self.storeinput(inspect.currentframe())
         LineDoublet1D.__init__(
             self,
             model,

timflow/steady/linesink.py

@@ -7,6 +7,7 @@
     River(ml, x1=-10, y1=0, x2=10, y2=0, hls=5, layers=0)
 """
 
+import inspect  # Used for storing the input
 import warnings
 
 import matplotlib.pyplot as plt
@@ -285,6 +286,7 @@ def __init__(
         label=None,
         addtomodel=True,
     ):
+        self.storeinput(inspect.currentframe())
         LineSinkBase.__init__(
             self,
             model,
@@ -929,6 +931,7 @@ def __init__(
         """Initialize a steady string of head-specified line-sinks."""
         if xy is None:
             xy = [(-1, 0), (1, 0)]
+        self.storeinput(inspect.currentframe())
         LineSinkStringBase2.__init__(
             self,
             model,
@@ -1081,6 +1084,7 @@ def __init__(
         """Initialize a steady string of discharge-specified line-sinks."""
         if xy is None:
             xy = [(-1, 0), (1, 0)]
+        self.storeinput(inspect.currentframe())
         RiverString.__init__(
             self,
             model,