How to use the torchkbnufft.kbmodule.KbModule function in torchkbnufft

import warnings

import numpy as np
import torch

from .functional.mrisensenufft import (AdjMriSenseNufftFunction,
                                       MriSenseNufftFunction)
from .kbmodule import KbModule
from .mri.sensenufft_functions import sense_toeplitz
from .nufft.utils import build_spmatrix, build_table, compute_scaling_coefs


class SenseNufftModule(KbModule):
    """Parent class for SENSE-NUFFT classes.

    This implementation collects all init functions into one place. It inherits
    from torch.nn.Module via torchkbnufft.KbModule.

    Args:
        smap (tensor): Sensitivity coils of size (batch_size, real/imag,) +
            im_size.
        im_size (int or tuple of ints): Size of base image.
        grid_size (int or tuple of ints, default=2*im_size): Size of the grid
            to interpolate to.
        numpoints (int or tuple of ints, default=6): Number of points to use
            for interpolation in each dimension. Default is six points in each
            direction.
        n_shift (int or tuple of ints, default=im_size//2): Number of points to
            shift for fftshifts.

'real_interp_mats' and 'imag_interp_mats', each key containing
                a list of interpolation matrices (see 
                mri.sparse_interp_mat.precomp_sparse_mats for construction).
                If None, then a standard interpolation is run.

        Returns:
            tensor: The image after adjoint NUFFT.
        """
        interpob = self._extract_nufft_interpob()

        x = AdjKbNufftFunction.apply(y, om, interpob, interp_mats)

        return x


class ToepNufft(KbModule):
    """Forward/backward NUFFT with Toeplitz embedding.

    This essentially is an torch.nn.Module wrapper for the
    torchkbnufft.nufft.fft_functions.fft_filter function.
    """

    def __init__(self):
        super(ToepNufft, self).__init__()

    def forward(self, x, kern, norm=None):
        """Toeplitz NUFFT forward function.

        Args:
            x (tensor): The image (or images) to apply the forward/backward
                Toeplitz-embedded NUFFT to.
            kern (tensor): The filter response taking into account Toeplitz

import warnings

import numpy as np
import torch

from .functional.kbnufft import AdjKbNufftFunction, KbNufftFunction
from .kbmodule import KbModule
from .nufft.fft_functions import fft_filter
from .nufft.utils import build_spmatrix, build_table, compute_scaling_coefs


class KbNufftModule(KbModule):
    """Parent class for KbNufft classes.

    This implementation collects all init functions into one place. It inherits
    from torch.nn.Module via torchkbnufft.kbmodule.KbModule.

    Args:
        im_size (int or tuple of ints): Size of base image.
        grid_size (int or tuple of ints, default=2*im_size): Size of the grid
            to interpolate from.
        numpoints (int or tuple of ints, default=6): Number of points to use
            for interpolation in each dimension. Default is six points in each
            direction.
        n_shift (int or tuple of ints, default=im_size//2): Number of points to
            shift for fftshifts.
        table_oversamp (int, default=2^10): Table oversampling factor.
        kbwidth (double, default=2.34): Kaiser-Bessel width parameter.

import warnings

import numpy as np
import torch

from .functional.kbinterp import AdjKbInterpFunction, KbInterpFunction
from .kbmodule import KbModule
from .nufft.utils import build_table


class KbInterpModule(KbModule):
    """Parent class for KbInterp classes.

    This implementation collects all init functions into one place. It inherits
    from torch.nn.Module via torchkbnufft.kbmodule.KbModule.

    Args:
        im_size (int or tuple of ints): Size of base image.
        grid_size (int or tuple of ints, default=2*im_size): Size of the grid
            to interpolate to.
        numpoints (int or tuple of ints, default=6): Number of points to use
            for interpolation in each dimension. Default is six points in each
            direction.
        n_shift (int or tuple of ints, default=im_size//2): Number of points to
            shift for fftshifts.
        table_oversamp (int, default=2^10): Table oversampling factor.
        kbwidth (double, default=2.34): Kaiser-Bessel width parameter.

instead of the one used on layer initialization.

        Returns:
            tensor: The image with an adjoint SENSE-NUFFT.
        """
        interpob = self._extract_sense_interpob()

        if smap is None:
            smap = self.smap_tensor

        x = AdjMriSenseNufftFunction.apply(y, smap, om, interpob, interp_mats)

        return x


class ToepSenseNufft(KbModule):
    """Forward/backward SENSE-NUFFT with Toeplitz embedding.

    This essentially is an torch.nn.Module wrapper for the
    mri.sensenufft_functions.sense_toeplitz function.

    Args:
        smap (tensor): Sensitivity coils of size (batch_size, real/imag,) +
            im_size.
    """

    def __init__(self, smap):
        super(ToepSenseNufft, self).__init__()

        self.smap_shape = smap.shape

        self.register_buffer('smap_tensor', smap)