utils.py

import tqdm
import numpy as np
import torch
from kornia.geometry.transform import remap
from scipy.interpolate import BSpline

def ifft(x):
    x = torch.fft.ifftshift(x, dim=(-2, -1))
    x = torch.fft.ifft2(x, dim=(-2, -1), norm='ortho')
    x = torch.fft.fftshift(x, dim=(-2, -1))
    return x

# Centered, orthogonal fft in torch >= 1.7
def fft(x):
    x = torch.fft.fftshift(x, dim=(-2, -1))
    x = torch.fft.fft2(x, dim=(-2, -1), norm='ortho')
    x = torch.fft.ifftshift(x, dim=(-2, -1))
    return x

#----------------------------------------------------------------------------
# Wrapper for torch.Generator that allows specifying a different random seed
# for each sample in a minibatch.

class StackedRandomGenerator:
    def __init__(self, device, seeds):
        super().__init__()
        self.generators = [torch.Generator(device).manual_seed(int(seed) % (1 << 32)) for seed in seeds]

    def randn(self, size, **kwargs):
        assert size[0] == len(self.generators)
        return torch.stack([torch.randn(size[1:], generator=gen, **kwargs) for gen in self.generators])

    def randn_like(self, input):
        return self.randn(input.shape, dtype=input.dtype, layout=input.layout, device=input.device)

    def randint(self, *args, size, **kwargs):
        assert size[0] == len(self.generators)
        return torch.stack([torch.randint(*args, size=size[1:], generator=gen, **kwargs) for gen in self.generators])

def mult_sample(x_0, mask):
    """
    Computes F^{-1} mask F x_0.
    """
    img = x_0[:, 0] + x_0[:, 1] * 1j
    img = fft(img)
    img = mask * img
    img = ifft(img)
    img = img[:, None, :, :]
    x_t = torch.cat([img.real, img.imag], 1)
    return x_t

def calc_basis(device, n = 3, delta = 16):
    """
    Compute B-spline basis functions. n = spline order (default: cubic)
    """
    spl = BSpline.basis_element([-1,-0.5,0,0.5,1], extrapolate=False)
    def b(n, u):
        if u >= 1 or u <= -1:
            return 0
        if n == 1:
            return 1 - abs(u)
        elif n == 3:
            return spl(u)
        
    spline_kernel = torch.zeros((2 * delta, 2 * delta)).to(device)
    for x in range(0, 2 * delta):
        for y in range(0, 2 * delta):
            spline_kernel[x,y] = b(n, y / delta - 1) * b(n, x / delta - 1)
    
    return spline_kernel

