From abc53b113b2cf64efa644eff1f52f9255151f71e Mon Sep 17 00:00:00 2001 From: Avinash Anand <36325275+anand-avinash@users.noreply.github.com> Date: Mon, 3 Jun 2024 10:35:50 +0200 Subject: [PATCH] fixed the tests and circular import case --- .github/workflows/tests.yaml | 2 +- brahmap/__init__.py | 15 ++++++------ brahmap/_extensions/compute_weights.cpp | 2 -- brahmap/interfaces/linearoperators.py | 28 +++++++++++++---------- brahmap/mapmakers/GLS.py | 9 +++++--- brahmap/mapmakers/lbsim_mapmakers.py | 10 +++++--- brahmap/{utilities => }/mpi.py | 4 ++-- brahmap/utilities/__init__.py | 3 --- brahmap/utilities/process_time_samples.py | 16 ++++++------- tests/test_BlkDiagPrecondLO.py | 2 +- tests/test_InvNoiseCov_tools_cpp.py | 2 +- tests/test_PointingLO.py | 6 ----- 12 files changed, 50 insertions(+), 49 deletions(-) rename brahmap/{utilities => }/mpi.py (100%) diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml index 63e9614..f199bda 100644 --- a/.github/workflows/tests.yaml +++ b/.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 }} diff --git a/brahmap/__init__.py b/brahmap/__init__.py index 12e048b..2e7eb42 100644 --- a/brahmap/__init__.py +++ b/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", ] diff --git a/brahmap/_extensions/compute_weights.cpp b/brahmap/_extensions/compute_weights.cpp index 54a4396..e5d83f9 100644 --- a/brahmap/_extensions/compute_weights.cpp +++ b/brahmap/_extensions/compute_weights.cpp @@ -3,8 +3,6 @@ #include #include -#include - #include "mpi_utils.hpp" namespace py = pybind11; diff --git a/brahmap/interfaces/linearoperators.py b/brahmap/interfaces/linearoperators.py index 59f3592..5b63eb8 100644 --- a/brahmap/interfaces/linearoperators.py +++ b/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}", diff --git a/brahmap/mapmakers/GLS.py b/brahmap/mapmakers/GLS.py index 44d1efb..51d5420 100644 --- a/brahmap/mapmakers/GLS.py +++ b/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}", diff --git a/brahmap/mapmakers/lbsim_mapmakers.py b/brahmap/mapmakers/lbsim_mapmakers.py index 7e1c799..0ab5cd0 100644 --- a/brahmap/mapmakers/lbsim_mapmakers.py +++ b/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)] diff --git a/brahmap/utilities/mpi.py b/brahmap/mpi.py similarity index 100% rename from brahmap/utilities/mpi.py rename to brahmap/mpi.py index 3e08237..068d919 100644 --- a/brahmap/utilities/mpi.py +++ b/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: diff --git a/brahmap/utilities/__init__.py b/brahmap/utilities/__init__.py index acaaced..19168ae 100644 --- a/brahmap/utilities/__init__.py +++ b/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", ] diff --git a/brahmap/utilities/process_time_samples.py b/brahmap/utilities/process_time_samples.py index 9eaeae6..0b573ae 100644 --- a/brahmap/utilities/process_time_samples.py +++ b/brahmap/utilities/process_time_samples.py @@ -1,8 +1,8 @@ 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 @@ -10,8 +10,8 @@ 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}", diff --git a/tests/test_BlkDiagPrecondLO.py b/tests/test_BlkDiagPrecondLO.py index b7825be..d34390f 100644 --- a/tests/test_BlkDiagPrecondLO.py +++ b/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 diff --git a/tests/test_InvNoiseCov_tools_cpp.py b/tests/test_InvNoiseCov_tools_cpp.py index fa7fee4..f3320ad 100644 --- a/tests/test_InvNoiseCov_tools_cpp.py +++ b/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 diff --git a/tests/test_PointingLO.py b/tests/test_PointingLO.py index facb7a0..440325a 100644 --- a/tests/test_PointingLO.py +++ b/tests/test_PointingLO.py @@ -225,8 +225,6 @@ def test_I(self, initint, initfloat, rtol): # Test for P.T * 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 * @@ -273,8 +271,6 @@ def test_QU(self, initint, initfloat, rtol): # Test for P.T * 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 * 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])