Source code for corrct.solvers

#!/usr/bin/env python3
"""
Solvers for the tomographic reconstruction problem.

@author: Nicola VIGANÒ, Computational Imaging group, CWI, The Netherlands,
and ESRF - The European Synchrotron, Grenoble, France
"""

import copy as cp
from abc import ABC, abstractmethod
from typing import Any, Callable, Optional, Union
from collections.abc import Sequence

import numpy as np
from numpy.typing import DTypeLike, NDArray
from tqdm.auto import tqdm

from . import data_terms, filters, operators, projectors, regularizers

eps = np.finfo(np.float32).eps


NDArrayFloat = NDArray[np.floating]



[docs]
class SolutionInfo:
    """Reconstruction info."""

    method: str
    iterations: int
    max_iterations: int

    residual0: Union[float, np.floating]
    residual0_cv: Union[float, np.floating]

    residuals: NDArrayFloat
    residuals_cv: NDArrayFloat
    tolerance: Union[float, np.floating, None]

    def __init__(
        self,
        method: str,
        max_iterations: int,
        tolerance: Union[float, np.floating, None],
        residual0: float = np.inf,
        residual0_cv: float = np.inf,
    ) -> None:
        self.method = method
        self.max_iterations = max_iterations
        self.tolerance = tolerance

        self.residual0 = residual0
        self.residual0_cv = residual0_cv

        self.residuals = np.zeros(max_iterations)
        self.residuals_cv = np.zeros(max_iterations)

        self.iterations = 0

    @property
    def residuals_rel(self) -> NDArrayFloat:
        return self.residuals / self.residual0

    @property
    def residuals_cv_rel(self) -> NDArrayFloat:
        return self.residuals_cv / self.residual0_cv




[docs]
class Solver(ABC):
    """
    Initialize the base solver class.

    Parameters
    ----------
    verbose : bool, optional
        Turn on verbose output. The default is False.
    tolerance : Optional[float], optional
        Tolerance on the data residual for computing when to stop iterations.
        The default is None.
    relaxation : float, optional
        The relaxation length. The default is 1.0.
    data_term : Union[str, data_terms.DataFidelityBase], optional
        Data fidelity term for computing the data residual. The default is "l2".
    data_term_test : Optional[data_terms.DataFidelityBase], optional
        The data fidelity to be used for the test set.
        If None, it will use the same as for the rest of the data.
        The default is None.
    """

    def __init__(
        self,
        verbose: bool = False,
        leave_progress: bool = True,
        relaxation: float = 1.0,
        tolerance: Optional[float] = None,
        data_term: Union[str, data_terms.DataFidelityBase] = "l2",
        data_term_test: Union[str, data_terms.DataFidelityBase, None] = None,
    ):
        self.verbose = verbose
        self.leave_progress = leave_progress
        self.relaxation = relaxation
        self.tolerance = tolerance

        self.data_term = self._initialize_data_fidelity_function(data_term)
        if data_term_test is None:
            data_term_test = self.data_term
        else:
            data_term_test = self._initialize_data_fidelity_function(data_term_test)
        self.data_term_test = cp.deepcopy(data_term_test)


[docs]
    def info(self) -> str:
        """
        Return the solver info.

        Returns
        -------
        str
            Solver info string.
        """
        return type(self).__name__



[docs]
    def upper(self) -> str:
        """
        Return the upper case name of the solver.

        Returns
        -------
        str
            Upper case string name of the solver.
        """
        return type(self).__name__.upper()



[docs]
    def lower(self) -> str:
        """
        Return the lower case name of the solver.

        Returns
        -------
        str
            Lower case string name of the solver.
        """
        return type(self).__name__.lower()



[docs]
    @abstractmethod
    def __call__(
        self, A: operators.BaseTransform, b: NDArrayFloat, *args: Any, **kwds: Any
    ) -> tuple[NDArrayFloat, SolutionInfo]:
        """Execute the reconstruction of the data.

        Parameters
        ----------
        A : operators.BaseTransform
            The projection operator.
        b : NDArrayFloat
            The data to be reconstructed.

        Returns
        -------
        Tuple[NDArrayFloat, SolutionInfo]
            The reconstruction and related information.
        """



[docs]
    @staticmethod
    def _initialize_data_fidelity_function(data_term: Union[str, data_terms.DataFidelityBase]) -> data_terms.DataFidelityBase:
        if isinstance(data_term, str):
            if data_term.lower() == "l2":
                return data_terms.DataFidelity_l2()
            elif data_term.lower() == "kl":
                return data_terms.DataFidelity_KL()
            else:
                raise ValueError(f"Unknown data term: '{data_term}', only accepted terms are: 'l2' | 'kl'.")
        elif isinstance(data_term, (data_terms.DataFidelity_l2, data_terms.DataFidelity_KL)):
            return cp.deepcopy(data_term)
        else:
            raise ValueError(f"Unsupported data term: '{data_term.info()}', only accepted terms are 'kl' and 'l2'-based.")



