Multi Session Temporal Registration
# Configuration file for the Sphinx documentation builder.
#
# For the full list of built-in configuration values, see the documentation:
# https://www.sphinx-doc.org/en/master/usage/configuration.html

import os
import sys
import toml

sys.path.insert(0, os.path.abspath(".."))

# get version from pyproject.toml
with open("../pyproject.toml") as f:
    pyproject = toml.load(f)

# -- Project information -----------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information

project = 'pymuster'
copyright = '2023, Edvard Grødem'
author = 'Edvard Grødem'
release = pyproject["tool"]["poetry"]["version"]

# -- General configuration ---------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration
extensions = [
    "sphinx.ext.autodoc",
    "sphinx.ext.napoleon",
    "sphinx.ext.coverage",
]

templates_path = ["_templates"]
exclude_patterns = []
master_doc = "index"
exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store']

# -- Options for HTML output -------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output

html_theme = 'sphinx_rtd_theme'
html_static_path = ['_static']

docs/conf.py
CRAI-OUS/MUSTER
import nox


@nox.session(python=["3.11"])
def tests(session):
    session.install("poetry")
    session.run("poetry", "install", external=True)
    session.run("poetry", "run", "pytest", "tests", external=True)

noxfile.py
CRAI-OUS/MUSTER
from .muster import Registration

src/pymuster/__init__.py
CRAI-OUS/MUSTER
import click


@click.command()
@click.option('--count', default=1, help='Number of greetings.')
@click.option('--name', default="Stupid", help='The person to greet.')
def main(count, name):
    """Simple program that greets NAME for a total of COUNT times."""
    for _ in range(count):
        click.echo(f"Hello, {name}!")

src/pymuster/cli/cli.py
CRAI-OUS/MUSTER
""" Containing useful functions for field calculus
"""

import torch
import torch.nn as nn
import torch.nn.functional as F

import numpy as np


def jacobian(x, pix_dim=(1.0, 1.0, 1.0), filter=False, midpoint=False):
    """
    Calculates the Jacobian of a 3D field
    Args:
        x: tensor of shape (batch_size, channels, x, y, z)
    Returns:
        J: jacobian of shape (batch_size, 3, channels, x, y, z)
    """
    m = x.shape[1]
    # Assert that filter only is used with midpoint
    assert not filter or midpoint, "Filter only works with midpoint"

    if filter:
        h = torch.tensor([[[1, 2, 1]]], dtype=torch.float32) / 4
    else:
        h = torch.tensor([[[0, 1, 0]]], dtype=torch.float32)

    if midpoint:
        h_m = torch.tensor([[[1, 0, -1]]], dtype=torch.float32) / 2
    else:
        h_m = torch.tensor([[[0, 1, -1]]], dtype=torch.float32)

    s3_x = torch.einsum('ijk, jki, kij -> kij', h_m / pix_dim[0], h, h)
    s3_y = torch.einsum('ijk, jki, kij -> kij', h, h_m / pix_dim[1], h)
    s3_z = torch.einsum('ijk, jki, kij -> kij', h, h, h_m / pix_dim[2])

    # s3_x = torch.einsum('ijk, jki, kij -> kij', h_m, h, h)
    # s3_y = torch.einsum('ijk, jki, kij -> kij', h, h_m, h)
    # s3_z = torch.einsum('ijk, jki, kij -> kij', h, h, h_m)

    s_filter = torch.cat((s3_x[None, ...], s3_y[None, ...], s3_z[None, ...]), dim=0)

    s_filter = s_filter.repeat(m, 1, 1, 1)[:, None, ...].to(x.device)

    # Use convolutions to calculate the jacobian
    J = F.conv3d(x, s_filter, padding=1, groups=m)
    J = J.reshape(J.shape[0], 3, m, J.shape[2], J.shape[3], J.shape[4])

    return J


def curl(J):
    """
    Returns the curl of a 3D field
    Args:
        J: Jacobian of shape (batch_size, 3, 3, x, y, z)
    """
    C = torch.zeros(J.shape[0], 3, J.shape[3], J.shape[4], J.shape[5]).to(J.device)
    C[:, 0, :, :, :] = J[:, 2, 1, :, :, :] - J[:, 1, 2, :, :, :]
    C[:, 1, :, :, :] = J[:, 0, 2, :, :, :] - J[:, 2, 0, :, :, :]
    C[:, 2, :, :, :] = J[:, 1, 0, :, :, :] - J[:, 0, 1, :, :, :]
    return C


def divergence(J):
    """
    Returns the divergence of a 3D field
    Args:
        J: Jacobian of shape (batch_size, 3, 3, x, y, z)
    """
    return J[:, 0, 0, :, :, :] + J[:, 1, 1, :, :, :] + J[:, 2, 2, :, :, :]


def laplacian(J, pix_dim=(1.0, 1.0, 1.0)):
    """
    Return the second order jacobian of a field
    """
    channels = J.shape[2]
    L = jacobian(J.reshape(J.shape[0], 3 * channels, J.shape[3], J.shape[4], J.shape[5]), pix_dim=pix_dim)
    L = L.reshape(J.shape[0], channels, 3, 3, J.shape[3], J.shape[4], J.shape[5])
    return L


def determinant(J):
    """
    Computes the determinant of a 3x3 matrix
    """
    J = J.permute(0, 3, 4, 5, 1, 2)
    return torch.det(J)


def integrate_forward_flow(v, sp_tr, n_iter=5, dt=1):
    """
    implements step (4): Calculation of forward flows
    @return:
    """
    image_size = v.shape[-3:]
    T = v.shape[0]
    # create flow
    Phi0 = torch.zeros(v.shape, device=v.device)

    # Phi0_0 is the identity mapping
    # create sampling grid
    vectors = [torch.arange(0, s) for s in image_size]
    grids = torch.meshgrid(vectors)
    grid = torch.stack(grids)
    grid = torch.unsqueeze(grid, 0)
    print(f"Grid shape: {grid.shape}")

    # Phi0_0 is the identity mapping
    Phi0[0] = grid[0]

    for t in range(0, T - 1):
        alpha = forward_alpha(v[[t]], sp_tr, n_iter, dt)
        Phi0[t + 1] = sp_tr(Phi0[[t]], -alpha)

    return Phi0


def forward_alpha(v_t, sp_tr, n_iter, dt=1):
    """
    helper function for step (4): Calculation of forward flows
    @param v_t: the velocity field
    @param sp_tr: the spatial transformer
    @return:
    """
    alpha = torch.zeros(v_t.shape, device=v_t.device)
    for i in range(n_iter):
        alpha = dt * sp_tr(v_t, -0.5 * alpha)
    return alpha

src/pymuster/core/field_calculus.py
CRAI-OUS/MUSTER
import torch
import torch.nn.functional as F

import numpy as np
import math

from . import field_calculus


class NCC:
    """Normalized Local Cross Correlation.
    
    Args:
        kernal_size (int): The size of the kernal used for calculating the local mean
    """

    def __init__(self, kernal_size=7):
        self.kernal_size = kernal_size
        self.n = kernal_size**3
        self.padding = kernal_size // 2

    def loss(self, y_true, y_pred, mask):
        """
        Args:
            y_true (torch.Tensor): The true field
            y_pred (torch.Tensor): The predicted field
            mask (torch.Tensor): The mask used for calculating the loss
        """
        vox_img = torch.prod(torch.tensor(y_true.shape[2:]))

        y_true_means = torch.nn.functional.avg_pool3d(y_true, self.kernal_size, stride=1, padding=self.padding)
        y_pred_means = torch.nn.functional.avg_pool3d(y_pred, self.kernal_size, stride=1, padding=self.padding)
        diff_true = y_true - y_true_means
        diff_pred = y_pred - y_pred_means

        diff_true_squared = torch.nn.functional.avg_pool3d(
            diff_true * diff_true, self.kernal_size, stride=1, padding=self.padding)
        diff_pred_squared = torch.nn.functional.avg_pool3d(
            diff_pred * diff_pred, self.kernal_size, stride=1, padding=self.padding)

        if mask is not None:
            diff_true = diff_true * mask
            diff_pred = diff_pred * mask
            diff_true_squared = diff_true_squared * mask
            diff_pred_squared = diff_pred_squared * mask

        # Equvalent to the NCC see: B. Panm (2011) Recent Progress in Digital Image Correlation
        return 1 / 2.0 * torch.sum(
            (diff_true / torch.sqrt(diff_true_squared + 0.01) - diff_pred / torch.sqrt(diff_pred_squared + 0.01))**
            2) / torch.sum(mask)


