Merge pull request #10 from anand-avinash/toeplitz_precond

Adding preconditioners for the inverse Toeplitz operators

Author: anand-avinash

brahmap/base/linop.py

@@ -39,7 +39,7 @@
 
 # Default (null) logger.
 null_log = logging.getLogger("linop")
-null_log.setLevel(logging.INFO)
+null_log.setLevel(logging.WARNING)
 null_log.addHandler(logging.NullHandler())
 
 

brahmap/base/noise_ops.py

@@ -1,4 +1,5 @@
 import numpy as np
+from typing import Literal
 
 from ..base import LinearOperator, BlockDiagonalLinearOperator
 
@@ -12,7 +13,7 @@ def __init__(
         self,
         nargin: int,
         matvec: int,
-        input_type: str = "covariance",
+        input_type: Literal["covariance", "power_spectrum"] = "covariance",
         dtype: DTypeFloat = np.float64,
         **kwargs,
     ):
@@ -54,7 +55,7 @@ def __init__(
         self,
         nargin: int,
         matvec: int,
-        input_type: str = "covariance",
+        input_type: Literal["covariance", "power_spectrum"] = "covariance",
         dtype: DTypeFloat = np.float64,
         **kwargs,
     ):

brahmap/core/noise_ops_block_diag.py

@@ -1,5 +1,5 @@
 import numpy as np
