diff --git a/docs/steady/03xsections/cross_section_models.ipynb b/docs/steady/03xsections/cross_section_models.ipynb index 61c4a591..29fefe50 100644 --- a/docs/steady/03xsections/cross_section_models.ipynb +++ b/docs/steady/03xsections/cross_section_models.ipynb @@ -206,7 +206,7 @@ "\n", "ml.solve()\n", "\n", - "ml.plots.xsection(xy=[(-100, 0), (100, 0)]);" + "ax = ml.plots.xsection(xy=[(-100, 0), (100, 0)])" ] }, { @@ -314,6 +314,41 @@ "ld1.plot(ax); # plot wall" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plot velocity quiver near the wall." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "x = np.linspace(-5, 5, 30)\n", + "y = np.array([0])\n", + "z = np.array([[4.5, 3.5, 2.5, 1.5, 0.5]]).T\n", + "fig, ax = plt.subplots(1, 1, figsize=(10, 3))\n", + "ml.plots.xsection(xy=[(x[0], 0), (x[-1], 0)], horizontal_axis=\"x\", labels=False, ax=ax)\n", + "\n", + "npts = 31\n", + "ml.plots.tracelines(\n", + " -5 * np.ones(npts),\n", + " np.zeros(npts),\n", + " np.linspace(0, 5, npts),\n", + " ax=ax,\n", + " linewidth=1,\n", + " hstepmax=1.0,\n", + " vstepfrac=0.1,\n", + " orientation=\"ver\",\n", + " silent=True,\n", + ")\n", + "ld1.plot(ax) # plot wall\n", + "ml.plots.quiver_z(x, y, z, ax=ax, normalize=True, scale=30, zorder=10);" + ] + }, { "cell_type": "markdown", "metadata": {}, @@ -506,13 +541,6 @@ "ml.solve()" ] }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [] - }, { "cell_type": "code", "execution_count": null, diff --git a/docs/transient/02examples/horizontal_well.ipynb b/docs/transient/02examples/horizontal_well.ipynb index 19f59830..1c7cbb43 100644 --- a/docs/transient/02examples/horizontal_well.ipynb +++ b/docs/transient/02examples/horizontal_well.ipynb @@ -80,7 +80,6 @@ " t=[10, 20],\n", " layers=range(7),\n", " parallel=True,\n", - " show_progress=True,\n", ")" ] }, diff --git a/pyproject.toml b/pyproject.toml index 633493be..c7bed1e1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,7 +32,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Hydrology", ] requires-python = ">=3.11" -dependencies = ["numpy", "scipy", "numba", "matplotlib", "pandas"] +dependencies = ["numpy", "scipy", "numba", "matplotlib", "pandas", "tqdm"] [project.urls] homepage = "https://github.com/timflow-org/timflow" @@ -41,8 +41,8 @@ documentation = "https://timflow.readthedocs.io/en/latest/" [project.optional-dependencies] lint = ["ruff"] -parallel = ["tqdm"] -optional = ["timflow[parallel]", "lmfit"] +lmfit = ["lmfit"] +optional = ["timflow[lmfit]"] ci = [ "timflow[lint,optional]", "pytest", @@ -61,7 +61,7 @@ docs = [ "myst_nb", "sphinxcontrib-bibtex", "sphinx-autoapi", - "jupyter-cache" + "jupyter-cache", ] dev = ["timflow[docs]"] diff --git a/timflow/plots/plots.py b/timflow/plots/plots.py index 569fda8f..f16a7a4f 100644 --- a/timflow/plots/plots.py +++ b/timflow/plots/plots.py @@ -934,15 +934,20 @@ def vcontour_array( x = np.sqrt((x - x[0]) ** 2 + (y - y[0]) ** 2) else: raise ValueError("horizontal_axis must be 'x', 'y', or 's'") + # find aquifer for z coordinates + if self._ml.name == "ModelXsection": + aq = self._ml.aq.find_aquifer_data(x[0], y[0]) # use aquifer at first coord + else: + aq = self._ml.aq if vinterp: - z = 0.5 * (self._ml.aq.zaqbot + self._ml.aq.zaqtop) - z = np.hstack((self._ml.aq.zaqtop[0], z, self._ml.aq.zaqbot[-1])) + z = 0.5 * (aq.zaqbot + aq.zaqtop) + z = np.hstack((aq.zaqtop[0], z, aq.zaqbot[-1])) arr = np.vstack((arr[0], arr, arr[-1])) else: - z = np.empty(2 * self._ml.aq.naq) - for i in range(self._ml.aq.naq): - z[2 * i] = self._ml.aq.zaqtop[i] - z[2 * i + 1] = self._ml.aq.zaqbot[i] + z = np.empty(2 * aq.naq) + for i in range(aq.naq): + z[2 * i] = aq.zaqtop[i] + z[2 * i + 1] = aq.zaqbot[i] arr = np.repeat(arr, 2, 0) if ax is None: _, ax = plt.subplots(figsize=figsize) diff --git a/timflow/steady/inhomogeneity1d.py b/timflow/steady/inhomogeneity1d.py index 6baa2e0c..e15687dc 100644 --- a/timflow/steady/inhomogeneity1d.py +++ b/timflow/steady/inhomogeneity1d.py @@ -186,7 +186,7 @@ def plot( x1 = kwargs.pop("x1") if np.isfinite(self.x1): x1 = max(x1, self.x1) - else: + elif not np.isfinite(x1): x1 = self.x2 - 100.0 elif np.isfinite(self.x1): x1 = self.x1 @@ -197,7 +197,7 @@ def plot( x2 = kwargs.pop("x2") if np.isfinite(self.x2): x2 = min(x2, self.x2) - else: + elif not np.isfinite(x2): x2 = self.x1 + 100.0 elif np.isfinite(self.x2): x2 = self.x2 diff --git a/timflow/steady/model.py b/timflow/steady/model.py index 30caf3f4..e8e02062 100644 --- a/timflow/steady/model.py +++ b/timflow/steady/model.py @@ -12,34 +12,47 @@ import multiprocessing as mp import warnings +from concurrent.futures import ProcessPoolExecutor +from itertools import repeat import numpy as np import pandas as pd from scipy.integrate import quad_vec +from tqdm import tqdm from timflow.steady.aquifer import Aquifer, SimpleAquifer from timflow.steady.aquifer_parameters import param_3d, param_maq from timflow.steady.constant import ConstantStar from timflow.steady.plots import PlotSteady -from timflow.version import check_tqdm_parallel __all__ = ["Model", "ModelMaq", "Model3D", "ModelXsection"] -def _compute_head_mp(args): +_WORKER_STATE = {"model": None} + + +def _init_worker(model): + """Initialize a single model instance per worker process.""" + _WORKER_STATE["model"] = model + + +def _compute_head_mp(xi, yi, layers): """Helper function for parallel computation of head_array.""" - model, xi, yi, layers, i = args - return i, model.head(xi, yi, layers=layers) + return _WORKER_STATE["model"].head(xi, yi, layers=layers) -def _compute_velocity_mp(args): +def _compute_velocity_mp(xi, yi, zi): """Helper function for parallel computation of velocity_array.""" - model, xi, yi, zi, i = args try: - vv = model.velocomp(xi, yi, zi) + vv = _WORKER_STATE["model"].velocomp(xi, yi, zi) except (ZeroDivisionError, ValueError): vv = np.full((3,), np.nan) - return i, vv + return vv + + +def _compute_disvec_mp(xi, yi, layers): + """Helper function for parallel computation of disvec_array.""" + return _WORKER_STATE["model"].disvec(xi, yi, layers=layers) class Model: @@ -115,7 +128,7 @@ def potential(self, x, y, aq=None): rv += aq.constantstar.potstar return rv - def disvec(self, x, y, aq=None): + def disvec(self, x, y, layers=None, aq=None): """Discharge vector at `x`, `y`. Returns @@ -130,7 +143,10 @@ def disvec(self, x, y, aq=None): for e in aq.elementlist: rv += e.disvec(x, y, aq) rv = np.sum(rv[:, np.newaxis, :] * aq.eigvec, 2) - return rv + if layers is None: + return rv + else: + return rv[:, np.atleast_1d(layers)] def normflux(self, x, y, theta): """Flux at point x, y in direction of angle theta. @@ -275,7 +291,7 @@ 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): + def head_array(self, x, y, layers=None, show_progress=True, parallel=False): """Head for array of points. Parameters @@ -287,8 +303,7 @@ def head_array(self, x, y, layers=None, show_progress=False, parallel=False): 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. + show computation progress, by default `True`. 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 @@ -299,7 +314,6 @@ def head_array(self, x, y, layers=None, show_progress=False, parallel=False): 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) @@ -310,31 +324,33 @@ def head_array(self, x, y, layers=None, show_progress=False, parallel=False): 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) - ): + for i in tqdm(range(npts), disable=not show_progress): 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, + nproc = mp.cpu_count() // 2 if parallel is True else int(parallel) + nproc = max(1, nproc) + responsive_progress = show_progress == "responsive" + chunksize = 1 if responsive_progress else max(1, npts // (4 * nproc)) + with ProcessPoolExecutor( max_workers=nproc, - chunksize=chunksize, - ) - - for i, result in results: - h[:, i] = result + initializer=_init_worker, + initargs=(self,), + ) as executor: + results = executor.map( + _compute_head_mp, + x, + y, + repeat(layers), + chunksize=chunksize, + ) + if show_progress: + results = tqdm(results, total=npts, desc="head array") + for i, result in enumerate(results): + h[:, i] = result return h def headgrid( - self, xg, yg, layers=None, printrow=False, show_progress=False, parallel=False + self, xg, yg, layers=None, printrow=False, show_progress=True, parallel=False ): """Grid of heads. @@ -347,8 +363,7 @@ def headgrid( 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. + show computation progress, by default `True`. 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 @@ -393,7 +408,7 @@ def headgrid2( y2, ny, layers=None, - show_progress=False, + show_progress=True, printrow=False, parallel=False, ): @@ -408,8 +423,7 @@ def headgrid2( 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. + show computation progress, by default `True`. 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 @@ -467,6 +481,65 @@ def headalongline(self, x, y, layers=None): h[:, i] = self.head(xg[i], yg[i], layers) return h + def disvec_array(self, x, y, layers=None, show_progress=True, parallel=False): + """Discharge vector 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 default `True`. + parallel : bool or int, optional + if `True`, computes disvec_array in parallel using multiprocessing, + by default `False`. If an integer is provided, it specifies the number of + processes to use. + + Returns + ------- + qx : array size `nlayers, npoints` + qy : array size `nlayers, npoints` + """ + 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)) + qx = np.empty((nlayers, npts)) + qy = np.empty((nlayers, npts)) + if not parallel: + for i in tqdm(range(npts), disable=not show_progress, desc="disvec array"): + qx[:, i], qy[:, i] = self.disvec(x[i], y[i], layers) + else: + nproc = mp.cpu_count() // 2 if parallel is True else int(parallel) + nproc = max(1, nproc) + responsive_progress = show_progress == "responsive" + chunksize = 1 if responsive_progress else max(1, npts // (4 * nproc)) + with ProcessPoolExecutor( + max_workers=nproc, + initializer=_init_worker, + initargs=(self,), + ) as executor: + results = executor.map( + _compute_disvec_mp, + x, + y, + repeat(layers), + chunksize=chunksize, + ) + if show_progress: + results = tqdm(results, total=npts, desc="disvec array") + for i, result in enumerate(results): + qx[:, i], qy[:, i] = result + return qx, qy + def disvecgrid( self, xg, @@ -486,11 +559,11 @@ def disvecgrid( 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 True. - parallel : bool, optional + show computation progress, by default `True`. + parallel : bool or int, optional if `True`, computes discharge vector grid in parallel using multiprocessing, - by default `False` + by default `False`. If an integer is provided, it is interpreted as the + number of processes to use. Returns ------- @@ -499,43 +572,18 @@ def disvecgrid( qy : array size (Nlayers, ny, nx) y component of discharge vector at each point in grid """ - 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)) - qx = np.empty((Nlayers, ny, nx)) - qy = 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): - qx[:, j, i], qy[:, j, i] = self.disvec(xg[i], yg[j], layers) - if show_progress: - print("", flush=True) - else: - - def compute(ij): - i, j = ij - return i, j, self.disvec(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="disvecgrid", - disable=not show_progress, - tqdm_class=tqdm, - ) - for i, j, result in results: - qx[:, j, i], qy[:, j, i] = result - - return qx, qy + x, y = np.meshgrid(xg, yg) + qx, qy = self.disvec_array( + x.ravel(), + y.ravel(), + layers=layers, + show_progress=show_progress, + parallel=parallel, + ) + nlayers = qx.shape[0] + return qx.reshape((nlayers, len(yg), len(xg))), qy.reshape( + (nlayers, len(yg), len(xg)) + ) def disvecalongline(self, x, y, layers=None): """Compute discharge vector along line. @@ -635,8 +683,7 @@ def velocity_array(self, x, y, z, show_progress=True, parallel=False): z : 1d-array z values show_progress : bool - show computation progress with tqdm progressbar if tqdm is installed. - Default is True. + show computation progress, by default `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 @@ -648,7 +695,6 @@ def velocity_array(self, x, y, z, show_progress=True, parallel=False): 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) @@ -656,30 +702,33 @@ def velocity_array(self, x, y, z, show_progress=True, parallel=False): 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) - ): + for i in tqdm(range(npts), disable=not show_progress): 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, + nproc = mp.cpu_count() // 2 if parallel is True else int(parallel) + nproc = max(1, nproc) + responsive_progress = show_progress == "responsive" + chunksize = 1 if responsive_progress else max(1, npts // (4 * nproc)) + with ProcessPoolExecutor( max_workers=nproc, - chunksize=chunksize, - ) - for i, result in results: - v[:, i] = result + initializer=_init_worker, + initargs=(self,), + ) as executor: + results = executor.map( + _compute_velocity_mp, + x, + y, + z, + chunksize=chunksize, + ) + if show_progress: + results = tqdm(results, total=npts, desc="velocity array") + for i, result in enumerate(results): + v[:, i] = result return v @@ -694,8 +743,8 @@ 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` + show_progress : bool + show computation progress, 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 diff --git a/timflow/steady/plots.py b/timflow/steady/plots.py index 5d9c5710..f73510e3 100644 --- a/timflow/steady/plots.py +++ b/timflow/steady/plots.py @@ -63,6 +63,7 @@ def contour( legend=True, return_contours=False, parallel=False, + show_progress=False, **kwargs, ): """Head contour plot. @@ -102,6 +103,8 @@ def contour( 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. + show_progress : bool + if True, show progress bar when computing headgrid, default is False. **kwargs additional keyword arguments passed to ax.contour() @@ -114,7 +117,9 @@ def contour( of contour sets for each contoured layer, only if return_contours=True """ xg, yg = self._get_xy_arrays(win, ngr) - h = self._ml.headgrid(xg, yg, np.atleast_1d(layers), parallel=parallel) + h = self._ml.headgrid( + xg, yg, np.atleast_1d(layers), show_progress=show_progress, parallel=parallel + ) return self.contour_array( xg, yg, diff --git a/timflow/transient/model.py b/timflow/transient/model.py index 1dd802aa..967aa5de 100644 --- a/timflow/transient/model.py +++ b/timflow/transient/model.py @@ -11,10 +11,13 @@ """ import multiprocessing as mp +from concurrent.futures import ProcessPoolExecutor +from itertools import repeat from warnings import warn import numpy as np import pandas as pd +from tqdm import tqdm from timflow.transient.aquifer import Aquifer, SimpleAquifer from timflow.transient.aquifer_parameters import param_3d, param_maq @@ -24,23 +27,32 @@ invlapcomp, ) from timflow.transient.plots import PlotTransient -from timflow.version import check_tqdm_parallel +_WORKER_STATE = {"model": None} -def _compute_head_mp(args): + +def _init_worker(model): + """Initialize a single model instance per worker process.""" + _WORKER_STATE["model"] = model + + +def _compute_head_mp(xi, yi, t, layers): """Helper function for parallel computation of head_array.""" - model, xi, yi, t, layers, i = args - return i, model.head(xi, yi, t, layers) + return _WORKER_STATE["model"].head(xi, yi, t, layers) -def _compute_velocity_mp(args): +def _compute_velocity_mp(xi, yi, zi, t): """Helper function for parallel computation of velocity_array.""" - model, xi, yi, zi, t, i = args try: - vv = model.velocomp(xi, yi, zi, t) + vv = _WORKER_STATE["model"].velocomp(xi, yi, zi, t) except (ZeroDivisionError, ValueError): vv = np.full((3,), np.nan) - return i, vv + return vv + + +def _compute_disvec_mp(xi, yi, t, layers): + """Helper function for parallel computation of disvec_array.""" + return _WORKER_STATE["model"].disvec(xi, yi, t, layers) class Model: @@ -479,8 +491,8 @@ def velocity_array(self, x, y, z, t, show_progress=True, parallel=False): 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` + show_progress : bool + show computation progress, 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 @@ -492,40 +504,40 @@ def velocity_array(self, x, y, z, t, show_progress=True, parallel=False): 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) - ): + for i in tqdm(range(npts), desc="velocity array", disable=not show_progress): 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, + nproc = mp.cpu_count() // 2 if parallel is True else int(parallel) + nproc = max(1, nproc) + responsive_progress = show_progress == "responsive" + chunksize = 1 if responsive_progress else max(1, npts // (4 * nproc)) + with ProcessPoolExecutor( max_workers=nproc, - chunksize=chunksize, - ) - for i, result in results: - v[:, i] = result + initializer=_init_worker, + initargs=(self,), + ) as executor: + results = executor.map( + _compute_velocity_mp, + x, + y, + z, + repeat(t), + chunksize=chunksize, + ) + if show_progress: + results = tqdm(results, total=npts, desc="velocity array") + for i, result in enumerate(results): + v[:, i] = result return v @@ -542,8 +554,8 @@ def velocity_grid(self, xg, yg, zg, t, show_progress=True, parallel=False): z values of grid t : float time for which grid is returned - show_progress : bool, optional - if `True`, shows progress bar when computing velocity grid, by default `True` + show_progress : bool + show computation progress, 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 @@ -647,7 +659,7 @@ 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): + def head_array(self, x, y, t, layers=None, show_progress=True, parallel=False): """Head for array of points. Parameters @@ -661,8 +673,7 @@ def head_array(self, x, y, t, layers=None, show_progress=False, parallel=False): 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. + show computation progress, by default `True`. 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 @@ -672,7 +683,6 @@ def head_array(self, x, y, t, layers=None, show_progress=False, parallel=False): ------- 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) @@ -685,29 +695,30 @@ def head_array(self, x, y, t, layers=None, show_progress=False, parallel=False): 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) - ): + for i in tqdm(range(npts), disable=not show_progress, desc="head array"): 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, + nproc = mp.cpu_count() // 2 if parallel is True else int(parallel) + nproc = max(1, nproc) + responsive_progress = show_progress == "responsive" + chunksize = 1 if responsive_progress else max(1, npts // (4 * nproc)) + with ProcessPoolExecutor( max_workers=nproc, - chunksize=chunksize, - ) - - for i, result in results: - h[:, :, i] = result + initializer=_init_worker, + initargs=(self,), + ) as executor: + results = executor.map( + _compute_head_mp, + x, + y, + repeat(t), + repeat(layers), + chunksize=chunksize, + ) + if show_progress: + results = tqdm(results, total=npts, desc="head array") + for i, result in enumerate(results): + h[:, :, i] = result return h def headgrid( @@ -717,7 +728,7 @@ def headgrid( t, layers=None, printrow=False, - show_progress=False, + show_progress=True, parallel=False, ): """Grid of heads. @@ -733,8 +744,7 @@ def headgrid( 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. + show computation progress, by default `True`. parallel : bool or int, optional if `True`, computes headgrid in parallel using multithreading, by default `False`. If an integer is provided, it specifies the number of @@ -783,7 +793,7 @@ def headgrid2( ny, t, layers=None, - show_progress=False, + show_progress=True, printrow=False, parallel=False, ): @@ -800,8 +810,7 @@ def headgrid2( 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. + show computation progress, by default `True`. 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 @@ -831,6 +840,71 @@ def headgrid2( parallel=parallel, ) + def disvec_array(self, x, y, t, layers=None, show_progress=True, parallel=False): + """Discharge vector 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 default `True`. + parallel : bool or int, optional + if `True`, computes disvec_array in parallel using multiprocessing, + by default `False`. If an integer is provided, it specifies the number of + processes to use. + + Returns + ------- + qx : array size `nlayers, ntimes, npoints` + qy : array size `nlayers, ntimes, npoints` + """ + 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)) + qx = np.empty((nlayers, ntimes, npts)) + qy = np.empty((nlayers, ntimes, npts)) + if not parallel: + for i in tqdm(range(npts), disable=not show_progress, desc="disvec array"): + qx[:, :, i], qy[:, :, i] = self.disvec(x[i], y[i], t, layers) + else: + nproc = mp.cpu_count() // 2 if parallel is True else int(parallel) + nproc = max(1, nproc) + responsive_progress = show_progress == "responsive" + chunksize = 1 if responsive_progress else max(1, npts // (4 * nproc)) + with ProcessPoolExecutor( + max_workers=nproc, + initializer=_init_worker, + initargs=(self,), + ) as executor: + results = executor.map( + _compute_disvec_mp, + x, + y, + repeat(t), + repeat(layers), + chunksize=chunksize, + ) + if show_progress: + results = tqdm(results, total=npts, desc="disvec array") + for i, (result_qx, result_qy) in enumerate(results): + qx[:, :, i] = result_qx + qy[:, :, i] = result_qy + return qx, qy + def disvecgrid( self, x, @@ -853,11 +927,11 @@ def disvecgrid( 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 True. - parallel : bool, optional + show computation progress, by default `True`. + parallel : bool or int, optional if `True`, computes discharge vector grid in parallel using multiprocessing, - by default `False` + by default `False`. If an integer is provided, it specifies the number of + processes to use. Returns ------- @@ -866,44 +940,20 @@ def disvecgrid( qy : array size (nlayers, ntimes, ny, nx) y component of discharge vector at each point in grid """ - parallel, thread_map, tqdm = check_tqdm_parallel(parallel) - - x = np.atleast_1d(x) - y = np.atleast_1d(y) - t = np.atleast_1d(t) - nx, ny = len(x), len(y) - 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)) - qx = np.empty((nlayers, ntimes, ny, nx)) - qy = 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): - qx[:, :, j, i], qy[:, :, j, i] = self.disvec(x[i], y[j], t, layers) - if show_progress: - print("", flush=True) - else: - - def compute(ij): - i, j = ij - return i, j, self.disvec(x[i], y[j], t, layers) - - results = thread_map( - compute, - [(i, j) for j in range(ny) for i in range(nx)], - total=nx * ny, - desc="disvecgrid", - disable=not show_progress, - tqdm_class=tqdm, - ) - for i, j, result in results: - qx[:, :, j, i], qy[:, :, j, i] = result - + xg = np.atleast_1d(x) + yg = np.atleast_1d(y) + nx, ny = len(xg), len(yg) + x, y = np.meshgrid(xg, yg) + qx, qy = self.disvec_array( + x.ravel(), + y.ravel(), + t, + layers=layers, + show_progress=show_progress, + parallel=parallel, + ) + qx = qx.reshape((qx.shape[0], qx.shape[1], ny, nx)) + qy = qy.reshape((qy.shape[0], qy.shape[1], ny, nx)) return qx, qy def inverseLapTran(self, pot, t): diff --git a/timflow/transient/plots.py b/timflow/transient/plots.py index 9483ec02..bc71e9a5 100644 --- a/timflow/transient/plots.py +++ b/timflow/transient/plots.py @@ -233,9 +233,8 @@ def contour( 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. - show_progress : bool, optional - if True, show progress bar when computing headgrid in parallel, - default is False. + show_progress : bool + if True, show progress bar when computing headgrid, default is False. **kwargs additional keyword arguments passed to ax.contour() diff --git a/timflow/version.py b/timflow/version.py index 407d69c0..19250bef 100644 --- a/timflow/version.py +++ b/timflow/version.py @@ -1,4 +1,3 @@ -import warnings from importlib import import_module, metadata from platform import python_version @@ -32,38 +31,3 @@ def show_versions(optional=True) -> None: msg += "Not Installed" print(msg) - - -def check_tqdm_parallel(parallel): - """Check if tqdm is installed when parallel processing is requested. - - Parameters - ---------- - parallel : bool - Whether parallel processing is requested. - - Returns - ------- - 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. - tqdm : class or None - The tqdm class from tqdm if parallel processing is available, else None. - """ - if not parallel: # short circuit when no parallel requested - return parallel, None, None - try: - from tqdm import tqdm - from tqdm.contrib.concurrent import process_map - except ImportError: - warnings.warn( - "Parallel requires 'tqdm'. Install 'timflow[parallel]' or 'tqdm' to" - " enable parallel execution. Falling back to serial execution.", - category=ImportWarning, - stacklevel=2, - ) - parallel = False - process_map = None - tqdm = None - return parallel, process_map, tqdm