class Fourier:
    """Spatial jacobian loss using the fourier space.
    
    Args:
        filter_omega (float): The cutoff frequency of the filter
        image_size (tuple): The size of the image
        pix_dim (tuple): The pixel dimensions of the image
    """

    def __init__(self, filter_omega, image_size, pix_dim):
        self.filter_omega = filter_omega  #/pix_dim[0]
        self.image_size = image_size
        self.pix_dim = pix_dim

    def loss(self, y_true, y_pred, mask=None):
        filtered_true = self.filter(y_true)
        filtered_pred = self.filter(y_pred)

        if mask is not None:
            filtered_true = filtered_true * mask
            filtered_pred = filtered_pred * mask

        dims = [-4, -3, -2, -1]
        diff_true = filtered_true  #- torch.mean(filtered_true, dims, keepdim=True)
        diff_pred = filtered_pred  #- torch.mean(filtered_pred, dims, keepdim=True)
        return 1 - torch.mean(
            torch.sum(diff_true * diff_pred, dims) /
            torch.sqrt(torch.sum(diff_true * diff_true, dims) * torch.sum(diff_pred * diff_pred, dims)))

    def filter(self, y):
        y_fft = torch.fft.fftn(y, dim=(-3, -2, -1))

        # Compute the derivative by multiplying with the fourier space coordinates
        freqsx = torch.fft.fftfreq(self.image_size[0], d=self.pix_dim[0], device=y.device)
        freqsy = torch.fft.fftfreq(self.image_size[1], d=self.pix_dim[1], device=y.device)
        freqsz = torch.fft.fftfreq(self.image_size[2], d=self.pix_dim[2], device=y.device)
        freqsx, freqsy, freqsz = torch.meshgrid(freqsx, freqsy, freqsz, indexing="ij")

        filter = torch.exp(-1 / 2 * (freqsx**2 + freqsy**2 + freqsz**2) / self.filter_omega**2)

        freqs = torch.stack([freqsx, freqsy, freqsz], dim=0).to(y.device)

        filter = filter * freqs * 1j

        y_fft = y_fft * filter[None, ...]
        y_fft = torch.fft.ifftn(y_fft, dim=(-3, -2, -1)).real
        return y_fft


class NCCS:
    """ NCC with Sobel filter.
    See: https://doi.org/10.3390/app12062828
    
    Args:
        kernal_size (int): The size of the kernal used for calculating the local mean
    """

    def __init__(self, kernal_size=13):
        self.kernal_size = kernal_size
        self.n = kernal_size**3
        self.padding = kernal_size // 2

    def loss(self, y_true, y_pred, mask=None):

        vox_img = torch.prod(torch.tensor(y_true.shape[2:]))

        y_true = field_calculus.jacobian(y_true, filter=True, midpoint=True)[:, :, 0, ...]
        y_pred = field_calculus.jacobian(y_pred, filter=True, midpoint=True)[:, :, 0, ...]

        y_true_means = torch.nn.functional.avg_pool3d(y_true, self.kernal_size, stride=1, padding=self.padding)
        y_pred_means = torch.nn.functional.avg_pool3d(y_pred, self.kernal_size, stride=1, padding=self.padding)
        diff_true = y_true - y_true_means
        diff_pred = y_pred - y_pred_means

        diff_true_squared = torch.nn.functional.avg_pool3d(
            diff_true * diff_true, self.kernal_size, stride=1, padding=self.padding)
        diff_pred_squared = torch.nn.functional.avg_pool3d(
            diff_pred * diff_pred, self.kernal_size, stride=1, padding=self.padding)

        diff_true_pred = torch.nn.functional.avg_pool3d(
            diff_true * diff_pred, self.kernal_size, stride=1, padding=self.padding)

        if mask is not None:
            diff_true_squared = diff_true_squared * mask
            diff_pred_squared = diff_pred_squared * mask
            diff_true_pred = diff_true_pred * mask

        return 1 - torch.sum(
            diff_true_pred / torch.sqrt(diff_true_squared * diff_pred_squared + 0.0001)) / torch.sum(mask)


class WNCC:

    def __init__(self, kernal_size=7):
        self.kernal_size = kernal_size
        self.n = kernal_size**3
        self.padding = kernal_size // 2

    def loss(self, y_true, y_pred, mask):
        vox_img = torch.prod(torch.tensor(y_true.shape[2:]))

        y_true = y_true / y_true.std(dim=[-3, -2, -1], keepdim=True)
        y_pred = y_pred / y_pred.std(dim=[-3, -2, -1], keepdim=True)

        y_true_means = torch.nn.functional.avg_pool3d(y_true, self.kernal_size, stride=1, padding=self.padding)
        y_pred_means = torch.nn.functional.avg_pool3d(y_pred, self.kernal_size, stride=1, padding=self.padding)
        diff_true = y_true - y_true_means
        diff_pred = y_pred - y_pred_means

        diff_true_squared = torch.nn.functional.avg_pool3d(
            diff_true * diff_true, self.kernal_size, stride=1, padding=self.padding)
        diff_pred_squared = torch.nn.functional.avg_pool3d(
            diff_pred * diff_pred, self.kernal_size, stride=1, padding=self.padding)

        # diff_true_pred = torch.nn.functional.avg_pool3d(
        #     diff_true * diff_pred, self.kernal_size, stride=1, padding=self.padding)

        if mask is not None:
            diff_true_squared = diff_true_squared * mask
            diff_pred_squared = diff_pred_squared * mask

        with torch.no_grad():
            cross_std = torch.sqrt(torch.sqrt(diff_true_squared * diff_pred_squared + 0.0001))
            tot_cross_std = torch.sum(cross_std, [-4, -3, -2, -1])

        return torch.mean(1 / tot_cross_std * torch.sum(
            cross_std * (diff_true / torch.sqrt(diff_true_squared + 0.00001) -
                         diff_pred / torch.sqrt(diff_pred_squared + 0.00001))**2, [-4, -3, -2, -1]))


class GaussNCC:

    def __init__(self, filter_omega, image_size, pix_dim):
        self.filter_omega = filter_omega  #/pix_dim[0]
        self.image_size = image_size
        self.pix_dim = pix_dim
        self.omega_f = 1 / (2 * torch.pi * filter_omega)

    def loss(self, y_true, y_pred, mask):

        mask = y_true != 0
        mask = mask.float()
        #mask = torch.ones_like(y_true)

        vox_img = torch.prod(torch.tensor(y_true.shape[2:]))

        y_true_f = torch.fft.fftn(y_true, dim=(-3, -2, -1))
        y_pred_f = torch.fft.fftn(y_pred, dim=(-3, -2, -1))

        freqsx = torch.fft.fftfreq(self.image_size[0], d=self.pix_dim[0], device=y_true.device)
        freqsy = torch.fft.fftfreq(self.image_size[1], d=self.pix_dim[1], device=y_true.device)
        freqsz = torch.fft.fftfreq(self.image_size[2], d=self.pix_dim[2], device=y_true.device)
        freqsx, freqsy, freqsz = torch.meshgrid(freqsx, freqsy, freqsz, indexing="ij")

        filter_response = torch.exp(-1 / 2 * ((freqsx**2 + freqsy**2 + freqsz**2) / self.omega_f**2))

        y_true_means_f = filter_response * y_true_f
        y_pred_means_f = filter_response * y_pred_f

        diff_true = y_true - torch.fft.ifftn(y_true_means_f, dim=(-3, -2, -1)).real
        diff_pred = y_pred - torch.fft.ifftn(y_pred_means_f, dim=(-3, -2, -1)).real

        diff_true_squared = torch.fft.ifftn(
            filter_response * torch.fft.fftn(diff_true * diff_true, dim=(-3, -2, -1)), dim=(-3, -2, -1))
        diff_pred_squared = torch.fft.ifftn(
            filter_response * torch.fft.fftn(diff_pred * diff_pred, dim=(-3, -2, -1)), dim=(-3, -2, -1))

        if mask is not None:
            diff_true = diff_true * mask
            diff_pred = diff_pred * mask
            diff_true_squared = diff_true_squared * mask
            diff_pred_squared = diff_pred_squared * mask

        # Convert to real
        diff_true_squared = diff_true_squared.real
        diff_pred_squared = diff_pred_squared.real

        # Equvalent to the NCC see: B. Panm (2011) Recent Progress in Digital Image Correlation
        return 1 / 2.0 * torch.sum(
            (diff_true / torch.sqrt(diff_true_squared + 0.0001) - diff_pred / torch.sqrt(diff_pred_squared + 0.0001))**
            2) / torch.sum(mask)


