Repository: Project-MONAI/MONAI
File: monai/networks/layers/conjugate_gradient.py (Python)

Code before:
        """
        super().__init__()

        self.A = linear_op
        self.num_iter = num_iter
        self.dbprint = dbprint

    def _zdot(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
        """
        Complex dot product between tensors x1 and x2: sum(x1.*x2)
        """
        if torch.is_complex(x1):
            assert torch.is_complex(x2), "x1 and x2 must both be complex"
            return torch.sum(x1.conj() * x2)
        else:
            return torch.sum(x1 * x2)

    def _zdot_single(self, x: torch.Tensor) -> torch.Tensor:
        """
        Complex dot product between tensor x and itself
        """
        res = self._zdot(x, x)
        if torch.is_complex(res):
            return res.real
        else:
            return res

    def _update(self, iter: int) -> Callable:
        """
        perform one iteration of the CG method. It takes the current solution x,
        the current search direction p, the current residual r, and the old
        residual norm rsold as inputs. Then it computes the new solution, search
        direction, residual, and residual norm, and returns them.
        """

        def update_fn(
            x: torch.Tensor, p: torch.Tensor, r: torch.Tensor, rsold: torch.Tensor
        ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
            dy = self.A(p)
            p_dot_dy = self._zdot(p, dy)
            alpha = rsold / p_dot_dy
            x = x + alpha * p
            r = r - alpha * dy
            rsnew = self._zdot_single(r)
            beta = rsnew / rsold
            rsold = rsnew
            p = beta * p + r

            # print residual
            if self.dbprint:
                print(f"CG Iteration {iter}: {rsnew}")

Review comment (refactor):
These two methods could be utility functions placed elsewhere since they don't rely on `self` internally. Moving them out would reduce the size of this class.
```suggestion
    def _zdot(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
        """
        Complex dot product between tensors x1 and x2: sum(x1.*x2)
        """
        if torch.is_complex(x1):
            if not torch.is_complex(x2):
                raise ValueError("x2 must be complex if x1 is complex")
            x1=x1.conj()
        
        return torch.sum(x1 * x2)

    def _zdot_single(self, x: torch.Tensor) -> torch.Tensor:
        """
        Complex dot product between tensor x and itself
        """
        res = self._zdot(x, x)
        return res.real if torch.is_complex(res) else res
```

Code after:
        return res


class ConjugateGradient(nn.Module):
    """
    Congugate Gradient (CG) solver for linear systems Ax = y.

    For linear_op that is positive definite and self-adjoint, CG is
    guaranteed to converge CG is often used to solve linear systems of the form
    Ax = y, where A is too large to store explicitly, but can be computed via a
    linear operator.

    As a result, here we won't set A explicitly as a matrix, but rather as a
    linear operator. For example, A could be a FFT/IFFT operation
    """

    def __init__(self, linear_op: Callable, num_iter: int):
        """
        Args:
            linear_op: Linear operator
            num_iter: Number of iterations to run CG
        """
        super().__init__()

        self.linear_op = linear_op
        self.num_iter = num_iter

    def update(
        self, x: torch.Tensor, p: torch.Tensor, r: torch.Tensor, rsold: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        perform one iteration of the CG method. It takes the current solution x,
        the current search direction p, the current residual r, and the old
        residual norm rsold as inputs. Then it computes the new solution, search
        direction, residual, and residual norm, and returns them.
        """

        dy = self.linear_op(p)
        p_dot_dy = _zdot(p, dy)
        alpha = rsold / p_dot_dy
        x = x + alpha * p
        r = r - alpha * dy
        rsnew = _zdot_single(r)
        beta = rsnew / rsold
        rsold = rsnew
        p = beta * p + r
        return x, p, r, rsold

    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        """
        run conjugate gradient for num_iter iterations to solve Ax = y