[docs]
    @staticmethod
    def _initialize_regularizer(
        regularizer: Union[regularizers.BaseRegularizer, None, Sequence[regularizers.BaseRegularizer]]
    ) -> Sequence[regularizers.BaseRegularizer]:
        if regularizer is None:
            return []
        elif isinstance(regularizer, regularizers.BaseRegularizer):
            return [regularizer]
        elif isinstance(regularizer, (list, tuple)):
            check_regs_ok = [isinstance(r, regularizers.BaseRegularizer) for r in regularizer]
            if not np.all(check_regs_ok):
                raise ValueError(
                    "The following regularizers are not derived from the regularizers.BaseRegularizer class: "
                    f"{np.array(np.arange(len(check_regs_ok))[np.array(check_regs_ok, dtype=bool)])}"
                )
            else:
                return list(regularizer)
        else:
            raise ValueError("Unknown regularizer type.")



[docs]
    @staticmethod
    def _initialize_b_masks(
        b: NDArrayFloat, b_mask: Optional[NDArrayFloat], b_test_mask: Optional[NDArrayFloat]
    ) -> tuple[Optional[NDArrayFloat], Optional[NDArrayFloat]]:
        if b_test_mask is not None:
            if b_mask is None:
                b_mask = np.ones_like(b)
            # As we are being passed a test residual pixel mask, we need
            # to make sure to mask those pixels out from the reconstruction.
            # At the same time, we need to remove any masked pixel from the test count.
            b_mask, b_test_mask = b_mask * (1 - b_test_mask), b_test_mask * b_mask
        return (b_mask, b_test_mask)





[docs]
class FBP(Solver):
    """Implementation of the Filtered Back-Projection (FBP) algorithm."""

    def __init__(
        self,
        verbose: bool = False,
        leave_progress: bool = False,
        regularizer: Union[Sequence[regularizers.BaseRegularizer], regularizers.BaseRegularizer, None] = None,
        data_term: Union[str, data_terms.DataFidelityBase] = "l2",
        fbp_filter: Union[str, NDArrayFloat, filters.Filter] = "ramp",
        pad_mode: str = "constant",
    ):
        """Initialize the Filtered Back-Projection (FBP) algorithm.

        Parameters
        ----------
        verbose : bool, optional
            Turn on verbose output. The default is False.
        leave_progress: bool, optional
            Leave the progress bar after the computation is finished. The default is True.
        regularizer : Sequence[regularizers.BaseRegularizer] | regularizers.BaseRegularizer | None, optional
            NOT USED, only exposed for compatibility reasons.
        data_term : Union[str, data_terms.DataFidelityBase], optional
            NOT USED, only exposed for compatibility reasons.
        fbp_filter : Union[str, NDArrayFloat], optional
            FBP filter to use. Either a string from scikit-image's list of `iradon` filters, or an array. The default is "ramp".
        pad_mode: str, optional
            The padding mode to use for the linear convolution. The default is "constant".
        """
        super().__init__(verbose=verbose)
        if isinstance(fbp_filter, str):
            fbp_filter = fbp_filter.lower()
        self.fbp_filter = fbp_filter
        self.pad_mode = pad_mode


[docs]
    def info(self) -> str:
        """
        Return the solver info.

        Returns
        -------
        str
            Solver info string.
        """
        if isinstance(self.fbp_filter, str):
            return super().info() + "(F:" + self.fbp_filter.upper() + ")"
        elif isinstance(self.fbp_filter, np.ndarray):
            return super().info() + "(F:" + filters.FilterCustom.__name__.upper() + ")"
        else:
            return super().info() + "(F:" + type(self.fbp_filter).__name__.upper() + ")"



[docs]
    def __call__(  # noqa: C901
        self,
        A: operators.BaseTransform,
        b: NDArrayFloat,
        iterations: int = 0,
        x0: Optional[NDArrayFloat] = None,
        lower_limit: Union[float, NDArrayFloat, None] = None,
        upper_limit: Union[float, NDArrayFloat, None] = None,
        x_mask: Optional[NDArrayFloat] = None,
        b_mask: Optional[NDArrayFloat] = None,
    ) -> tuple[NDArrayFloat, SolutionInfo]:
        """
        Reconstruct the data, using the FBP algorithm.

        Parameters
        ----------
        A : BaseTransform
            Projection operator.
        b : NDArrayFloat
            Data to reconstruct.
        iterations : int
            Number of iterations.
        x0 : Optional[NDArrayFloat], optional
            Initial solution. The default is None.
        lower_limit : Union[float, NDArrayFloat], optional
            Lower clipping value. The default is None.
        upper_limit : Union[float, NDArrayFloat], optional
            Upper clipping value. The default is None.
        x_mask : Optional[NDArrayFloat], optional
            Solution mask. The default is None.
        b_mask : Optional[NDArrayFloat], optional
            Data mask. The default is None.

        Raises
        ------
        ValueError
            In case the data is 1D.

        Returns
        -------
        Tuple[NDArrayFloat, SolutionInfo]
            The reconstruction, and None.
        """
        if len(b.shape) < 2:
            raise ValueError(f"Data should be at least 2-dimensional (b.shape = {b.shape})")

        info = SolutionInfo(self.info(), max_iterations=0, tolerance=0.0)

        if isinstance(self.fbp_filter, str):
            if self.fbp_filter in ("mr", "data"):
                local_filter = filters.FilterMR(projector=A)
            else:
                local_filter = filters.FilterFBP(filter_name=self.fbp_filter)
        elif isinstance(self.fbp_filter, np.ndarray):
            local_filter = filters.FilterCustom(self.fbp_filter)
        else:
            local_filter = self.fbp_filter
        local_filter.pad_mode = self.pad_mode

        if isinstance(A, operators.ProjectorOperator):
            pre_weights = A.get_pre_weights()
            if pre_weights is not None:
                b = b * pre_weights

        b_f = local_filter(b)

        x = A.T(b_f)

        if lower_limit is not None or upper_limit is not None:
            x = x.clip(lower_limit, upper_limit)
        if x_mask is not None:
            x *= x_mask

        return x, info





