solvers.py

"""Iterative solvers used by BASS.

This module currently contains a minimal conjugate-gradient (CG) solver for the
SENSE least-squares problem used throughout the project.
"""

import numpy as np

def solve_sense_cg(
    initial_guess: np.ndarray,
    E_op,
    m_under: np.ndarray,
    max_iter: int = 10,
    error_tol: float = 1e-6,
    verbose: bool = True,
) -> np.ndarray:
    """Solve a SENSE least-squares reconstruction using conjugate gradients.

    This is a direct translation of the legacy CG loop used in the
    reference implementation. It solves:

        argmin_x || E x - m_under ||_2^2

    where `E_op` is the encoding operator and `m_under` is the masked k-space.

    Args:
        initial_guess: Initial image estimate, shape (T, H, W).
        E_op: Encoding operator supporting `E_op * image` and `E_op.H * kspace`.
        m_under: Masked k-space measurements, shape (C, T, H, W).
        max_iter: Maximum CG iterations.
        error_tol: Stop early when 0.5*||m_under - E x||^2 falls below this threshold.
        verbose: If True, prints per-iteration cost/step sizes.

    Returns:
        Reconstructed image estimate, same shape as `initial_guess`.
    """
    if verbose:
        print("--- Starting Python SENSE CG Reconstruction ---")
    
    x = initial_guess.copy()
    E_op_adj = E_op.H

    Ex = E_op * x
    residual = m_under - Ex
    g = E_op_adj * residual
    d = g.copy()
    cost = 0.5 * np.real(np.vdot(residual, residual))
    e_old = np.real(np.vdot(g, g))

    for k in range(max_iter):
        if cost < error_tol:
            if verbose:
                print(
                    f"  Early stop at iter {k}/{max_iter}: cost={cost:.8e} < error_tol={error_tol:.8e}"
                )
            break
        Ad = E_op * d
        Ad_norm_sq = np.real(np.vdot(Ad, Ad))
        if Ad_norm_sq <= 0:
            if verbose:
                print("  Warning: Encountered zero-norm search direction; exiting CG early.")
            break
        alpha = e_old / Ad_norm_sq
        
        x = x + alpha * d
        residual = residual - alpha * Ad
        g_new = E_op_adj * residual
        cost = 0.5 * np.real(np.vdot(residual, residual))

        e_new = np.real(np.vdot(g_new, g_new))
        if e_old <= 0:
            if verbose:
                print("  Warning: Gradient norm vanished; exiting CG early.")
            break
        beta = e_new / e_old
        
        d = g_new + beta * d
        g = g_new
        e_old = e_new
        
        if verbose:
            print(
                f"  Iter {k+1}/{max_iter} | Cost: {cost:.8f} | "
                f"alpha: {alpha:.4f} | beta: {beta:.4f}"
            )
        
    if verbose:
        print("--- Python SENSE CG Finished ---")
    return x