Source code for imagine.tools.parallel_ops

"""
Interface module which allows automatically switching between the routines in
the :py:mod:`imagine.tools.mpi_helper` module and their:py:mod:`numpy`
or pure Python equivalents, depending on the contents of
:py:data:`imagine.rc['distributed_arrays']`
"""

# %% IMPORTS
# Package imports
from e13tools import add_to_all
import numpy as np

# IMAGINE imports
from imagine.tools import mpi_helper as m, rc

# All declaration
__all__ = []


# %% FUNCTION DEFINITIONS
[docs]@add_to_all def pshape(data): """ :py:func:`imagine.tools.mpi_helper.mpi_shape` or :py:meth:`numpy.ndarray.shape` depending on :py:data:`imagine.rc['distributed_arrays']`. """ if rc['distributed_arrays']: return m.mpi_shape(data) else: return data.shape
[docs]@add_to_all def prosecutor(data): """ :py:func:`imagine.tools.mpi_helper.mpi_prosecutor` or *nothing* depending on :py:data:`imagine.rc['distributed_arrays']`. """ if rc['distributed_arrays']: m.mpi_prosecutor(data)
[docs]@add_to_all def pmean(data): """ :py:func:`imagine.tools.mpi_helper.mpi_mean` or :py:func:`numpy.mean` depending on :py:data:`imagine.rc['distributed_arrays']`. """ if rc['distributed_arrays']: return m.mpi_mean(data) else: return (np.mean(data, axis=0)).reshape(1, -1)
[docs]@add_to_all def pvar(data): """ :py:func:`imagine.tools.mpi_helper.mpi_var` or :py:func:`numpy.var` depending on :py:data:`imagine.rc['distributed_arrays']`. """ if rc['distributed_arrays']: # This will be done later! raise NotImplementedError else: return data.var(axis=0)
[docs]@add_to_all def ptrans(data): """ :py:func:`imagine.tools.mpi_helper.mpi_mean` or :py:meth:`numpy.ndarray.T` depending on :py:data:`imagine.rc['distributed_arrays']`. """ if rc['distributed_arrays']: return m.mpi_trans(data) else: return data.T
[docs]@add_to_all def pmult(left, right): """ :py:func:`imagine.tools.mpi_helper.mpi_mult` or :py:meth:`numpy.matmul` depending on :py:data:`imagine.rc['distributed_arrays']`. """ if rc['distributed_arrays']: return m.mpi_mult(left, right) else: return left @ right
[docs]@add_to_all def ptrace(data): """ :py:func:`imagine.tools.mpi_helper.mpi_trace` or :py:func:`numpy.trace` depending on :py:data:`imagine.rc['distributed_arrays']`. """ if rc['distributed_arrays']: return m.mpi_trace(data) else: return np.trace(data)
[docs]@add_to_all def pdiag(data): """ :py:func:`imagine.tools.mpi_helper.mpi_diag` or :py:func:`numpy.diagonal` depending on :py:data:`imagine.rc['distributed_arrays']`. """ if rc['distributed_arrays']: return m.mpi_diag(data) else: return data.diagonal()
[docs]@add_to_all def pnewdiag(data): """ :py:func:`imagine.tools.mpi_helper.mpi_new_diag` or :py:func:`numpy.diag` depending on :py:data:`imagine.rc['distributed_arrays']`. """ if rc['distributed_arrays']: return m.mpi_new_diag(data) else: return np.diag(data)
[docs]@add_to_all def peye(size): """ :py:func:`imagine.tools.mpi_helper.mpi_eye` or :py:func:`numpy.eye` depending on :py:data:`imagine.rc['distributed_arrays']`. """ if rc['distributed_arrays']: return m.mpi_eye(size) else: return np.eye(size)
[docs]@add_to_all def distribute_matrix(full_matrix): """ :py:func:`imagine.tools.mpi_helper.mpi_distribute_matrix` or *nothing* depending on :py:data:`imagine.rc['distributed_arrays']`. """ if rc['distributed_arrays']: return m.mpi_distribute_matrix(full_matrix) else: return full_matrix
[docs]@add_to_all def plu_solve(operator, source): """ :py:func:`imagine.tools.mpi_helper.mpi_lu_solve` or :py:func:`numpy.linalg.solve` depending on :py:data:`imagine.rc['distributed_arrays']`. Notes ----- In the non-distributed case, the source is transposed before the calculation """ if rc['distributed_arrays']: return m.mpi_lu_solve(operator, source) else: return np.linalg.solve(operator, source.T)
[docs]@add_to_all def pslogdet(data): """ :py:func:`imagine.tools.mpi_helper.mpi_slogdet` or :py:func:`numpy.linalg.slogdet` depending on :py:data:`imagine.rc['distributed_arrays']`. """ if rc['distributed_arrays']: return m.mpi_slogdet(data) else: return np.linalg.slogdet(data)
[docs]@add_to_all def pglobal(data): """ :py:func:`imagine.tools.mpi_helper.mpi_global` or *nothing* depending on :py:data:`imagine.rc['distributed_arrays']`. """ if rc['distributed_arrays']: return m.mpi_global(data) else: return data
[docs]@add_to_all def plocal(data): """ :py:func:`imagine.tools.mpi_helper.mpi_local` or *nothing* depending on :py:data:`imagine.rc['distributed_arrays']`. """ if rc['distributed_arrays']: return m.mpi_local(data) else: return data