Source code for Muscat.Helpers.MPIInterface

# -*- coding: utf-8 -*-
#
# This file is subject to the terms and conditions defined in
# file 'LICENSE.txt', which is part of this source code package.
#




class MPIInterface():
    """
    Class to test if the mpi interface is available
    This class work even if the module mpi4py is not available

    more information on https://mpi4py.readthedocs.io/

    .. note::
        from (https://mpi4py.readthedocs.io/en/stable/overview.html)
        MPI_Init() or MPI_Init_thread() is actually called when you import the MPI module from the mpi4py package
        (this is done internally if mpi4py is installed), but only if MPI is not already initialized. In such case,
        calling Init or Init_thread from Python is expected to generate an MPI error, and in turn an exception will
        be raised.

        MPI_Finalize() is registered (by using Python C/API function Py_AtExit()) for being automatically called when
        Python processes exit, but only if mpi4py (or the Muscat MPIInterface) actually initialized MPI. Therefore,
        there is no need to call Finalize from Python to ensure MPI finalization.
    """

    def __init__(self) -> None:
        try:  # pragma: no coverage
            from mpi4py import MPI
            self.MPI = MPI
            self.mpiOK = True
            self.parallel = True
            self.comm = MPI.COMM_WORLD
            self.rank = self.comm.Get_rank()
            self.size = self.comm.Get_size()
            if self.size == 1:
                self.parallel = False
        except:
            self.MPI = None
            self.mpiOK = False
            self.parallel = False

            self.comm = None
            self.rank = 0
            self.size = 1



    @classmethod
    def IsParallel(cls) -> bool:
        """return True if the main program was launched in a mpi environment (mpirun with more than 1 core)

        Returns
        -------
        bool
            True there is more than one instance running this code (I.e. mpi with more than 1 core).
            False otherwise
        """
        return MPIInterface().parallel




    @classmethod
    def Rank(cls) -> int:
        """return the rank of the current process (0 if not in a mpi environment)

        Returns
        -------
        int
            Return the rank of the current process
        """
        return MPIInterface().rank




    @classmethod
    def Size(cls) -> int:
        """return the size of the mpi environment (0 if not in a mpi environment)

        Returns
        -------
        int
            The number of mpi processes running
        """
        return MPIInterface().size


    def __str__(self) -> str:
        """return a string representation of the class

        Returns
        -------
        str
            return a string representation of the class
        """
        res = "MPIInterface : " + ("ok" if self.mpiOK else "no") + "\n"
        res += 'Total size is ' + str(self.size) + "\n"
        res += 'My rank is ' + str(self.rank)
        return res





def CheckIntegrity(GUI:bool=False):
    obj = MPIInterface()
    assert (obj.IsParallel() == 0)
    assert (obj.Rank() == 0)
    assert (obj.Size() == 1)
    print(str(obj.__str__()))
    return "ok"



if __name__ == '__main__':
    print(CheckIntegrity(GUI=True))  # pragma: no cover