class MSE:
    """
    Mean squared error loss.
    """

    def loss(self, y_true, y_pred):
        return torch.mean((y_true - y_pred)**2)

src/pymuster/core/losses.py
CRAI-OUS/MUSTER
import torch
import torch.nn as nn
import torch.nn.functional as F

import numpy as np


def grid(grid_size, grid_interval):
    """
    Constructs a 3d image of a grid with a given size
    For visualization purposes
    Args:
        grid_size: (x, y, z)
        grid_interval: The interval between the grid lines
    """
    image = np.zeros(grid_size)
    # Draw vertical lines
    image[::grid_interval, :, :] = 1
    image[:, ::grid_interval, :] = 1
    image[:, :, ::grid_interval] = 1
    return image


class RunningMeanVar():
    """
    This class is used to calculate the mean and variance of the n last values
    Example:
        >>> rmv = RunningMeanVar(n=10)
        >>> for i in range(100):
        >>>     rmv.add(i)
        >>>     print(rmv.mean())
        
    Args:
        n: The number of values to use for calculating the mean and variance
    """

    def __init__(self, n=10):
        self.n = n
        self.values = np.zeros(n)
        self.last_index = 0

    def add(self, value):
        """ Add a value to the running mean and variance
        Args:
            value: The value to add
        
        """
        self.values[self.last_index] = value
        self.last_index = (self.last_index + 1) % self.n

    def mean(self):
        """
        Returns: The mean of the last n values
        """
        return np.mean(self.values)

    def var(self):
        """
        Returns: The variance of the last n values
        """
        return np.var(self.values)


def VectorPearsonCorrelation(x, y):
    """
    Calculates the Pearson correlation between two vectors sets
    Args:
        x: (num_samples, num_features)
        y: (num_samples, num_features)
    """
    mean_x = torch.mean(x, dim=0, keepdim=True)
    mean_y = torch.mean(y, dim=0, keepdim=True)

    diff_x = x - mean_x
    diff_y = y - mean_y

    sigma_x = torch.sqrt(torch.sum(diff_x**2))
    sigma_y = torch.sqrt(torch.sum(diff_y**2))

    pcc = torch.sum(diff_x * diff_y) / (sigma_x * sigma_y)
    return pcc


def get_dig_ind(N, offset):
    """Returns the indices of the diagonal elements of the offset diagonal matrices.
    """
    if offset == 0:
        return np.diag_indices(N)
    indices = np.array(np.diag_indices(N - np.abs(offset)))
    if offset > 0:
        indices[1] += offset
    else:
        indices[0] -= offset
    return (indices[0], indices[1])

src/pymuster/core/utils.py
CRAI-OUS/MUSTER
"""Multi Session Temporal Registration.

This module is the core of Multi Session Temporal Registration (MUSTER).

Classes:
    StageRegistration: Register a longitudinal series of images using a series of deformation fields using a single resolution.

"""

import torch
import torch.nn as nn
import torch.nn.functional as F

import numpy as np

from alive_progress import alive_bar

from .core import field_calculus
from .core import losses
from .core import utils

import warnings


