dubfi.linalg.mpipad.worker ========================== .. py:module:: dubfi.linalg.mpipad.worker .. autoapi-nested-parse:: MPI linear algebra worker script (private module). This script is run in parallel by multiple MPI processes (workers). Each worker has the same set of variables stored on the respective processor. The workers await and follow instructions from the parent. Instructions include creating objects, assigning variables, and doing linear algebra operations on these variables. This module shall be used jointly with mpi_parent.py and does not have a stable API by itself! .. codeauthor:: Valentin Bruch, DWD .. versionadded:: 0.1.0 (initial release) Attributes ---------- .. autoapisummary:: dubfi.linalg.mpipad.worker.comm dubfi.linalg.mpipad.worker.rank Classes ------- .. autoapisummary:: dubfi.linalg.mpipad.worker.MpiAbstractVector dubfi.linalg.mpipad.worker.MpiVectorWorker dubfi.linalg.mpipad.worker.MpiVectorWorkerZero dubfi.linalg.mpipad.worker.MpiAbstractOperator dubfi.linalg.mpipad.worker.MpiDenseBdMatrixWorker dubfi.linalg.mpipad.worker.MpiParamVectorWorker dubfi.linalg.mpipad.worker.MpiLinParamVectorWorker dubfi.linalg.mpipad.worker.MpiDensePostRWorker dubfi.linalg.mpipad.worker.MpiGradOpWorker dubfi.linalg.mpipad.worker.MpiTraceProdOpWorker Functions --------- .. autoapisummary:: dubfi.linalg.mpipad.worker.receive_array dubfi.linalg.mpipad.worker._trace_product_self dubfi.linalg.mpipad.worker._trace_product_generic Module Contents --------------- .. py:data:: comm :type: mpi4py.MPI.Intracomm | mpi4py.MPI.Intercomm MPI communication object .. py:data:: rank :type: int Rank (identifier) of this MPI process .. py:function:: receive_array() -> numpy.ndarray Receive numpy array of unknown shape from parent. .. py:class:: MpiAbstractVector Bases: :py:obj:`dubfi.linalg.types.AbstractVector` AbstractVector that sends reduce operations to MPI parent. .. py:method:: apply(other) Compute scalar product with other vector (along unbuffered segment) and send result to parent. .. py:method:: dot_array() Compute scalar product with array received from parent along some non-vector dimension. .. py:method:: tonumpy() Compute numpy array representation of unbuffered segment and send result to parent. .. py:method:: _apply(other: dubfi.linalg.types.AbstractVector) -> numpy.ndarray | float :abstractmethod: Compute scalar product with other vector (along unbuffered segment). .. py:class:: MpiVectorWorker Bases: :py:obj:`MpiAbstractVector` Distributed vector: Part of a vector handled by this worker process. Distributed AbstractVector. .. py:property:: shape Shape of this part of the vector. .. py:method:: fromparent() :classmethod: Create new instance and receive data from parent to initialize it. .. py:method:: fromdata(data: numpy.ndarray, other) :classmethod: Create new instance and initialize it with given data. :param data: array to construct vector, last dimension must be vector dimension :type data: np.ndarray :param other: some MPI worker object from which start_idx and end_idx can be copied :type other: (object providing start_idx and end_idx attributes) .. py:method:: iscompatible(other: dubfi.linalg.types.AbstractVector) -> bool Check for compatibility with other vector. .. py:property:: data Numpy array representation of self (buffered segment). .. py:method:: _apply(other: dubfi.linalg.types.AbstractVector) -> numpy.ndarray Compute scalar product with other vector (along unbuffered segment). .. py:method:: __add__(other) Add two vectors, applying broadcasting rules. .. py:method:: __iadd__(other) Add vector in-place. .. py:method:: __sub__(other) Subtract two vectors, applying broadcasting rules. .. py:method:: __isub__(other) Subtract vector in-place. .. py:method:: __mul__(other) Multiply by scalar or numpy array treated as scalar along vector space dimension. .. py:method:: mul_array() Multiply by array received from parent. .. py:method:: dot(other) Compute scalar product along some non-vector dimension. .. py:method:: iszero() Check if all elements are zero and send result to parent. .. py:class:: MpiVectorWorkerZero(shape, start_idx: int = -1, end_idx: int = -1) Bases: :py:obj:`MpiVectorWorker` Distributed all-zero vector. Distributed all-zero AbstractVector. .. py:property:: shape Shape of this part of the vector. .. py:property:: data Shape of this part of the vector. .. py:method:: iszero() Check if all elements are zero and send result to parent (always true). .. py:method:: tonumpy() Compute numpy array representation of unbuffered segment and send result to parent. .. py:class:: MpiAbstractOperator Bases: :py:obj:`dubfi.linalg.types.AbstractOperator` AbstractOperator that sends reduce operations to MPI parent. .. py:method:: trace() Compute trace of self and send result to parent. .. py:method:: trace_product(other) Compute trace of (self @ other) and send result to parent. .. py:method:: sandwich(vec) Compute vec @ self @ vec and send result to parent. .. py:method:: logdet() Compute log(det(self)) and send result to parent. This assumes that self is a positive definite, real-symmetric matrix. .. py:method:: dot_array() -> dubfi.linalg.types.AbstractOperator Compute scalar product with array received from parent along some non-vector dimension. .. py:method:: validate_uncertainty_statistics(vec: dubfi.linalg.types.AbstractVector) -> tuple[numpy.ndarray, numpy.ndarray] Indicate whether vec is a likely realization of a Gaussian random variable with variance self. Assume that self is the error covariance matrix of a Gaussian random variable X with mean 0. Provide indications whether X is a likely realization of X in the following form: Diagonalize self: self = V @ diag(D) @ V.T where D is the vector of eigenvalues of self and V is an orthogonal matrix. Return (V.T @ vec, D). One can expect that (V.T @ vec) / D is a Gaussian random variable with mean 0 and standard deviation 1. .. versionadded:: 0.1.3 .. py:class:: MpiDenseBdMatrixWorker Bases: :py:obj:`MpiAbstractOperator` Dense matrix representing part of distributed, approximately diagonal, linear operator. Dense matrix representation of segment of linear operator. .. py:property:: shape Shape of array (matrix) representation of self. .. py:property:: symmetric True if self equals its adjoint. .. py:method:: fromparent() :classmethod: Construct from parent by receiving all data. .. seealso:: :meth:`dubfi.linalg.mpipad.main.MpiDenseBdMatrix.fromarray` .. py:method:: fromdata(data: numpy.ndarray, other, symmetric: bool = False) :classmethod: Construct from given data. .. py:property:: data Numpy array representation of self. .. py:method:: iscompatible(other: dubfi.linalg.types.AbstractOperator) -> bool Check whether other is an MpiDenseBdMatrixWorker of compatible shape. .. py:method:: apply(vec) Apply operator on vector: self @ vec. .. py:method:: rapply(vec) Apply transpose operator on vector: self.T @ vec. .. py:method:: diagonal() Return diagonal of self as MPI vector worker. .. py:method:: __mul__(other) Multiply element-wise. .. py:method:: mul_array() Multiply element-wise by array received from parent. .. py:method:: inv() Compute (multiplicative) inverse operator. .. py:method:: solve(vec) Solve linear equation self @ x = vec for x. .. py:method:: _trace() -> numpy.ndarray | float Compute trace of self. .. py:method:: _trace_product(other: dubfi.linalg.types.AbstractOperator) -> numpy.ndarray Compute trace of (self @ other). .. py:method:: chain(other) Combine operators: self @ other. .. py:method:: _logdet() Compute log(det(self)) and send result to parent. This assumes that self is a positive definite, real-symmetric matrix. .. py:method:: dot(other: numpy.ndarray) Inner product along non-vector dimensions. .. py:method:: tonumpy() Compute numpy array (matrix) representation of unbuffered segment and send result to parent. .. py:method:: __add__(other) Add other operator. .. py:method:: __iadd__(other) Add other operator in-place. .. py:method:: __sub__(other) Subtract other operator. .. py:method:: __isub__(other) Subtract other operator in-place. .. py:method:: _sandwich(vec) Compute vec @ self @ vec and send result to parent. .. py:method:: _validate_uncertainty_statistics(vec) -> tuple[numpy.ndarray, numpy.ndarray] Indicate whether vec is a likely realization of a Gaussian random variable with variance self. Assume that self is the error covariance matrix of a Gaussian random variable X with mean 0. Provide indications whether X is a likely realization of X in the following form: Diagonalize self: self = V @ diag(D) @ V.T where D is the vector of eigenvalues of self and V is an orthogonal matrix. Return (V.T @ vec, D). One can expect that (V.T @ vec) / D is a Gaussian random variable with mean 0 and standard deviation 1. .. versionadded:: 0.1.3 .. py:class:: MpiParamVectorWorker Bases: :py:obj:`dubfi.linalg.generic.ParametrizedVector` Part of distributed, linearly parametrize vector. Worker of distributed linear parametrized vector. .. py:property:: size_s :type: int :abstractmethod: Size of expected input state. .. py:method:: call_mpi() Compute vector at parameters received from parent. .. py:method:: grad_mpi() Compute gradient of vector at parameters received from parent. .. py:method:: hess_mpi() Compute Hesse matrix of vector at parameters received from parent. .. py:class:: MpiLinParamVectorWorker Bases: :py:obj:`MpiParamVectorWorker` Part of distributed, linearly parametrize vector. Worker of distributed linear parametrized vector. .. py:property:: size_s :type: int Size of expected input state. .. py:property:: shape Shape of array representation of self. .. py:method:: fromparent() :classmethod: Create new instance and receive data from parent to initialize it. .. py:method:: __call__(s) Compute vector at given parameters. .. py:method:: grad(s) Compute gradient of vector at given parameters. .. py:method:: hess(s) Compute Hesse matrix of vector at given parameters. .. py:class:: MpiDensePostRWorker Bases: :py:obj:`dubfi.linalg.generic.ParametrizedOperator` Segment of distributed, parametrized linear operator constructed from flux ensemble. Dense matrix representation of distributed R matrix. .. py:method:: fromparent() :classmethod: Create new instance and receive data from parent to initialize it. .. py:method:: __call__(s) Compute operator at given parameters. .. py:method:: call_mpi() Compute operator at parameters received from parent. .. py:method:: grad(s) Compute gradient of operator at given parameters. .. py:method:: grad_dense(s: numpy.ndarray) -> MpiDenseBdMatrixWorker More flexible than :meth:`~grad`, but requires more memory. .. py:method:: grad_mpi() Compute gradient of operator at parameters received from parent. .. py:method:: hess(s) Compute Hesse matrix of operator at given parameters. .. py:method:: hess_dense(s: numpy.ndarray) -> MpiDenseBdMatrixWorker More flexible than :meth:`~hess`, but requires more memory. .. py:method:: hess_mpi() Compute Hesse matrix of operator at parameters received from parent. .. py:method:: diagonal_obs() -> MpiVectorWorker Observation and regularization contribution to diagonal. .. py:method:: diagonal_bc() -> MpiVectorWorker Boundary condition uncertainty contribution to diagonal. .. py:method:: diagonal_bg() -> MpiVectorWorker Background meteorology uncertainty contribution to diagonal. .. py:class:: MpiGradOpWorker(other, x, chain: None | MpiDenseBdMatrixWorker = None) Bases: :py:obj:`MpiAbstractOperator` Segment of distributed gradient operator, used for gradient of :class:`MpiDensePostRWorker`. Construct auxiliary operator for computing gradient of R matrix. Define operator G of shape (state, obs, obs) by:: G[s, i, j] = sum H[m, s, i] X[m, j] L[i, j] + (i <-> j) m if chain is None or:: G[s, i, j] = sum ( H[m, s, i] X[m, k] L[i, k] + (i <-> k) ) C[k, j] m,k where C=chain, X=x, H=other.h_ens, L=other.loc and it is assumed that L == L.T .. py:property:: shape Shape of array (matrix) representation of self. .. py:method:: apply(vec) Apply operator on vector: self @ vec. .. py:method:: chain(other) Combine operators: self @ other. .. py:method:: tonumpy() :abstractmethod: Numpy array (matrix) representation of self. .. py:method:: solve(vec) :abstractmethod: Solve linear equation self @ x = vec for x. .. py:method:: __mul__(other) :abstractmethod: Multiply element-wise. .. py:method:: _trace() Compute trace of self and send result to parent. .. py:method:: _trace_product(other) Compute trace(self @ other) and send result to parent. Note: This is usually by far the slowest part of the inversion. .. py:method:: dot(other) Inner product along non-vector dimensions. .. py:function:: _trace_product_self(s: int, e: int, hens: numpy.ndarray, x: numpy.ndarray, loc: numpy.ndarray, chainT: numpy.ndarray) -> numpy.ndarray Compute part of trace product for :class:`MpiGradOpWorker` (special case). Helper for :meth:`MpiGradOpWorker._trace_product` in commonly used case. Runtime: O(n_state * n_obs**3) + O(n_state * n_ens * n_obs**2) + O(n_state**2 * n_ens * n_obs) Required memory (theoretically): 16 * n_obs**2 bytes per numba thread :param s: start index in reduced observation dimension :type s: int :param e: end index in reduced observation dimension :type e: int :param hens: centered ensemble observation operator, shape (n_ens, n_state, n_obs) :type hens: array :param x: ensemble of states, shape (n_ens, n_obs) :type x: array :param loc: localization, real-symmetric matrix, shape (n_obs, n_obs) :type loc: array :param chainT: chained array, real-symmetric matrix, shape (n_obs, n_obs) :type chainT: array .. py:function:: _trace_product_generic(s: int, e: int, hens: numpy.ndarray, x: numpy.ndarray, loc: numpy.ndarray, chain: numpy.ndarray, ohens: numpy.ndarray, ox: numpy.ndarray, oloc: numpy.ndarray, ochain: numpy.ndarray) -> numpy.ndarray Compute part of trace product for :class:`MpiGradOpWorker` (generic). Helper for :meth:`MpiGradOpWorker._trace_product` in general case. .. py:class:: MpiTraceProdOpWorker(ensemble: numpy.ndarray, loc: numpy.ndarray, other) Bases: :py:obj:`MpiAbstractOperator` Segment of distributed trace product operator. Linear operator of limited functionality to be used in trace products. .. py:property:: shape Shape of array (matrix) representation of self. .. py:method:: _trace_product(other) Compute trace of (self @ other). .. py:method:: _trace() Compute trace of self. .. py:method:: _sandwich(vec) Compute vec @ self @ vec. .. py:method:: apply(other) :abstractmethod: Apply operator on vector: self @ vec (not implemented). .. py:method:: rapply(other) :abstractmethod: Apply transpose operator on vector: self.T @ vec (not implemented). .. py:method:: __mul__(other) :abstractmethod: Multiply element-wise (not implemented). .. py:method:: solve(vec) :abstractmethod: Solve linear equation self @ x = vec for x (not implemented). .. py:method:: tonumpy() :abstractmethod: Numpy array (matrix) representation of self (not implemented). .. py:method:: dot(other) Inner product along non-vector dimensions.