[docs]
class SART(Solver):
    """Solver class implementing the Simultaneous Algebraic Reconstruction Technique (SART) algorithm."""


[docs]
    def compute_residual(
        self,
        A: Callable,
        b: NDArrayFloat,
        x: NDArrayFloat,
        A_num_rows: int,
        b_mask: Optional[NDArrayFloat],
    ) -> NDArrayFloat:
        """Compute the solution residual.

        Parameters
        ----------
        A : Callable
            The forward projector.
        b : NDArrayFloat
            The detector data.
        x : NDArrayFloat
            The current solution
        A_num_rows : int
            The number of projections.
        b_mask : Optional[NDArrayFloat]
            The mask to apply

        Returns
        -------
        NDArrayFloat
            The residual.
        """
        fp = np.stack([A(x, ii) for ii in range(A_num_rows)], axis=-1)
        fp = np.ascontiguousarray(fp, dtype=b.dtype)
        res = fp - b
        if b_mask is not None:
            res *= b_mask
        return res



[docs]
    def __call__(  # noqa: C901
        self,
        A: Union[Callable[[NDArray, int], NDArray], projectors.ProjectorUncorrected],
        b: NDArrayFloat,
        iterations: int,
        A_num_rows: Optional[int] = None,
        At: Optional[Callable] = None,
        x0: Optional[NDArrayFloat] = None,
        lower_limit: Union[float, NDArrayFloat, None] = None,
        upper_limit: Union[float, NDArrayFloat, None] = None,
        x_mask: Optional[NDArrayFloat] = None,
        b_mask: Optional[NDArrayFloat] = None,
    ) -> tuple[NDArrayFloat, SolutionInfo]:
        """
        Reconstruct the data, using the SART algorithm.

        Parameters
        ----------
        A : Union[Callable, BaseTransform]
            Projection operator.
        b : NDArrayFloat
            Data to reconstruct.
        iterations : int
            Number of iterations.
        A_num_rows : int
            Number of projections.
        x0 : Optional[NDArrayFloat], optional
            Initial solution. The default is None.
        At : Callable, optional
            The back-projection operator. This is only needed if the projection operator does not have an adjoint.
            The default is None.
        lower_limit : Union[float, NDArrayFloat], optional
            Lower clipping value. The default is None.
        upper_limit : Union[float, NDArrayFloat], optional
            Upper clipping value. The default is None.
        x_mask : Optional[NDArrayFloat], optional
            Solution mask. The default is None.
        b_mask : Optional[NDArrayFloat], optional
            Data mask. The default is None.

        Returns
        -------
        Tuple[NDArrayFloat, SolutionInfo]
            The reconstruction, and the residuals.
        """
        if isinstance(A, projectors.ProjectorUncorrected):
            p = A

            if not p.projector_backend.has_individual_projs:
                raise ValueError("The projector needs to have enabled single projections.")

            A = lambda x, ii: p.fp_angle(x, ii)  # noqa: E731
            if isinstance(p, projectors.ProjectorAttenuationXRF):
                At = lambda y, ii: p.bp_angle(y, ii, single_line=True)  # noqa: E731
            else:
                At = lambda y, ii: p.bp_angle(y, ii)  # noqa: E731
            A_num_rows = len(p.angles_rot_rad)
        elif At is None:
            raise ValueError("Parameter `At` is required, if `A` is not a projector.")
        elif A_num_rows is None:
            raise ValueError("Parameter `A_num_rows` is required, if `A` is not a projector.")

        # Back-projection diagonal re-scaling
        b_ones = np.ones_like(b)
        if b_mask is not None:
            b_ones *= b_mask
        tau = [At(b_ones[..., ii, :], ii) for ii in range(A_num_rows)]
        tau = np.abs(np.stack(tau, axis=-2))
        tau[(tau / np.max(tau)) < 1e-5] = 1
        tau = self.relaxation / tau

        # Forward-projection diagonal re-scaling
        x_ones = np.ones([*tau.shape[:-2], tau.shape[-1]], dtype=tau.dtype)
        if x_mask is not None:
            x_ones *= x_mask
        sigma = [A(x_ones, ii) for ii in range(A_num_rows)]
        sigma = np.abs(np.stack(sigma, axis=-2))
        sigma[(sigma / np.max(sigma)) < 1e-5] = 1
        sigma = 1 / sigma

        if x0 is None:
            x0 = np.zeros_like(x_ones)
        else:
            x0 = np.array(x0).copy()
        x = x0

        info = SolutionInfo(self.info(), max_iterations=iterations, tolerance=self.tolerance)

        if self.tolerance is not None:
            res = self.compute_residual(A, b, x, A_num_rows=A_num_rows, b_mask=b_mask)
            info.residual0 = np.linalg.norm(res.flatten())

        rows_sequence = np.random.permutation(A_num_rows)

        algo_info = f"- Performing {self.upper()} iterations: "

        for ii in tqdm(range(iterations), desc=algo_info, disable=(not self.verbose), leave=self.leave_progress):
            info.iterations += 1

            for ii_a in rows_sequence:
                res = A(x, ii_a) - b[..., ii_a, :]
                if b_mask is not None:
                    res *= b_mask[..., ii_a, :]

                x -= At(res * sigma[..., ii_a, :], ii_a) * tau[..., ii_a, :]

                if lower_limit is not None:
                    x = np.fmax(x, lower_limit)
                if upper_limit is not None:
                    x = np.fmin(x, upper_limit)
                if x_mask is not None:
                    x *= x_mask

            if self.tolerance is not None:
                res = self.compute_residual(A, b, x, A_num_rows=A_num_rows, b_mask=b_mask)
                info.residuals[ii] = np.linalg.norm(res)

                if self.tolerance > info.residuals[ii]:
                    break

        return x, info