def create_shells(H, W, alpha=1.5, one_D=False):
    """
    Compute coarse to fine shells. Specify rate of increase with alpha, and 1D vs 2D sampling.
    """
    # coarse to fine diffusion shells
    grid = torch.linspace(0,1,steps=300+1)
    std_devs = 5.0 * ((5.0*grid).exp())
    scale = alpha * grid
    scale[scale > 1] = 1

    all_shells = torch.ones((300 + 1,H,W))
    for i in range(H):
        for j in range(W):
            if one_D:
                all_shells[:,i,j] = 1 - scale * np.exp(-((i - H//2)**2) / (2 * std_devs**2))
            else:
                all_shells[:,i,j] = 1 - scale * np.exp(-((i - H//2)**2 + (j - W//2)**2) / (2 * std_devs**2))

    all_shells[all_shells > 1] = 1
    all_shells[all_shells < 0] = 0
    all_shells[0] = 1
    all_shells[300] = 0
    return all_shells

def ODE_motion_sampler(net, y, all_A, maps, latents, img_l_ss=1.0, motion_l_ss=1.0, second_order=False,
                     class_labels=None, randn_like=torch.randn_like, num_steps=100,
                      sigma_min=0.002, sigma_max=50, rho=7, S_churn=0,S_min=0,  known=False,
                      S_max=float('inf'), S_noise=1, verbose=True, motion_est=1, coords=None,
                      gt_theta=None, gt_dx=None, gt_dy=None,group_ETL=True, posterior=False):
    """
    MI-PS Code from Levac et al. (2023), left as unmodified as possible.
    """
    # Adjust noise levels based on what's supported by the network.
    sigma_min = max(sigma_min, net.sigma_min)
    sigma_max = min(sigma_max, net.sigma_max)
    posy, posx = torch.from_numpy(np.mgrid[:all_A[0].shape[0], :all_A[0].shape[1]]).to(latents.device)    
    
    # Time step discretization.
    step_indices = torch.arange(num_steps, dtype=torch.float64, device=latents.device)
    t_steps = (sigma_max ** (1 / rho) + step_indices / (num_steps - 1) * (sigma_min ** (1 / rho) - sigma_max ** (1 / rho))) ** rho
    t_steps = torch.cat([net.round_sigma(t_steps), torch.zeros_like(t_steps[:1])]) # t_N = 0

    # initialize motion estimates
    if known:
        est_theta = gt_theta
        est_dx    = gt_dx
        est_dy    = gt_dy
    else:
        est_theta = torch.zeros_like(gt_theta)
        est_dx    = torch.zeros_like(gt_dx)
        est_dy    = torch.zeros_like(gt_dy)

    # Main sampling loop.
    x_next = latents.to(torch.float64) * t_steps[0]

    deform_x, deform_y = None, None
    for i, (t_cur, t_next) in enumerate(tqdm.tqdm(zip(t_steps[:-1], t_steps[1:]))): # 0, ..., N-1
        x_cur = x_next
        x_hat = x_cur
        x_hat = x_hat.requires_grad_() #starting grad tracking with the noised img

        # Euler step.
        denoised = net(x_hat, t_cur, class_labels).to(torch.float64)
        d_cur = (x_hat - denoised)/t_cur
        # take step over prior score and add noise
        x_next = x_hat + (t_next - t_cur) * d_cur #+ ((2*t_cur)*(t_cur-t_next))**0.5 * randn_like(x_cur)

        if posterior:
            sse = 0
            for n, Phixi in enumerate(range(len(y))):
                img = torch.view_as_complex(denoised[0].permute(1, 2, 0).contiguous())
                ksp = fft(maps * img)
                residual = y[n] - all_A[n] * ksp
                sse += torch.norm(residual)**2

            likelihood_score = torch.autograd.grad(outputs=sse, inputs=x_hat)[0]
            x_next = x_next - img_l_ss * likelihood_score

        # take step on motion parameters
        if motion_est:
            est_theta = est_theta.requires_grad_()
            est_dy    = est_dy.requires_grad_()
            est_dx    = est_dx.requires_grad_()

            Phix_list, deform_x, deform_y = rigid_deform(image=denoised, coords=coords,
                            angles=est_theta, dx=est_dx, dy=est_dy, posx=posx, posy=posy, 
                            device=denoised.device)
                
            sse_m = 0
            for n, Phixi in enumerate(Phix_list[1:]):
                img = torch.view_as_complex(Phixi.permute(1, 2, 0).contiguous())
                ksp = fft(maps * img)
                residual = y[n+1] - all_A[n+1] * ksp
                sse_m += torch.norm(residual)**2


            meas_grad_motion = torch.autograd.grad(outputs = sse_m, inputs = (est_theta, est_dx, est_dy), create_graph = not True)
            norm = torch.sqrt(sse_m)
            est_theta = est_theta - (motion_l_ss/norm)*meas_grad_motion[0]
            est_dx = est_dx - (motion_l_ss/norm)*meas_grad_motion[1]
            est_dy = est_dy - (motion_l_ss/norm)*meas_grad_motion[2]

        x_next = x_next.detach()
        x_hat = x_hat.requires_grad_(False)

    return x_next, deform_x, deform_y# , deform_x, deform_y, all_images

def general_sampler(
    net, y, all_A, maps, latents, spline_kernels, all_shells,
    img_l_ss=1.0, num_steps=100, sigma_min=0.002, sigma_max=80.0, rho=7,
    posterior=False, update_motion=[], what_n=[5]
):
    """
    Perform coarse-to-fine sampling using a diffusion model with deformation updates.
    
    Args:
        net: Diffusion model.
        y: List of (num_coils, H, W) measurements for each motion state.
        all_A: List of (H, W) sampling masks for each motion state.
        maps: Sensitivity maps (num_coils, H, W).
        latents: Initial latent (1, 2, H, W).
        spline_kernels: Dictionary of spline kernels.
        all_shells: Precomputed noise shells (num_steps+1, H, W).
        img_l_ss: Step size for gradient descent.
        num_steps: Number of diffusion steps.
        sigma_min: Minimum noise level.
        sigma_max: Maximum noise level.
        rho: Noise schedule exponent.
        posterior: Whether to apply posterior score gradient.
        update_motion: Iteration indices at which to update deformation.
        what_n: Which spline kernels to use (customizable; tradeoff between accuracy and efficiency)

    Returns:
        Final sampled image, x- and y-deformations, and list of deformed images.
    """
    device = latents.device
    H, W = latents.shape[2], latents.shape[3]
    posy, posx = torch.from_numpy(np.mgrid[:H, :W]).to(device)
    skip = int(300 / num_steps)

    sigma_min = max(sigma_min, net.sigma_min)
    sigma_max = min(sigma_max, net.sigma_max)

    step_indices = torch.arange(num_steps, dtype=torch.float64, device=device)
    sigma_steps = (sigma_max ** (1 / rho) + step_indices / (num_steps - 1) * (sigma_min ** (1 / rho) - sigma_max ** (1 / rho))) ** rho
    sigma_steps = torch.cat([net.round_sigma(sigma_steps), torch.zeros_like(sigma_steps[:1])])

    est_phi = torch.zeros((len(y), 2, H, W), dtype=torch.float64, device=device) 
    x_next = latents.to(torch.float64) * sigma_steps[0]

    counter = 0
    Phix_list = deform_x = deform_y = None

    for i, (sigma_cur, sigma_next) in enumerate(tqdm.auto.tqdm(zip(sigma_steps[:-1], sigma_steps[1:]), desc="Reverse diffusion", total = num_steps, position=0, dynamic_ncols=True)):
        x_hat = x_next.requires_grad_()

        if i in update_motion:
            counter += 1
            x_next1 = x_next.clone().detach()

            for j, (sc1, sn1) in enumerate(tqdm.auto.tqdm(zip(sigma_steps[i:-1], sigma_steps[i+1:]), desc=f"Sampling clean image xbar ({counter}/{len(update_motion)})", total = len(sigma_steps[i:-1]), position=1, dynamic_ncols=True, leave=False)):
                x_hat1 = x_next1.requires_grad_()
                denoised1 = net(x_hat1, sc1, None).to(torch.float64)
                d_cur1 = denoised1 - x_hat1

                mask1 = (sc1 * all_shells[(i + j) * skip] - sn1 * all_shells[(i + j + 1) * skip]) / (sc1 * all_shells[(i + j) * skip])
                mask1 = torch.nan_to_num(mask1)

                d_update1 = mult_sample(d_cur1, mask1)
                x_next1 = x_next1 + d_update1

                img = torch.view_as_complex(denoised1[0].permute(1, 2, 0).contiguous())
                ksp = fft(maps * img)
                residual = all_A[0] * y[0] - all_A[0] * ksp
                sse = torch.norm(residual)**2

                grads = torch.autograd.grad(sse, x_hat1, retain_graph=False)[0]
                x_next1 = x_next1 - img_l_ss * grads # grad_update

                x_next1 = x_next1.detach()
            est_phi = learn_phi(x_next1, est_phi, y, all_A, spline_kernels, posx, posy, device, what_n=what_n, H=H, W=W, maps=maps,counter=counter,update_motion=update_motion)

        # Euler update step
        denoised = net(x_hat, sigma_cur, None).to(torch.float64)
        d_cur = denoised - x_hat

        mask = (sigma_cur * all_shells[i * skip] - sigma_next * all_shells[(i + 1) * skip]) / (sigma_cur * all_shells[i * skip])
        mask = torch.nan_to_num(mask)
        d_update = mult_sample(d_cur, mask)
        x_next = x_next + d_update

        if posterior:
            Phix_list, deform_x, deform_y = randn_deform(denoised[0], est_phi, posx, posy, device)
            sse = 0

            for n, Phixi in enumerate(Phix_list):
                img = torch.view_as_complex(
                    (denoised[0] if n == 0 else Phixi).permute(1, 2, 0).contiguous()
                )
                ksp = fft(maps * img)
                residual = all_A[n] * y[n] - all_A[n] * ksp
                sse += torch.norm(residual)**2

            grads = torch.autograd.grad(sse, x_hat, retain_graph=True)[0]
            x_next = x_next - img_l_ss * grads / (len(y))

        x_next = x_next.detach()

    return x_next, deform_x, deform_y, Phix_list

def rigid_deform(image, coords, angles, dx, dy, posx, posy, device):
    """
    Applies a set of rigid deforms to an input image.
    - coords: (H, W, 2). For kornia.remap. Pre-computed in the notebooks for efficiency.
    - angles: list of rotation angles.
    - dx: list of x translations.
    - dy: list of y translations.
    - posx, posy: (H, W)
    """
    Phix_list = []
    deform_x = []
    deform_y = []

    for i in range(len(angles)):
        rot_mat = torch.zeros(2,2, dtype=torch.double).to(device)
        theta = torch.tensor(np.pi) * angles[i:i+1]/180

        rot_mat[0,0] = torch.cos(theta)
        rot_mat[0,1] = -1*torch.sin(theta)
        rot_mat[1,0] = torch.sin(theta)
        rot_mat[1,1] =torch.cos(theta)

        deform_coords = coords @ rot_mat
        deform_coords = deform_coords + coords.shape[0] // 2
        deform_coords[:,:,0] += dx[i:i+1]
        deform_coords[:,:,1] += dy[i:i+1]
    
        deformed_i = remap(image, deform_coords[:,:,0][None], deform_coords[:,:,1][None], padding_mode='zeros')
        Phix_list.append(deformed_i[0])
        deform_x.append(deform_coords[:,:,0] - posx)
        deform_y.append(deform_coords[:,:,1] - posy)

    return Phix_list,deform_x,deform_y

def randn_deform(image, d_field, posx, posy, device):
    """
    Deform an image with given d_field (pixel-wise vector field).
    - image: (2, H, W)
    - d_field: (num_images, 2, H, W)
    - posx, posy: (H, W)
    """
    flow_x = posx + d_field[:,0]
    flow_y = posy + d_field[:,1]
    images = image[None].repeat(len(d_field), 1, 1, 1)

    return remap(images, flow_x, flow_y, padding_mode='zeros'), d_field[:,0], d_field[:,1]

def spline_deform(image, d_field, phi_list, spline_kernel, posx, posy, n_x, n_y, delta, device, H=256, W=256):
    """
    Deform an image by given d_field plus an additional B-spline deform with coefficients specified by phi_list and basis specified by spline_kernel.
    - image: (2, H, W)
    - d_field: (num_images, 2, H, W)
    - phi_list: list of length num_images, each of shape (n_x, n_y, 2)
    - spline_kernel: (2*delta, 2*delta)
    - posx, posy: (H, W)
    """
    num_images = len(phi_list)
    pad_H, pad_W = H + 2 * delta, W + 2 * delta
    deform_field = torch.zeros((num_images, pad_H, pad_W, 2), dtype=torch.float64, device=device)

    # Precompute spline patch and reshape for broadcasting
    kernel = spline_kernel[None, None, :, :, None]  # shape: (1, 1, 2*delta, 2*delta, 1)

    # Expand phi to shape: (num_images, n_x, n_y, 1, 1, 2) -> to broadcast over spline_kernel
    phi_tensor = torch.stack(phi_list, dim=0).unsqueeze(3).unsqueeze(4)  # (N, n_x, n_y, 1, 1, 2)
    patches = kernel * phi_tensor  # (N, n_x, n_y, 2*delta, 2*delta, 2)

    for i in range(n_x):
        for j in range(n_y):
            x_start, x_end = (i * delta), (i + 2) * delta
            y_start, y_end = (j * delta), (j + 2) * delta
            deform_field[:, x_start:x_end, y_start:y_end, :] += patches[:, i, j]

    deform_field = deform_field[:, delta:-delta, delta:-delta]  # crop padding
    remap_field = deform_field.permute(3, 0, 1, 2)  # (2, N, H, W)

    flow_x = posx + remap_field[0] + d_field[:, 0]
    flow_y = posy + remap_field[1] + d_field[:, 1]

    images = image[None].repeat(num_images, 1, 1, 1)
    return remap(images, flow_x, flow_y, padding_mode='zeros'), remap_field[0] + d_field[:,0], remap_field[1] + d_field[:,1]

def learn_phi(x_next, est_phi, all_Y, all_A, spline_kernels, posx, posy, device, what_n=[5, 12], H=256, W=256, maps=None, counter=0, update_motion=[]):
    """
    Optimize deformation fields phi using spline-based and random perturbation methods (multicoil only).

    Args:
        x_next (torch.Tensor): Input image, shape (1, 2, H, W).
        est_phi (torch.Tensor): Initial deformation field estimate, shape (N, 2, H, W).
        all_Y (List[torch.Tensor]): Measured k-space data per state, each (C, H, W).
        all_A (List[torch.Tensor]): Sampling mask per state, each (H, W).
        spline_kernels (Dict[int, torch.Tensor]): Map from n_star -> kernel (2*delta, 2*delta).
        posx (torch.Tensor): X grid, shape (H, W).
        posy (torch.Tensor): Y grid, shape (H, W).
        device (torch.device): CUDA or CPU device.
        what_n (List[int]): Spline control point grid sizes.
        H (int): Image height.
        W (int): Image width.
        maps (torch.Tensor): Sensitivity maps (C, H, W).

    Returns:
        torch.Tensor: Final deformation field (N, 2, H, W).
    """
    num_images = len(all_Y)
    num_gradient_iters = 1000
    num_spline_iters = 50

    for n_star in tqdm.auto.tqdm(what_n, desc=f"Fitting B-spline ({counter}/{len(update_motion)})...", leave=False):
        n_x = n_star
        n_y = n_star
        delta = H // (n_star - 1)
        spline_coefs = [torch.zeros((n_x, n_y, 2), dtype=torch.float64, device=device) for _ in range(num_images)]
        for n in range(num_images):
            spline_coefs[n] = spline_coefs[n].requires_grad_()
        
        for j in range(num_spline_iters):
            Phix_list, deform_x, deform_y = spline_deform(x_next[0], est_phi, spline_coefs, spline_kernels[n_star], posx, posy, n_x, n_y, delta, all_Y[0].device, H, W)
            
            all_sse = 0
            for n, Phixi in enumerate(Phix_list[1:]):
                img = torch.view_as_complex(Phixi.permute(1, 2, 0).contiguous())
                ksp = fft(maps * img)
                residual = all_A[n+1] * all_Y[n+1] - all_A[n+1] * ksp
                all_sse += torch.norm(residual)**2

            grads = torch.autograd.grad(outputs=all_sse, inputs=(spline_coefs), retain_graph=False)
            norm = 50
            for j in range(1, len(spline_coefs)):
                spline_coefs[j] = spline_coefs[j] - (1.0/norm) *grads[j]

        est_phi[:,0] = deform_x
        est_phi[:,1] = deform_y

    for _ in tqdm.auto.tqdm(range(num_gradient_iters), desc=f"Fitting vector field ({counter}/{len(update_motion)})...", leave=False):
        est_phi = est_phi.requires_grad_() # num images, 2, H, W
        Phix_list, deform_x, deform_y = randn_deform(x_next[0], est_phi, posx, posy, all_Y[0].device)

        all_sse = 0
        for n, Phixi in enumerate(Phix_list[1:]):
            img = torch.view_as_complex(Phixi.permute(1, 2, 0).contiguous())
            ksp = fft(maps * img)
            residual = all_A[n+1] * all_Y[n+1] - all_A[n+1] * ksp
            all_sse += torch.norm(residual)**2

        reg_grad_x = est_phi[:,:,:-1,:] - est_phi[:,:,1:,:]
        reg_grad_y = est_phi[:,:,:,:-1] - est_phi[:,:,:,1:]
        all_sse += 0.5 * torch.norm(reg_grad_x)**2
        all_sse += 0.5 * torch.norm(reg_grad_y)**2
        grads = torch.autograd.grad(outputs=all_sse, inputs=(est_phi), retain_graph=True)
        norm = 50
        est_phi = est_phi - (1.0/norm)*grads[0]
        # print(grads[0])


    return est_phi.detach()