class StageRegistration:
    """Register a longitudinal series of images using a series of deformation fields using a single resolution.

    The images must be in the same space and must be pre-registered using an rigid or affine transformation.

    The deformation between two consectutive timepoints is represented by a deformation field. These fields are 
    generated by integrating a vector field, also known as the deformation flow. Various methods for 
    this integration are supported:
        * ``integration_method="ss"``: Use scaling and squaring to integrate stationary vector fields. (Default)
        * ``integration_method=euler``: Use euler integration for stationary vector fields.
        * ``integration_method=rk4``: Use Runge-Kutta 4 integration for stationary vector field.
        * ``integration_method=euler_time``: Use euler integration for time varying vector fields.
        * ``integration_method=rk4_time``: Use Runge-Kutta 4 integration for time varying vector fields.

    For non-consecutive timepoints, the deformation field can be composed in two ways:
        * ``field_composition_method="interpolate"``: Use interpolation to compose the intermediate deformation 
            fields. (Default)
        * ``field_composition_method="flow_continuation"``: Continuate the deformation field by integrating the 
            intermediate deformation flows.

    Supported image similarity metrics include:
        * ``img_similarity_metric="NCC"``: Normalized local cross correlation between the two images as loss. 
            The local neighborhood is a cube of size with side length ``img_similarity_spatial_size``.
        * ``img_similarity_metric="L2"``: L2 norm of the difference between the two images as loss. 
        * ``img_similarity_metric="NCCS"``: Use a sobel filter to compute the gradient of the two images, and use the 
            normalized local cross correlation between the two gradients as loss. The local window
            is a cube of size with side length ``3``.
        * ``img_similarity_metric="WNCC"``: Normalized local cross correlation between the two images as loss where each 
            local neighborhood is weighted by the cross standard deviation of the two images.
            The local neighborhood is a cube of size with side length ``img_similarity_spatial_size``.
        * ``img_similarity_metric="GaussNCC"``: Normalized local cross correlation between the two images as loss where
            each local neighborhood is weighted by a gaussian filter. The standard deviation of the gaussian filter is 
            ``img_similarity_spatial_size``.
        * ``img_similarity_metric="Fourier"``: Use the Fourier tranform to compute a filtered gradient of each images,
            and compute the global normalized cross correlation between the two filtered gradients as loss. 
            The standard deviation of the gaussian filter is ``img_similarity_spatial_size``.


    Args:
        image_size (tuple): The size of the input images in the format ``(x, y, z)``.
        pix_dim (tuple): The pixel dimensions of the input images in the format ``(dx, dy, dz)``. Default is ``(1, 1, 1)``.
        deform_res_scale (int): The resolution scale of the deformation field. The deformation field will have a 
            resolution of image_size/deform_res_scale. Default is ``1``.
        device (str): The device to use for computation. ('cpu' or 'cuda:n' or torch.device object). Default is ``"cpu"``.
        num_iterations (int): Number of optimization iterations. Default is ``100``.
        spatial_smoothness_penalty (float): Weight for spatial smoothness in loss function. Default is ``1.0``.
        temporal_smoothness_penalty (float): Weight for temporal smoothness in loss function. Default is ``1.0``.
        invertability_penalty (float): Weight for invertibility penalty in loss function. Default is ``1.0``.
        l2_penalty (float): Weight for L2 penalty on the deformation flow. Default is ``0.0``.
        smoothing_sigma (float): Sigma for Gaussian smoothing of deformation flows. Default is ``1.0``.
        mode (str): The mode for interpolation of the images. ('nearest', 'bilinear', or 'bicubic'). Default is ``"bilinear"``.
        intergration_steps (int): The number of integration steps for integrating deformation flow. Default is ``7``.
        integration_method (str): The method for integrating deformation field.
            ('ss', 'euler', 'rk4', 'euler_time', or 'rk4'). Default is ``"ss"``.
        field_composition_method (str): The method for composing deformation fields.
            ('interpolate' or 'flow_continuation'). Default is ``"interpolate"``.
        rigid (bool): Flag for including rigid registration. Default is ``False``.
        learning_rate (float): Learning rate for optimizer. Default is ``1e-3``.
        betas (tuple): The beta coefficients for Adam optimizer. Default is ``(0.9, 0.999)``.
        tol (float): Tolerance for optimization convergence. Default is ``1e-4``.
        img_similarity_metric (str): The image similarity metric ('NCC', 'L2', 'NCCS', 'WNCC', 'GaussNCC', or 'Fourier').
            Default is ``"NCC"``.
        img_similarity_spatial_size (int/float): Size or standard deviation of the local neighborhood for the image 
            similarity metric. Default is ``3``.
        verbose (bool): Flag for printing progress during optimization. Default is ``True``.
    """

    def __init__(
        self,
        image_size: tuple,
        pix_dim: tuple = (1, 1, 1),
        deform_res_scale: int = 1,
        device: str = "cpu",
        num_iterations: int = 100,
        spatial_smoothness_penalty: float = 1.0,
        temporal_smoothness_penalty: float = 1.0,
        invertability_penalty: float = 1.0,
        l2_penalty: float = 0.0,
        smoothing_sigma: float = 1.0,
        mode: str = "bilinear",
        intergration_steps: int = 7,
        integration_method: str = "ss",
        field_composition_method: str = "interpolate",
        rigid: bool = False,
        learning_rate: float = 1e-3,
        betas: tuple = (0.9, 0.999),
        tol: float = 1e-4,
        img_similarity_metric: str = "NCC",
        img_similarity_spatial_size: int or float = 3,
        verbose: bool = True,
    ):
        # Valid values for integration_method
        valid_integration_methods = ["ss", "euler", "rk4", "euler_time", "rk4_time"]
        if integration_method not in valid_integration_methods:
            raise ValueError(f"integration_method must be one of {valid_integration_methods}")

        # Valid values for field_composition_method
        valid_field_composition_methods = ["interpolate", "flow_continuation"]
        if field_composition_method not in valid_field_composition_methods:
            raise ValueError(f"field_composition_method must be one of {valid_field_composition_methods}")

        # flow_continuation can not be combined with ss
        if field_composition_method == "flow_continuation" and integration_method == "ss":
            raise ValueError("flow_continuation can not be combined with ss")

        # Valid values for img_similarity_metric
        valid_img_similarity_metrics = ["NCC", "L2", "NCCS", "Fourier", "WNCC", "GaussNCC"]
        if img_similarity_metric not in valid_img_similarity_metrics:
            raise ValueError(f"img_similarity_metric must be one of {valid_img_similarity_metrics}")

        # Valid values for mode
        valid_modes = ["nearest", "bilinear", "bicubic"]
        if mode not in valid_modes:
            raise ValueError(f"mode must be one of {valid_modes}")

        # Check the dimensions of the image
        if len(image_size) != 3:
            raise ValueError("image_size must be a tuple of length 3")

        # Check the dimensions of the pixel dimensions
        if len(pix_dim) != 3:
            raise ValueError("pix_dim must be a tuple of length 3")

        self.device = torch.device(device)

        self.num_iterations = num_iterations
        self.spatial_smoothness_penalty = spatial_smoothness_penalty
        self.temporal_smoothness_penalty = temporal_smoothness_penalty
        self.l2_penalty = l2_penalty
        self.smoothing_sigma = smoothing_sigma
        self.intergration_steps = intergration_steps
        self.integration_method = integration_method
        self.field_composition_method = field_composition_method

        self.rigid = rigid
        self.betas = betas
        self.learning_rate = learning_rate
        self.image_size = image_size
        self.deform_res_scale = deform_res_scale
        self.verbose = verbose
        self.tol = tol
        self.img_similarity_metric = img_similarity_metric
        self.pix_dim = pix_dim
        self.invertability_penalty = invertability_penalty
        self.img_similarity_spatial_size = img_similarity_spatial_size

        self.deform_size = [int(s // self.deform_res_scale) for s in self.image_size]
        self.deform_pix_dim = [s * self.deform_res_scale for s in self.pix_dim]
        self.mode = mode

        self.sp_tr = SpatialTransformer(
            self.deform_size,
            self.deform_pix_dim,
            mode=mode,
            padding_mode="border",
        ).to(self.device)

        self.sp_tr_img = SpatialTransformer(self.image_size, pix_dim, mode=mode, padding_mode="border").to(self.device)

    def fit(
        self,
        images,
        inital_deform_flow=None,
        timepoints=None,
        initial_rotations=None,
        initial_translations=None,
        masks=None,
    ):
        """Fits the deformation flow to the images.

        Args:
            images (torch.Tensor or np.ndarray): shape ``(N, channels, x, y, z)``
            inital_deform_flow (torch.Tensor or np.ndarray): shape ``(N-1, 3, x, y, z)``
            timepoints (torch.Tensor or np.ndarray): shape ``(N,)``
                The timepoints of the images. Used to adjust the temporal
                penalties. May be None, in which case the timepoints are
                assumed to be equally spaced.

        Returns:
            out (dict): Dictionary containing the deformation flow and optionally the rotations and translations.
        """

        self.n_img = images.shape[0]  # Number of images
        if self.integration_method in ['ss', 'euler', 'rk4']:
            self.n_flow = self.n_img - 1  # Number of deformations flow fields
        else:
            self.n_flow = self.n_img

        # Set device
        device = self.device

        images = self._convert_to_torch(images)

        if masks is None:
            masks = torch.ones_like(images)

        # The time difference between the images
        if timepoints is None:
            delta_times = torch.ones((self.n_img - 1, 1, 1, 1, 1), device=device)
        else:
            timepoints = self._convert_to_torch(timepoints)
            delta_times = timepoints[1:] - timepoints[:-1]
            delta_times = delta_times[:, None, None, None, None]

        # The integral deformation flow
        if inital_deform_flow is None:
            deform_flow = nn.parameter.Parameter(torch.zeros((self.n_flow, 3, *self.deform_size), device=device))
        else:
            inital_deform_flow = self._convert_to_torch(inital_deform_flow)
            deform_flow = nn.parameter.Parameter(inital_deform_flow.clone().to(device))

        param_groups = [{"params": [deform_flow], "lr": self.learning_rate}]

        # Parameters for rigid ajustment of the images
        if self.rigid:
            if initial_rotations is None:
                rotation_params = nn.Parameter(torch.zeros((self.n_img - 1, 3), device=device))
            else:
                initial_rotations = self._convert_to_torch(initial_rotations[1:])
                rotation_params = nn.Parameter(initial_rotations.clone().to(device))

            if initial_translations is None:
                translation_params = nn.Parameter(torch.zeros((self.n_img - 1, 3), device=device))
            else:
                initial_translations = self._convert_to_torch(initial_translations[1:])
                translation_params = nn.Parameter(initial_translations.clone().to(device))

            param_groups.append({
                "params": [rotation_params, translation_params],
                "lr": self.learning_rate * 0.0001,
            })

        optimizer = torch.optim.Adam(param_groups, lr=self.learning_rate, betas=self.betas)
        scaler = torch.cuda.amp.GradScaler()

        # Cosine annealing with linear warmup
        warmup_steps = int(self.num_iterations * 0.2)
        scheduler1 = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda step: 1 / (warmup_steps) * (step + 1))
        scheduler2 = torch.optim.lr_scheduler.CosineAnnealingLR(
            optimizer, T_max=int(self.num_iterations - warmup_steps))
        scheduler = torch.optim.lr_scheduler.SequentialLR(
            optimizer, schedulers=[scheduler1, scheduler2], milestones=[warmup_steps])

        if self.img_similarity_metric == "MSE" or self.img_similarity_metric == 'L2':
            image_img_similarity_metric = losses.MSE()
        elif self.img_similarity_metric == "NCC":
            image_img_similarity_metric = losses.NCC(self.img_similarity_spatial_size)
        elif self.img_similarity_metric == "WNCC":
            image_img_similarity_metric = losses.WNCC(self.img_similarity_spatial_size)
        elif self.img_similarity_metric == "GaussNCC":
            image_img_similarity_metric = losses.GaussNCC(self.img_similarity_spatial_size, self.image_size,
                                                          self.pix_dim)
        elif self.img_similarity_metric == "NCCS":
            image_img_similarity_metric = losses.NCCS(self.img_similarity_spatial_size)
        elif self.img_similarity_metric == "Fourier":
            image_img_similarity_metric = losses.Fourier(self.img_similarity_spatial_size, self.image_size,
                                                         self.pix_dim)

        # Object for keeping track of convergence of the optimization process
        running_mean_var = utils.RunningMeanVar()

        use_autocast = torch.cuda.is_available() and (str(device) != "cpu")

        with alive_bar(
                self.num_iterations,
                force_tty=True,
                max_cols=130,
                dual_line=True,
                title_length=50,
                elapsed="({elapsed})",
                stats="(eta: {eta})",
                disable=not self.verbose,
        ) as bar:
            for iteration_idx in range(self.num_iterations):
                optimizer.zero_grad()

                with torch.cuda.amp.autocast(use_autocast):
                    if self.rigid:
                        # Add the identity transformation to first timepoint
                        rots = torch.cat((torch.zeros((1, 3), device=device), rotation_params), dim=0)
                        trans = torch.cat((torch.zeros((1, 3), device=device), translation_params), dim=0)
                        tran_grid = self._get_affine_grid(rots, trans)
                        batch = self.sp_tr_img(images, tran_grid, displacement=False)
                    else:
                        batch = images

                    loss = 0
                    image_loss = 0
                    inv_loss = 0

                    # The cumulative deformation fields
                    cum_deform_field_fwd = self._intergate_flow(deform_flow, dir="fwd")  # mm/step
                    cum_deform_field_bwd = self._intergate_flow(deform_flow, dir="bwd")  # mm/step

                    # The consecutive deformation fields
                    deform_field_fwd = cum_deform_field_fwd
                    deform_field_bwd = cum_deform_field_bwd

                    # Regularize the deformation field
                    loss += self.spatial_smoothness_penalty * self._spatial_continuity_loss(deform_flow)
                    loss += self.l2_penalty * torch.mean(deform_flow**2)
                    if self.n_img > 2:
                        loss += (self.temporal_smoothness_penalty * self._temperal_continuity_loss(deform_flow))

                    # Loop over the seperation between the images starting with consecutive images and ending with the
                    # first and last image
                    for delta_timestep in range(self.n_img - 1):
                        if delta_timestep > 0:
                            (cum_deform_field_fwd, cum_deform_field_bwd) = self._commute_deform_fields(
                                cum_deform_field_fwd,
                                cum_deform_field_bwd,
                                deform_flow,
                                delta_timestep,
                                deform_field_fwd,
                                deform_field_bwd,
                            )

                        # Interpolate the cum_deform_fields to the resolution of the images
                        if np.any(np.array(self.image_size) != np.array(self.deform_size)):
                            cum_deform_field_fwd_interp = nn.functional.interpolate(
                                cum_deform_field_fwd,
                                size=self.image_size,
                                mode="trilinear",
                                align_corners=True,
                            )
                            cum_deform_field_bwd_interp = nn.functional.interpolate(
                                cum_deform_field_bwd,
                                size=self.image_size,
                                mode="trilinear",
                                align_corners=True,
                            )
                        else:
                            cum_deform_field_fwd_interp = cum_deform_field_fwd
                            cum_deform_field_bwd_interp = cum_deform_field_bwd

                        sample = self.sp_tr_img(batch[:-delta_timestep - 1, ...], cum_deform_field_fwd_interp)

                        image_loss += image_img_similarity_metric.loss(
                            sample,
                            batch[delta_timestep + 1:, ...],
                            mask=masks[delta_timestep + 1:, ...],
                        ) * (
                            self.n_img - 1 - delta_timestep)

                        sample = self.sp_tr_img(
                            batch[delta_timestep + 1:, ...],
                            cum_deform_field_bwd_interp,
                        )

                        image_loss += image_img_similarity_metric.loss(
                            sample,
                            batch[:-delta_timestep - 1, ...],
                            mask=masks[:-delta_timestep - 1, ...],
                        ) * (
                            self.n_img - 1 - delta_timestep)

                        if self.invertability_penalty > 0:
                            inv_loss += self.invertability_penalty * self._invertability_loss(
                                cum_deform_field_fwd, cum_deform_field_bwd) * (
                                    self.n_img - 1 - delta_timestep)

                    # Normalize with respect to the number of images
                    image_loss /= (self.n_img - 1) * self.n_img
                    inv_loss /= (self.n_img - 1) * self.n_img

                    # Combine the losses
                    loss += image_loss + inv_loss

                if use_autocast:
                    # Propagate the loss and update the parameters
                    scaler.scale(loss).backward()
                    scaler.step(optimizer)
                    scaler.update()
                else:
                    loss.backward()
                    optimizer.step()

                # Ignore the warning from the scheduler
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore")
                    scheduler.step()

                bar()
                bar.title(f"\tLoss: {loss.item():.4e}, Image Loss: {image_loss.item():.4e}")
                bar.title_length = 200

                running_mean_var.add(loss.item())

                if self.verbose:
                    if iteration_idx == self.num_iterations - 1:
                        bar.title(f"\tMax iterations reached")

                if (iteration_idx > 20 and np.sqrt(running_mean_var.var()) / running_mean_var.mean() < self.tol):
                    bar.title(f"\tConvergence reached")
                    break

        optimizer.zero_grad()

        return_dict = {
            "deform_flow": deform_flow.detach().cpu().numpy(),
        }
        if self.rigid:
            return_dict["rotations"] = rotation_params.detach().cpu().numpy()
            return_dict["translations"] = translation_params.detach().cpu().numpy()

        return return_dict

    def _get_homogeneous_transformation_matrix(self, rotation, translation):
        """Computes the affine matrix from rotation and translation
        """
        Rx = torch.eye(3, device=self.device).repeat(rotation.shape[0], 1, 1)
        Rx[:, [1, 2], [1, 2]] = torch.cos(rotation[:, 0])[:, None]
        Rx[:, 1, 2] = -torch.sin(rotation[:, 0])
        Rx[:, 2, 1] = torch.sin(rotation[:, 0])

        Ry = torch.eye(3, device=self.device).repeat(rotation.shape[0], 1, 1)
        Ry[:, [0, 2], [0, 2]] = torch.cos(rotation[:, 1])[:, None]
        Ry[:, 0, 2] = torch.sin(rotation[:, 1])
        Ry[:, 2, 0] = -torch.sin(rotation[:, 1])

        Rz = torch.eye(3, device=self.device).repeat(rotation.shape[0], 1, 1)
        Rz[:, [0, 1], [0, 1]] = torch.cos(rotation[:, 2])[:, None]
        Rz[:, 0, 1] = -torch.sin(rotation[:, 2])
        Rz[:, 1, 0] = torch.sin(rotation[:, 2])

        R = torch.matmul(Rz, torch.matmul(Ry, Rx))
        matrix = torch.cat([R, translation[:, :, None]], dim=2)
        return matrix

    def _get_affine_grid(self, rotation=None, translation=None, matrix=None):
        """
        Computes the affine grid from rotation and translation
        """

        if matrix is None:
            matrix = self._get_homogeneous_transformation_matrix(rotation, translation)

        grid = self.sp_tr_img.grid

        # Grid is in the shape (N, 3, H, W, D)
        # Matrix is in the shape (N, 3, 4)
        grid = (torch.einsum("nij, njhwd -> nihwd", matrix[:, :3, :3], grid) + matrix[:, :3, 3, None, None, None])
        return grid

    def _invertability_loss(self, deform_field_fwd, deform_field_bwd):
        method = "ICON"  # 'ICON' or 'gradICON'

        grid = self.sp_tr.grid

        fwd_loss = (self.sp_tr(deform_field_bwd + grid, deform_field_fwd) - grid)
        bwd_loss = (self.sp_tr(deform_field_fwd + grid, deform_field_bwd) - grid)

        if method == "ICON":
            return torch.mean(fwd_loss**2 + bwd_loss**2)
        elif method == "gradICON":
            J = field_calculus.jacobian(torch.cat([fwd_loss, bwd_loss], dim=0), self.deform_pix_dim)
            margin = 2
            J = J[:, :, :, margin:-margin, margin:-margin, margin:-margin]
            return torch.mean((J)**2)

    def _spatial_continuity_loss(self, deform_flow, omega_1=0.1, omega_2=1.0, omega_3=0.0):
        J = field_calculus.jacobian(deform_flow, pix_dim=self.deform_pix_dim)
        loss = (
            omega_1 * torch.mean(J**2) + omega_2 * torch.mean(field_calculus.divergence(J)**2) +
            torch.mean(field_calculus.curl(J)**2) * omega_3)
        return loss

    def _temperal_continuity_loss(self, deform_flow):
        return torch.mean(((deform_flow[1:, ...] - deform_flow[:-1, ...])**2))

    def _intergate_flow(self, deform_flow, dir, target_grid=None):
        """
        Integrates a vector field via scaling and squaring.
        Args:
            deform_flow: The vector field to integrate. Shape (N, 3, x, y, z)
            dir: The direction to integrate in. ('fwd' or 'bwd')
            target_grid: The grid to start the integration from. If None, the grid will be set to the identity transformation.
        """
        super().__init__()
        nsteps = self.intergration_steps
        method = self.integration_method

        if target_grid is None:
            if method in ["ss", "euler", "rk4"]:
                target_grid = (self.sp_tr.grid.detach().clone().repeat(deform_flow.shape[0], 1, 1, 1, 1))
            else:
                target_grid = (self.sp_tr.grid.detach().clone().repeat(deform_flow.shape[0] - 1, 1, 1, 1, 1))

        if self.smoothing_sigma > 0:
            deform_flow = self._gaussian_smooth(deform_flow, self.smoothing_sigma)

        if method == "ss":
            # Integrate a constant vector field with scaling and squaring
            scale = 1.0 / (2**nsteps)
            if dir in ["backward", "bwd"]:
                scale = -scale
            vec = deform_flow * scale
            for _ in range(nsteps):
                vec = vec + self.sp_tr(vec, vec)
            return vec

        elif method == "euler":
            # Integrate a constant vector field with euler integration
            vec = torch.zeros_like(deform_flow)
            dt = 1.0 / nsteps
            if dir in ["backward", "bwd"]:
                dt = -dt
            for i in range(nsteps):
                vec = vec + self.sp_tr(deform_flow * dt, vec + target_grid, displacement=False)
            return vec
        elif method == "euler_time":
            # Integrate a constant vector field with euler integration
            vec = torch.zeros(
                (deform_flow.shape[0] - 1, 3, *self.deform_size),
                device=self.device,
            )
            dt = 1.0 / nsteps
            if dir in ["backward", "bwd"]:
                dt = -dt
            for i in range(nsteps):
                if dir in ["backward", "bwd"]:
                    flow = (1 - i / nsteps) * deform_flow[1:, ...] + (i / nsteps) * deform_flow[:-1, ...]
                else:
                    flow = (1 - i / nsteps) * deform_flow[:-1, ...] + (i / nsteps) * deform_flow[1:, ...]

                vec = vec + self.sp_tr(flow * dt, vec + target_grid, displacement=False)
            return vec

        elif method == "rk4":
            vec = torch.zeros_like(deform_flow)
            dt = 1.0 / nsteps
            if dir in ["backward", "bwd"]:
                dt = -dt
            for i in range(nsteps):
                k1 = self.sp_tr(deform_flow * dt, vec + target_grid, displacement=False)
                k2 = self.sp_tr(
                    deform_flow * dt,
                    vec + target_grid + k1 * 0.5,
                    displacement=False,
                )
                k3 = self.sp_tr(
                    deform_flow * dt,
                    vec + target_grid + k2 * 0.5,
                    displacement=False,
                )
                k4 = self.sp_tr(deform_flow * dt, vec + target_grid + k3, displacement=False)
                vec = vec + (k1 + 2 * k2 + 2 * k3 + k4) / 6.0
            return vec
        elif method == "rk4_time":
            vec = torch.zeros(
                (deform_flow.shape[0] - 1, 3, *self.deform_size),
                device=self.device,
            )
            dt = 1.0 / nsteps
            if dir in ["backward", "bwd"]:
                dt = -dt
            for i in range(nsteps):
                if dir in ["backward", "bwd"]:
                    flow = (1 - i / nsteps) * deform_flow[1:, ...] + (i / nsteps) * deform_flow[:-1, ...]
                else:
                    flow = (1 - i / nsteps) * deform_flow[:-1, ...] + (i / nsteps) * deform_flow[1:, ...]
                k1 = self.sp_tr(flow * dt, vec + target_grid, displacement=False)
                k2 = self.sp_tr(flow * dt, vec + target_grid + k1 * 0.5, displacement=False)
                k3 = self.sp_tr(flow * dt, vec + target_grid + k2 * 0.5, displacement=False)
                k4 = self.sp_tr(flow * dt, vec + target_grid + k3, displacement=False)
                vec = vec + (k1 + 2 * k2 + 2 * k3 + k4) / 6.0
            return vec
        elif method == "trapz":
            # A secret method that also works
            # Find the vector field that transforms the midpoint grid to the target grid
            k = 0.5
            midpoint_grid = target_grid
            landing_grid = target_grid
            if dir in ["backward", "bwd"]:
                dir = -1
            elif dir in ["forward", "fwd"]:
                dir = 1

            for i in range(nsteps):
                midpoint_grid = midpoint_grid - \
                    k * (landing_grid - target_grid)
                # The boundary condition is that the deformation field continue outside the grid,
                landing_grid = (self.sp_tr(-deform_flow * 0.5 * dir, midpoint_grid, displacement=False) + midpoint_grid)

            return dir * self.sp_tr(
                deform_flow,
                midpoint_grid - k * (landing_grid - target_grid),
                displacement=False,
            )

    def _commute_deform_fields(
        self,
        cum_deform_field_fwd,
        cum_deform_field_bwd,
        deform_flow,
        delta_timestep,
        deform_field_fwd,
        deform_field_bwd,
    ):
        """Adds the integral of the deformation flow to the cumulative deformation fields.
        
            If ``field_composition_method`` is set to ``interpolate``, the composition is done by interpolation of 
            the consecutive deformation fields with the cumulative deformation fields.
            If ``field_composition_method`` is set to ``flow_continuation``, the composition is done by integrating 
            the deformation flow from the enpoint of the cumulative deformation fields.
            
            ``delta_timestep`` gives the number of times the cumulative deformation fields has been integrated.
            
            Args:
                cum_deform_field_fwd (torch.Tensor): Cumulative forward deformation field.
                cum_deform_field_bwd (torch.Tensor): Cumulative backward deformation field.
                deform_flow (torch.Tensor): Deformation flow.
                delta_timestep (int): The number of timesteps the cumulative deformation fields has been integrated.
                deform_field_fwd (torch.Tensor): Consecutive forward deformation field.
                deform_field_bwd (torch.Tensor): Consecutive backward deformation field.

            Returns:
                tuple: A tuple containing:

                    - **cum_deform_field_fwd** (torch.Tensor): Cumulative forward deformation field.
                    - **cum_deform_field_bwd** (torch.Tensor): Cumulative backward deformation field.
            """
        Nd = deform_field_fwd.shape[0]

        method = self.field_composition_method

        fwd_grid = cum_deform_field_fwd[:-1, ...] + self.sp_tr.grid
        bwd_grid = cum_deform_field_bwd[1:, ...] + self.sp_tr.grid

        if method == "interpolate":
            cum_deform_field_fwd = cum_deform_field_fwd[:-1, ...] + self.sp_tr(
                deform_field_fwd[delta_timestep:, ...],
                fwd_grid,
                displacement=False,
            )
            cum_deform_field_bwd = cum_deform_field_bwd[1:, ...] + self.sp_tr(
                deform_field_bwd[:Nd - delta_timestep, ...],
                bwd_grid,
                displacement=False,
            )
        elif method == "flow_continuation":
            Nd = deform_flow.shape[0]
            cum_deform_field_fwd = cum_deform_field_fwd[:-1, ...] + self._intergate_flow(
                deform_flow[delta_timestep:, ...], dir="fwd", target_grid=fwd_grid)
            cum_deform_field_bwd = cum_deform_field_bwd[1:, ...] + self._intergate_flow(
                deform_flow[:Nd - delta_timestep, ...],
                dir="bwd",
                target_grid=bwd_grid,
            )
        else:
            raise ValueError("field_composition_method must be either 'interpolation' or 'flow_continuity'")
        return cum_deform_field_fwd, cum_deform_field_bwd

    def _gaussian_smooth(self, field, sigma):
        """Smooths the deformation field with a Gaussian filter in the Fourier domain.

        Args:
            field (torch.Tensor): The deformation field to be smoothed.
            sigma (float): The standard deviation of the Gaussian filter.

        Returns:
            torch.Tensor: The smoothed deformation field.
        """
        v_fft = torch.fft.fftn(field, dim=[2, 3, 4])

        freqx = torch.fft.fftfreq(field.shape[2], d=self.deform_pix_dim[0]).to(self.device)
        freqy = torch.fft.fftfreq(field.shape[3], d=self.deform_pix_dim[1]).to(self.device)
        freqz = torch.fft.fftfreq(field.shape[4], d=self.deform_pix_dim[2]).to(self.device)

        freqx, freqy, freqz = torch.meshgrid(freqx, freqy, freqz, indexing="ij")
        omega_f = 1 / (2 * torch.pi * sigma)
        filter_response = torch.exp(-1 / 2 * ((freqx**2 + freqy**2 + freqz**2) / omega_f**2))
        v_fft = filter_response * v_fft

        field_filtered = torch.real(torch.fft.ifftn(v_fft, dim=[2, 3, 4]))
        return field_filtered

    def get_deform_fields(self, deform_flow):
        """Computes the all deformation fields from the deformation flow.
        
        The deformations are organized in a matrix such that the deformation from the deformation that distortes 
        image_i to image_j is given by deform_matrix[j, i]. Another way of viewing the matrix is that the 
        the displacment at time t_j of a particle starting at the identity grid in image_i is given with
        deform_matrix[i, j]. 
        
        Args:
            deform_flow (torch.Tensor or np.ndarray): The deformation flow. Shape ``(N, 3, x, y, z)``
            
        Returns:
            torch.Tensor or np.ndarray: The deformation fields. Shape ``(N, N, 3, x, y, z)``
        """
        # Get the type of the output
        if isinstance(deform_flow, np.ndarray):
            dtype = "np"
        elif isinstance(deform_flow, torch.Tensor):
            dtype = "torch"
        else:
            raise TypeError("deform_flow must be either numpy array or torch tensor")

        deform_flow = self._convert_to_torch(deform_flow)

        with torch.no_grad():
            deform_matrix = torch.zeros((self.n_img, self.n_img, 3, *self.deform_size), device=self.device)

            deform_field_fwd = self._intergate_flow(deform_flow, "fwd")
            deform_field_bwd = self._intergate_flow(deform_flow, "bwd")

            deform_matrix[utils.get_dig_ind(self.n_img, 1)] = deform_field_fwd
            deform_matrix[utils.get_dig_ind(self.n_img, -1)] = deform_field_bwd

            cum_deform_field_fwd = deform_field_fwd
            cum_deform_field_bwd = deform_field_bwd

            for delta_timestep in range(1, self.n_img - 1):
                (
                    cum_deform_field_fwd,
                    cum_deform_field_bwd,
                ) = self._commute_deform_fields(cum_deform_field_fwd, cum_deform_field_bwd, deform_flow, delta_timestep,
                                                deform_field_fwd, deform_field_bwd)
                deform_matrix[utils.get_dig_ind(self.n_img, delta_timestep + 1)] = cum_deform_field_fwd
                deform_matrix[utils.get_dig_ind(self.n_img, -delta_timestep - 1)] = cum_deform_field_bwd

        if dtype == "np":
            return deform_matrix.detach().cpu().numpy()
        elif dtype == "torch":
            return deform_matrix

    def deform(self, images, deform_field, mode="bilinear", padding_mode="zeros", displacement=True):
        """Transforms the batch of images according to the deformation field.
        
        Args:
            images (torch.Tensor or np.ndarray): The images to deform. Shape ``(N, C, x, y, z)``
            deform_field (torch.Tensor or np.ndarray): The deformation field. Shape ``(N, 3, x, y, z)``
            mode (str, optional): The interpolation mode. Either 'bilinear' or 'nearest'.
            padding_mode (str, optional): The padding mode. Either 'zeros' or 'border'.
            displacement (bool, optional): Whether the deformation field is a displacement field or a deformation field.
            
        Returns:
            torch.Tensor or np.ndarray: The deformed images. Shape ``(N, C, x, y, z)``
        """
        if isinstance(images, np.ndarray):
            dtype = "np"
        elif isinstance(images, torch.Tensor):
            dtype = "torch"
        else:
            raise TypeError("images must be either numpy array or torch tensor")

        images = self._convert_to_torch(images)
        deform_field = self._convert_to_torch(deform_field)

        if ((mode == self.mode) and np.all(np.array(self.image_size) == np.array(images.shape[2:]))):
            with torch.no_grad():
                def_images = self.sp_tr_img(images, deform_field)
        else:
            sp_tr = SpatialTransformer(
                images.shape[2:],
                self.pix_dim,
                mode=mode,
                padding_mode=padding_mode,
            ).to(self.device)
            def_images = sp_tr(images, deform_field, displacement=displacement)

        if dtype == "np":
            return def_images.detach().cpu().numpy()
        elif dtype == "torch":
            return def_images

    def rigid_transform(self, images, rotations, translations):
        grid = self._get_affine_grid(rotations, translations)
        return self.sp_tr_img(images, grid, displacement=False)

    def linear_transform(self, images, matrix):
        grid = self._get_affine_grid(matrix=matrix)
        return self.sp_tr_img(images, grid, displacement=False)

    def add_linear_transform(self, displacements, matrix):
        if isinstance(displacements, np.ndarray):
            dtype = "np"
        elif isinstance(displacements, torch.Tensor):
            dtype = "torch"
        else:
            raise TypeError("displacements must be either numpy array or torch tensor")

        displacements = self._convert_to_torch(displacements)
        matrix = self._convert_to_torch(matrix)

        deform_field = displacements + self._get_affine_grid(matrix=matrix)
        # deform_field = torch.cat([deform_field, torch.ones_like(deform_field[:, :1, ...])], dim=1)
        # deform_field = torch.einsum("nji, nihwd -> njhwd", matrix, deform_field)

        deform_field = deform_field[:, :3, ...] - self.sp_tr_img.grid
        if dtype == "np":
            return deform_field.detach().cpu().numpy()
        elif dtype == "torch":
            return deform_field

    def _convert_to_torch(self, data):
        """
        Converts the data to torch tensor.
        """
        if isinstance(data, np.ndarray):
            return torch.from_numpy(data).to(dtype=torch.float32, device=self.device)
        elif isinstance(data, torch.Tensor):
            return data.to(dtype=torch.float32, device=self.device)
        elif isinstance(data, list):
            return torch.tensor(data, dtype=torch.float32, device=self.device)
        else:
            raise TypeError("Data must be either list, numpy array or torch tensor")


class Registration:
    """Multi Stage Temporal Registration.
    Registrates a series of images. The images are first registrated at a low resolution and then the
    resolution is increased and the images are registrated again. At each resolution the deformation field is
    has a resolution relative to the resampled image resolution set by ``deform_res_scale``.
    
    Accepts same arguments as :class:`StageRegistration` but :attr:`iterations` is replaced by 
    :attr:`stages_iterations`, attr:`deform_res_scale` is replaced by :attr:`stages_deform_scales`, and
    :attr:`image_size` is replaced by :attr:`stages_img_scales`.
    
    Example:
        >>> from muster import Registration
        >>> deform_reg = Registration(
        >>>     stages_iterations=[500, 250, 100],
        >>>     stages_img_scales=[4, 2, 1],
        >>>     stages_deform_scales=[4, 2, 2],
        >>>     image_size=[128, 128, 128],
        >>>     pix_dim=[1, 1, 1],
        >>>     device="cuda:0",
        >>> )
        >>> out = deform_reg.fit(images)
    This example will registrate the image first at a resolution of ``[32, 32, 32]`` with a deformation grid of 
    ``[8, 8, 8]``. Then the resolution is increased to ``[64, 64, 64]`` with a deformation grid of ``[32, 32, 32]``.
    Finally the resolution is increased to ``[128, 128, 128]`` with a deformation grid of ``[64, 64, 64]``.
    
    See :class:`StageRegistration` for more information about the arguments.
    
    Args:
        stages_iterations (list of int): Number of iterations for each stage.
        stages_img_scales (list of int): Image rescaling factors for each stage.
        stages_deform_scales (list of int): Deformation rescaling factors for each stage relative to the image rescaling factor.
        """

    def __init__(self, stages_iterations: list, stages_img_scales: list, stages_deform_scales: list, **kwargs):

        # Make sure the arguments are the same length
        if len(stages_iterations) != len(stages_img_scales) or len(stages_img_scales) != len(stages_deform_scales):
            raise ValueError("stages_iterations, stages_img_scales and stages_deform_scales must have the same length")

        self.stages_iterations = stages_iterations
        self.device = kwargs.get("device", "cpu")
        self.image_sizes = []
        self.deform_scales = stages_deform_scales
        self.image_size = kwargs["image_size"]
        self.pix_dim = kwargs["pix_dim"]
        self.deform_sizes = []
        self.pix_dims = []
        self.deform_pix_dims = []
        self.verbose = kwargs.get("verbose", False)
        self.rigid = kwargs.get("rigid", False)

        for stage in range(len(stages_iterations)):
            img_size = np.array(kwargs["image_size"]) // stages_img_scales[stage]
            self.image_sizes.append(img_size.tolist())
            self.pix_dims.append(np.array(kwargs["pix_dim"]) * stages_img_scales[stage])

            self.deform_sizes.append((np.array(self.image_sizes[stage]) // self.deform_scales[stage]).tolist())
            self.deform_pix_dims.append(np.array(self.pix_dims[stage]) * self.deform_scales[stage])

        # Initialize the deformable registration model of the last stage
        self.kwargs = kwargs
        stage_args = self.kwargs
        stage_args["deform_res_scale"] = self.deform_scales[-1]
        stage_args["num_iterations"] = self.stages_iterations[-1]
        stage_args["image_size"] = self.image_sizes[-1]
        stage_args["pix_dim"] = self.pix_dims[-1]

        self.dr = StageRegistration(**stage_args)

    def fit(self, images, timepoints=None, masks=None):
        """
        Fits the registration model to the images
        Args:.
            - images (numpy.ndarray or torch.Tensor): The images to registrate. Shape (N, C, x, y, z)
            - timepoints (numpy.ndarray or torch.Tensor, optional): The timepoints of the images. Shape (N,)
            - masks (numpy.ndarray or torch.Tensor, optional): The masks of the images. Shape (N, C, x, y, z)

        Returns:
            dict: A dictionary containing the deformation field, rotations and translations.
        
        """

        images = self._convert_to_torch(images)

        # Register the images to the current stage resolution
        deform_flow = None
        rotation = None
        translation = None

        for stage in range(len(self.stages_iterations)):

            # Reinitalize the deformable registration model for the current stage
            stage_args = self.kwargs
            stage_args["deform_res_scale"] = self.deform_scales[stage]
            stage_args["num_iterations"] = self.stages_iterations[stage]
            stage_args["image_size"] = self.image_sizes[stage]
            stage_args["pix_dim"] = self.pix_dims[stage]

            self.dr = None
            torch.cuda.empty_cache()
            self.dr = StageRegistration(**stage_args)

            if self.verbose:
                print(
                    f"Stage {stage+1}/{len(self.stages_iterations)} at image resolution {self.image_sizes[stage]}, deformation resolution {self.deform_sizes[stage]}"
                )
            # Resample the images to the current stage resolution
            img_resampled = nn.functional.interpolate(
                images,
                size=self.image_sizes[stage],
                mode="trilinear",
                align_corners=True,
            )
            if masks is not None:
                masks_resampled = nn.functional.interpolate(masks, size=self.image_sizes[stage], mode="nearest")
            else:
                masks_resampled = None

            out = self.dr.fit(
                img_resampled,
                deform_flow,
                timepoints,
                rotation,
                translation,
                masks_resampled,
            )

            deform_flow = torch.tensor(out["deform_flow"], device=self.device)
            if self.rigid:
                rotation = torch.tensor(out["rotations"], device=self.device)
                translation = torch.tensor(out["translations"], device=self.device)

            out["deform_flow_not_resampled"] = (deform_flow.detach().cpu().numpy())

            if stage < len(self.stages_iterations) - 1:
                deform_flow = nn.functional.interpolate(
                    deform_flow,
                    size=self.deform_sizes[stage + 1],
                    mode="trilinear",
                    align_corners=True,
                )
        if not self.rigid:
            out["rotations"] = torch.zeros((images.shape[0], 3), device=self.device)
            out["translations"] = torch.zeros((images.shape[0], 3), device=self.device)

        # Add the deform fields
        out["deform_field"] = self.get_deform_fields(deform_flow).detach().cpu().numpy()

        return out

    def rigid_transform(self, images, rotations, translations):
        """
        Applies the rigid transformation to the images.
        """
        rotations = self._convert_to_torch(rotations)
        translations = self._convert_to_torch(translations)
        return self.dr.rigid_transform(images, rotations, translations)

    def deform(self, images, deform_field, mode="bilinear", padding_mode="zeros", displacement=True):
        """Transforms the batch of images according to the deformation field.
        
        Args:
            images (torch.Tensor or np.ndarray): The images to deform. Shape ``(N, C, x, y, z)``
            deform_field (torch.Tensor or np.ndarray): The deformation field. Shape ``(N, 3, x, y, z)``
            mode (str, optional): The interpolation mode. Either 'bilinear' or 'nearest'.
            padding_mode (str, optional): The padding mode. Either 'zeros' or 'border'.
            displacement (bool, optional): Whether the deformation field is a displacement field or a deformation field.
            
        Returns:
            torch.Tensor or np.ndarray: The deformed images. Shape ``(N, C, x, y, z)``
        """
        return self.dr.deform(images, deform_field, mode, padding_mode, displacement)

    def get_deform_fields(self, deform_flow):
        """Computes the all deformation fields from the deformation flow.
        
        The deformations are organized in a matrix such that the deformation from the deformation that distortes 
        image_i to image_j is given by deform_matrix[j, i]. Another way of viewing the matrix is that the 
        the displacment at time t_j of a particle starting at the identity grid in image_i is given with
        deform_matrix[i, j]. 
        
        Args:
            deform_flow (torch.Tensor or np.ndarray): The deformation flow. Shape ``(N, 3, x, y, z)``
            
        Returns:
            torch.Tensor or np.ndarray: The deformation fields. Shape ``(N, N, 3, x, y, z)``
        """

        if isinstance(deform_flow, np.ndarray):
            dtype = "np"
        elif isinstance(deform_flow, torch.Tensor):
            dtype = "torch"
        else:
            raise TypeError("deform_field must be either numpy array or torch tensor")

        deform_flow = self._convert_to_torch(deform_flow)
        if self.deform_scales[-1] != 1:
            deform_matrix = self.dr.get_deform_fields(deform_flow)
            n_img = deform_matrix.shape[0]
            deform_matrix = nn.functional.interpolate(
                deform_matrix.reshape((-1, 3, *self.deform_sizes[-1])),
                size=self.image_size,
                mode="trilinear",
                align_corners=True,
            )
            deform_matrix = deform_matrix.reshape((
                n_img,
                n_img,
                3,
                *self.image_size,
            ))
        else:
            deform_matrix = self.dr.get_deform_fields(deform_flow)

        if dtype == "np":
            return deform_matrix.detach().cpu().numpy()
        elif dtype == "torch":
            return deform_matrix

    def _convert_to_torch(self, data):
        """Converts the data to torch tensor.
        
        Args:
            data (numpy.ndarray or torch.Tensor): The data to convert.
            
        Returns:
            torch.Tensor: The data converted to torch tensor.
        """
        if isinstance(data, np.ndarray):
            return torch.from_numpy(data).to(dtype=torch.float32, device=self.device)
        elif isinstance(data, torch.Tensor):
            return data.to(dtype=torch.float32, device=self.device)
        else:
            raise TypeError("Data must be either numpy array or torch tensor")

    def get_identity_grid(self):
        return self.dr.sp_tr_img.grid.detach().cpu().numpy()


class SpatialTransformer(nn.Module):
    r""" Spatial Transformer for performing grid pull operations on spaces.
    Based on https://github.com/voxelmorph/voxelmorph. The spatial transformer uses a deformation field to deform
    the input space. The deformation field is given in mm/step. 
    
    Args:
        size (tuple): the target size of the output tensor. shape ``(x, y, z)``
        pix_dim (tuple): the pixel spacing (mm/px) of the output tensor ``(dx, dy, dz)``
        mode ('bilinear' or 'nearest'): interpolation mode
        padding_mode ('zeros' or 'border'): padding mode
    """

    def __init__(self, size, pix_dim=(1, 1, 1), mode="bilinear", padding_mode="zeros"):
        super().__init__()

        self.mode = mode
        self.padding_mode = padding_mode
        self.pix_dim = pix_dim
        self.size = size
        self.size_mm = [(s - 1) * p for s, p in zip(size, pix_dim)]

        # Create sampling grid
        vectors = [torch.linspace(0, self.size_mm[i], steps=self.size[i]) for i in range(len(self.size))]
        grids = torch.meshgrid(vectors, indexing='ij')
        grid = torch.stack(grids)
        grid = torch.unsqueeze(grid, 0)
        grid = grid.type(torch.float)
        self.register_buffer("grid", grid)

    def forward(self, field, transformation_field, displacement=True):
        r"""Deforms the field image with the transformation field with units in mm/step.

            Args:
                field: the source space to be transformed
                transformation_field: the transformation field to be applied to the source image, with units in mm/step
                displacement: whether the transformation field is a displacement field or a deformation field
                
            Returns:
                torch.Tensor: the transformed source image
            """

        if displacement:
            new_locs = self.grid + transformation_field
        else:
            new_locs = transformation_field.clone()

        # need to normalize grid values to [-1, 1] for resampler
        for axis_idx in range(len(self.size_mm)):
            new_locs[:, axis_idx, ...] = 2 * (new_locs[:, axis_idx, ...] / (self.size_mm[axis_idx]) - 0.5)

        if len(self.size_mm) == 2:
            new_locs = new_locs.permute(0, 2, 3, 1)
            new_locs = new_locs[..., [1, 0]]
        elif len(self.size_mm) == 3:
            new_locs = new_locs.permute(0, 2, 3, 4, 1)
            new_locs = new_locs[..., [2, 1, 0]]

        return F.grid_sample(
            field,
            new_locs,
            align_corners=True,
            mode=self.mode,
            padding_mode=self.padding_mode,
        )

src/pymuster/muster.py
CRAI-OUS/MUSTER
import pymuster
import numpy as np


def test_registration():
    # Create a smal 3d image time series
    img = np.random.rand(4, 1, 10, 10, 10)

    deform_reg = pymuster.Registration(
        stages_iterations=[5],
        stages_img_scales=[2],
        stages_deform_scales=[2],
        image_size=[10, 10, 10],
        pix_dim=[1, 1, 1],
        device="cpu",
        verbose=False,
    )

    out = deform_reg.fit(img)
    deform_matrix = out["deform_field"]
    assert deform_matrix.shape == (4, 4, 3, 10, 10, 10)
    assert deform_matrix.dtype == np.float32


if __name__ == "__main__":
    test_registration()

tests/test_registration.py
CRAI-OUS/MUSTER