diff --git a/timflow/bessel/besselnumba.py b/timflow/bessel/besselnumba.py index 994d030b..955f2483 100644 --- a/timflow/bessel/besselnumba.py +++ b/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 diff --git a/timflow/steady/aquifer.py b/timflow/steady/aquifer.py index 9bcdec7d..89713d22 100644 --- a/timflow/steady/aquifer.py +++ b/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]: diff --git a/timflow/steady/circareasink.py b/timflow/steady/circareasink.py index 4b92e85e..0cf7c252 100644 --- a/timflow/steady/circareasink.py +++ b/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 ) diff --git a/timflow/steady/constant.py b/timflow/steady/constant.py index dde0baba..89d662bc 100644 --- a/timflow/steady/constant.py +++ b/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 ) diff --git a/timflow/steady/element.py b/timflow/steady/element.py index 8095f98e..bd93a5cf 100644 --- a/timflow/steady/element.py +++ b/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 diff --git a/timflow/steady/inhomogeneity.py b/timflow/steady/inhomogeneity.py index 9f94158e..b5caa554 100644 --- a/timflow/steady/inhomogeneity.py +++ b/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 ) diff --git a/timflow/steady/inhomogeneity1d.py b/timflow/steady/inhomogeneity1d.py index 0131ed88..74665f3c 100644 --- a/timflow/steady/inhomogeneity1d.py +++ b/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, diff --git a/timflow/steady/linedoublet.py b/timflow/steady/linedoublet.py index 29ad827b..b9997c14 100644 --- a/timflow/steady/linedoublet.py +++ b/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, diff --git a/timflow/steady/linedoublet1d.py b/timflow/steady/linedoublet1d.py index 3fcd5e62..e08f8a56 100644 --- a/timflow/steady/linedoublet1d.py +++ b/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, diff --git a/timflow/steady/linesink.py b/timflow/steady/linesink.py index 145accd9..ade4a925 100644 --- a/timflow/steady/linesink.py +++ b/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, diff --git a/timflow/steady/model.py b/timflow/steady/model.py index d013f897..79ae990d 100644 --- a/timflow/steady/model.py +++ b/timflow/steady/model.py @@ -10,6 +10,7 @@ ml.solve() """ +import inspect # Used for storing the input import multiprocessing as mp import warnings @@ -26,22 +27,6 @@ __all__ = ["Model", "ModelMaq", "Model3D", "ModelXsection"] -def _compute_head_mp(args): - """Helper function for parallel computation of head_array.""" - model, xi, yi, layers, i = args - return i, model.head(xi, yi, layers=layers) - - -def _compute_velocity_mp(args): - """Helper function for parallel computation of velocity_array.""" - model, xi, yi, zi, i = args - try: - vv = model.velocomp(xi, yi, zi) - except (ZeroDivisionError, ValueError): - vv = np.full((3,), np.nan) - return i, vv - - class Model: """Create a model consisting of an arbitrary sequence of aquifers and leaky layers. @@ -75,6 +60,7 @@ def __init__(self, kaq, z, c, npor, ltype, model3d=False): self.elementlist = [] self.elementdict = {} # only elements that have a label self.aq = Aquifer(self, kaq, c, z, npor, ltype, model3d=model3d) + self.modelname = "ml" # Used for writing out input self.name = "Model" self.model_type = "steady" # Model type for plotting and other purposes @@ -98,6 +84,9 @@ def remove_element(self, e): self.elementdict.pop(e.label) self.elementlist.remove(e) + def storeinput(self, frame): + self.inputargs, _, _, self.inputvalues = inspect.getargvalues(frame) + def potential(self, x, y, aq=None): if aq is None: aq = self.aq.find_aquifer_data(x, y) @@ -270,64 +259,6 @@ def head(self, x, y, layers=None, aq=None): else: return rv[layers] - def head_array(self, x, y, layers=None, show_progress=False, parallel=False): - """Head for array of points. - - Parameters - ---------- - x : 1D array or list - x values of points - y : 1D array or list - y values of points - layers : integer, list or array, optional - layers for which grid is returned - show_progress : bool - show computation progress, by printing dots per row or with tqdm progressbar - when parallel is True. Default is False. - parallel : bool or int, optional - if `True`, computes head_array in parallel using multiprocessing, - by default `False`. If an integer is provided, it is interpreted as the - number of processes to use. - - Returns - ------- - h : array - heads array with size (nlayers, npoints) - """ - parallel, process_map, tqdm = check_tqdm_parallel(parallel) - x = np.atleast_1d(x) - y = np.atleast_1d(y) - npts = len(x) - assert npts == len(y), "x and y must have the same length" - if layers is None: - nlayers = self.aq.find_aquifer_data(x[0], y[0]).naq - else: - nlayers = len(np.atleast_1d(layers)) - h = np.empty((nlayers, npts)) - if not parallel: - for i in ( - tqdm(range(npts), disable=not show_progress) if tqdm else range(npts) - ): - h[:, i] = self.head(x[i], y[i], layers) - else: - nproc = mp.cpu_count() if parallel is True else int(parallel) - chunksize = max(1, npts // (4 * nproc)) if nproc > 0 else 1 - tasks = [(self, x[i], y[i], layers, i) for i in range(npts)] - results = process_map( - _compute_head_mp, - tasks, - total=npts, - desc="head array", - disable=not show_progress, - tqdm_class=tqdm, - max_workers=nproc, - chunksize=chunksize, - ) - - for i, result in results: - h[:, i] = result - return h - def headgrid( self, xg, yg, layers=None, printrow=False, show_progress=False, parallel=False ): @@ -344,10 +275,9 @@ def headgrid( show_progress : bool show computation progress, by printing dots per row or with tqdm progressbar when parallel is True. Default is False. - parallel : bool or int, optional - if `True`, computes headgrid in parallel using multiprocessing, - by default `False`. If an integer is provided, it is interpreted as the - number of processes to use. + parallel : bool, optional + if `True`, computes headgrid in parallel using multi threading, + by default `False` printrow : bool, optional .. deprecated:: 0.2.0 @@ -359,7 +289,6 @@ def headgrid( See Also -------- - :func:`~timflow.steady.Model.head_array` :func:`~timflow.steady.Model.headgrid2` """ if printrow: @@ -369,15 +298,43 @@ def headgrid( stacklevel=2, ) show_progress = printrow - x, y = np.meshgrid(xg, yg) - h = self.head_array( - x.ravel(), - y.ravel(), - layers=layers, - show_progress=show_progress, - parallel=parallel, - ) - return h.reshape((h.shape[0], len(yg), len(xg))) + + parallel, thread_map, tqdm = check_tqdm_parallel(parallel) + + xg = np.atleast_1d(xg) + yg = np.atleast_1d(yg) + nx, ny = len(xg), len(yg) + if layers is None: + Nlayers = self.aq.find_aquifer_data(xg[0], yg[0]).naq + else: + Nlayers = len(np.atleast_1d(layers)) + h = np.empty((Nlayers, ny, nx)) + if not parallel: + for j in range(ny): + if show_progress: + print(".", end="", flush=True) + for i in range(nx): + h[:, j, i] = self.head(xg[i], yg[j], layers) + if show_progress: + print("", flush=True) + else: + + def compute(ij): + i, j = ij + return i, j, self.head(xg[i], yg[j], layers) + + results = thread_map( + compute, + [(i, j) for j in range(ny) for i in range(nx)], + total=nx * ny, + desc="headgrid", + disable=not show_progress, + tqdm_class=tqdm, + ) + for i, j, result in results: + h[:, j, i] = result + + return h def headgrid2( self, @@ -405,10 +362,6 @@ def headgrid2( show_progress : bool show computation progress, by printing dots per row or with tqdm progressbar when parallel is True. Default is False. - parallel : bool or int, optional - if `True`, computes headgrid in parallel using multiprocessing, - by default `False`. If an integer is provided, it is interpreted as the - number of processes to use. printrow : boolean, optional .. deprecated:: 0.2.0 @@ -420,7 +373,6 @@ def headgrid2( See Also -------- - :func:`~timflow.steady.Model.head_array` :func:`~timflow.steady.Model.headgrid` """ xg, yg = np.linspace(x1, x2, nx), np.linspace(y1, y2, ny) @@ -484,7 +436,7 @@ def disvecgrid( show computation progress, by printing dots per row or with tqdm progressbar when parallel is True. Default is True. parallel : bool, optional - if `True`, computes discharge vector grid in parallel using multiprocessing, + if `True`, computes discharge vector grid in parallel using multi threading, by default `False` Returns @@ -618,68 +570,8 @@ def velocity(self, x, y, z): """ return self.velocomp(x, y, z) - def velocity_array(self, x, y, z, show_progress=True, parallel=False): - """Compute velocity grid. - - Parameters - ---------- - x : 1d-array - x values - y : 1d-array - y values - z : 1d-array - z values - show_progress : bool - show computation progress with tqdm progressbar if tqdm is installed. - Default is True. - parallel : bool or int, optional - if `True`, computes velocity grid in parallel using multi processing, - by default `False`. If an integer is provided, it is interpreted - as the number of processes to use. - - Returns - ------- - velocity : array - velocity vector (vx, vy, vz) at each point in grid, - size (3, len(x)) - """ - parallel, process_map, tqdm = check_tqdm_parallel(parallel) - x = np.atleast_1d(x) - y = np.atleast_1d(y) - z = np.atleast_1d(z) - npts = len(x) - assert npts == len(y) == len(z), "x, y and z must have the same length" - v = np.empty((3, npts)) - if not parallel: - for i in ( - tqdm(range(npts), disable=not show_progress) if tqdm else range(npts) - ): - try: - vv = self.velocomp(x[i], y[i], z[i]) - except (ZeroDivisionError, ValueError): - vv = np.full((3,), np.nan) - v[:, i] = vv - else: - nproc = mp.cpu_count() if parallel is True else int(parallel) - chunksize = max(1, npts // (4 * nproc)) if nproc > 0 else 1 - tasks = [(self, x[i], y[i], z[i], i) for i in range(npts)] - results = process_map( - _compute_velocity_mp, - tasks, - total=npts, - desc="velocity array", - disable=not show_progress, - tqdm_class=tqdm, - max_workers=nproc, - chunksize=chunksize, - ) - for i, result in results: - v[:, i] = result - - return v - def velocity_grid(self, xg, yg, zg, show_progress=True, parallel=False): - """Compute velocities for an array of points. + """Compute velocity grid. Parameters ---------- @@ -689,31 +581,57 @@ def velocity_grid(self, xg, yg, zg, show_progress=True, parallel=False): y values of grid zg : 1d-array z values of grid - show_progress : bool, optional - if `True`, shows progress bar when computing velocity grid, by default `True` - parallel : bool or int, optional - if `True`, computes velocity grid in parallel using multiprocessing, - by default `False`. If an integer is provided, it specifies the number of - processes to use. + parallel : bool, optional + if `True`, computes velocity grid in parallel using multi threading, + by default `False` Returns ------- velocity : array - velocity vector (vx, vy, vz) at each point in grid, - size (3, len(zg), len(yg), len(xg)) + velocity vector (vz, vy, vx) at each point in grid, + size (3, len(z), len(y), len(x)) """ + parallel, thread_map, tqdm = check_tqdm_parallel(parallel) + + def compute(kij): + k, i, j = kij + try: + vv = self.velocomp(xg[j], yg[i], zg[k]) + except ZeroDivisionError: + vv = np.full((3,), np.nan) + return k, i, j, vv + xg = np.atleast_1d(xg) yg = np.atleast_1d(yg) zg = np.atleast_1d(zg) - x, y, z = np.meshgrid(xg, yg, zg) - v = self.velocity_array( - x.ravel(), - y.ravel(), - z.ravel(), - show_progress=show_progress, - parallel=parallel, - ) - return v.reshape((3, len(yg), len(xg), len(zg))).transpose((0, 3, 1, 2)) + nz, ny, nx = len(zg), len(yg), len(xg) + v = np.empty((3, nz, ny, nx)) + if not parallel: + for k in range(nz): + if show_progress: + print(".", end="", flush=True) + for i in range(ny): + for j in range(nx): + try: + vv = self.velocomp(xg[j], yg[i], zg[k]) + except ZeroDivisionError: + vv = np.full((3,), np.nan) + v[:, k, i, j] = vv + if show_progress: + print("", flush=True) + else: + results = thread_map( + compute, + [(k, i, j) for k in range(nz) for i in range(ny) for j in range(nx)], + total=nz * nx * ny, + desc="velocity grid", + disable=not show_progress, + tqdm_class=tqdm, + ) + for k, i, j, result in results: + v[:, k, i, j] = result + + return v def velocomp(self, x, y, z, aq=None, layer_ltype=None): if aq is None: @@ -874,6 +792,32 @@ def solve_mp(self, nproc=4, printmat=0, sendback=0, silent=False): return sol return + def write(self): + rv = self.modelname + " = " + self.name + "(\n" + for key in self.inputargs[1:]: # The first argument (self) is 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 + + def writemodel(self, fname): + self.initialize() # So that the model can be written without solving first + f = open(fname, "w") + f.write("from timflow.steady import *\n") + f.write(self.write()) + for e in self.elementlist: + f.write(e.write()) + f.close() + def aquifer_summary(self): """Return DataFrame with summary of aquifer(s) parameters in model. @@ -981,6 +925,7 @@ def __init__(self, kaq=1, z=None, c=None, npor=0.3, topboundary="conf", hstar=No c = [] if z is None: z = [1, 0] + self.storeinput(inspect.currentframe()) kaq, c, npor, ltype = param_maq(kaq, z, c, npor, topboundary) super().__init__(kaq=kaq, z=z, c=c, npor=npor, ltype=ltype) self.name = "ModelMaq" @@ -1059,6 +1004,7 @@ def __init__( """ if z is None: z = [1, 0] + self.storeinput(inspect.currentframe()) kaq, kzoverkh, c, npor, ltype = param_3d( kaq, z, kzoverkh, npor, topboundary, topres ) @@ -1095,6 +1041,7 @@ def __init__(self, naq=1): self.elementlist = [] self.elementdict = {} # only elements that have a label self.aq = SimpleAquifer(naq) + self.modelname = "ml" # Used for writing out input self.plots = PlotSteady(self) self.name = "ModelXsection" diff --git a/timflow/steady/plots.py b/timflow/steady/plots.py index 4f9dcd8b..3db4baa3 100644 --- a/timflow/steady/plots.py +++ b/timflow/steady/plots.py @@ -98,10 +98,9 @@ def contour( if list of strings: use strings as names in legend return_contours : bool, optional if True, return list of contour sets for each contoured layer - parallel : bool or int, optional - if True, compute head grid in parallel using multiprocessing, - default is False. If int is provided, it is interpreted as the number of - processes to use. + parallel : bool, optional + if True, compute headgrid in parallel using multiple threads, + default is False **kwargs additional keyword arguments passed to ax.contour() @@ -784,7 +783,7 @@ def quiver_xy( figsize : tuple of 2 values size of figure parallel : bool - if True, compute velocity grid in parallel using multiprocessing, + if True, compute velocity grid in parallel using multiple threads, default is False **kwargs additional keyword arguments passed to ax.quiver() @@ -820,7 +819,7 @@ def quiver_z( figsize : tuple of 2 values size of figure parallel : bool - if True, compute velocity grid in parallel using multiprocessing, + if True, compute velocity grid in parallel using multiple threads, default is False **kwargs additional keyword arguments passed to ax.quiver() diff --git a/timflow/steady/uflow.py b/timflow/steady/uflow.py index 1242426e..83b36d58 100644 --- a/timflow/steady/uflow.py +++ b/timflow/steady/uflow.py @@ -7,6 +7,8 @@ Uflow(ml, slope=1e-4, angle=0.0) """ +import inspect # Used for storing the input + import numpy as np from timflow.steady.element import Element @@ -40,6 +42,7 @@ def __init__(self, model, slope, angle, label=None): "timflow error: Uflow can only be added to " "model with background confined aquifer" ) + self.storeinput(inspect.currentframe()) Element.__init__( self, model, nparam=2, nunknowns=0, layers=0, name="Uflow", label=label ) diff --git a/timflow/steady/well.py b/timflow/steady/well.py index a5cd27df..50df3c9f 100644 --- a/timflow/steady/well.py +++ b/timflow/steady/well.py @@ -8,6 +8,7 @@ Well(ml, xw=100, yw=200, Qw=1000, layers=[0]) """ +import inspect # Used for storing the input import warnings import matplotlib.pyplot as plt @@ -683,6 +684,7 @@ def __init__( xc=None, yc=None, ): + self.storeinput(inspect.currentframe()) WellBase.__init__( self, model, diff --git a/timflow/transient/__init__.py b/timflow/transient/__init__.py index 239a98d6..6327fc0f 100644 --- a/timflow/transient/__init__.py +++ b/timflow/transient/__init__.py @@ -7,6 +7,7 @@ """ # ruff : noqa: F401 +from timflow.transient import parallel from timflow.transient.circareasink import CircAreaSink from timflow.transient.fit import Calibrate from timflow.transient.inhom1d import Xsection3D, XsectionMaq diff --git a/timflow/transient/aquifer.py b/timflow/transient/aquifer.py index fa3d8115..92b028e0 100644 --- a/timflow/transient/aquifer.py +++ b/timflow/transient/aquifer.py @@ -11,6 +11,8 @@ import numpy as np import pandas as pd +from timflow.transient.parallel.dtypes import AquiferTuple + class AquiferData: def __init__( @@ -365,6 +367,16 @@ def add_inhom(self, inhom): self.inhomdict[inhom.name] = inhom return inhom_number + def to_numba_tuple(self): + """Write aquifer data to a numba tuple for use in parallel computations.""" + return AquiferTuple( + naq=self.naq, + lab2=self.lab2, + lababs=self.lababs, + eigvec=self.eigvec, + Tcol=self.Tcol, + ) + class SimpleAquifer(Aquifer): def __init__(self, naq): diff --git a/timflow/transient/circareasink.py b/timflow/transient/circareasink.py index 2b55039d..83797c25 100644 --- a/timflow/transient/circareasink.py +++ b/timflow/transient/circareasink.py @@ -8,6 +8,8 @@ CircAreaSink(ml, xc=0, yc=0, R=50, tsandN=[(0, 0.001)], layers=0) """ +import inspect # Used for storing the input + import matplotlib.pyplot as plt import numpy as np from scipy.special import iv, kv @@ -38,6 +40,7 @@ class CircAreaSink(Element): def __init__( self, model, xc=0, yc=0, R=0.1, tsandN=[(0, 1)], name="CircAreaSink", label=None ): + self.storeinput(inspect.currentframe()) Element.__init__( self, model, diff --git a/timflow/transient/element.py b/timflow/transient/element.py index f9f33baa..1f09a7ce 100644 --- a/timflow/transient/element.py +++ b/timflow/transient/element.py @@ -9,13 +9,48 @@ def initialize(self): pass """ +import inspect # Used for storing the input from abc import ABC, abstractmethod +from enum import IntEnum import numpy as np from timflow.transient.invlapnumba import invlapcomp +class BCType(IntEnum): + """Map boundary condition types to integers for use in numba.""" + + ZERO = 0 + GIVEN = 1 + VARIABLE = 2 + + @classmethod + def from_str(cls, char: str | bytes) -> "BCType": + """Maps boundary condition string flags to their corresponding integer.""" + # Convert bytes to string if needed (handles b'g', b'v', etc.) + if isinstance(char, bytes): + char = char.decode("utf-8") + + mapping = { + "z": cls.ZERO, + "g": cls.GIVEN, + "v": cls.VARIABLE, + } + return mapping[char.lower()] + + +class ElementType(IntEnum): + """Map element types to integers for use in numba.""" + + LINESINK = 1 + LINESINKSTRING = 2 + LINEDOUBLET = 3 + LINEDOUBLETSTRING = 4 + WELL = 5 + WELLSTRING = 6 + + class Element(ABC): def __init__( self, @@ -323,6 +358,26 @@ def headinside(self, t): print("This function not implemented for this element") return + def storeinput(self, frame): + self.inputargs, _, _, self.inputvalues = inspect.getargvalues(frame) + + 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 + def run_after_solve(self): # noqa: B027 """Function to run after a solution is completed. diff --git a/timflow/transient/linedoublet.py b/timflow/transient/linedoublet.py index 414fff5b..29c22345 100644 --- a/timflow/transient/linedoublet.py +++ b/timflow/transient/linedoublet.py @@ -7,6 +7,7 @@ LeakyLineDoublet(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 @@ -242,6 +243,7 @@ def __init__( label=None, addtomodel=True, ): + self.storeinput(inspect.currentframe()) super().__init__( model, x1=x1, @@ -297,6 +299,7 @@ class LeakyWallString(Element, LeakyWallEquation): def __init__( self, model, xy=[(-1, 0), (1, 0)], res="imp", order=0, layers=0, label=None ): + self.storeinput(inspect.currentframe()) Element.__init__( self, model, diff --git a/timflow/transient/linesink.py b/timflow/transient/linesink.py index 98bdfc5a..96f1e2d3 100644 --- a/timflow/transient/linesink.py +++ b/timflow/transient/linesink.py @@ -8,19 +8,21 @@ River(ml, x1=-10, y1=0, x2=10, y2=0, tsandbc=[(0, 5)], layers=0) """ +import inspect # Used for storing the input import warnings import matplotlib.pyplot as plt import numpy as np from timflow.bessel import besselnumba -from timflow.transient.element import Element +from timflow.transient.element import BCType, Element, ElementType from timflow.transient.equation import ( HeadEquation, HeadEquationNores, MscreenDitchEquation, MscreenEquation, ) +from timflow.transient.parallel.dtypes import ElementTuple, line_element_dtype class LineSinkBase(Element): @@ -200,6 +202,28 @@ def plot(self, ax=None, layer=None): if layer is None or np.isin(self.layers, layer).any(): ax.plot([self.x1, self.x2], [self.y1, self.y2], "k") + def to_numba_tuple(self): + """Write element data to a numba tuple for use in parallel computations.""" + struc_array = np.empty(1, dtype=line_element_dtype) + struc_array["etype"] = ElementType.LINESINK + struc_array["bctype"] = BCType.from_str(self.type) + struc_array["aq_id"] = 0 + struc_array["nparam"] = self.nparam + struc_array["z1"] = self.z1 + struc_array["z2"] = self.z2 + struc_array["L"] = self.L + struc_array["order"] = self.order + struc_array["rzero"] = self.rzero + struc_array["p0"] = 0 + struc_array["p1"] = self.nparam + return ElementTuple( + meta=struc_array, + layers=self.layers, + layer_ptr=np.array([0, len(self.layers)], dtype=np.int64), + term2=self.term2, + parameters=self.parameters, + ) + class LineSink(LineSinkBase): """LineSink with non-zero and potentially variable discharge through time. @@ -243,6 +267,7 @@ def __init__( #'''HeadLineSink that remains zero and constant through time''' # def __init__(self, model, x1=-1, y1=0, x2=1, y2=0, res=0.0, wh='H', \ # layers=0, label=None, addtomodel=True): +# self.storeinput(inspect.currentframe()) # LineSinkBase.__init__(self, model, x1=x1, y1=y1, x2=x2, y2=y2, \ # tsandbc=[(0, 0)], res=res, wh=wh, layers=layers,\ # type='z', name='ZeroHeadLineSink', label=label, \ @@ -502,6 +527,49 @@ def discharge_list(self, t, derivative=0): rv[i, :, :] = self.lslist[i].discharge(t, derivative=derivative) return rv + def to_numba_tuple(self): + """Write element data to a numba tuple for use in parallel computations.""" + struc_array = np.empty(self.nls, dtype=line_element_dtype) + _, naq, nint, npint = self.lslist[0].term2.shape + term2 = np.empty( + (self.nparam, naq, nint, npint), + dtype=np.complex128, + ) + parameters = np.empty( + (self.model.ngvbc, self.nparam, self.model.npval), + dtype=np.complex128, + ) + layers = [] + layer_ptr = np.empty(self.nls + 1, dtype=np.int64) + layer_ptr[0] = 0 + p = 0 + for i, ils in enumerate(self.lslist): + n = ils.nparam + struc_array[i]["etype"] = ElementType.LINESINK + struc_array[i]["bctype"] = BCType.from_str(ils.type) + struc_array[i]["aq_id"] = 0 + struc_array[i]["nparam"] = ils.nparam + struc_array[i]["z1"] = ils.z1 + struc_array[i]["z2"] = ils.z2 + struc_array[i]["L"] = ils.L + struc_array[i]["order"] = 0 + struc_array[i]["rzero"] = ils.rzero + struc_array[i]["p0"] = p + struc_array[i]["p1"] = p + n + term2[p : p + n, :, :, :] = ils.term2 + parameters[:, p : p + n, :] = ils.parameters + layers.extend(ils.layers.tolist()) + layer_ptr[i + 1] = len(layers) + p += n + + return ElementTuple( + meta=struc_array, + layers=np.asarray(layers, dtype=np.int64), + layer_ptr=layer_ptr, + term2=term2, + parameters=parameters, + ) + class RiverString(LineSinkStringBase, HeadEquation): r"""String of head-specified line-sinks with optional width and resistance. @@ -631,6 +699,7 @@ def __init__( ): """Initialize a transient multi-screen line-sink.""" # assert len(layers) > 1, "number of layers must be at least 2" + self.storeinput(inspect.currentframe()) super().__init__( model, x1=x1, @@ -814,6 +883,7 @@ def __init__( Astorage=None, label=None, ): + self.storeinput(inspect.currentframe()) super().__init__( model, tsandbc=tsandQ, diff --git a/timflow/transient/model.py b/timflow/transient/model.py index acb806c1..524bd4a5 100644 --- a/timflow/transient/model.py +++ b/timflow/transient/model.py @@ -10,7 +10,8 @@ ml.solve() """ -import multiprocessing as mp +import inspect # Used for storing the input +import os from warnings import warn import numpy as np @@ -23,26 +24,11 @@ invlap, invlapcomp, ) +from timflow.transient.parallel.dtypes import ModelTuple from timflow.transient.plots import PlotTransient from timflow.version import check_tqdm_parallel -def _compute_head_mp(args): - """Helper function for parallel computation of head_array.""" - model, xi, yi, t, layers, i = args - return i, model.head(xi, yi, t, layers) - - -def _compute_velocity_mp(args): - """Helper function for parallel computation of velocity_array.""" - model, xi, yi, zi, t, i = args - try: - vv = model.velocomp(xi, yi, zi, t) - except (ZeroDivisionError, ValueError): - vv = np.full((3,), np.nan) - return i, vv - - class Model: def __init__( self, @@ -96,7 +82,8 @@ def __init__( model3d, ) self.compute_laplace_parameters() - self.name = "Model" + self.name = "TimModel" + self.modelname = "ml" # Used for writing out input self.model_type = "transient" # Model type for plotting and other purposes self.steady = steady if self.steady is not None: @@ -461,69 +448,6 @@ def velocomp(self, x, y, z, t, aq=None, layer_ltype=None): return velo - def velocity_array(self, x, y, z, t, show_progress=True, parallel=False): - """Compute velocity for an array of points. - - Parameters - ---------- - x : 1d-array - x values - y : 1d-array - y values - z : 1d-array - z values - t : float - time at which velocity computed - show_progress : bool, optional - if `True`, shows progress bar when computing velocity grid, by default `True` - parallel : bool or int, optional - if `True`, computes velocity grid in parallel using multiprocessing, - by default `False`. If an integer is provided, it specifies the number of - processes to use. - - Returns - ------- - velocity : array - velocity vector (vx, vy, vz) at each point in grid, - size (3, len(x)) - """ - parallel, thread_map, tqdm = check_tqdm_parallel(parallel) - - x = np.atleast_1d(x) - y = np.atleast_1d(y) - z = np.atleast_1d(z) - npts = len(x) - v = np.empty((3, npts)) - if not parallel: - for i in ( - tqdm(range(npts), desc="velocity array", disable=not show_progress) - if tqdm - else range(npts) - ): - try: - vv = self.velocomp(x[i], y[i], z[i], t) - except ZeroDivisionError: - vv = np.full((3,), np.nan) - v[:, i] = vv - else: - nproc = mp.cpu_count() if parallel is True else int(parallel) - chunksize = max(1, npts // (4 * nproc)) if nproc > 0 else 1 - tasks = [(self, x[i], y[i], z[i], t, i) for i in range(npts)] - results = thread_map( - _compute_velocity_mp, - tasks, - total=npts, - desc="velocity array", - disable=not show_progress, - tqdm_class=tqdm, - max_workers=nproc, - chunksize=chunksize, - ) - for i, result in results: - v[:, i] = result - - return v - def velocity_grid(self, xg, yg, zg, t, show_progress=True, parallel=False): """Compute velocity grid. @@ -539,30 +463,58 @@ def velocity_grid(self, xg, yg, zg, t, show_progress=True, parallel=False): time for which grid is returned show_progress : bool, optional if `True`, shows progress bar when computing velocity grid, by default `True` - parallel : bool or int, optional - if `True`, computes velocity grid in parallel using multiprocessing, - by default `False`. If an integer is provided, it specifies the number of - processes to use. + parallel : bool, optional + if `True`, computes velocity grid in parallel using multi threading, + by default `False` Returns ------- velocity : array - velocity vector (vx, vy, vz) at each point in grid, + velocity vector (vz, vy, vx) at each point in grid, size (3, len(z), len(y), len(x)) """ + parallel, thread_map, tqdm = check_tqdm_parallel(parallel) + + def compute(kij): + k, i, j = kij + try: + vv = self.velocomp(xg[j], yg[i], zg[k], t) + except ZeroDivisionError: + vv = np.full((3,), np.nan) + return k, i, j, vv + xg = np.atleast_1d(xg) yg = np.atleast_1d(yg) zg = np.atleast_1d(zg) - x, y, z = np.meshgrid(xg, yg, zg) - v = self.velocity_array( - x.ravel(), - y.ravel(), - z.ravel(), - t, - show_progress=show_progress, - parallel=parallel, - ) - return v.reshape((3, len(yg), len(xg), len(zg))).transpose((0, 3, 1, 2)) + nz, ny, nx = len(zg), len(yg), len(xg) + v = np.empty((3, nz, ny, nx)) + if not parallel: + for k in range(nz): + if show_progress: + print(".", end="", flush=True) + for i in range(ny): + for j in range(nx): + try: + vv = self.velocomp(xg[j], yg[i], zg[k], t) + except ZeroDivisionError: + vv = np.full((3,), np.nan) + v[:, k, i, j] = vv + if show_progress: + print("", flush=True) + else: + results = thread_map( + compute, + [(k, i, j) for k in range(nz) for i in range(ny) for j in range(nx)], + total=nz * nx * ny, + desc="velocity grid", + disable=not show_progress, + tqdm_class=tqdm, + max_workers=os.cpu_count(), + ) + for k, i, j, result in results: + v[:, k, i, j] = result + + return v def velo_one(self, x, y, z, t, aq=None, layer_ltype=[0, 0]): # implemented for one layer and one time @@ -642,78 +594,8 @@ def disvecalongline(self, x, y, t, layers=None): qx[:, :, i], qy[:, :, i] = self.disvec(xg[i], yg[i], t, layers) return qx, qy - def head_array(self, x, y, t, layers=None, show_progress=False, parallel=False): - """Head for array of points. - - Parameters - ---------- - x : 1D array or list - x values of points - y : 1D array or list - y values of points - t : float or 1D array or list - times for which grid is returned - layers : integer, list or array, optional - layers for which grid is returned - show_progress : bool - show computation progress, by printing dots per row or with tqdm progressbar - when parallel is True. Default is False. - parallel : bool or int, optional - if `True`, computes head_array in parallel using multiprocessing, - by default `False`. If an integer is provided, it specifies the number of - processes to use. - - Returns - ------- - h : array size `nlayers, ntimes, npoints` - """ - parallel, process_map, tqdm = check_tqdm_parallel(parallel) - x = np.atleast_1d(x) - y = np.atleast_1d(y) - t = np.atleast_1d(t) - npts = len(x) - assert npts == len(y), "x and y must have the same length" - ntimes = len(t) - if layers is None: - nlayers = self.aq.find_aquifer_data(x[0], y[0]).naq - else: - nlayers = len(np.atleast_1d(layers)) - h = np.empty((nlayers, ntimes, npts)) - if not parallel: - for i in ( - tqdm(range(npts), disable=not show_progress, desc="head array") - if tqdm - else range(npts) - ): - h[:, :, i] = self.head(x[i], y[i], t, layers) - else: - nproc = mp.cpu_count() if parallel is True else int(parallel) - chunksize = max(1, npts // (4 * nproc)) if nproc > 0 else 1 - tasks = [(self, x[i], y[i], t, layers, i) for i in range(npts)] - results = process_map( - _compute_head_mp, - tasks, - total=npts, - desc="head array", - disable=not show_progress, - tqdm_class=tqdm, - max_workers=nproc, - chunksize=chunksize, - ) - - for i, result in results: - h[:, :, i] = result - return h - def headgrid( - self, - xg, - yg, - t, - layers=None, - printrow=False, - show_progress=False, - parallel=False, + self, xg, yg, t, layers=None, printrow=False, show_progress=False, parallel=False ): """Grid of heads. @@ -730,10 +612,9 @@ def headgrid( show_progress : bool show computation progress, by printing dots per row or with tqdm progressbar when parallel is True. Default is False. - parallel : bool or int, optional + parallel : bool, optional if `True`, computes headgrid in parallel using multithreading, - by default `False`. If an integer is provided, it specifies the number of - processes to use. + by default `False` printrow : bool, optional .. deprecated:: 0.2.0 @@ -755,18 +636,45 @@ def headgrid( ) show_progress = printrow + parallel, thread_map, tqdm = check_tqdm_parallel(parallel) + + xg = np.atleast_1d(xg) + yg = np.atleast_1d(yg) + t = np.atleast_1d(t) nx = len(xg) ny = len(yg) - x, y = np.meshgrid(xg, yg) - h = self.head_array( - x.ravel(), - y.ravel(), - t, - layers=layers, - show_progress=show_progress, - parallel=parallel, - ) - return h.reshape((h.shape[0], h.shape[1], ny, nx)) + ntimes = len(t) + if layers is None: + nlayers = self.aq.find_aquifer_data(xg[0], yg[0]).naq + else: + nlayers = len(np.atleast_1d(layers)) + t = np.atleast_1d(t) + h = np.empty((nlayers, ntimes, ny, nx)) + if not parallel: + for j in range(ny): + if show_progress: + print(".", end="", flush=True) + for i in range(nx): + h[:, :, j, i] = self.head(xg[i], yg[j], t, layers) + else: + + def compute(ij): + i, j = ij + return i, j, self.head(xg[i], yg[j], t, layers) + + results = thread_map( + compute, + [(i, j) for j in range(ny) for i in range(nx)], + total=nx * ny, + desc="headgrid", + disable=not show_progress, + tqdm_class=tqdm, + max_workers=os.cpu_count(), + ) + + for i, j, result in results: + h[:, :, j, i] = result + return h def headgrid2( self, @@ -797,10 +705,9 @@ def headgrid2( show_progress : bool show computation progress, by printing dots per row or with tqdm progressbar when parallel is True. Default is False. - parallel : bool or int, optional - if `True`, computes headgrid in parallel using multiprocessing, - by default `False`. If an integer is provided, it specifies the number of - processes to use. + parallel : bool, optional + if `True`, computes headgrid in parallel using multi threading, + by default `False` printrow : boolean, optional .. deprecated:: 0.2.0 @@ -851,7 +758,7 @@ def disvecgrid( show computation progress, by printing dots per row or with tqdm progressbar when parallel is True. Default is True. parallel : bool, optional - if `True`, computes discharge vector grid in parallel using multiprocessing, + if `True`, computes discharge vector grid in parallel using multi threading, by default `False` Returns @@ -895,6 +802,7 @@ def compute(ij): desc="disvecgrid", disable=not show_progress, tqdm_class=tqdm, + max_workers=os.cpu_count(), ) for i, j, result in results: qx[:, :, j, i], qy[:, :, j, i] = result @@ -970,6 +878,35 @@ def solve(self, printmat=0, sendback=0, silent=False): return sol return + def storeinput(self, frame): + self.inputargs, _, _, self.inputvalues = inspect.getargvalues(frame) + + def write(self): + rv = self.modelname + " = " + self.name + "(\n" + for key in self.inputargs[1:]: # The first argument (self) is 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 + + def writemodel(self, fname): + self.initialize() # So that model can be written without solving first + f = open(fname, "w") + f.write("from timflow.transient import *\n") + f.write(self.write()) + for e in self.elementlist: + f.write(e.write()) + f.close() + def aquifer_summary(self): """Return DataFrame with summary of aquifer(s) parameters in model. @@ -990,6 +927,22 @@ def aquifer_summary(self): aqs[name] = iaq.summary() return pd.concat(aqs, axis=0) + def to_numba_tuple(self): + """Write model data to a numba tuple for use in parallel computations.""" + return ModelTuple( + nint=self.nint, + npint=self.npint, + npval=self.npval, + ngvbc=self.ngvbc, + M=self.M, + tstart=self.tstart, + p=self.p, + tintervals=self.tintervals, + enumber=self.enumber, + etstart=self.etstart, + ebc=self.ebc, + ) + class ModelMaq(Model): """Create model specifying a multi-aquifer sequence of aquifer-leakylayer-etc. @@ -1063,6 +1016,7 @@ def __init__( M=10, steady=None, ): + self.storeinput(inspect.currentframe()) if phreatictop is None: phreatictop = False if topboundary[:3] == "phr": @@ -1178,6 +1132,7 @@ def __init__( steady=None, ): """Z must have the length of the number of layers + 1.""" + self.storeinput(inspect.currentframe()) if phreatictop is None: phreatictop = False if topboundary[:3] == "phr": @@ -1277,7 +1232,8 @@ def __init__( self.M = M self.aq = SimpleAquifer(naq) self.compute_laplace_parameters() - self.name = "ModelXsection" + self.name = "TimModel" + self.modelname = "ml" # Used for writing out input self.steady = steady if self.steady is not None: self.steady.solve() diff --git a/timflow/transient/parallel/__init__.py b/timflow/transient/parallel/__init__.py new file mode 100644 index 00000000..f59780d6 --- /dev/null +++ b/timflow/transient/parallel/__init__.py @@ -0,0 +1,4 @@ +"""Optimized numba code for parallel computations.""" + +# ruff : noqa: F401 +from timflow.transient.parallel.core import headgrid diff --git a/timflow/transient/parallel/core.py b/timflow/transient/parallel/core.py new file mode 100644 index 00000000..28f3bf37 --- /dev/null +++ b/timflow/transient/parallel/core.py @@ -0,0 +1,579 @@ +# %% +import numba as nb +import numpy as np + +from timflow.bessel.besselnumba import besselk0, bessellsv2 +from timflow.transient.element import ElementType +from timflow.transient.invlapnumba import invlapcomp +from timflow.transient.parallel.dtypes import ( + AquiferTuple, + ModelTuple, + line_element_dtype, + well_element_dtype, +) + +FASTMATH = False + +# %% + + +@nb.njit(nogil=True, cache=True, fastmath=FASTMATH) +def potinf_linesink( + x: float, + y: float, + nint: int, + naq: int, + lab2: np.ndarray, + meta: np.ndarray, + term2: np.ndarray, + out4d: np.ndarray, +) -> None: + """Evaluate analytic line element contributions into thread memory. + + Parameters + ---------- + x : float + X-coordinate of evaluation point. + y : float + Y-coordinate of evaluation point. + nint : int + Number of time intervals. + naq : int + Number of aquifer layers. + lab2 : np.ndarray + Array of eigenvalues of the aquifer layers. + meta : np.ndarray + Structured array containing line sink element metadata. + term2 : np.ndarray + Term 2 values for the element. + out4d : np.ndarray + Output array to store the computed potential contributions. + + Returns + ------- + None + """ + nls = meta.shape[0] + for i in range(nls): + m = meta[i] + p0, p1 = m["p0"], m["p1"] + z1, z2, L, order, rzero = m["z1"], m["z2"], m["L"], m["order"], m["rzero"] + + for a in range(naq): + for j in range(nint): + out4d[p0:p1, a, j, :] = ( + term2[p0:p1, a, j, :] + * bessellsv2(x, y, z1, z2, lab2[a, j, :], order, rzero) + / L + ) + + +@nb.njit(nogil=True, cache=True, fastmath=FASTMATH) +def potinf_well( + x: float, + y: float, + nint: int, + npint: int, + naq: int, + lab2: np.ndarray, + meta: np.ndarray, + term2: np.ndarray, + out4d: np.ndarray, +) -> None: + """Evaluate analytic well element contributions into thread memory. + + Parameters + ---------- + x : float + X-coordinate of evaluation point. + y : float + Y-coordinate of evaluation point. + nint : int + Number of time intervals. + npint : int + Number of Laplace parameters per interval. + naq : int + Number of aquifer layers. + lab2 : np.ndarray + Array of eigenvalues of the aquifer layers. + meta : np.ndarray + Structured array containing well element metadata. + term2 : np.ndarray + Term 2 values for the element. + out4d : np.ndarray + Output array to store the computed potential contributions. + + Returns + ------- + None + """ + nwells = meta.shape[0] + for i in range(nwells): + m = meta[i] + p0, p1 = m["p0"], m["p1"] + xw, yw, rw, rzero = m["xw"], m["yw"], m["rw"], m["rzero"] + + dx = x - xw + dy = y - yw + r = np.sqrt(dx**2 + dy**2) + if r < rw: + dx = rw + dy = 0.0 + r = rw + + for a in range(naq): + for j in range(nint): + if r / abs(lab2[a, j, 0]) < rzero: + for k in range(npint): + out4d[p0:p1, a, j, k] = term2[p0:p1, a, j, k] * besselk0( + dx, dy, lab2[a, j, k] + ) + else: + # Explicitly zero out + out4d[p0:p1, a, j, :] = 0.0 + 0.0j + + +@nb.njit(nogil=True, cache=True, fastmath=FASTMATH) +def pot_linesink( + x: float, + y: float, + nint: int, + npint: int, + naq: int, + ngvbc: int, + lab2: np.ndarray, + meta: np.ndarray, + term2: np.ndarray, + params: np.ndarray, + pot_view: np.ndarray, +) -> None: + """Evaluate and accumulate analytic line elements into thread memory. + + Parameters + ---------- + x : float + X-coordinate of evaluation point. + y : float + Y-coordinate of evaluation point. + nint : int + Number of time intervals. + npint : int + Number of Laplace parameters per interval. + naq : int + Number of aquifer layers. + ngvbc : int + Number of boundary conditions or groups. + lab2 : np.ndarray + Array of eigenvalues of the aquifer layers. + meta : np.ndarray + Structured array containing line sink element metadata. + term2 : np.ndarray + Term 2 values for the element. + params : np.ndarray + Array of element parameters. + pot_view : np.ndarray + Array view to accumulate the computed potentials. + + Returns + ------- + None + """ + nls = meta.shape[0] + for i in range(nls): + m = meta[i] + p0, p1 = m["p0"], m["p1"] + z1, z2, L, order, rzero = m["z1"], m["z2"], m["L"], m["order"], m["rzero"] + bessel_vals = np.empty((order + 1, npint), dtype=np.complex128) + for a in range(naq): + for j in range(nint): + bessellsv2(x, y, z1, z2, lab2[a, j, :], order, rzero, out=bessel_vals) + bessel_vals /= L + for g in range(ngvbc): + for ip in range(p0, p1): + order_idx = ip - p0 + params_slice = params[g, ip, :] + for k in range(npint): + v = j * npint + k + t_val = term2[ip, a, j, k] * bessel_vals[order_idx, k] + pot_view[g, a, v] += params_slice[v] * t_val + + +@nb.njit(nogil=True, cache=True, fastmath=FASTMATH) +def pot_well( + x: float, + y: float, + nint: int, + npint: int, + naq: int, + ngvbc: int, + lab2: np.ndarray, + meta: np.ndarray, + term2: np.ndarray, + params: np.ndarray, + pot_view: np.ndarray, +) -> None: + """Evaluate and accumulate analytic well elements into thread memory. + + Parameters + ---------- + x : float + X-coordinate of evaluation point. + y : float + Y-coordinate of evaluation point. + nint : int + Number of time intervals. + npint : int + Number of Laplace parameters per interval. + naq : int + Number of aquifer layers. + ngvbc : int + Number of boundary conditions or groups. + lab2 : np.ndarray + Array of eigenvalues of the aquifer layers. + meta : np.ndarray + Structured array containing well element metadata. + term2 : np.ndarray + Term 2 values for the element. + params : np.ndarray + Array of element parameters. + pot_view : np.ndarray + Array view to accumulate the computed potentials. + + Returns + ------- + None + """ + nwells = meta.shape[0] + for i in range(nwells): + m = meta[i] + p0, p1 = m["p0"], m["p1"] + xw, yw, rw, rzero = m["xw"], m["yw"], m["rw"], m["rzero"] + + dx = x - xw + dy = y - yw + r = np.sqrt(dx**2 + dy**2) + if r < rw: + dx = rw + dy = 0.0 + r = rw + + for a in range(naq): + for j in range(nint): + if r / abs(lab2[a, j, 0]) < rzero: + for k in range(npint): + b_val = besselk0(dx, dy, lab2[a, j, k]) + v = j * npint + k + for ip in range(p0, p1): + t_val = term2[ip, a, j, k] * b_val + for g in range(ngvbc): + pot_view[g, a, v] += params[g, ip, v] * t_val + + +def get_element_data(ml): + """Prepare element data for numba processing. + + Parameters + ---------- + ml : Model3D or ModelMaq + The model object containing the elements. + + Returns + ------- + dict + Dictionary mapping ElementType keys to lists of element data tuples. + """ + edict = {i: [] for i in ElementType} + for e in ml.elementlist: + if not hasattr(e, "to_numba_tuple"): + raise NotImplementedError( + f"Element of type {type(e)} does not have a to_numba_tuple() method." + ) + edata = e.to_numba_tuple() + edict[edata.meta["etype"][0]].append(edata) + + return edict + + +@nb.njit(nogil=True, cache=True) +def elements_to_numba_arrays(etuples, element_dtype, mtuple, aqtuple): + """Convert element data tuples to a structured numpy array. + + Parameters + ---------- + etuples : list of tuple + List of element data tuples. + element_dtype : np.dtype + The structured array schema. + mtuple : ModelTuple + Model configuration parameters. + aqtuple : AquiferTuple + Aquifer parameters. + + Returns + ------- + meta : np.ndarray + Structured numpy array containing element metadata. + term2 : np.ndarray + 4D array of term2 values for each element. + params : np.ndarray + 3D array of parameters for each element. + """ + # dimensions + nint, npint, npval, ngvbc = mtuple.nint, mtuple.npint, mtuple.npval, mtuple.ngvbc + naq = aqtuple.naq + + # no. of elements and parameters + n_params = 0 + n_elements = 0 + for e in etuples: + n_elements += e.meta.shape[0] + n_params += np.sum(e.meta["nparam"]) + + # build metadata, term2 and param arrays + meta = np.empty(n_elements, dtype=element_dtype) + term2 = np.empty((n_params, naq, nint, npint), dtype=np.complex128) + params = np.empty((ngvbc, n_params, npval), dtype=np.complex128) + + ie = 0 + ip = 0 + for e in etuples: + for i in range(e.meta.shape[0]): + imeta = e.meta[i] + meta[ie] = imeta + p0, p1 = int(imeta["p0"]), int(imeta["p1"]) + n = p1 - p0 + meta[ie]["p0"] = ip + meta[ie]["p1"] = ip + n + term2[ip : ip + n] = e.term2[p0:p1] + params[:, ip : ip + n, :] = e.parameters[:, p0:p1, :] + ip += n + ie += 1 + + return meta, term2, params + + +def prepare_element_data(ml, mtuple, aqtuple): + """Prepare and partition elements into line and well arrays. + + Parameters + ---------- + ml : Model3D or ModelMaq + The model object containing elements. + mtuple : ModelTuple + Model configuration parameters. + aqtuple : AquiferTuple + Aquifer parameters. + + Returns + ------- + line_meta : np.ndarray + Structured metadata array for line elements. + line_term2 : np.ndarray + 4D term2 values for line elements. + line_params : np.ndarray + 3D parameter values for line elements. + well_meta : np.ndarray + Structured metadata array for well elements. + well_term2 : np.ndarray + 4D term2 values for well elements. + well_params : np.ndarray + 3D parameter values for well elements. + """ + edict = get_element_data(ml) + + # 1. Gather and build line arrays + ltuples = tuple(edict[ElementType.LINESINK] + edict[ElementType.LINESINKSTRING]) + line_meta, line_term2, line_params = elements_to_numba_arrays( + ltuples, line_element_dtype, mtuple, aqtuple + ) + + # 2. Gather and build well arrays + wtuples = tuple(edict[ElementType.WELL] + edict[ElementType.WELLSTRING]) + well_meta, well_term2, well_params = elements_to_numba_arrays( + wtuples, well_element_dtype, mtuple, aqtuple + ) + + return line_meta, line_term2, line_params, well_meta, well_term2, well_params + + +@nb.njit(nogil=True, cache=True, parallel=True) +def _headgrid_numba( + x: np.ndarray, + y: np.ndarray, + t: np.ndarray, + mtuple: ModelTuple, + aqtuple: AquiferTuple, + line_meta: np.ndarray, + line_term2: np.ndarray, + line_params: np.ndarray, + well_meta: np.ndarray, + well_term2: np.ndarray, + well_params: np.ndarray, +): + """Compute transient heads on a grid using parallelized numba. + + Parameters + ---------- + x : np.ndarray + X-coordinates of evaluation points. + y : np.ndarray + Y-coordinates of evaluation points. + t : np.ndarray + Times at which to evaluate heads. + mtuple : ModelTuple + Model configuration parameters. + aqtuple : AquiferTuple + Aquifer parameters. + line_meta : np.ndarray + Structured metadata array for line elements. + line_term2 : np.ndarray + 4D term2 values for line elements. + line_params : np.ndarray + 3D parameter values for line elements. + well_meta : np.ndarray + Structured metadata array for well elements. + well_term2 : np.ndarray + 4D term2 values for well elements. + well_params : np.ndarray + 3D parameter values for well elements. + + Returns + ------- + out : np.ndarray + Computed transient heads array with shape (naq, ntimes, npts). + """ + # get model and aquifer data + nint, npint, npval, ngvbc = (mtuple.nint, mtuple.npint, mtuple.npval, mtuple.ngvbc) + naq = aqtuple.naq + lab2 = aqtuple.lab2 + npts = len(x) + + # create output array and compute delta time + out = np.empty((naq, len(t), npts), dtype=np.float64) + time = np.atleast_1d(t) - mtuple.tstart + + # allocate computation arrays (with dimension per thread) + num_threads = nb.config.NUMBA_NUM_THREADS + pot_scratch = np.empty((num_threads, ngvbc, naq, npval), dtype=np.complex128) + pot_tx_scratch = np.empty((num_threads, ngvbc, naq, npval), dtype=np.complex128) + + # ensure parameter and eigvec arrays are contiguous + eigvec_contiguous = np.ascontiguousarray(aqtuple.eigvec) + + # parallel loop + for i in nb.prange(npts): + xi, yi = x[i], y[i] + tid = nb.get_thread_id() + + # Zero out the reused array views for this iteration + pot_view = pot_scratch[tid] + pot_tx_view = pot_tx_scratch[tid] + pot_view[:, :, :] = 0.0 + 0.0j + pot_tx_view[:, :, :] = 0.0 + 0.0j + + # Compute potential line sinks + if line_meta.shape[0] > 0: + pot_linesink( + xi, + yi, + nint, + npint, + naq, + ngvbc, + lab2, + line_meta, + line_term2, + line_params, + pot_view, + ) + + # Compute potential wells + if well_meta.shape[0] > 0: + pot_well( + xi, + yi, + nint, + npint, + naq, + ngvbc, + lab2, + well_meta, + well_term2, + well_params, + pot_view, + ) + + # compute potentials + for g in range(ngvbc): + for a_out in range(naq): + for a_in in range(naq): + for v in range(npval): + pot_tx_view[g, a_out, v] += ( + pot_view[g, a_in, v] * eigvec_contiguous[a_out, a_in, v] + ) + + # inverse laplace transform + rv_inv = invlapcomp( + time, + pot_tx_view, + npint, + mtuple.M, + mtuple.tintervals, + mtuple.enumber, + mtuple.etstart, + mtuple.ebc, + naq, + ) + out[:, :, i] = rv_inv / aqtuple.Tcol + + return out + + +def headgrid( + ml, + x: np.ndarray, + y: np.ndarray, + t: np.ndarray, +): + """Compute transient heads for specified coordinates and times. + + Parameters + ---------- + ml : Model3D or ModelMaq + The model object. + x : np.ndarray or float + X-coordinates of evaluation points. + y : np.ndarray or float + Y-coordinates of evaluation points. + t : np.ndarray or float + Times at which to evaluate heads. + + Returns + ------- + np.ndarray + Computed transient heads array with shape (naq, ntimes, npts). + """ + mtuple = ml.to_numba_tuple() + aqtuple = ml.aq.to_numba_tuple() + + lmeta, lterm2, lparams, wmeta, wterm2, wparams = prepare_element_data( + ml, mtuple, aqtuple + ) + + x = np.atleast_1d(x) + y = np.atleast_1d(y) + t = np.atleast_1d(t) + + return _headgrid_numba( + x, + y, + t, + mtuple, + aqtuple, + lmeta, + lterm2, + lparams, + wmeta, + wterm2, + wparams, + ) diff --git a/timflow/transient/parallel/dtypes.py b/timflow/transient/parallel/dtypes.py new file mode 100644 index 00000000..a6951dcf --- /dev/null +++ b/timflow/transient/parallel/dtypes.py @@ -0,0 +1,72 @@ +# %% +from typing import NamedTuple + +import numpy as np + +# %% +line_element_dtype = np.dtype( + [ + ("etype", np.int64), + ("bctype", np.int8), + ("aq_id", np.int64), + ("nparam", np.int64), + ("z1", np.complex128), + ("z2", np.complex128), + ("L", np.float64), + ("order", np.int64), + ("rzero", np.float64), + ("p0", np.int64), + ("p1", np.int64), + ] +) + +well_element_dtype = np.dtype( + [ + ("etype", np.int64), + ("bctype", np.int8), + ("aq_id", np.int64), + ("nparam", np.int64), + ("xw", np.float64), + ("yw", np.float64), + ("rw", np.float64), + ("rzero", np.float64), + ("p0", np.int64), + ("p1", np.int64), + ] +) + + +class ModelTuple(NamedTuple): + """Data structure for timflow Model parameters and arrays.""" + + nint: np.int64 + npint: np.int64 + npval: np.int64 + ngvbc: np.int64 + M: np.int64 + tstart: np.float64 + p: np.ndarray[np.complex128] + tintervals: np.ndarray[np.float64] + enumber: np.ndarray[np.int64] + etstart: np.ndarray[np.float64] + ebc: np.ndarray[np.float64] + + +class AquiferTuple(NamedTuple): + """Data structure for Aquifer parameters and arrays.""" + + naq: np.int64 + lab2: np.ndarray[np.complex128] + lababs: np.ndarray[np.float64] + eigvec: np.ndarray[np.complex128] + Tcol: np.ndarray[np.float64] + + +class ElementTuple(NamedTuple): + """Unified element tuple for all element types (lines, wells, etc.).""" + + meta: np.ndarray + layers: np.ndarray[np.int32] + layer_ptr: np.ndarray[np.int32] + term2: np.ndarray[np.complex128] + parameters: np.ndarray[np.complex128] diff --git a/timflow/transient/plots.py b/timflow/transient/plots.py index 9483ec02..370fffb8 100644 --- a/timflow/transient/plots.py +++ b/timflow/transient/plots.py @@ -229,10 +229,9 @@ def contour( if list of strings: use strings as names in legend return_contours : bool, optional if True, return list of contour sets for each contoured layer - parallel : bool or int, optional - if True, compute headgrid in parallel using multiprocessing, - default is False. If int is provided, it is interpreted as the number of - processes to use. + parallel : bool, optional + if True, compute headgrid in parallel using multiple threads, + default is False show_progress : bool, optional if True, show progress bar when computing headgrid in parallel, default is False. @@ -374,7 +373,7 @@ def quiver_xy( figsize : tuple of 2 values size of figure parallel : bool - if True, compute velocity grid in parallel using multiprocessing, + if True, compute velocity grid in parallel using multiple threads, default is False **kwargs additional keyword arguments passed to ax.quiver() @@ -412,7 +411,7 @@ def quiver_z( figsize : tuple of 2 values size of figure parallel : bool - if True, compute velocity grid in parallel using multiprocessing, + if True, compute velocity grid in parallel using multiple threads, default is False **kwargs additional keyword arguments passed to ax.quiver() diff --git a/timflow/transient/well.py b/timflow/transient/well.py index 8449bb4a..6422a406 100644 --- a/timflow/transient/well.py +++ b/timflow/transient/well.py @@ -8,15 +8,18 @@ Well(ml, xw=100, yw=200, tsandbc=[(0, 1000)], layers=[0]) """ +import inspect # Used for storing the input + import matplotlib.pyplot as plt import numpy as np # from scipy.special import iv # Needed for K1 in Well class, and in CircInhom from scipy.special import kv -from timflow.transient.element import Element +from timflow.transient.element import BCType, Element, ElementType from timflow.transient.equation import HeadEquation, WellBoreStorageEquation from timflow.transient.invlapnumba import invlapcomp +from timflow.transient.parallel.dtypes import ElementTuple, well_element_dtype class WellBase(Element): @@ -210,6 +213,27 @@ def changetrace( message = "reached element of type well: " + str(self) return changed, terminate, xyztnew, message + def to_numba_tuple(self): + """Write element data to a numba tuple for use in parallel computations.""" + struc_array = np.empty(1, dtype=well_element_dtype) + struc_array["etype"] = ElementType.WELL + struc_array["bctype"] = BCType.from_str(self.type) + struc_array["aq_id"] = 0 + struc_array["nparam"] = self.nparam + struc_array["xw"] = self.xw + struc_array["yw"] = self.yw + struc_array["rw"] = self.rw + struc_array["rzero"] = self.rzero + struc_array["p0"] = 0 + struc_array["p1"] = self.nparam + return ElementTuple( + meta=struc_array, + layers=self.layers, + layer_ptr=np.array([0], dtype=np.int32), + term2=self.term2, + parameters=self.parameters, + ) + class DischargeWell(WellBase): r"""Well with a specified discharge for each layer that the well is screened in. @@ -256,6 +280,7 @@ class DischargeWell(WellBase): def __init__( self, model, xw=0, yw=0, tsandQ=[(0, 1)], rw=0.1, res=0, layers=0, label=None ): + self.storeinput(inspect.currentframe()) super().__init__( model, xw, @@ -438,6 +463,7 @@ def __init__( label=None, fp=None, ): + self.storeinput(inspect.currentframe()) super().__init__( model, xw, diff --git a/timflow/version.py b/timflow/version.py index 8a2deb81..deabe242 100644 --- a/timflow/version.py +++ b/timflow/version.py @@ -45,8 +45,8 @@ def check_tqdm_parallel(parallel): ------- parallel : bool Whether parallel processing was requested and can be used. - process_map : function or None - The process_map function from tqdm if parallel processing is available, else None. + thread_map : function or None + The thread_map function from tqdm if parallel processing is available, else None. tqdm : class or None The tqdm class from tqdm if parallel processing is available, else None. """ @@ -54,7 +54,7 @@ def check_tqdm_parallel(parallel): return parallel, None, None try: from tqdm import tqdm - from tqdm.contrib.concurrent import process_map + from tqdm.contrib.concurrent import thread_map except ImportError: warnings.warn( "Parallel requires 'tqdm'. Install 'timflow[parallel]' or 'tqdm' to" @@ -63,6 +63,6 @@ def check_tqdm_parallel(parallel): stacklevel=2, ) parallel = False - process_map = None + thread_map = None tqdm = None - return parallel, process_map, tqdm + return parallel, thread_map, tqdm