fixed the tests and circular import case

.github/workflows/tests.yaml

@@ -21,7 +21,7 @@ jobs:
         with:
           submodules: true
 
-      - name: Set up MPI ${{ matrix.mpi }}
+      - name: Install MPI ${{ matrix.mpi }}
         uses: mpi4py/setup-mpi@v1
         with:
           mpi: ${{ matrix.mpi }}

brahmap/__init__.py

@@ -7,18 +7,19 @@
 if MPI.Is_initialized() is False:
     MPI.Init_thread(required=MPI.THREAD_FUNNELED)
 
-from . import interfaces, utilities, linop, mapmakers, _extensions  # noqa: E402
+bMPI = None
 
-from .utilities import Initialize, MPI_RAISE_EXCEPTION  # noqa: E402
+from .mpi import Initialize, MPI_RAISE_EXCEPTION  # noqa: E402
 
-bMPI = None
+from . import linop, _extensions, interfaces, utilities, mapmakers  # noqa: E402
 
 __all__ = [
+    "bMPI",
+    "Initialize",
+    "MPI_RAISE_EXCEPTION",
+    "linop",
+    "_extensions",
     "interfaces",
     "utilities",
-    "linop",
     "mapmakers",
-    "_extensions",
-    "Initialize",
-    "MPI_RAISE_EXCEPTION",
 ]

brahmap/_extensions/compute_weights.cpp

@@ -3,8 +3,6 @@
 #include <pybind11/numpy.h>
 #include <pybind11/pybind11.h>
 
-#include <omp.h>
-
 #include "mpi_utils.hpp"
 
 namespace py = pybind11;

brahmap/interfaces/linearoperators.py

@@ -1,15 +1,19 @@
 import numpy as np
 import warnings
 
-import brahmap
+
 from brahmap.linop import linop as lp
 from brahmap.linop import blkop as blk
+
 from brahmap.utilities import ProcessTimeSamples, TypeChangeWarning
 
 from brahmap._extensions import PointingLO_tools
 from brahmap._extensions import BlkDiagPrecondLO_tools
 from brahmap._extensions import InvNoiseCov_tools
 
+from brahmap import MPI_RAISE_EXCEPTION
+import brahmap
+
 
 class PointingLO(lp.LinearOperator):
     r"""Derived class from the one from the  :class:`LinearOperator` in :mod:`linop`.
@@ -84,7 +88,7 @@ def _mult_I(self, vec: np.ndarray):
 
         """
 