[docs]
class MLEM(Solver):
    """
    Initialize the MLEM solver class.

    This class implements the Maximul Likelihood Expectation Maximization (MLEM) algorithm.

    Parameters
    ----------
    verbose : bool, optional
        Turn on verbose output. The default is False.
    leave_progress: bool, optional
        Leave the progress bar after the computation is finished. The default is True.
    tolerance : Optional[float], optional
        Tolerance on the data residual for computing when to stop iterations.
        The default is None.
    regularizer : Sequence[regularizers.BaseRegularizer] | regularizers.BaseRegularizer | None, optional
        Regularizer to be used. The default is None.
    data_term : Union[str, data_terms.DataFidelityBase], optional
        Data fidelity term for computing the data residual. The default is "l2".
    data_term_test : Optional[data_terms.DataFidelityBase], optional
        The data fidelity to be used for the test set.
        If None, it will use the same as for the rest of the data.
        The default is None.
    """

    def __init__(
        self,
        verbose: bool = False,
        leave_progress: bool = True,
        tolerance: Optional[float] = None,
        regularizer: Union[Sequence[regularizers.BaseRegularizer], regularizers.BaseRegularizer, None] = None,
        data_term: Union[str, data_terms.DataFidelityBase] = "kl",
        data_term_test: Union[str, data_terms.DataFidelityBase, None] = None,
    ):
        super().__init__(
            verbose=verbose,
            leave_progress=leave_progress,
            tolerance=tolerance,
            data_term=data_term,
            data_term_test=data_term_test,
        )
        self.regularizer = self._initialize_regularizer(regularizer)


[docs]
    def info(self) -> str:
        """
        Return the MLEM info.

        Returns
        -------
        str
            info string.
        """
        return Solver.info(self) + f"(B:{self.data_term.background:g})" if self.data_term.background is not None else ""



