diff --git a/docs/source/api/index.rst b/docs/source/api/index.rst index f8e8df12..8084d114 100755 --- a/docs/source/api/index.rst +++ b/docs/source/api/index.rst @@ -28,6 +28,7 @@ Templates LinearOperator FunctionOperator MemoizeOperator + MultiOperator PyTensorOperator TorchOperator JaxOperator diff --git a/pylops/__init__.py b/pylops/__init__.py index 7bf8bca1..94b38128 100755 --- a/pylops/__init__.py +++ b/pylops/__init__.py @@ -51,6 +51,7 @@ from .config import * from .linearoperator import * +from .multioperator import * from .torchoperator import * from .pytensoroperator import * from .jaxoperator import * diff --git a/pylops/_multioperator.py b/pylops/_multioperator.py new file mode 100644 index 00000000..b3a90edf --- /dev/null +++ b/pylops/_multioperator.py @@ -0,0 +1,18 @@ +import threading +from collections.abc import Callable + +from pylops.utils.typing import NDArray + + +def _matvec_rmatvec_map(op: Callable[[NDArray], NDArray], x: NDArray) -> NDArray: + """matvec/rmatvec for multiprocessing / multithreading""" + return op(x).squeeze() + + +def _matvec_rmatvec_map_mt( + op: Callable[[NDArray], NDArray], x: NDArray, y: NDArray, lock: threading.Lock +) -> None: + """rmatvec for multithreading with lock""" + ylocal = op(x).squeeze() + with lock: + y[:] += ylocal diff --git a/pylops/basicoperators/blockdiag.py b/pylops/basicoperators/blockdiag.py index 7d98026a..80953738 100644 --- a/pylops/basicoperators/blockdiag.py +++ b/pylops/basicoperators/blockdiag.py @@ -1,8 +1,5 @@ __all__ = ["BlockDiag"] -import concurrent.futures as mt -import multiprocessing as mp - import numpy as np import scipy as sp @@ -22,18 +19,14 @@ from collections.abc import Sequence -from pylops import LinearOperator +from pylops import LinearOperator, MultiOperator +from pylops._multioperator import _matvec_rmatvec_map from pylops.basicoperators import MatrixMult from pylops.utils.backend import get_array_module, get_module, inplace_set from pylops.utils.typing import DTypeLike, NDArray, Tinoutengine, Tparallel_kind -def _matvec_rmatvec_map(op, x: NDArray) -> NDArray: - """matvec/rmatvec for multiprocessing""" - return op(x).squeeze() - - -class BlockDiag(LinearOperator): +class BlockDiag(MultiOperator): r"""Block-diagonal operator. Create a block-diagonal operator from N linear operators. @@ -149,9 +142,6 @@ def __init__( parallel_kind: Tparallel_kind = "multiproc", dtype: DTypeLike | None = None, ) -> None: - if parallel_kind not in ["multiproc", "multithread"]: - msg = "parallel_kind must be 'multiproc' or 'multithread'" - raise ValueError(msg) # identify dimensions self.ops = ops mops = np.zeros(len(ops), dtype=int) @@ -181,15 +171,10 @@ def __init__( else: dimsd = (self.nops,) forceflat = True + # create pool for multithreading / multiprocessing - self.parallel_kind = parallel_kind - self._nproc = nproc - self.pool: mp.pool.Pool | None = None - if self.nproc > 1: - if self.parallel_kind == "multiproc": - self.pool = mp.Pool(processes=nproc) - else: - self.pool = mt.ThreadPoolExecutor(max_workers=nproc) + self._setup_pool(nproc, parallel_kind=parallel_kind) + self.inoutengine = inoutengine dtype = _get_dtype(ops) if dtype is None else np.dtype(dtype) clinear = all([getattr(oper, "clinear", True) for oper in self.ops]) @@ -201,25 +186,6 @@ def __init__( forceflat=forceflat, ) - @property - def nproc(self) -> int: - return self._nproc - - @nproc.setter - def nproc(self, nprocnew: int) -> None: - if self._nproc > 1 and self.pool is not None: - if self.parallel_kind == "multiproc": - self.pool.close() - self.pool.join() - else: - self.pool.shutdown() - if nprocnew > 1: - if self.parallel_kind == "multiproc": - self.pool = mp.Pool(processes=nprocnew) - else: - self.pool = mt.ThreadPoolExecutor(max_workers=nprocnew) - self._nproc = nprocnew - def _matvec_serial(self, x: NDArray) -> NDArray: ncp = ( get_array_module(x) @@ -297,35 +263,3 @@ def _rmatvec_multithread(self, x: NDArray) -> NDArray: ) y = np.hstack(ys) return y - - def _matvec(self, x: NDArray) -> NDArray: - if self.nproc == 1: - y = self._matvec_serial(x) - else: - if self.parallel_kind == "multiproc": - y = self._matvec_multiproc(x) - else: - y = self._matvec_multithread(x) - return y - - def _rmatvec(self, x: NDArray) -> NDArray: - if self.nproc == 1: - y = self._rmatvec_serial(x) - else: - if self.parallel_kind == "multiproc": - y = self._rmatvec_multiproc(x) - else: - y = self._rmatvec_multithread(x) - return y - - def close(self): - """Close the pool of workers used for multiprocessing - / multithreading. - """ - if self.pool is not None: - if self.parallel_kind == "multiproc": - self.pool.close() - self.pool.join() - else: - self.pool.shutdown() - self.pool = None diff --git a/pylops/basicoperators/hstack.py b/pylops/basicoperators/hstack.py index ca817e9f..218145db 100644 --- a/pylops/basicoperators/hstack.py +++ b/pylops/basicoperators/hstack.py @@ -1,9 +1,5 @@ __all__ = ["HStack"] -import concurrent.futures as mt -import multiprocessing as mp -import threading - import numpy as np import scipy as sp @@ -19,27 +15,16 @@ ) from scipy.sparse.linalg._interface import _get_dtype -from collections.abc import Callable, Sequence +from collections.abc import Sequence -from pylops import LinearOperator +from pylops import LinearOperator, MultiOperator +from pylops._multioperator import _matvec_rmatvec_map, _matvec_rmatvec_map_mt from pylops.basicoperators import MatrixMult, Zero from pylops.utils.backend import get_array_module, get_module, inplace_add, inplace_set from pylops.utils.typing import NDArray, Tinoutengine, Tparallel_kind -def _matvec_rmatvec_map(op, x: NDArray) -> NDArray: - """matvec/rmatvec for multiprocessing""" - return op(x).squeeze() - - -def _matvec_map_mt(op: Callable, x: NDArray, y: NDArray, lock: threading.Lock) -> None: - """rmatvec for multithreading with lock""" - ylocal = op(x).squeeze() - with lock: - y[:] += ylocal - - -class HStack(LinearOperator): +class HStack(MultiOperator): r"""Horizontal stacking. Stack a set of N linear operators horizontally. Note that in case @@ -157,9 +142,6 @@ def __init__( parallel_kind: Tparallel_kind = "multiproc", dtype: str | None = None, ) -> None: - if parallel_kind not in ["multiproc", "multithread"]: - msg = "parallel_kind must be 'multiproc' or 'multithread'" - raise ValueError(msg) # identify dimensions self.ops = ops mops = np.zeros(len(ops), dtype=int) @@ -182,16 +164,10 @@ def __init__( else: dimsd = (self.nops,) forceflat = True + # create pool for multithreading / multiprocessing - self.parallel_kind = parallel_kind - self._nproc = nproc - self.pool = None - if self.nproc > 1: - if self.parallel_kind == "multiproc": - self.pool = mp.Pool(processes=nproc) - else: - self.pool = mt.ThreadPoolExecutor(max_workers=nproc) - self.lock = threading.Lock() + self._setup_pool(nproc, parallel_kind=parallel_kind) + self.inoutengine = inoutengine dtype = _get_dtype(self.ops) if dtype is None else np.dtype(dtype) clinear = all([getattr(oper, "clinear", True) for oper in self.ops]) @@ -203,25 +179,6 @@ def __init__( forceflat=forceflat, ) - @property - def nproc(self) -> int: - return self._nproc - - @nproc.setter - def nproc(self, nprocnew: int): - if self._nproc > 1 and self.pool is not None: - if self.parallel_kind == "multiproc": - self.pool.close() - self.pool.join() - else: - self.pool.shutdown() - if nprocnew > 1: - if self.parallel_kind == "multiproc": - self.pool = mp.Pool(processes=nprocnew) - else: - self.pool = mt.ThreadPoolExecutor(max_workers=nprocnew) - self._nproc = nprocnew - def _matvec_serial(self, x: NDArray) -> NDArray: ncp = ( get_array_module(x) @@ -277,7 +234,7 @@ def _matvec_multithread(self, x: NDArray) -> NDArray: y = np.zeros(self.nops, dtype=self.dtype) list( self.pool.map( - lambda args: _matvec_map_mt(*args), + lambda args: _matvec_rmatvec_map_mt(*args), [ ( oper._matvec, @@ -300,35 +257,3 @@ def _rmatvec_multithread(self, x: NDArray) -> NDArray: ) y = np.hstack(ys) return y - - def _matvec(self, x: NDArray) -> NDArray: - if self.nproc == 1: - y = self._matvec_serial(x) - else: - if self.parallel_kind == "multiproc": - y = self._matvec_multiproc(x) - else: - y = self._matvec_multithread(x) - return y - - def _rmatvec(self, x: NDArray) -> NDArray: - if self.nproc == 1: - y = self._rmatvec_serial(x) - else: - if self.parallel_kind == "multiproc": - y = self._rmatvec_multiproc(x) - else: - y = self._rmatvec_multithread(x) - return y - - def close(self): - """Close the pool of workers used for multiprocessing / - multithreading. - """ - if self.pool is not None: - if self.parallel_kind == "multiproc": - self.pool.close() - self.pool.join() - else: - self.pool.shutdown() - self.pool = None diff --git a/pylops/basicoperators/kronecker.py b/pylops/basicoperators/kronecker.py index b342fabf..cc8221cf 100644 --- a/pylops/basicoperators/kronecker.py +++ b/pylops/basicoperators/kronecker.py @@ -1,12 +1,17 @@ __all__ = ["Kronecker"] +from typing import TYPE_CHECKING + import numpy as np -from pylops import LinearOperator -from pylops.utils.typing import DTypeLike, NDArray +from pylops import MultiOperator +from pylops.utils.typing import DTypeLike, NDArray, Tparallel_kind + +if TYPE_CHECKING: + from pylops.linearoperator import LinearOperator -class Kronecker(LinearOperator): +class Kronecker(MultiOperator): r"""Kronecker operator. Perform Kronecker product of two operators. Note that the combined operator @@ -22,6 +27,18 @@ class Kronecker(LinearOperator): Second operator dtype : :obj:`str`, optional Type of elements in input array. + nproc : :obj:`int`, optional + .. versionadded:: 2.9.0 + + Number of processes/threads used to evaluate the N operators in parallel + using ``multiprocessing``/``concurrent.futures``. If ``nproc=1``, work in serial mode. + parallel_kind : :obj:`str`, optional + .. versionadded:: 2.9.0 + + Parallelism kind when ``nproc>1``. Can be ``multiproc`` (using + :mod:`multiprocessing`) or ``multithread`` (using + :class:`concurrent.futures.ThreadPoolExecutor`). Defaults + to ``multiproc``. name : :obj:`str`, optional .. versionadded:: 2.0.0 @@ -65,8 +82,10 @@ class Kronecker(LinearOperator): def __init__( self, - Op1: LinearOperator, - Op2: LinearOperator, + Op1: "LinearOperator", + Op2: "LinearOperator", + nproc: int = 1, + parallel_kind: Tparallel_kind = "multiproc", dtype: DTypeLike = "float64", name: str = "K", ) -> None: @@ -74,6 +93,10 @@ def __init__( self.Op2 = Op2 self.Op1H = self.Op1.H self.Op2H = self.Op2.H + + # create pool for multithreading / multiprocessing + self._setup_pool(nproc, parallel_kind=parallel_kind) + shape = ( self.Op1.shape[0] * self.Op2.shape[0], self.Op1.shape[1] * self.Op2.shape[1], @@ -82,12 +105,12 @@ def __init__( def _matvec(self, x: NDArray) -> NDArray: x = x.reshape(self.Op1.shape[1], self.Op2.shape[1]) - y = self.Op2.matmat(x.T).T - y = self.Op1.matmat(y).ravel() + y = self.Op2.matmat(x.T, pool=self.pool).T + y = self.Op1.matmat(y, pool=self.pool).ravel() return y def _rmatvec(self, x: NDArray) -> NDArray: x = x.reshape(self.Op1.shape[0], self.Op2.shape[0]) - y = self.Op2H.matmat(x.T).T - y = self.Op1H.matmat(y).ravel() + y = self.Op2H.matmat(x.T, pool=self.pool).T + y = self.Op1H.matmat(y, pool=self.pool).ravel() return y diff --git a/pylops/basicoperators/vstack.py b/pylops/basicoperators/vstack.py index 50d75609..3de7bdf5 100644 --- a/pylops/basicoperators/vstack.py +++ b/pylops/basicoperators/vstack.py @@ -1,9 +1,5 @@ __all__ = ["VStack"] -import concurrent.futures as mt -import multiprocessing as mp -import threading - import numpy as np import scipy as sp @@ -19,27 +15,16 @@ ) from scipy.sparse.linalg._interface import _get_dtype -from collections.abc import Callable, Sequence +from collections.abc import Sequence -from pylops import LinearOperator +from pylops import LinearOperator, MultiOperator +from pylops._multioperator import _matvec_rmatvec_map, _matvec_rmatvec_map_mt from pylops.basicoperators import MatrixMult, Zero from pylops.utils.backend import get_array_module, get_module, inplace_add, inplace_set from pylops.utils.typing import DTypeLike, NDArray, Tinoutengine, Tparallel_kind -def _matvec_rmatvec_map(op: Callable, x: NDArray) -> NDArray: - """matvec/rmatvec for multiprocessing / multithreading""" - return op(x).squeeze() - - -def _rmatvec_map_mt(op: Callable, x: NDArray, y: NDArray, lock: threading.Lock) -> None: - """rmatvec for multithreading with lock""" - ylocal = op(x).squeeze() - with lock: - y[:] += ylocal - - -class VStack(LinearOperator): +class VStack(MultiOperator): r"""Vertical stacking. Stack a set of N linear operators vertically. Note that in case @@ -156,9 +141,6 @@ def __init__( parallel_kind: Tparallel_kind = "multiproc", dtype: DTypeLike | None = None, ) -> None: - if parallel_kind not in ["multiproc", "multithread"]: - msg = "parallel_kind must be 'multiproc' or 'multithread'" - raise ValueError(msg) # identify dimensions self.ops = ops nops = np.zeros(len(self.ops), dtype=int) @@ -181,16 +163,10 @@ def __init__( else: dims = (self.mops,) forceflat = True + # create pool for multithreading / multiprocessing - self.parallel_kind = parallel_kind - self._nproc = nproc - self.pool = None - if self.nproc > 1: - if self.parallel_kind == "multiproc": - self.pool = mp.Pool(processes=nproc) - else: - self.pool = mt.ThreadPoolExecutor(max_workers=nproc) - self.lock = threading.Lock() + self._setup_pool(nproc, parallel_kind=parallel_kind) + self.inoutengine = inoutengine dtype = _get_dtype(self.ops) if dtype is None else np.dtype(dtype) clinear = all([getattr(oper, "clinear", True) for oper in self.ops]) @@ -202,25 +178,6 @@ def __init__( forceflat=forceflat, ) - @property - def nproc(self) -> int: - return self._nproc - - @nproc.setter - def nproc(self, nprocnew: int): - if self._nproc > 1 and self.pool is not None: - if self.parallel_kind == "multiproc": - self.pool.close() - self.pool.join() - else: - self.pool.shutdown() - if nprocnew > 1: - if self.parallel_kind == "multiproc": - self.pool = mp.Pool(processes=nprocnew) - else: - self.pool = mt.ThreadPoolExecutor(max_workers=nprocnew) - self._nproc = nprocnew - def _matvec_serial(self, x: NDArray) -> NDArray: ncp = ( get_array_module(x) @@ -286,7 +243,7 @@ def _rmatvec_multithread(self, x: NDArray) -> NDArray: y = np.zeros(self.mops, dtype=self.dtype) list( self.pool.map( - lambda args: _rmatvec_map_mt(*args), + lambda args: _matvec_rmatvec_map_mt(*args), [ ( oper._rmatvec, @@ -299,35 +256,3 @@ def _rmatvec_multithread(self, x: NDArray) -> NDArray: ) ) return y - - def _matvec(self, x: NDArray) -> NDArray: - if self.nproc == 1: - y = self._matvec_serial(x) - else: - if self.parallel_kind == "multiproc": - y = self._matvec_multiproc(x) - else: - y = self._matvec_multithread(x) - return y - - def _rmatvec(self, x: NDArray) -> NDArray: - if self.nproc == 1: - y = self._rmatvec_serial(x) - else: - if self.parallel_kind == "multiproc": - y = self._rmatvec_multiproc(x) - else: - y = self._rmatvec_multithread(x) - return y - - def close(self): - """Close the pool of workers used for multiprocessing / - multithreading. - """ - if self.pool is not None: - if self.parallel_kind == "multiproc": - self.pool.close() - self.pool.join() - else: - self.pool.shutdown() - self.pool = None diff --git a/pylops/linearoperator.py b/pylops/linearoperator.py index 87934a54..efac4100 100644 --- a/pylops/linearoperator.py +++ b/pylops/linearoperator.py @@ -6,6 +6,10 @@ ] from abc import ABC, abstractmethod +from collections.abc import Callable, Sequence +from concurrent.futures import ThreadPoolExecutor +from multiprocessing import Pool +from multiprocessing.pool import Pool as PoolClass import numpy as np import scipy as sp @@ -27,14 +31,13 @@ else: from scipy.sparse._sputils import isintlike, isshape -from collections.abc import Callable, Sequence - from pylops import get_ndarray_multiplication +from pylops._multioperator import _matvec_rmatvec_map from pylops.optimization.basic import cgls from pylops.utils.backend import get_array_module, get_module, get_sparse_eye from pylops.utils.decorators import count from pylops.utils.estimators import trace_hutchinson, trace_hutchpp, trace_nahutchpp -from pylops.utils.typing import DTypeLike, InputDimsLike, NDArray, ShapeLike +from pylops.utils.typing import DTypeLike, InputDimsLike, NDArray, ShapeLike, Tpool class _LinearOperator(ABC): @@ -458,31 +461,101 @@ def _rmatvec(self, x: NDArray) -> NDArray: if self.Op is not None: return self.Op._rmatvec(x) - def _matmat(self, X: NDArray) -> NDArray: + def _matmat_serial(self, X: NDArray) -> NDArray: + """Matrix-matrix multiplication (serial version)""" + ncp = get_array_module(X) + if sp.sparse.issparse(X): + y = ncp.vstack([self._matvec(col.toarray().reshape(-1)) for col in X.T]).T + else: + y = ncp.vstack([self._matvec(col.reshape(-1)) for col in X.T]).T + return y + + def _matmat_multithread(self, X: NDArray, pool: ThreadPoolExecutor) -> NDArray: + """Matrix-matrix multiplication (multithreaded version)""" + ys = list( + pool.map( + lambda args: _matvec_rmatvec_map(*args), + [(self._matvec, col.reshape(-1)) for col in X.T], + ) + ) + y = np.vstack(ys).T + return y + + def _matmat_multiproc(self, X: NDArray, pool: Pool) -> NDArray: + """Matrix-matrix multiplication (multiprocess version)""" + ys = pool.starmap( + _matvec_rmatvec_map, + [(self._matvec, col.reshape(-1)) for col in X.T], + ) + y = np.vstack(ys).T + return y + + def _matmat(self, X: NDArray, pool: Tpool | None = None) -> NDArray: """Matrix-matrix multiplication handler. Modified version of scipy _matmat to avoid having trailing dimension in col when provided to matvec """ + if pool is None: + # serial + return self._matmat_serial(X) + elif isinstance(pool, ThreadPoolExecutor): + # multithread + return self._matmat_multithread(X, pool) + elif isinstance(pool, PoolClass): + # multiprocess + return self._matmat_multiproc(X, pool) + else: + msg = f"Received pool of unsupported type ({type(pool)})" + raise NotImplementedError(msg) + + def _rmatmat_serial(self, X: NDArray) -> NDArray: + """Matrix-matrix adjoint multiplication (serial version)""" ncp = get_array_module(X) if sp.sparse.issparse(X): - y = ncp.vstack([self.matvec(col.toarray().reshape(-1)) for col in X.T]).T + y = ncp.vstack([self._rmatvec(col.toarray().reshape(-1)) for col in X.T]).T else: - y = ncp.vstack([self.matvec(col.reshape(-1)) for col in X.T]).T + y = ncp.vstack([self._rmatvec(col.reshape(-1)) for col in X.T]).T + return y + + def _rmatmat_multithread(self, X: NDArray, pool: ThreadPoolExecutor) -> NDArray: + """Matrix-matrix adjoint multiplication (multithreaded version)""" + ys = list( + pool.map( + lambda args: _matvec_rmatvec_map(*args), + [(self._rmatvec, col.reshape(-1)) for col in X.T], + ) + ) + y = np.vstack(ys).T return y - def _rmatmat(self, X: NDArray) -> NDArray: + def _rmatmat_multiproc(self, X: NDArray, pool: Pool) -> NDArray: + """Matrix-matrix adjoint multiplication (multiprocess version)""" + ys = pool.starmap( + _matvec_rmatvec_map, + [(self._rmatvec, col.reshape(-1)) for col in X.T], + ) + y = np.vstack(ys).T + return y + + def _rmatmat(self, X: NDArray, pool: Tpool | None = None) -> NDArray: """Matrix-matrix adjoint multiplication handler. Modified version of scipy _rmatmat to avoid having trailing dimension in col when provided to rmatvec """ - ncp = get_array_module(X) - if sp.sparse.issparse(X): - y = ncp.vstack([self.rmatvec(col.toarray().reshape(-1)) for col in X.T]).T + if pool is None: + # serial + return self._rmatmat_serial(X) + elif isinstance(pool, ThreadPoolExecutor): + # multithread + return self._rmatmat_multithread(X, pool) + elif isinstance(pool, PoolClass): + # multiprocess + return self._rmatmat_multiproc(X, pool) else: - y = ncp.vstack([self.rmatvec(col.reshape(-1)) for col in X.T]).T - return y + msg = f"Received pool of unsupported type ({type(pool)})" + raise NotImplementedError(msg) def _adjoint(self) -> LinearOperator: Op = _AdjointLinearOperator(self) @@ -583,7 +656,7 @@ def rmatvec(self, x: NDArray) -> NDArray: return y @count(forward=True, matmat=True) - def matmat(self, X: NDArray) -> NDArray: + def matmat(self, X: NDArray, pool: Tpool | None = None) -> NDArray: """Matrix-matrix multiplication. Modified version of scipy matmat which does not consider the case @@ -594,6 +667,8 @@ def matmat(self, X: NDArray) -> NDArray: ---------- x : :obj:`numpy.ndarray` Input array of shape (N,K) + pool : :obj:`multiprocessing.Pool` or :obj:`concurrent.futures.ThreadPoolExecutor` or :obj:`None` + Pool of workers used to evaluate the operator on different columns of ``X`` in parallel. Returns ------- @@ -607,11 +682,11 @@ def matmat(self, X: NDArray) -> NDArray: if X.shape[0] != self.shape[1]: msg = f"Dimension mismatch: {self.shape}, {X.shape}" raise ValueError(msg) - Y = self._matmat(X) + Y = self._matmat(X, pool=pool) return Y @count(forward=False, matmat=True) - def rmatmat(self, X: NDArray) -> NDArray: + def rmatmat(self, X: NDArray, pool: Tpool | None = None) -> NDArray: """Matrix-matrix multiplication. Modified version of scipy rmatmat which does not consider the case @@ -622,6 +697,8 @@ def rmatmat(self, X: NDArray) -> NDArray: ---------- x : :obj:`numpy.ndarray` Input array of shape (M,K) + pool : :obj:`multiprocessing.Pool` or :obj:`concurrent.futures.ThreadPoolExecutor` or :obj:`None` + Pool of workers used to evaluate the operator on different columns of ``X`` in parallel. Returns ------- @@ -633,9 +710,9 @@ def rmatmat(self, X: NDArray) -> NDArray: msg = f"Expected 2-d ndarray or matrix, not {X.ndim}-d ndarray" raise ValueError(msg) if X.shape[0] != self.shape[0]: - f"Dimension mismatch: {self.shape}, {X.shape}" + msg = f"Dimension mismatch: {self.shape}, {X.shape}" raise ValueError(msg) - Y = self._rmatmat(X) + Y = self._rmatmat(X, pool=pool) return Y def dot(self, x: NDArray) -> NDArray: @@ -1312,11 +1389,11 @@ def _matvec(self, x: NDArray) -> NDArray: def _rmatvec(self, x: NDArray) -> NDArray: return np.conj(self.args[1]) * self.args[0].rmatvec(x) - def _rmatmat(self, x: NDArray) -> NDArray: - return np.conj(self.args[1]) * self.args[0].rmatmat(x) + def _matmat(self, x: NDArray, pool: Tpool | None = None) -> NDArray: + return self.args[1] * self.args[0]._matmat(x, pool=pool) - def _matmat(self, x: NDArray) -> NDArray: - return self.args[1] * self.args[0].matmat(x) + def _rmatmat(self, x: NDArray, pool: Tpool | None = None) -> NDArray: + return np.conj(self.args[1]) * self.args[0]._rmatmat(x, pool=pool) def _adjoint(self) -> LinearOperator: A, alpha = self.args @@ -1400,11 +1477,11 @@ def _matvec(self, x: NDArray) -> NDArray: def _rmatvec(self, x: NDArray) -> NDArray: return self.A._matvec(x) - def _matmat(self, X: NDArray) -> NDArray: - return self.A._rmatmat(X) + def _matmat(self, X: NDArray, pool: Tpool | None = None) -> NDArray: + return self.A._rmatmat(X, pool=pool) - def _rmatmat(self, X: NDArray) -> NDArray: - return self.A._matmat(X) + def _rmatmat(self, X: NDArray, pool: Tpool | None = None) -> NDArray: + return self.A._matmat(X, pool=pool) class _TransposedLinearOperator(LinearOperator): @@ -1422,11 +1499,11 @@ def _matvec(self, x: NDArray) -> NDArray: def _rmatvec(self, x: NDArray) -> NDArray: return np.conj(self.A._matvec(np.conj(x))) - def _matmat(self, X: NDArray) -> NDArray: - return np.conj(self.A._rmatmat(np.conj(X))) + def _matmat(self, X: NDArray, pool: Tpool | None = None) -> NDArray: + return np.conj(self.A._rmatmat(np.conj(X), pool=pool)) - def _rmatmat(self, X: NDArray) -> NDArray: - return np.conj(self.A._matmat(np.conj(X))) + def _rmatmat(self, X: NDArray, pool: Tpool | None = None) -> NDArray: + return np.conj(self.A._matmat(np.conj(X), pool=pool)) class _ProductLinearOperator(LinearOperator): @@ -1445,16 +1522,16 @@ def __init__(self, A: LinearOperator, B: LinearOperator): self.args = (A, B) def _matvec(self, x: NDArray) -> NDArray: - return self.args[0].matvec(self.args[1].matvec(x)) + return self.args[0]._matvec(self.args[1]._matvec(x)) def _rmatvec(self, x: NDArray) -> NDArray: - return self.args[1].rmatvec(self.args[0].rmatvec(x)) + return self.args[1]._rmatvec(self.args[0]._rmatvec(x)) - def _rmatmat(self, X: NDArray) -> NDArray: - return self.args[1].rmatmat(self.args[0].rmatmat(X)) + def _matmat(self, X: NDArray, pool: Tpool | None = None) -> NDArray: + return self.args[0]._matmat(self.args[1]._matmat(X), pool=pool) - def _matmat(self, X: NDArray) -> NDArray: - return self.args[0].matmat(self.args[1].matmat(X)) + def _rmatmat(self, X: NDArray, pool: Tpool | None = None) -> NDArray: + return self.args[1]._rmatmat(self.args[0]._rmatmat(X), pool=pool) def _adjoint(self): A, B = self.args @@ -1477,16 +1554,16 @@ def __init__( super().__init__(dtype=_get_dtype([A, B]), shape=A.shape) def _matvec(self, x: NDArray) -> NDArray: - return self.args[0].matvec(x) + self.args[1].matvec(x) + return self.args[0]._matvec(x) + self.args[1]._matvec(x) def _rmatvec(self, x: NDArray) -> NDArray: - return self.args[0].rmatvec(x) + self.args[1].rmatvec(x) + return self.args[0]._rmatvec(x) + self.args[1]._rmatvec(x) - def _rmatmat(self, x: NDArray) -> NDArray: - return self.args[0].rmatmat(x) + self.args[1].rmatmat(x) + def _matmat(self, x: NDArray, pool: Tpool | None = None) -> NDArray: + return self.args[0]._matmat(x, pool=pool) + self.args[1]._matmat(x, pool=pool) - def _matmat(self, x: NDArray) -> NDArray: - return self.args[0].matmat(x) + self.args[1].matmat(x) + def _rmatmat(self, x: NDArray, pool: Tpool | None = None) -> NDArray: + return self.args[0]._rmatmat(x, pool=pool) + self.args[1]._rmatmat(x, pool=pool) def _adjoint(self) -> LinearOperator: A, B = self.args @@ -1508,10 +1585,10 @@ def __init__(self, A: LinearOperator, p: int) -> None: super().__init__(dtype=A.dtype, shape=A.shape) self.args = (A, p) - def _power(self, fun: Callable, x: NDArray) -> NDArray: + def _power(self, fun: Callable, x: NDArray, **kwargs) -> NDArray: res = x.copy() for _ in range(self.args[1]): - res = fun(res) + res = fun(res, **kwargs) return res def _matvec(self, x: NDArray) -> NDArray: @@ -1520,11 +1597,11 @@ def _matvec(self, x: NDArray) -> NDArray: def _rmatvec(self, x: NDArray) -> NDArray: return self._power(self.args[0].rmatvec, x) - def _rmatmat(self, x: NDArray) -> NDArray: - return self._power(self.args[0].rmatmat, x) + def _matmat(self, x: NDArray, pool: Tpool | None = None) -> NDArray: + return self._power(self.args[0]._matmat, x, pool=pool) - def _matmat(self, x: NDArray) -> NDArray: - return self._power(self.args[0].matmat, x) + def _rmatmat(self, x: NDArray, pool: Tpool | None = None) -> NDArray: + return self._power(self.args[0]._rmatmat, x, pool=pool) def _adjoint(self) -> LinearOperator: A, p = self.args diff --git a/pylops/multioperator.py b/pylops/multioperator.py new file mode 100644 index 00000000..28cb94a6 --- /dev/null +++ b/pylops/multioperator.py @@ -0,0 +1,108 @@ +from __future__ import annotations + +__all__ = ["MultiOperator"] + +import concurrent.futures as mt +import multiprocessing as mp +import threading + +from pylops import LinearOperator +from pylops.utils.typing import NDArray, Tparallel_kind, Tpool + + +class MultiOperator(LinearOperator): + """Multiprocess/threading operator + + This class acts as a base class for all operators that wish to + support multiprocessing/multithreading in their ``matvec``/``rmatvec`` + methods. + + Implements basic methods to instantiate and tear down a pool or workers, + and ``_matvec``/``_rmatvec`` interfaces that dispatch to the actual + implementations for serial/multiprocess/multithread, namely: + + - ``_matvec_serial`` / ``_rmatvec_serial``: serial implementation + - ``_matvec_multiproc`` / ``_rmatvec_multiproc``: multiprocess implementation + - ``_matvec_multithread`` / ``_matvec_multithread``: multithreading implementation + + Developers are in charge of implementing these methods for specific operators or + overwriting ``_matvec``/``_rmatvec`` if not all of the implementations are available. + + .. note:: End users of PyLops should not use this class directly but simply + use operators that are already implemented. This class is meant for + developers and it has to be used as the parent class of any new operator + with multiprocess/multithreading capabilities developed within PyLops. + + """ + + def _setup_pool( + self, + nproc: int = 1, + parallel_kind: Tparallel_kind = "multiproc", + ) -> None: + """Setup pool for multiprocessing/multithreading""" + if parallel_kind not in ["multiproc", "multithread"]: + msg = "parallel_kind must be 'multiproc' or 'multithread'" + raise ValueError(msg) + + # create pool for multithreading / multiprocessing + self.parallel_kind = parallel_kind + self._nproc = nproc + self.pool: Tpool | None = None + if self.nproc > 1: + if self.parallel_kind == "multiproc": + self.pool = mp.Pool(processes=nproc) + else: + self.pool = mt.ThreadPoolExecutor(max_workers=nproc) + self.lock = threading.Lock() + + @property + def nproc(self) -> int: + return self._nproc + + @nproc.setter + def nproc(self, nprocnew: int) -> None: + if self._nproc > 1 and self.pool is not None: + if self.parallel_kind == "multiproc": + self.pool.close() + self.pool.join() + else: + self.pool.shutdown() + if nprocnew > 1: + if self.parallel_kind == "multiproc": + self.pool = mp.Pool(processes=nprocnew) + else: + self.pool = mt.ThreadPoolExecutor(max_workers=nprocnew) + self._nproc = nprocnew + + def _matvec(self, x: NDArray) -> NDArray: + if self.nproc == 1: + y = self._matvec_serial(x) + else: + if self.parallel_kind == "multiproc": + y = self._matvec_multiproc(x) + else: + y = self._matvec_multithread(x) + return y + + def _rmatvec(self, x: NDArray) -> NDArray: + if self.nproc == 1: + y = self._rmatvec_serial(x) + else: + if self.parallel_kind == "multiproc": + y = self._rmatvec_multiproc(x) + else: + y = self._rmatvec_multithread(x) + return y + + def close(self) -> None: + """Close the pool of workers used for multiprocessing + / multithreading. + """ + if self.pool is not None: + if self.parallel_kind == "multiproc": + self.pool.close() + self.pool.join() + else: + self.pool.shutdown() + self.pool = None diff --git a/pylops/utils/decorators.py b/pylops/utils/decorators.py index 4c61f14e..8dd5797f 100644 --- a/pylops/utils/decorators.py +++ b/pylops/utils/decorators.py @@ -189,20 +189,18 @@ def decorator(f): mat = matmat @wraps(f) - def wrapper(self, x): + def wrapper(self, x, *args, **kwargs): # perform operation - y = f(self, x) + y = f(self, x, *args, **kwargs) # increase count of the associated operation if fwd: if mat: self.matmat_count += 1 - self.matvec_count -= x.shape[-1] else: self.matvec_count += 1 else: if mat: self.rmatmat_count += 1 - self.rmatvec_count -= x.shape[-1] else: self.rmatvec_count += 1 return y diff --git a/pylops/utils/typing.py b/pylops/utils/typing.py index b7f85436..db36847b 100644 --- a/pylops/utils/typing.py +++ b/pylops/utils/typing.py @@ -10,6 +10,8 @@ ] from collections.abc import Sequence +from concurrent.futures import ThreadPoolExecutor +from multiprocessing.pool import Pool from typing import Literal import numpy as np @@ -75,3 +77,5 @@ Tpwdsmoothing = Literal["triangle", "boxcar"] Tsampler = Literal["gaussian", "rayleigh", "rademacher", "unitvector"] Tsampler2 = Literal["gaussian", "rayleigh", "rademacher"] + +Tpool = Pool | ThreadPoolExecutor diff --git a/pytests/test_combine.py b/pytests/test_combine.py index 75b8e79b..88978833 100644 --- a/pytests/test_combine.py +++ b/pytests/test_combine.py @@ -467,7 +467,7 @@ def test_BlockDiag(par): ) @pytest.mark.parametrize("par", [(par1), (par2), (par1j), (par2j)]) def test_VStack_multiproc_multithread(par): - """Single and multiprocess/multithreading consistentcy for VStack operator""" + """Single and multiprocess/multithreading consistency for VStack operator""" for parallel_kind in ["multiproc", "multithread"]: np.random.seed(0) nproc = 2 @@ -503,7 +503,7 @@ def test_VStack_multiproc_multithread(par): ) @pytest.mark.parametrize("par", [(par2), (par2j)]) def test_HStack_multiproc_multithread(par): - """Single and multiprocess/multithreading consistentcy for HStack operator""" + """Single and multiprocess/multithreading consistency for HStack operator""" for parallel_kind in ["multiproc", "multithread"]: np.random.seed(0) nproc = 2 @@ -539,7 +539,7 @@ def test_HStack_multiproc_multithread(par): ) @pytest.mark.parametrize("par", [(par1), (par2), (par1j), (par2j)]) def test_Block_multiproc_multithread(par): - """Single and multiprocess/multithreading consistentcy for Block operator""" + """Single and multiprocess/multithreading consistency for Block operator""" for parallel_kind in ["multiproc", "multithread"]: np.random.seed(0) nproc = 2 @@ -574,7 +574,7 @@ def test_Block_multiproc_multithread(par): ) @pytest.mark.parametrize("par", [(par1), (par2), (par1j), (par2j)]) def test_BlockDiag_multiproc_multithread(par): - """Single and multiprocess/multithreading consistentcy for BlockDiag operator""" + """Single and multiprocess/multithreading consistency for BlockDiag operator""" for parallel_kind in ["multiproc", "multithread"]: np.random.seed(0) nproc = 2 diff --git a/pytests/test_kronecker.py b/pytests/test_kronecker.py index ed188aa4..48199784 100644 --- a/pytests/test_kronecker.py +++ b/pytests/test_kronecker.py @@ -103,3 +103,62 @@ def test_Kroneker_Derivative(par): y = D2op * x.ravel() assert_array_equal(y, yk) + + +@pytest.mark.skipif( + int(os.environ.get("TEST_CUPY_PYLOPS", 0)) == 1, reason="Not CuPy enabled" +) +@pytest.mark.parametrize("par", [(par1), (par2), (par1s), (par2s), (par1j), (par2j)]) +def test_Kroneker_multiproc_multithread(par): + """Single and multiprocess/multithreading consistency for Kroneker operator""" + for parallel_kind in ["multiproc", "multithread"]: + np.random.seed(10) + nproc = 2 + + dtype = np.empty(0, dtype=par["dtype"]).real.dtype + + G1 = np.random.normal(0, 10, (par["ny"], par["nx"])).astype(dtype) + par[ + "imag" + ] * np.random.normal(0, 10, (par["ny"], par["nx"])).astype(dtype) + G2 = np.random.normal(0, 10, (par["ny"], par["nx"])).astype(dtype) + par[ + "imag" + ] * np.random.normal(0, 10, (par["ny"], par["nx"])).astype(dtype) + x = np.ones((par["nx"] ** 2, 4), dtype=dtype) + par["imag"] * np.ones( + (par["nx"] ** 2, 4), dtype=dtype + ) + y = np.ones((par["ny"] ** 2, 4), dtype=dtype) + par["imag"] * np.ones( + (par["ny"] ** 2, 4), dtype=dtype + ) + + Kop = Kronecker( + MatrixMult(G1, dtype=par["dtype"]), + MatrixMult(G2, dtype=par["dtype"]), + dtype=par["dtype"], + ) + Kmultiop = Kronecker( + MatrixMult(G1, dtype=par["dtype"]), + MatrixMult(G2, dtype=par["dtype"]), + nproc=nproc, + parallel_kind=parallel_kind, + dtype=par["dtype"], + ) + assert dottest( + Kmultiop, + par["ny"] ** 2, + par["nx"] ** 2, + complexflag=0 if par["imag"] == 0 else 3, + rtol=1e-4 if dtype == np.float32 else 1e-6, + backend=backend, + ) + + # forward + assert_array_almost_equal( + Kop * x, Kmultiop * x, decimal=3 if dtype == np.float32 else 8 + ) + # adjoint + assert_array_almost_equal( + Kop.H * y, Kmultiop.H * y, decimal=3 if dtype == np.float32 else 8 + ) + + # close pool + Kmultiop.close() diff --git a/pytests/test_linearoperator.py b/pytests/test_linearoperator.py index e413ee38..f42be2e9 100644 --- a/pytests/test_linearoperator.py +++ b/pytests/test_linearoperator.py @@ -1,3 +1,5 @@ +import concurrent.futures as mt +import multiprocessing as mp import os if int(os.environ.get("TEST_CUPY_PYLOPS", 0)): @@ -386,3 +388,54 @@ def test_counts(par): assert Aop.rmatvec_count == 0 assert Aop.matmat_count == 0 assert Aop.rmatmat_count == 0 + + +@pytest.mark.skipif( + int(os.environ.get("TEST_CUPY_PYLOPS", 0)) == 1, reason="Not CuPy enabled" +) +@pytest.mark.parametrize("par", [(par1), (par1j), (par2), (par2j)]) +def test_matmul_multiproc_multithread(par): + """Single and multiprocess/multithreading consistency for matmat/rmatmat""" + dtype = np.empty(0, dtype=par["dtype"]).real.dtype + + for parallel_kind in ["multiproc", "multithread"]: + np.random.seed(0) + nproc = 2 + + # create pool + if parallel_kind == "multiproc": + pool = mp.Pool(processes=nproc) + else: + pool = mt.ThreadPoolExecutor(max_workers=nproc) + + M = np.random.normal(0, 10, (par["ny"], par["nx"])).astype(dtype) + par[ + "imag" + ] * np.random.normal(0, 10, (par["ny"], par["nx"])).astype(dtype) + x = np.ones((par["nx"], 4), dtype=dtype) + par["imag"] * np.ones( + (par["nx"], 4), dtype=dtype + ) + y = np.ones((par["ny"], 4), dtype=dtype) + par["imag"] * np.ones( + (par["ny"], 4), dtype=dtype + ) + + Mop = MatrixMult(M, dtype=par["dtype"]) + + # forward + assert_array_almost_equal( + Mop.matmat(x), + Mop.matmat(x, pool=pool), + decimal=3 if dtype == np.float32 else 8, + ) + # adjoint + assert_array_almost_equal( + Mop.rmatmat(y), + Mop.rmatmat(y, pool=pool), + decimal=3 if dtype == np.float32 else 8, + ) + + # close pool + if parallel_kind == "multiproc": + pool.close() + pool.join() + else: + pool.shutdown()