-        brahmap.MPI_RAISE_EXCEPTION(
+        MPI_RAISE_EXCEPTION(
             condition=(len(vec) != self.ncols),
             exception=ValueError,
             message=f"Dimensions of `vec` is not compatible with the dimension of this `PointingLO` instance.\nShape of `PointingLO` instance: {self.shape}\nShape of `vec`: {vec.shape}",
@@ -116,7 +120,7 @@ def _rmult_I(self, vec: np.ndarray):
 
         """
 
-        brahmap.MPI_RAISE_EXCEPTION(
+        MPI_RAISE_EXCEPTION(
             condition=(len(vec) != self.nrows),
             exception=ValueError,
             message=f"Dimensions of `vec` is not compatible with the dimension of this `PointingLO` instance.\nShape of `PointingLO` instance: {self.shape}\nShape of `vec`: {vec.shape}",
@@ -153,7 +157,7 @@ def _mult_QU(self, vec: np.ndarray):
             d_t=  Q_p \cos(2\phi_t)+ U_p \sin(2\phi_t).
         """
 
-        brahmap.MPI_RAISE_EXCEPTION(
+        MPI_RAISE_EXCEPTION(
             condition=(len(vec) != self.ncols),
             exception=ValueError,
             message=f"Dimensions of `vec` is not compatible with the dimension of this `PointingLO` instance.\nShape of `PointingLO` instance: {self.shape}\nShape of `vec`: {vec.shape}",
@@ -186,7 +190,7 @@ def _rmult_QU(self, vec: np.ndarray):
         Performs :math:`A^T * v`. The output vector will be a QU-map-like array.
         """
 
-        brahmap.MPI_RAISE_EXCEPTION(
+        MPI_RAISE_EXCEPTION(
             condition=(len(vec) != self.nrows),
             exception=ValueError,
             message=f"Dimensions of `vec` is not compatible with the dimension of this `PointingLO` instance.\nShape of `PointingLO` instance: {self.shape}\nShape of `vec`: {vec.shape}",
@@ -232,7 +236,7 @@ def _mult_IQU(self, vec: np.ndarray):
             :math:`\phi_t`.
         """
 
-        brahmap.MPI_RAISE_EXCEPTION(
+        MPI_RAISE_EXCEPTION(
             condition=(len(vec) != self.ncols),
             exception=ValueError,
             message=f"Dimensions of `vec` is not compatible with the dimension of this `PointingLO` instance.\nShape of `PointingLO` instance: {self.shape}\nShape of `vec`: {vec.shape}",
@@ -268,7 +272,7 @@ def _rmult_IQU(self, vec: np.ndarray):
 
         """
 
-        brahmap.MPI_RAISE_EXCEPTION(
+        MPI_RAISE_EXCEPTION(
             condition=(len(vec) != self.nrows),
             exception=ValueError,
             message=f"Dimensions of `vec` is not compatible with the dimension of this `PointingLO` instance.\nShape of `PointingLO` instance: {self.shape}\nShape of `vec`: {vec.shape}",
@@ -373,7 +377,7 @@ def __init__(self, diag: np.ndarray, dtype=None):
             dtype = np.float64
             self.diag = np.asarray(diag, dtype=dtype)
 
-        brahmap.MPI_RAISE_EXCEPTION(
+        MPI_RAISE_EXCEPTION(
             condition=(self.diag.ndim != 1),
             exception=ValueError,
             message="The `diag` array must be a 1-d vector",
@@ -389,7 +393,7 @@ def __init__(self, diag: np.ndarray, dtype=None):
         )
 
     def _mult(self, vec: np.ndarray):
-        brahmap.MPI_RAISE_EXCEPTION(
+        MPI_RAISE_EXCEPTION(
             condition=(len(vec) != self.diag.shape[0]),
             exception=ValueError,
             message=f"Dimensions of `vec` is not compatible with the dimensions of this `InvNoiseCovLO_Uncorrelated` instance.\nShape of `InvNoiseCovLO_Uncorrelated` instance: {self.shape}\nShape of `vec`: {vec.shape}",
@@ -564,7 +568,7 @@ def _mult_I(self, vec: np.ndarray):
         where :math:`x` is   an :math:`n_{pix}` array.
         """
 
-        brahmap.MPI_RAISE_EXCEPTION(
+        MPI_RAISE_EXCEPTION(
             condition=(len(vec) != self.size),
             exception=ValueError,
             message=f"Dimenstions of `vec` is not compatible with the dimension of this `BlockDiagonalPreconditionerLO` instance.\nShape of `BlockDiagonalPreconditionerLO` instance: {self.shape}\nShape of `vec`: {vec.shape}",
@@ -588,7 +592,7 @@ def _mult_QU(self, vec: np.ndarray):
         where :math:`x` is   an :math:`n_{pix}` array.
         """
 
-        brahmap.MPI_RAISE_EXCEPTION(
+        MPI_RAISE_EXCEPTION(
             condition=(len(vec) != self.size),
             exception=ValueError,
             message=f"Dimenstions of `vec` is not compatible with the dimension of this `BlockDiagonalPreconditionerLO` instance.\nShape of `BlockDiagonalPreconditionerLO` instance: {self.shape}\nShape of `vec`: {vec.shape}",
@@ -622,7 +626,7 @@ def _mult_IQU(self, vec: np.ndarray):
         where :math:`x` is   an :math:`n_{pix}` array.
         """
 
-        brahmap.MPI_RAISE_EXCEPTION(
+        MPI_RAISE_EXCEPTION(
             condition=(len(vec) != self.size),
             exception=ValueError,
             message=f"Dimenstions of `vec` is not compatible with the dimension of this `BlockDiagonalPreconditionerLO` instance.\nShape of `BlockDiagonalPreconditionerLO` instance: {self.shape}\nShape of `vec`: {vec.shape}",

brahmap/mapmakers/GLS.py

@@ -2,9 +2,12 @@
 import scipy
 from dataclasses import dataclass
 
-import brahmap
+from brahmap import MPI_RAISE_EXCEPTION
+
 from brahmap.linop import DiagonalOperator
+
 from brahmap.utilities import ProcessTimeSamples, SolverType
+
 from brahmap.interfaces import (
     PointingLO,
     ToeplitzLO,
@@ -51,7 +54,7 @@ def compute_GLS_maps(
     update_pointings_inplace: bool = True,
     GLSParameters: GLSParameters = GLSParameters(),
 ) -> GLSResult | tuple[ProcessTimeSamples, GLSResult]:
-    brahmap.MPI_RAISE_EXCEPTION(
+    MPI_RAISE_EXCEPTION(
         condition=(len(pointings) != len(time_ordered_data)),
         exception=ValueError,
         message=f"Size of `pointings` must be equal to the size of `time_ordered_data` array:\nlen(pointings) = {len(pointings)}\nlen(time_ordered_data) = {len(time_ordered_data)}",
@@ -68,7 +71,7 @@ def compute_GLS_maps(
             diag=np.ones(len(pointings)), dtype=dtype_float
         )
     else:
-        brahmap.MPI_RAISE_EXCEPTION(
+        MPI_RAISE_EXCEPTION(
             condition=(inv_noise_cov_operator.shape[0] != len(time_ordered_data)),
             exception=ValueError,
             message=f"The shape of `inv_noise_cov_operator` must be same as `(len(time_ordered_data), len(time_ordered_data))`:\nlen(time_ordered_data) = {len(time_ordered_data)}\ninv_noise_cov_operator.shape = {inv_noise_cov_operator.shape}",

brahmap/mapmakers/lbsim_mapmakers.py

@@ -4,10 +4,14 @@
 import litebird_sim as lbs
 from typing import List
 
-import brahmap
-from brahmap.mapmakers import GLSParameters, GLSResult, compute_GLS_maps
+from brahmap import MPI_RAISE_EXCEPTION
+
 from brahmap.linop import DiagonalOperator
+
+from brahmap.mapmakers import GLSParameters, GLSResult, compute_GLS_maps
+
 from brahmap.interfaces import ToeplitzLO, BlockLO, InvNoiseCovLO_Uncorrelated
+
 from brahmap.utilities import ProcessTimeSamples
 
 
@@ -91,7 +95,7 @@ def __init__(
                     idx = det_list.index(detector)
                     noise_variance[detector] = np.ones(tod_len[idx])
 
-                brahmap.MPI_RAISE_EXCEPTION(
+                MPI_RAISE_EXCEPTION(
                     condition=(
                         len(noise_variance[detector])
                         != tod_len[det_list.index(detector)]

brahmap/utilities/mpi.py → brahmap/mpi.py

@@ -1,8 +1,8 @@
 import os
-import brahmap
-
 from mpi4py import MPI
 
+import brahmap
+
 
 def Initialize(communicator=None, raise_exception_per_process: bool = True):
     if brahmap.bMPI is None:

brahmap/utilities/__init__.py

@@ -22,7 +22,6 @@
 
 from .process_time_samples import ProcessTimeSamples, SolverType
 
-from .mpi import Initialize, MPI_RAISE_EXCEPTION
 
 __all__ = [
     "is_sorted",
@@ -40,6 +39,4 @@
     "ProcessTimeSamples",
     "SolverType",
     "TypeChangeWarning",
-    "Initialize",
-    "MPI_RAISE_EXCEPTION",
 ]

brahmap/utilities/process_time_samples.py

@@ -1,17 +1,17 @@
 from enum import IntEnum
 import numpy as np
 import warnings
+from mpi4py import MPI
 
-import brahmap
 
 from brahmap.utilities.tools import TypeChangeWarning
 from brahmap.utilities import bash_colors
 
 from brahmap._extensions import compute_weights
 from brahmap._extensions import repixelize
 
-
-from mpi4py import MPI
+from brahmap import Initialize, MPI_RAISE_EXCEPTION
+import brahmap
 
 
 class SolverType(IntEnum):
@@ -34,7 +34,7 @@ def __init__(
         update_pointings_inplace: bool = True,
     ):
         if brahmap.bMPI is None:
-            brahmap.Initialize()
+            Initialize()
 
         self.npix = npix
         self.nsamples = len(pointings)
@@ -50,7 +50,7 @@ def __init__(
         if self.pointings_flag is None:
             self.pointings_flag = np.ones(self.nsamples, dtype=bool)
 
-        brahmap.MPI_RAISE_EXCEPTION(
+        MPI_RAISE_EXCEPTION(
             condition=(len(self.pointings_flag) != self.nsamples),
             exception=AssertionError,
             message=f"Size of `pointings_flag` must be equal to the size of `pointings` array:\nlen(pointings_flag) = {len(pointings_flag)}\nlen(pointings) = {self.nsamples}",
@@ -59,7 +59,7 @@ def __init__(
         self.threshold = threshold
         self.solver_type = solver_type
 
-        brahmap.MPI_RAISE_EXCEPTION(
+        MPI_RAISE_EXCEPTION(
             condition=(self.solver_type not in [1, 2, 3]),
             exception=ValueError,
             message="Invalid `solver_type`!!!\n`solver_type` must be either SolverType.I, SolverType.QU or SolverType.IQU (equivalently 1, 2 or 3).",
@@ -81,7 +81,7 @@ def __init__(
         if noise_weights is None:
             noise_weights = np.ones(self.nsamples, dtype=self.dtype_float)
 
-        brahmap.MPI_RAISE_EXCEPTION(
+        MPI_RAISE_EXCEPTION(
             condition=(len(noise_weights) != self.nsamples),
             exception=AssertionError,
             message=f"Size of `noise_weights` must be equal to the size of `pointings` array:\nlen(noise_weigths) = {len(noise_weights)}\nlen(pointings) = {self.nsamples}",
@@ -96,7 +96,7 @@ def __init__(
             noise_weights = noise_weights.astype(dtype=self.dtype_float, copy=False)
 
         if self.solver_type != 1:
-            brahmap.MPI_RAISE_EXCEPTION(
+            MPI_RAISE_EXCEPTION(
                 condition=(len(pol_angles) != self.nsamples),
                 exception=AssertionError,
                 message=f"Size of `pol_angles` must be equal to the size of `pointings` array:\nlen(pol_angles) = {len(pol_angles)}\nlen(pointings) = {self.nsamples}",

tests/test_BlkDiagPrecondLO.py

@@ -1,7 +1,7 @@
 import pytest
 import numpy as np
-import brahmap
 
+import brahmap
 import helper_BlkDiagPrecondLO as bdplo
 import helper_ProcessTimeSamples as hpts
 

tests/test_InvNoiseCov_tools_cpp.py

@@ -1,7 +1,7 @@
 import pytest
 import numpy as np
-from brahmap._extensions import InvNoiseCov_tools
 
+from brahmap._extensions import InvNoiseCov_tools
 from brahmap.interfaces import InvNoiseCovLO_Uncorrelated
 
 

tests/test_PointingLO.py

@@ -225,8 +225,6 @@ def test_I(self, initint, initfloat, rtol):
         # Test for P.T * <vector>
         weights = P.T * initfloat.noise_weights
 
-        # brahmap.bMPI.comm.Allreduce(MPI.IN_PLACE, weights, MPI.SUM)
-
         np.testing.assert_allclose(PTS.weighted_counts, weights)
 
         # Test for P * <vector>
@@ -273,8 +271,6 @@ def test_QU(self, initint, initfloat, rtol):
         # Test for P.T * <vector>
         weights = P.T * initfloat.noise_weights
 
-        # brahmap.bMPI.comm.Allreduce(MPI.IN_PLACE, weights, MPI.SUM)
-
         weighted_sin = np.zeros(PTS.new_npix, dtype=initfloat.dtype)
         weighted_cos = np.zeros(PTS.new_npix, dtype=initfloat.dtype)
 
@@ -337,8 +333,6 @@ def test_IQU(self, initint, initfloat, rtol):
         # Test for P.T * <vector>
         weights = P.T * initfloat.noise_weights
 
-        # brahmap.bMPI.comm.Allreduce(MPI.IN_PLACE, weights, MPI.SUM)
-
         np.testing.assert_allclose(PTS.weighted_counts, weights[0::3])
         np.testing.assert_allclose(PTS.weighted_cos, weights[1::3])
         np.testing.assert_allclose(PTS.weighted_sin, weights[2::3])