[docs]
    def __call__(  # noqa: C901
        self,
        A: operators.BaseTransform,
        b: NDArrayFloat,
        iterations: int,
        x0: Optional[NDArrayFloat] = None,
        lower_limit: Union[float, NDArrayFloat, None] = None,
        upper_limit: Union[float, NDArrayFloat, None] = None,
        x_mask: Optional[NDArrayFloat] = None,
        b_mask: Optional[NDArrayFloat] = None,
        b_test_mask: Optional[NDArrayFloat] = None,
    ) -> tuple[NDArrayFloat, SolutionInfo]:
        """
        Reconstruct the data, using the MLEM algorithm.

        Parameters
        ----------
        A : BaseTransform
            Projection operator.
        b : NDArrayFloat
            Data to reconstruct.
        iterations : int
            Number of iterations.
        x0 : Optional[NDArrayFloat], optional
            Initial solution. The default is None.
        lower_limit : Union[float, NDArrayFloat], optional
            Lower clipping value. The default is None.
        upper_limit : Union[float, NDArrayFloat], optional
            Upper clipping value. The default is None.
        x_mask : Optional[NDArrayFloat], optional
            Solution mask. The default is None.
        b_mask : Optional[NDArrayFloat], optional
            Data mask. The default is None.
        b_test_mask : Optional[NDArrayFloat], optional
            Test data mask. The default is None.

        Returns
        -------
        Tuple[NDArrayFloat, SolutionInfo]
            The reconstruction, and the residuals.
        """
        b = np.array(b)

        (b_mask, b_test_mask) = self._initialize_b_masks(b, b_mask, b_test_mask)

        # Back-projection diagonal re-scaling
        b_ones = np.ones_like(b)
        if b_mask is not None:
            b_ones *= b_mask
        tau = A.T(b_ones)

        # Forward-projection diagonal re-scaling
        x_ones = np.ones_like(tau)
        if x_mask is not None:
            x_ones *= x_mask
        sigma = np.abs(A(x_ones))
        sigma[(sigma / np.max(sigma)) < 1e-5] = 1
        sigma = 1 / sigma

        if x0 is None:
            x = np.ones_like(tau)
        else:
            x = np.array(x0).copy()
        if x_mask is not None:
            x *= x_mask

        self.data_term.assign_data(b)

        info = SolutionInfo(self.info(), max_iterations=iterations, tolerance=self.tolerance)

        if b_test_mask is not None or self.tolerance is not None:
            Ax = A(x)

            if b_test_mask is not None:
                if self.data_term_test.background != self.data_term.background:
                    print("WARNING - the data_term and and data_term_test should have the same background. Making them equal.")
                    self.data_term_test.background = self.data_term.background
                self.data_term_test.assign_data(b)

                res_test_0 = self.data_term_test.compute_residual(Ax, mask=b_test_mask)
                info.residual0_cv = self.data_term_test.compute_residual_norm(res_test_0)

            if self.tolerance is not None:
                res_0 = self.data_term.compute_residual(Ax, mask=b_mask)
                info.residual0 = self.data_term.compute_residual_norm(res_0)

        reg_info = "".join(["-" + r.info().upper() for r in self.regularizer])
        algo_info = f"- Performing {self.upper()}-{self.data_term.upper()}{reg_info} iterations: "

        for ii in tqdm(range(iterations), desc=algo_info, disable=(not self.verbose), leave=self.leave_progress):
            info.iterations += 1

            # The MLEM update
            Ax = A(x)

            if b_test_mask is not None:
                res_test = self.data_term_test.compute_residual(Ax, mask=b_test_mask)
                info.residuals_cv[ii] = self.data_term_test.compute_residual_norm(res_test)

            if self.tolerance is not None:
                res = self.data_term.compute_residual(Ax, mask=b_mask)
                info.residuals[ii] = self.data_term.compute_residual_norm(res)
                if self.tolerance > info.residuals[ii]:
                    break

            if self.data_term.background is not None:
                Ax = Ax + self.data_term.background
            Ax = Ax.clip(eps, None)

            upd = A.T(b / Ax)
            x *= upd / tau

            if lower_limit is not None or upper_limit is not None:
                x = x.clip(lower_limit, upper_limit)
            if x_mask is not None:
                x *= x_mask

        return x, info





[docs]
class SIRT(Solver):
    """
    Initialize the SIRT solver class.

    This class implements the Simultaneous Iterative Reconstruction Technique (SIRT) algorithm.

    Parameters
    ----------
    verbose : bool, optional
        Turn on verbose output. The default is False.
    leave_progress: bool, optional
        Leave the progress bar after the computation is finished. The default is True.
    tolerance : Optional[float], optional
        Tolerance on the data residual for computing when to stop iterations.
        The default is None.
    relaxation : float, optional
        The relaxation length. The default is 1.95.
    regularizer : Sequence[regularizers.BaseRegularizer] | regularizers.BaseRegularizer | None, optional
        Regularizer to be used. The default is None.
    data_term : Union[str, data_terms.DataFidelityBase], optional
        Data fidelity term for computing the data residual. The default is "l2".
    data_term_test : Optional[data_terms.DataFidelityBase], optional
        The data fidelity to be used for the test set.
        If None, it will use the same as for the rest of the data.
        The default is None.
    """

    def __init__(
        self,
        verbose: bool = False,
        leave_progress: bool = True,
        relaxation: float = 1.95,
        tolerance: Optional[float] = None,
        regularizer: Union[Sequence[regularizers.BaseRegularizer], regularizers.BaseRegularizer, None] = None,
        data_term: Union[str, data_terms.DataFidelityBase] = "l2",
        data_term_test: Union[str, data_terms.DataFidelityBase, None] = None,
    ):
        super().__init__(
            verbose=verbose,
            leave_progress=leave_progress,
            relaxation=relaxation,
            tolerance=tolerance,
            data_term=data_term,
            data_term_test=data_term_test,
        )
        self.regularizer = self._initialize_regularizer(regularizer)


