Memory efficient inverse noise covariance operators for litebird_sim

brahmap/base/noise_ops.py

@@ -39,6 +39,8 @@ def __init__(
             message="Please provide only one of `covariance` or `power_spectrum`",
         )
 
+        self.__size = nargin
+
         super(NoiseCovLinearOperator, self).__init__(
             nargin=nargin,
             nargout=nargin,
@@ -48,7 +50,9 @@ def __init__(
             **kwargs,
         )
 
-        self.size = nargin
+    @property
+    def size(self) -> int:
+        return self.__size
 
     @property
     def diag(self) -> np.ndarray:
@@ -123,7 +127,10 @@ def __init__(
             exception=ValueError,
             message="The noise (inv-)covariance operators must be symmetric",
         )
-        self.size = sum(self.col_size)
+
+    @property
+    def size(self) -> int:
+        return sum(self.col_size)
 
     @property
     def diag(self) -> np.ndarray:

brahmap/core/noise_ops_block_diag.py

@@ -18,7 +18,7 @@ class BlockDiagNoiseCovLO(BaseBlockDiagNoiseCovLinearOperator):
         _description_
     block_size : Union[np.ndarray, List]
         _description_
-    block_input : List[Union[np.ndarray, List]]
+    block_input : Union[List, Dict]
         _description_
     input_type : Literal["covariance", "power_spectrum"], optional
         _description_, by default "power_spectrum"
@@ -32,36 +32,55 @@ def __init__(
         self,
         operator,
         block_size: Union[np.ndarray, List],
-        block_input: List[Union[np.ndarray, List]],
+        block_input: Union[List, Dict],
         input_type: Literal["covariance", "power_spectrum"] = "power_spectrum",
         dtype: DTypeFloat = np.float64,
         extra_kwargs: Dict[str, Any] = {},
     ):
-        MPI_RAISE_EXCEPTION(
-            condition=(len(block_size) != len(block_input)),
-            exception=ValueError,
-            message="The number of blocks listed in `block_size` is different"
-            "from the number of blocks listed in `block_input`",
-        )
+        if isinstance(block_input, list):
+            MPI_RAISE_EXCEPTION(
+                condition=(len(block_size) != len(block_input)),
+                exception=ValueError,
+                message="The number of blocks listed in `block_size` is different"
+                " from the number of blocks provided in `block_input`",
+            )
 
-        block_list = self.__build_blocks(
-            operator=operator,
-            block_size=block_size,
-            block_input=block_input,
-            input_type=input_type,
-            dtype=dtype,
-            extra_kwargs=extra_kwargs,
-        )
+            block_list = self.__build_blocks_from_list(
+                operator=operator,
+                block_size=block_size,
+                block_input=block_input,
+                input_type=input_type,
+                dtype=dtype,
+                extra_kwargs=extra_kwargs,
+            )
+
+        elif isinstance(block_input, dict):
+            block_list = self.__build_blocks_from_dict(
+                operator=operator,
+                block_size=block_size,
+                block_input=block_input,
+                input_type=input_type,
+                dtype=dtype,
+                extra_kwargs=extra_kwargs,
+            )
+
+        else:
+            MPI_RAISE_EXCEPTION(
+                condition=True,
+                exception=ValueError,
+                message="`block_input` must be either a list of arrays or list"
+                " OR a dictionary that maps operator size to an array or a list",
+            )
 
         super(BlockDiagNoiseCovLO, self).__init__(
             block_list=block_list,
         )
 
-    def __build_blocks(
+    def __build_blocks_from_list(
         self,
         operator,
-        block_input,
-        block_size,
+        block_input: List,
+        block_size: Union[np.ndarray, List],
         input_type,
         dtype,
         extra_kwargs,
@@ -76,6 +95,39 @@ def __build_blocks(
                 **extra_kwargs,
             )
             block_list.append(block_op)
+
+        return block_list
+
+    def __build_blocks_from_dict(
+        self,
+        operator,
+        block_input: Dict,
+        block_size: Union[np.ndarray, List],
+        input_type,
+        dtype,
+        extra_kwargs,
+    ):
+        op_dict = {}
+        for shape in block_input.keys():
+            op_dict[shape] = operator(
+                size=shape,
+                input=block_input[shape],
+                input_type=input_type,
+                dtype=dtype,
+                **extra_kwargs,
+            )
+
+        block_list = []
+        for shape in block_size:
+            if shape in op_dict.keys():
+                block_list.append(op_dict[shape])
+            else:
+                MPI_RAISE_EXCEPTION(
+                    condition=True,
+                    exception=ValueError,
+                    message=f"Operator for shape {shape} is missing from the input dictionary",
+                )
+
         return block_list
 
 
@@ -88,7 +140,7 @@ class BlockDiagInvNoiseCovLO(BlockDiagNoiseCovLO):
         _description_
     block_size : Union[np.ndarray, List]
         _description_
-    block_input : List[Union[np.ndarray, List]]
+    block_input : Union[List, Dict]
         _description_
     input_type : Literal["covariance", "power_spectrum"], optional
         _description_, by default "power_spectrum"
@@ -102,7 +154,7 @@ def __init__(
         self,
         operator,
         block_size: Union[np.ndarray, List],
-        block_input: List[Union[np.ndarray, List]],
+        block_input: Union[List, Dict],
         input_type: Literal["covariance", "power_spectrum"] = "power_spectrum",
         dtype: DTypeFloat = np.float64,
         extra_kwargs: Dict[str, Any] = {},

brahmap/core/noise_ops_toeplitz.py

@@ -133,8 +133,6 @@ class InvNoiseCovLO_Toeplitz01(InvNoiseCovLinearOperator):
         _description_, by default None
     precond_maxiter : int, optional
         _description_, by default 50
-    precond_rtol : float, optional
-        _description_, by default 1.0e-10
     precond_atol : float, optional
         _description_, by default 1.0e-10
     precond_callback : Callable, optional
@@ -152,7 +150,6 @@ def __init__(
             LinearOperator, Literal[None, "Strang", "TChan", "RChan", "KK2"]
         ] = None,
         precond_maxiter: int = 50,
-        precond_rtol: float = 1.0e-10,
         precond_atol: float = 1.0e-10,
         precond_callback: Callable = None,
         dtype: DTypeFloat = np.float64,
@@ -164,11 +161,12 @@ def __init__(
             dtype=dtype,
         )
 
-        self.__precond_rtol = precond_rtol
         self.__precond_atol = precond_atol
         self.__precond_maxiter = precond_maxiter
         self.__precond_callback = precond_callback
 
+        self.__last_num_iterations = 0
+
         if precond_op is None:
             self.__precond_op = None
         elif isinstance(precond_op, LinearOperator) or isinstance(
@@ -231,16 +229,27 @@ def diag(self) -> np.ndarray:
         factor = 1.0
         return factor * np.ones(self.size, dtype=self.dtype)
 
+    @property
+    def get_last_num_iterations(self) -> int:
+        return self.__last_num_iterations
+
     def get_inverse(self):
         return self.__toeplitz_op
 
+    def __callback(self, x, r, norm_residual):
+        self.__last_num_iterations += 1
+        if self.__precond_callback is not None:
+            self.__precond_callback(x, r, norm_residual)
+
     def _mult(self, vec: np.ndarray):
         MPI_RAISE_EXCEPTION(
             condition=(len(vec) != self.shape[0]),
             exception=ValueError,
             message=f"Dimensions of `vec` is not compatible with the dimensions of this `InvNoiseCovLO_Toeplitz` instance.\nShape of `InvNoiseCovLO_Toeplitz` instance: {self.shape}\nShape of `vec`: {vec.shape}",
         )
 
+        self.__last_num_iterations = 0
+
         if vec.dtype != self.dtype:
             if MPI_UTILS.rank == 0:
                 warnings.warn(
@@ -252,11 +261,10 @@ def _mult(self, vec: np.ndarray):
         prod, _ = cg(
             A=self.__toeplitz_op,
             b=vec,
-            rtol=self.__precond_rtol,
             atol=self.__precond_atol,
             maxiter=self.__precond_maxiter,
             M=self.__precond_op,
-            callback=self.__precond_callback,
+            callback=self.__callback,
             parallel=False,
         )