Trajectree 0.0.0__py3-none-any.whl → 0.0.1__py3-none-any.whl

trajectree/quimb/quimb/linalg/mpi_launcher.py

@@ -0,0 +1,397 @@
+"""Manages the spawning of mpi processes to send to the various solvers.
+"""
+
+import os
+import functools
+
+import numpy as np
+
+from .slepc_linalg import (
+    eigs_slepc,
+    svds_slepc,
+    mfn_multiply_slepc,
+    ssolve_slepc,
+)
+from ..core import _NUM_THREAD_WORKERS
+
+# Work out if already running as mpi
+if (
+    ("OMPI_COMM_WORLD_SIZE" in os.environ)  # OpenMPI
+    or ("PMI_SIZE" in os.environ)  # MPICH
+):
+    QUIMB_MPI_LAUNCHED = "_QUIMB_MPI_LAUNCHED" in os.environ
+    ALREADY_RUNNING_AS_MPI = True
+    USE_SYNCRO = "QUIMB_SYNCRO_MPI" in os.environ
+else:
+    QUIMB_MPI_LAUNCHED = False
+    ALREADY_RUNNING_AS_MPI = False
+    USE_SYNCRO = False
+
+# default to not allowing mpi spawning capabilities
+ALLOW_SPAWN = {
+    "TRUE": True,
+    "ON": True,
+    "FALSE": False,
+    "OFF": False,
+}[os.environ.get("QUIMB_MPI_SPAWN", "False").upper()]
+
+# Work out the desired total number of workers
+for _NUM_MPI_WORKERS_VAR in (
+    "QUIMB_NUM_MPI_WORKERS",
+    "QUIMB_NUM_PROCS",
+    "OMPI_COMM_WORLD_SIZE",
+    "PMI_SIZE",
+    "OMP_NUM_THREADS",
+):
+    if _NUM_MPI_WORKERS_VAR in os.environ:
+        NUM_MPI_WORKERS = int(os.environ[_NUM_MPI_WORKERS_VAR])
+        break
+else:
+    import psutil
+
+    _NUM_MPI_WORKERS_VAR = "psutil"
+    NUM_MPI_WORKERS = psutil.cpu_count(logical=False)
+
+
+def can_use_mpi_pool():
+    """Function to determine whether we are allowed to call `get_mpi_pool`."""
+    return ALLOW_SPAWN or ALREADY_RUNNING_AS_MPI
+
+
+def bcast(result, comm, result_rank):
+    """Broadcast a result to all workers, dispatching to proper MPI (rather
+    than pickled) communication if the result is a numpy array.
+    """
+    rank = comm.Get_rank()
+
+    # make sure all workers know if result is an array or not
+    if rank == result_rank:
+        is_ndarray = isinstance(result, np.ndarray)
+    else:
+        is_ndarray = None
+    is_ndarray = comm.bcast(is_ndarray, root=result_rank)
+
+    # standard (pickle) bcast if not array
+    if not is_ndarray:
+        return comm.bcast(result, root=result_rank)
+
+    # make sure all workers have shape and dtype
+    if rank == result_rank:
+        shape_dtype = result.shape, str(result.dtype)
+    else:
+        shape_dtype = None
+
+    shape_dtype = comm.bcast(shape_dtype, root=result_rank)
+    shape, dtype = shape_dtype
+
+    # allocate data space
+    if rank != result_rank:
+        result = np.empty(shape, dtype=dtype)
+
+    # use fast communication for main array
+    comm.Bcast(result, root=result_rank)
+
+    return result
+
+
+class SyncroFuture:
+    def __init__(self, result, result_rank, comm):
+        self._result = result
+        self.result_rank = result_rank
+        self.comm = comm
+
+    def result(self):
+        rank = self.comm.Get_rank()
+
+        if rank == self.result_rank:
+            should_it = isinstance(self._result, tuple) and any(
+                isinstance(x, np.ndarray) for x in self._result
+            )
+            if should_it:
+                iterate_over = len(self._result)
+            else:
+                iterate_over = 0
+        else:
+            iterate_over = None
+        iterate_over = self.comm.bcast(iterate_over, root=self.result_rank)
+
+        if iterate_over:
+            if rank != self.result_rank:
+                self._result = (None,) * iterate_over
+
+            result = tuple(
+                bcast(x, self.comm, self.result_rank) for x in self._result
+            )
+        else:
+            result = bcast(self._result, self.comm, self.result_rank)
+
+        return result
+
+    @staticmethod
+    def cancel():
+        raise ValueError(
+            "SyncroFutures cannot be cancelled - they are "
+            "submitted in a parallel round-robin fasion where "
+            "each worker immediately computes all its results."
+        )
+
+
+class SynchroMPIPool:
+    """An object that looks like a ``concurrent.futures`` executor but actually
+    distributes tasks in a round-robin fashion based to MPI workers, before
+    broadcasting the results to each other.
+    """
+
+    def __init__(self):
+        import itertools
+        from mpi4py import MPI
+
+        self.comm = MPI.COMM_WORLD
+        self.size = self.comm.Get_size()
+        self.rank = self.comm.Get_rank()
+        self.counter = itertools.cycle(range(0, self.size))
+        self._max_workers = self.size
+
+    def submit(self, fn, *args, **kwargs):
+        # round robin iterate through ranks
+        current_counter = next(self.counter)
+
+        # accept job and compute if have the same rank, else do nothing
+        if current_counter == self.rank:
+            res = fn(*args, **kwargs)
+        else:
+            res = None
+
+        # wrap the result in a SyncroFuture, that will broadcast result
+        return SyncroFuture(res, current_counter, self.comm)
+
+    def shutdown(self):
+        pass
+
+
+class CachedPoolWithShutdown:
+    """Decorator for caching the mpi pool when called with the equivalent args,
+    and shutting down previous ones when not needed.
+    """
+
+    def __init__(self, pool_fn):
+        self._settings = "__UNINITIALIZED__"
+        self._pool_fn = pool_fn
+
+    def __call__(self, num_workers=None, num_threads=1):
+        # convert None to default so the cache the same
+        if num_workers is None:
+            num_workers = NUM_MPI_WORKERS
+        elif ALREADY_RUNNING_AS_MPI and (num_workers != NUM_MPI_WORKERS):
+            raise ValueError(
+                "Can't specify number of processes when running "
+                "under MPI rather than spawning processes."
+            )
+
+        # first call
+        if self._settings == "__UNINITIALIZED__":
+            self._pool = self._pool_fn(num_workers, num_threads)
+            self._settings = (num_workers, num_threads)
+        # new type of pool requested
+        elif self._settings != (num_workers, num_threads):
+            self._pool.shutdown()
+            self._pool = self._pool_fn(num_workers, num_threads)
+            self._settings = (num_workers, num_threads)
+        return self._pool
+
+
+@CachedPoolWithShutdown
+def get_mpi_pool(num_workers=None, num_threads=1):
+    """Get the MPI executor pool, with specified number of processes and
+    threads per process.
+    """
+    if (num_workers == 1) and (num_threads == _NUM_THREAD_WORKERS):
+        from concurrent.futures import ProcessPoolExecutor
+
+        return ProcessPoolExecutor(1)
+
+    if not QUIMB_MPI_LAUNCHED:
+        raise RuntimeError(
+            "For the moment, quimb programs using `get_mpi_pool` need to be "
+            "explicitly launched using `quimb-mpi-python`."
+        )
+
+    if USE_SYNCRO:
+        return SynchroMPIPool()
+
+    if not can_use_mpi_pool():
+        raise RuntimeError(
+            "`get_mpi_pool()` cannot be explicitly called unless already "
+            "running under MPI, or you set the environment variable "
+            "`QUIMB_MPI_SPAWN=True`."
+        )
+
+    from mpi4py.futures import MPIPoolExecutor
+
+    return MPIPoolExecutor(
+        num_workers,
+        main=False,
+        env={
+            "OMP_NUM_THREADS": str(num_threads),
+            "QUIMB_NUM_MPI_WORKERS": str(num_workers),
+            "_QUIMB_MPI_LAUNCHED": "SPAWNED",
+        },
+    )
+
+
+class GetMPIBeforeCall(object):
+    """Wrap a function to automatically get the correct communicator before
+    its called, and to set the `comm_self` kwarg to allow forced self mode.
+
+    This is called by every mpi process before the function evaluation.
+    """
+
+    def __init__(self, fn):
+        self.fn = fn
+
+    def __call__(
+        self, *args, comm_self=False, wait_for_workers=None, **kwargs
+    ):
+        """
+        Parameters
+        ----------
+        args
+            Supplied to self.fn
+        comm_self : bool, optional
+            Whether to force use of MPI.COMM_SELF
+        wait_for_workers : int, optional
+            If set, wait for the communicator to have this many workers, this
+            can help to catch some errors regarding expected worker numbers.
+        kwargs
+            Supplied to self.fn
+        """
+        from mpi4py import MPI
+
+        if not comm_self:
+            comm = MPI.COMM_WORLD
+        else:
+            comm = MPI.COMM_SELF
+
+        if wait_for_workers is not None:
+            from time import time
+
+            t0 = time()
+            while comm.Get_size() != wait_for_workers:
+                if time() - t0 > 2:
+                    raise RuntimeError(
+                        f"Timeout while waiting for {wait_for_workers} "
+                        f"workers to join comm {comm}."
+                    )
+
+        res = self.fn(*args, comm=comm, **kwargs)
+        return res
+
+
+class SpawnMPIProcessesFunc(object):
+    """Automatically wrap a function to be executed in parallel by a
+    pool of mpi workers.
+
+    This is only called by the master mpi process in manual mode, only by
+    the (non-mpi) spawning process in automatic mode, or by all processes in
+    syncro mode.
+    """
+
+    def __init__(self, fn):
+        self.fn = fn
+
+    def __call__(
+        self,
+        *args,
+        num_workers=None,
+        num_threads=1,
+        mpi_pool=None,
+        spawn_all=USE_SYNCRO or (not ALREADY_RUNNING_AS_MPI),
+        **kwargs,
+    ):
+        """
+        Parameters
+        ----------
+        args
+            Supplied to `self.fn`.
+        num_workers : int, optional
+            How many total process should run function in parallel.
+        num_threads : int, optional
+            How many (OMP) threads each process should use
+        mpi_pool : pool-like, optional
+            If not None (default), submit function to this pool.
+        spawn_all : bool, optional
+            Whether all the parallel processes should be spawned (True), or
+            num_workers - 1, so that the current process can also do work.
+        kwargs
+            Supplied to `self.fn`.
+
+        Returns
+        -------
+        `fn` output from the master process.
+        """
+        if num_workers is None:
+            num_workers = NUM_MPI_WORKERS
+
+        if (
+            # use must explicitly run program as
+            (not can_use_mpi_pool())
+            or
+            # no pool or communicator needed
+            (num_workers == 1)
+        ):
+            return self.fn(*args, comm_self=True, **kwargs)
+
+        kwargs["wait_for_workers"] = num_workers
+
+        if mpi_pool is not None:
+            pool = mpi_pool
+        else:
+            pool = get_mpi_pool(num_workers, num_threads)
+
+        # the (non mpi) main process is idle while the workers compute.
+        if spawn_all:
+            futures = [
+                pool.submit(self.fn, *args, **kwargs)
+                for _ in range(num_workers)
+            ]
+            results = [f.result() for f in futures]
+
+        # the master process is the master mpi process and contributes
+        else:
+            futures = [
+                pool.submit(self.fn, *args, **kwargs)
+                for _ in range(num_workers - 1)
+            ]
+            results = [self.fn(*args, **kwargs)] + [
+                f.result() for f in futures
+            ]
+
+        # Get master result, (not always first submitted)
+        return next(r for r in results if r is not None)
+
+
+# ---------------------------------- SLEPC ---------------------------------- #
+
+eigs_slepc_mpi = functools.wraps(eigs_slepc)(GetMPIBeforeCall(eigs_slepc))
+eigs_slepc_spawn = functools.wraps(eigs_slepc)(
+    SpawnMPIProcessesFunc(eigs_slepc_mpi)
+)
+
+svds_slepc_mpi = functools.wraps(svds_slepc)(GetMPIBeforeCall(svds_slepc))
+svds_slepc_spawn = functools.wraps(svds_slepc)(
+    SpawnMPIProcessesFunc(svds_slepc_mpi)
+)
+
+mfn_multiply_slepc_mpi = functools.wraps(mfn_multiply_slepc)(
+    GetMPIBeforeCall(mfn_multiply_slepc)
+)
+mfn_multiply_slepc_spawn = functools.wraps(mfn_multiply_slepc)(
+    SpawnMPIProcessesFunc(mfn_multiply_slepc_mpi)
+)
+
+ssolve_slepc_mpi = functools.wraps(ssolve_slepc)(
+    GetMPIBeforeCall(ssolve_slepc)
+)
+ssolve_slepc_spawn = functools.wraps(ssolve_slepc)(
+    SpawnMPIProcessesFunc(ssolve_slepc_mpi)
+)