[docs]
    def info(self) -> str:
        """
        Return the SIRT info.

        Returns
        -------
        str
            SIRT info string.
        """
        reg_info = "".join(["-" + r.info().upper() for r in self.regularizer])
        return Solver.info(self) + "-" + self.data_term.info() + reg_info



[docs]
    def __call__(  # noqa: C901
        self,
        A: operators.BaseTransform,
        b: NDArrayFloat,
        iterations: int,
        x0: Optional[NDArrayFloat] = None,
        lower_limit: Union[float, NDArrayFloat, None] = None,
        upper_limit: Union[float, NDArrayFloat, None] = None,
        x_mask: Optional[NDArrayFloat] = None,
        b_mask: Optional[NDArrayFloat] = None,
        b_test_mask: Optional[NDArrayFloat] = None,
    ) -> tuple[NDArrayFloat, SolutionInfo]:
        """
        Reconstruct the data, using the SIRT algorithm.

        Parameters
        ----------
        A : BaseTransform
            Projection operator.
        b : NDArrayFloat
            Data to reconstruct.
        iterations : int
            Number of iterations.
        x0 : Optional[NDArrayFloat], optional
            Initial solution. The default is None.
        lower_limit : Union[float, NDArrayFloat], optional
            Lower clipping value. The default is None.
        upper_limit : Union[float, NDArrayFloat], optional
            Upper clipping value. The default is None.
        x_mask : Optional[NDArrayFloat], optional
            Solution mask. The default is None.
        b_mask : Optional[NDArrayFloat], optional
            Data mask. The default is None.
        b_test_mask : Optional[NDArrayFloat], optional
            Test data mask. The default is None.

        Returns
        -------
        Tuple[NDArrayFloat, SolutionInfo]
            The reconstruction, and the residuals.
        """
        b = np.array(b)

        (b_mask, b_test_mask) = self._initialize_b_masks(b, b_mask, b_test_mask)

        # Back-projection diagonal re-scaling
        b_ones = np.ones_like(b)
        if b_mask is not None:
            b_ones *= b_mask
        tau = np.abs(A.T(b_ones))
        for reg in self.regularizer:
            tau += reg.initialize_sigma_tau(tau)
        tau[(tau / np.max(tau)) < 1e-5] = 1
        tau = self.relaxation / tau

        # Forward-projection diagonal re-scaling
        x_ones = np.ones_like(tau)
        if x_mask is not None:
            x_ones *= x_mask
        sigma = np.abs(A(x_ones))
        sigma[(sigma / np.max(sigma)) < 1e-5] = 1
        sigma = 1 / sigma

        if x0 is None:
            x = np.zeros_like(x_ones)
        else:
            x = np.array(x0).copy()

        self.data_term.assign_data(b, sigma)

        info = SolutionInfo(self.info(), max_iterations=iterations, tolerance=self.tolerance)

        if b_test_mask is not None or self.tolerance is not None:
            Ax = A(x)

            res_0 = self.data_term.compute_residual(Ax, mask=b_mask)
            info.residual0 = self.data_term.compute_residual_norm(res_0)

            if b_test_mask is not None:
                if self.data_term_test.background != self.data_term.background:
                    print("WARNING - the data_term and and data_term_test should have the same background. Making them equal.")
                    self.data_term_test.background = self.data_term.background
                self.data_term_test.assign_data(b, sigma)

                res_test_0 = self.data_term_test.compute_residual(Ax, mask=b_test_mask)
                info.residual0_cv = self.data_term_test.compute_residual_norm(res_test_0)

        reg_info = "".join(["-" + r.info().upper() for r in self.regularizer])
        algo_info = f"- Performing {self.upper()}-{self.data_term.upper()}{reg_info} iterations: "

        for ii in tqdm(range(iterations), desc=algo_info, disable=(not self.verbose), leave=self.leave_progress):
            info.iterations += 1

            Ax = A(x)
            res = self.data_term.compute_residual(Ax, mask=b_mask)

            if b_test_mask is not None or self.tolerance is not None:
                info.residuals[ii] = self.data_term.compute_residual_norm(res)

                if b_test_mask is not None:
                    res_test = self.data_term_test.compute_residual(Ax, mask=b_test_mask)
                    info.residuals_cv[ii] = self.data_term_test.compute_residual_norm(res_test)

                if self.tolerance is not None and self.tolerance > info.residuals[ii]:
                    if self.verbose:
                        print(f"Residual reached the desired tolerance of {self.tolerance}. Ending iterations..")
                    break

            q = [reg.initialize_dual() for reg in self.regularizer]
            for q_r, reg in zip(q, self.regularizer):
                reg.update_dual(q_r, x)
                reg.apply_proximal(q_r)

            upd = A.T(res * sigma)
            for q_r, reg in zip(q, self.regularizer):
                upd -= reg.compute_update_primal(q_r)
            x += upd * tau

            if lower_limit is not None or upper_limit is not None:
                x = x.clip(lower_limit, upper_limit)
            if x_mask is not None:
                x *= x_mask

        return x, info