-from typing import List, Union
+from typing import List, Union, Literal, Dict, Any
 
 from ..base import (
     BaseBlockDiagNoiseCovLinearOperator,
@@ -15,8 +15,9 @@ def __init__(
         operator,
         block_size: Union[np.ndarray, List],
         block_input: List[Union[np.ndarray, List]],
-        input_type: str = "power_spectrum",
+        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)),
@@ -31,6 +32,7 @@ def __init__(
             block_input=block_input,
             input_type=input_type,
             dtype=dtype,
+            extra_kwargs=extra_kwargs,
         )
 
         super(BlockDiagNoiseCovLO, self).__init__(
@@ -44,6 +46,7 @@ def __build_blocks(
         block_size,
         input_type,
         dtype,
+        extra_kwargs,
     ):
         block_list = []
         for idx, input in enumerate(block_input):
@@ -52,6 +55,7 @@ def __build_blocks(
                 input=input,
                 input_type=input_type,
                 dtype=dtype,
+                **extra_kwargs,
             )
             block_list.append(block_op)
         return block_list
@@ -63,13 +67,15 @@ def __init__(
         operator,
         block_size: Union[np.ndarray, List],
         block_input: List[Union[np.ndarray, List]],
-        input_type: str = "power_spectrum",
+        input_type: Literal["covariance", "power_spectrum"] = "power_spectrum",
         dtype: DTypeFloat = np.float64,
+        extra_kwargs: Dict[str, Any] = {},
     ):
         super(BlockDiagInvNoiseCovLO, self).__init__(
             operator,
             block_size,
             block_input,
             input_type,
             dtype,
+            extra_kwargs,
         )

brahmap/core/noise_ops_circulant.py

@@ -1,6 +1,6 @@
 import numpy as np
 import warnings
-from typing import List, Union
+from typing import List, Union, Literal
 
 from ..utilities import TypeChangeWarning
 from ..base import NoiseCovLinearOperator, InvNoiseCovLinearOperator
@@ -15,7 +15,7 @@ def __init__(
         self,
         size: int,
         input: Union[np.ndarray, List],
-        input_type: str = "power_spectrum",
+        input_type: Literal["covariance", "power_spectrum"] = "power_spectrum",
         dtype: DTypeFloat = np.float64,
     ):
         input = np.asarray(a=input, dtype=dtype)
@@ -84,7 +84,7 @@ def __init__(
         self,
         size: int,
         input: Union[np.ndarray, List],
-        input_type: str = "power_spectrum",
+        input_type: Literal["covariance", "power_spectrum"] = "power_spectrum",
         dtype: DTypeFloat = np.float64,
     ):
         input = np.asarray(a=input, dtype=dtype)

brahmap/core/noise_ops_diagonal.py

@@ -1,7 +1,7 @@
 import numpy as np
 import warnings
 from numbers import Number
-from typing import List, Union
+from typing import List, Union, Literal
 
 
 from ..utilities import TypeChangeWarning
@@ -20,7 +20,7 @@ def __init__(
         self,
         size: int,
         input: Union[np.ndarray, List, DTypeFloat] = 1.0,
-        input_type="covariance",
+        input_type: Literal["covariance", "power_spectrum"] = "covariance",
         dtype: DTypeFloat = np.float64,
     ):
         if isinstance(input, Number) and input_type == "covariance":
@@ -95,7 +95,7 @@ def __init__(
         self,
         size: int,
         input: Union[np.ndarray, List, DTypeFloat] = 1.0,
-        input_type="covariance",
+        input_type: Literal["covariance", "power_spectrum"] = "covariance",
         dtype: DTypeFloat = np.float64,
     ):
         if isinstance(input, Number) and input_type == "covariance":

brahmap/core/noise_ops_toeplitz.py

@@ -1,13 +1,12 @@
 import numpy as np
 import warnings
-from typing import List, Union
+from typing import List, Union, Literal
 
 from ..utilities import TypeChangeWarning
-from ..base import NoiseCovLinearOperator, InvNoiseCovLinearOperator
-from ..math import DTypeFloat
+from ..base import LinearOperator, NoiseCovLinearOperator, InvNoiseCovLinearOperator
+from ..math import DTypeFloat, cg
 from ..mpi import MPI_RAISE_EXCEPTION
-
-import scipy.sparse.linalg
+from ..core import InvNoiseCovLO_Circulant
 
 from brahmap import MPI_UTILS
 
@@ -19,7 +18,7 @@ def __init__(
         self,
         size: int,
         input: Union[np.ndarray, List],
-        input_type: str = "power_spectrum",
+        input_type: Literal["covariance", "power_spectrum"] = "power_spectrum",
         dtype: DTypeFloat = np.float64,
     ):
         input = np.asarray(a=input, dtype=dtype)
@@ -108,9 +107,12 @@ def __init__(
         self,
         size: int,
         input: Union[np.ndarray, List],
-        input_type: str = "power_spectrum",
-        precond_op=None,
+        input_type: Literal["covariance", "power_spectrum"] = "power_spectrum",
+        precond_op: Union[
+            LinearOperator, Literal[None, "Strang", "TChan", "RChan", "KK2"]
+        ] = None,
         precond_maxiter=50,
+        precond_rtol=1.0e-10,
         precond_atol=1.0e-10,
         precond_callback=None,
         dtype: DTypeFloat = np.float64,
@@ -122,11 +124,61 @@ def __init__(
             dtype=dtype,
         )
 
+        self.__precond_rtol = precond_rtol
         self.__precond_atol = precond_atol
         self.__precond_maxiter = precond_maxiter
-        self.__precond_op = precond_op
         self.__precond_callback = precond_callback
 
+        if precond_op is None:
+            self.__precond_op = None
+        elif isinstance(precond_op, LinearOperator) or isinstance(
+            precond_op, np.ndarray
+        ):
+            self.__precond_op = precond_op
+        elif precond_op in ["Strang", "TChan", "RChan", "KK2"]:
+            if input_type == "power_spectrum":
+                cov = np.fft.ifft(input).real[:size]
+            else:
+                cov = input[:size]
+
+            if precond_op == "Strang":
+                temp_size = int(np.floor(cov.size / 2))
+                if cov.size % 2 == 0:
+                    new_cov = np.concatenate(
+                        [cov[:temp_size], cov[1 : temp_size + 1][::-1]]
+                    )
+                else:
+                    new_cov = np.concatenate(
+                        [cov[: temp_size + 1], cov[1 : temp_size + 1][::-1]]
+                    )
+            elif precond_op == "TChan":
+                new_cov = np.empty_like(cov)
+                new_cov[0] = cov[0]
+                n = cov.size
+                for idx in range(1, n):
+                    new_cov[idx] = ((n - idx) * cov[idx] + idx * cov[n - idx]) / n
+            elif precond_op == "RChan":
+                new_cov = np.roll(np.flip(cov), 1)
+                new_cov += cov
+                new_cov[0] = cov[0]
+            elif precond_op == "KK2":  # Circulant but not symmetric
+                new_cov = np.roll(np.flip(cov), 1)
+                new_cov[0] = 0
+                new_cov = cov - new_cov
+
+            self.__precond_op = InvNoiseCovLO_Circulant(
+                size=size,
+                input=new_cov,
+                input_type="covariance",
+                dtype=dtype,
+            )
+        else:
+            MPI_RAISE_EXCEPTION(
+                condition=True,
+                exception=ValueError,
+                message="Invalid preconditioner operator provided!",
+            )
+
         super(InvNoiseCovLO_Toeplitz01, self).__init__(
             nargin=size,
             matvec=self._mult,
@@ -157,13 +209,15 @@ def _mult(self, vec: np.ndarray):
                 )
             vec = vec.astype(dtype=self.dtype, copy=False)
 
-        prod, _ = scipy.sparse.linalg.gmres(
+        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,
+            parallel=False,
         )
 
         return prod

brahmap/lbsim/lbsim_noise_operators.py

@@ -1,4 +1,4 @@
-from typing import List, Union
+from typing import List, Union, Literal, Dict, Any
 
 import numpy as np
 import litebird_sim as lbs
@@ -63,7 +63,7 @@ def __init__(
         self,
         obs: Union[lbs.Observation, List[lbs.Observation]],
         input: Union[dict, Union[np.ndarray, List]],
-        input_type: str = "power_spectrum",
+        input_type: Literal["covariance", "power_spectrum"] = "power_spectrum",
         dtype=np.float64,
     ):
         if isinstance(obs, lbs.Observation):
@@ -142,9 +142,10 @@ def __init__(
         self,
         obs: Union[lbs.Observation, List[lbs.Observation]],
         input: Union[dict, Union[np.ndarray, List]],
-        input_type: str = "power_spectrum",
+        input_type: Literal["covariance", "power_spectrum"] = "power_spectrum",
         operator=InvNoiseCovLO_Toeplitz01,
         dtype=np.float64,
+        extra_kwargs: Dict[str, Any] = {},
     ):
         if isinstance(obs, lbs.Observation):
             obs_list = [obs]
@@ -184,6 +185,7 @@ def __init__(
             block_input=block_input,
             input_type=input_type,
             dtype=dtype,
+            extra_kwargs=extra_kwargs,
         )
 
     def _resize_input(self, new_size, input, input_type, dtype):

brahmap/math/linalg.py

@@ -1,3 +1,4 @@
+from typing import Callable
 import numpy as np
 import scipy
 import scipy.sparse
@@ -13,10 +14,37 @@ def parallel_norm(x: np.ndarray):
     return ret
 
 
-def cg(A, b, x0=None, atol=1.0e-12, maxiter=100, M=None, callback=None):
-    A, M, x, b, postprocess = scipy.sparse.linalg._isolve.utils.make_system(A, M, x0, b)
+def cg(
+    A,
+    b,
+    x0=None,
+    rtol=1.0e-12,
+    atol=1.0e-12,
+    maxiter=100,
+    M=None,
+    callback=None,
+    parallel=True,
+):
+    A, M, x, b, postprocess = scipy.sparse.linalg._isolve.utils.make_system(
+        A,
+        M,
+        x0,
+        b,
+    )
+
+    if parallel:
+        norm_function: Callable = parallel_norm
+    else:
+        norm_function: Callable = np.linalg.norm
+
+    b_norm = norm_function(b)
 
-    b_norm = parallel_norm(b)
+    atol, _ = scipy.sparse.linalg._isolve.iterative._get_atol_rtol(
+        "cg",
+        b_norm,
+        atol,
+        rtol,
+    )
 
     if b_norm == 0:
         return postprocess(b), 0
@@ -51,7 +79,7 @@ def cg(A, b, x0=None, atol=1.0e-12, maxiter=100, M=None, callback=None):
         r -= alpha * q
         rho_prev = rho_cur
 
-        norm_residual = parallel_norm(r) / b_norm
+        norm_residual = norm_function(r) / b_norm
 
         if callback:
             callback(x, r, norm_residual)