trajectree/quimb/quimb/linalg/numpy_linalg.py

@@ -0,0 +1,244 @@
+"""Numpy base linear algebra."""
+
+import numpy as np
+import numpy.linalg as nla
+import scipy.linalg as scla
+
+import quimb as qu
+from .autoblock import eigensystem_autoblocked
+
+
+_NUMPY_EIG_FUNCS = {
+    (True, True): nla.eigh,
+    (True, False): nla.eig,
+    (False, True): nla.eigvalsh,
+    (False, False): nla.eigvals,
+}
+
+
+def eig_numpy(A, sort=True, isherm=True, return_vecs=True, autoblock=False):
+    """Numpy based dense eigensolve.
+
+    Parameters
+    ----------
+    A : array_like
+        The operator to decompose.
+    sort : bool, optional
+        Whether to sort into ascending order.
+    isherm : bool, optional
+        Whether ``A`` is hermitian.
+    return_vecs : bool, optional
+        Whether to return the eigenvectors.
+    autoblock : bool, optional
+        If true, automatically identify and exploit symmetries appearing in the
+        current basis as block diagonals formed via permutation of rows and
+        columns.
+
+    Returns
+    -------
+    evals : 1D-array
+        The eigenvalues.
+    evecs : qarray
+        If ``return_vecs=True``, the eigenvectors.
+    """
+    if autoblock:
+        return eigensystem_autoblocked(
+            A, sort=sort, isherm=isherm, return_vecs=return_vecs
+        )
+
+    evals = _NUMPY_EIG_FUNCS[return_vecs, isherm](A)
+
+    if return_vecs:
+        evals, evecs = evals
+
+        if sort:
+            sortinds = np.argsort(evals)
+            evals, evecs = evals[sortinds], evecs[:, sortinds]
+
+        return evals, qu.qarray(evecs)
+
+    if sort:
+        evals.sort()
+
+    return evals
+
+
+def sort_inds(a, method, sigma=None):
+    """Return the sorting inds of a list
+
+    Parameters
+    ----------
+        a : array_like
+            List to base sort on.
+        method : str
+            Method of sorting list, one of
+                * "LM" - Largest magnitude first
+                * "SM" - Smallest magnitude first
+                * "SA" - Smallest algebraic first
+                * "SR" - Smallest real part first
+                * "SI" - Smallest imaginary part first
+                * "LA" - Largest algebraic first
+                * "LR" - Largest real part first
+                * "LI" - Largest imaginary part first
+                * "TM" - Magnitude closest to target sigma first
+                * "TR" - Real part closest to target sigma first
+                * "TI" - Imaginary part closest to target sigma first
+        sigma : float, optional
+            The target if method={"TM", "TR", or "TI"}.
+
+    Returns
+    -------
+        inds : array of int
+            Indices that would sort `a` based on `method`
+    """
+    _SORT_FUNCS = {
+        "LM": lambda a: -abs(a),
+        "SM": lambda a: -abs(1 / a),
+        "SA": lambda a: a,
+        "SR": lambda a: a.real,
+        "SI": lambda a: a.imag,
+        "LA": lambda a: -a,
+        "LR": lambda a: -a.real,
+        "LI": lambda a: -a.imag,
+        "TM": lambda a: -1 / abs(abs(a) - sigma),
+        "TR": lambda a: -1 / abs(a.real - sigma),
+        "TI": lambda a: -1 / abs(a.imag - sigma),
+    }
+    return np.argsort(_SORT_FUNCS[method.upper()](a))
+
+
+_DENSE_EIG_METHODS = {
+    (True, True, False): nla.eigh,
+    (True, False, False): nla.eigvalsh,
+    (False, True, False): nla.eig,
+    (False, False, False): nla.eigvals,
+    (True, True, True): scla.eigh,
+    (True, False, True): scla.eigvalsh,
+    (False, True, True): scla.eig,
+    (False, False, True): scla.eigvals,
+}
+
+
+def eigs_numpy(
+    A,
+    k,
+    B=None,
+    which=None,
+    return_vecs=True,
+    sigma=None,
+    isherm=True,
+    P=None,
+    sort=True,
+    **eig_opts,
+):
+    """Partial eigen-decomposition using numpy's dense linear algebra.
+
+    Parameters
+    ----------
+    A : array_like or quimb.Lazy
+        Operator to partially eigen-decompose.
+    k : int
+        Number of eigenpairs to return.
+    B : array_like or quimb.Lazy
+        If given, the RHS operator defining a generalized eigen problem.
+    which : str, optional
+        Which part of the spectrum to target.
+    return_vecs : bool, optional
+        Whether to return eigenvectors.
+    sigma : None or float, optional
+        Target eigenvalue.
+    isherm : bool, optional
+        Whether `a` is hermitian.
+    P : array_like or quimb.Lazy
+        Perform the eigensolve in the subspace defined by this projector.
+    sort : bool, optional
+        Whether to sort reduced list of eigenpairs into ascending order.
+    eig_opts
+        Settings to pass to numpy.eig... functions.
+
+    Returns
+    -------
+        lk, (vk): k eigenvalues (and eigenvectors) sorted according to which
+    """
+    if isinstance(A, qu.Lazy):
+        A = A()
+    if isinstance(B, qu.Lazy):
+        B = B()
+    if isinstance(P, qu.Lazy):
+        P = P()
+
+    # project into subspace
+    if P is not None:
+        A = qu.dag(P) @ (A @ P)
+
+    generalized = B is not None
+
+    eig_fn = _DENSE_EIG_METHODS[(isherm, return_vecs, generalized)]
+
+    if generalized:
+        eig_opts["b"] = B
+
+    # these might be given for partial eigsys but not relevant for numpy
+    eig_opts.pop("ncv", None)
+    eig_opts.pop("v0", None)
+    eig_opts.pop("tol", None)
+    eig_opts.pop("maxiter", None)
+    eig_opts.pop("EPSType", None)
+
+    if return_vecs:
+        # get all eigenpairs
+        lk, vk = eig_fn(A.toarray() if qu.issparse(A) else A, **eig_opts)
+
+        # sort and trim according to which k we want
+        sk = sort_inds(lk, method=which, sigma=sigma)[:k]
+        lk, vk = lk[sk], vk[:, sk]
+
+        # also potentially sort into ascending order
+        if sort:
+            so = np.argsort(lk)
+            lk, vk = lk[so], vk[:, so]
+
+        # map eigenvectors out of subspace
+        if P is not None:
+            vk = P @ vk
+
+        return lk, qu.qarray(vk)
+
+    else:
+        # get all eigenvalues
+        lk = eig_fn(A.toarray() if qu.issparse(A) else A, **eig_opts)
+
+        # sort and trim according to which k we want
+        sk = sort_inds(lk, method=which, sigma=sigma)[:k]
+        lk = lk[sk]
+
+        # also potentially sort into ascending order
+        return np.sort(lk) if sort else lk
+
+
+def svds_numpy(a, k, return_vecs=True, **_):
+    """Partial singular value decomposition using numpys (full) singular value
+    decomposition.
+
+    Parameters
+    ----------
+    a : array_like
+        Operator to decompose.
+    k : int, optional
+        Number of singular value triplets to retrieve.
+    return_vecs : bool, optional
+        whether to return the computed vecs or values only
+
+    Returns
+    -------
+    (uk,) sk (, vkt) :
+        Singlar value triplets.
+    """
+    if return_vecs:
+        uk, sk, vkt = nla.svd(
+            a.toarray() if qu.issparse(a) else a, compute_uv=True
+        )
+        return qu.qarray(uk[:, :k]), sk[:k], qu.qarray(vkt[:k, :])
+    else:
+        sk = nla.svd(a.toarray() if qu.issparse(a) else a, compute_uv=False)
+        return sk[:k]