[docs]
class PDHG(Solver):
    """
    Initialize the PDHG solver class.

    PDHG stands for primal-dual hybrid gradient algorithm from Chambolle and Pock.

    Parameters
    ----------
    verbose : bool, optional
        Turn on verbose output. The default is False.
    leave_progress: bool, optional
        Leave the progress bar after the computation is finished. The default is True.
    tolerance : Optional[float], optional
        Tolerance on the data residual for computing when to stop iterations.
        The default is None.
    relaxation : float, optional
        The relaxation length. The default is 0.95.
    regularizer : Sequence[regularizers.BaseRegularizer] | regularizers.BaseRegularizer | None, optional
        Regularizer to be used. The default is None.
    data_term : Union[str, data_terms.DataFidelityBase], optional
        Data fidelity term for computing the data residual. The default is "l2".
    data_term_test : Optional[data_terms.DataFidelityBase], optional
        The data fidelity to be used for the test set.
        If None, it will use the same as for the rest of the data.
        The default is None.
    """

    def __init__(
        self,
        verbose: bool = False,
        leave_progress: bool = True,
        tolerance: Optional[float] = None,
        relaxation: float = 0.95,
        regularizer: Union[Sequence[regularizers.BaseRegularizer], regularizers.BaseRegularizer, None] = None,
        data_term: Union[str, data_terms.DataFidelityBase] = "l2",
        data_term_test: Union[str, data_terms.DataFidelityBase, None] = None,
    ):
        super().__init__(
            verbose=verbose,
            leave_progress=leave_progress,
            relaxation=relaxation,
            tolerance=tolerance,
            data_term=data_term,
            data_term_test=data_term_test,
        )
        self.regularizer = self._initialize_regularizer(regularizer)


[docs]
    def info(self) -> str:
        """
        Return the PDHG info.

        Returns
        -------
        str
            PDHG info string.
        """
        reg_info = "".join(["-" + r.info().upper() for r in self.regularizer])
        return Solver.info(self) + "-" + self.data_term.info() + reg_info



[docs]
    @staticmethod
    def _initialize_data_fidelity_function(data_term: Union[str, data_terms.DataFidelityBase]):
        if isinstance(data_term, str):
            if data_term.lower() == "l2":
                return data_terms.DataFidelity_l2()
            if data_term.lower() == "l1":
                return data_terms.DataFidelity_l1()
            if data_term.lower() == "kl":
                return data_terms.DataFidelity_KL()
            else:
                raise ValueError(f'Unknown data term: "{data_term}", accepted terms are: "l2" | "l1" | "kl".')
        else:
            return cp.deepcopy(data_term)



[docs]
    def power_method(
        self, A: operators.BaseTransform, b: NDArrayFloat, iterations: int = 5
    ) -> tuple[np.floating, Sequence[int], DTypeLike]:
        """
        Compute the l2-norm of the operator A, with the power method.

        Parameters
        ----------
        A : BaseTransform
            Operator whose l2-norm needs to be computed.
        b : NDArrayFloat
            The data vector.
        iterations : int, optional
            Number of power method iterations. The default is 5.

        Returns
        -------
        Tuple[float, Tuple[int], DTypeLike]
            The l2-norm of A, and the shape and type of the solution.
        """
        x: NDArrayFloat = np.array(np.random.rand(*b.shape))
        x = x.astype(b.dtype)
        x /= np.linalg.norm(x)
        x = A.T(x)

        x_norm = np.linalg.norm(x)
        L = x_norm

        for _ in range(iterations):
            x /= x_norm
            x = A.T(A(x))

            x_norm = np.linalg.norm(x)
            L = np.sqrt(x_norm)

        return (L, x.shape, x.dtype)



[docs]
    def _get_data_sigma_tau_unpreconditioned(self, A: operators.BaseTransform, b: NDArrayFloat):
        (L, x_shape, x_dtype) = self.power_method(A, b)
        tau = L

        dummy_x = np.empty(x_shape, dtype=x_dtype)
        for reg in self.regularizer:
            tau += reg.initialize_sigma_tau(dummy_x)

        tau = self.relaxation / tau
        sigma = self.relaxation / L
        return (x_shape, x_dtype, sigma, tau)



[docs]
    def __call__(  # noqa: C901
        self,
        A: operators.BaseTransform,
        b: NDArrayFloat,
        iterations: int,
        x0: Optional[NDArrayFloat] = None,
        lower_limit: Union[float, NDArrayFloat, None] = None,
        upper_limit: Union[float, NDArrayFloat, None] = None,
        x_mask: Optional[NDArrayFloat] = None,
        b_mask: Optional[NDArrayFloat] = None,
        b_test_mask: Optional[NDArrayFloat] = None,
        precondition: bool = True,
    ) -> tuple[NDArrayFloat, SolutionInfo]:
        """
        Reconstruct the data, using the PDHG algorithm.

        Parameters
        ----------
        A : BaseTransform
            Projection operator.
        b : NDArrayFloat
            Data to reconstruct.
        iterations : int
            Number of iterations.
        x0 : Optional[NDArrayFloat], optional
            Initial solution. The default is None.
        lower_limit : Union[float, NDArrayFloat], optional
            Lower clipping value. The default is None.
        upper_limit : Union[float, NDArrayFloat], optional
            Upper clipping value. The default is None.
        x_mask : Optional[NDArrayFloat], optional
            Solution mask. The default is None.
        b_mask : Optional[NDArrayFloat], optional
            Data mask. The default is None.
        b_test_mask : Optional[NDArrayFloat], optional
            Test data mask. The default is None.
        precondition : bool, optional
            Whether to use the preconditioned version of the algorithm. The default is True.

        Returns
        -------
        Tuple[NDArrayFloat, SolutionInfo]
            The reconstruction, and the residuals.
        """
        b = np.array(b)

        if precondition:
            try:
                At_abs = A.T.absolute()
                A_abs = A.absolute()
            except AttributeError:
                print(A)
                print("WARNING: Turning off preconditioning because system matrix does not support absolute")
                precondition = False

        (b_mask, b_test_mask) = self._initialize_b_masks(b, b_mask, b_test_mask)

        if precondition:
            tau = np.ones_like(b)
            if b_mask is not None:
                tau *= b_mask
            tau = np.abs(At_abs(tau))
            for reg in self.regularizer:
                tau += reg.initialize_sigma_tau(tau)
            tau[(tau / np.max(tau)) < 1e-5] = 1
            tau = self.relaxation / tau

            x_shape = tau.shape
            x_dtype = tau.dtype

            sigma = np.ones_like(tau)
            if x_mask is not None:
                sigma *= x_mask
            sigma = np.abs(A_abs(sigma))
            sigma[(sigma / np.max(sigma)) < 1e-5] = 1
            sigma = self.relaxation / sigma
        else:
            (x_shape, x_dtype, sigma, tau) = self._get_data_sigma_tau_unpreconditioned(A, b)

        if x0 is None:
            x0 = np.zeros(x_shape, dtype=x_dtype)
        else:
            x0 = np.array(x0).copy()
        x = x0
        x_relax = x.copy()

        self.data_term.assign_data(b, sigma)
        p = self.data_term.initialize_dual()

        q = [reg.initialize_dual() for reg in self.regularizer]

        info = SolutionInfo(self.info(), max_iterations=iterations, tolerance=self.tolerance)

        if b_test_mask is not None or self.tolerance is not None:
            Ax = A(x)

            res_0 = self.data_term.compute_residual(Ax, mask=b_mask)
            info.residual0 = self.data_term.compute_residual_norm(res_0)

            if b_test_mask is not None:
                if self.data_term_test.background != self.data_term.background:
                    print("WARNING - the data_term and and data_term_test should have the same background. Making them equal.")
                    self.data_term_test.background = self.data_term.background
                self.data_term_test.assign_data(b, sigma)

                res_test_0 = self.data_term_test.compute_residual(Ax, mask=b_test_mask)
                info.residual0_cv = self.data_term_test.compute_residual_norm(res_test_0)

        reg_info = "".join(["-" + r.info().upper() for r in self.regularizer])
        algo_info = f"- Performing {self.upper()}-{self.data_term.upper()}{reg_info} iterations: "

        for ii in tqdm(range(iterations), desc=algo_info, disable=(not self.verbose), leave=self.leave_progress):
            info.iterations += 1

            Ax_rlx = A(x_relax)
            self.data_term.update_dual(p, Ax_rlx)
            self.data_term.apply_proximal(p)

            if b_mask is not None:
                p *= b_mask

            for q_r, reg in zip(q, self.regularizer):
                reg.update_dual(q_r, x_relax)
                reg.apply_proximal(q_r)

            upd = A.T(p)
            for q_r, reg in zip(q, self.regularizer):
                upd += reg.compute_update_primal(q_r)
            x_new = x - upd * tau

            if lower_limit is not None or upper_limit is not None:
                x_new = x_new.clip(lower_limit, upper_limit)
            if x_mask is not None:
                x_new *= x_mask

            x_relax = x_new + (x_new - x)
            x = x_new

            if b_test_mask is not None or self.tolerance is not None:
                Ax = A(x)
                res = self.data_term.compute_residual(Ax, mask=b_mask)
                info.residuals[ii] = self.data_term.compute_residual_norm(res)

                if b_test_mask is not None:
                    res_test = self.data_term_test.compute_residual(Ax, mask=b_test_mask)
                    info.residuals_cv[ii] = self.data_term_test.compute_residual_norm(res_test)

                if self.tolerance is not None and self.tolerance > info.residuals[ii]:
                    if self.verbose:
                        print(f"Residual reached the desired tolerance of {self.tolerance}. Ending iterations..")
                    break

        return x, info