__init__.py
"""
NUFFT class
=======================================
"""
from __future__ import absolute_import
import numpy
import warnings
import scipy.sparse
import numpy.fft
#import scipy.signal
import scipy.linalg
import scipy.special
from functools import wraps as _wraps
# from ..linalg.nufft_cpu import NUFFT_cpu
# from ..linalg.nufft_hsa import NUFFT_hsa
def push_cuda_context(hsa_method):
"""
Decorator: Push cude context to the top of the stack for current use
Add @push_cuda_context before the methods of NUFFT_device()
"""
@_wraps(hsa_method)
def wrapper(*args, **kwargs):
try:
args[0].thr._context.push()
except:
pass
return hsa_method(*args, **kwargs)
return wrapper
class NUFFT:
"""
NUFFT class
=======================================
A super class of cpu and gpu NUFFT functions.
Note: NUFFT does not inherit NUFFT_cpu (deprecated) and NUFFT_hsa (deprecated).
"""
#import cpu codes
from ._nufft_class_methods_cpu import _init__cpu, _plan_cpu, _precompute_sp_cpu, _solve_cpu, _forward_cpu, _adjoint_cpu, _selfadjoint_cpu, _selfadjoint2_cpu, _x2xx_cpu, _xx2k_cpu, _xx2k_one2one_cpu, _k2vec_cpu, _vec2y_cpu, _k2y_cpu, _y2vec_cpu, _vec2k_cpu, _y2k_cpu, _k2xx_cpu, _k2xx_one2one_cpu, _xx2x_cpu, _k2y2k_cpu
# import host codes
from ._nufft_class_methods_cpu import _forward_host, _adjoint_host, _selfadjoint_host, _solve_host, _xx2k_host, _k2xx_host, _x2xx_host, _xx2x_host, _k2y_host, _y2k_host
# import device codes
from ._nufft_class_methods_device import _init__device, _plan_device, _set_wavefront_device, _offload_device, to_device, to_host, _x2xx_device, _xx2k_device, _k2y_device, _y2k_device, _k2xx_device, _xx2x_device, _selfadjoint_device, _forward_device, _adjoint_device, release, _solve_device, release
# legacy codes (csr format for device)
from ._nufft_class_methods_device import _y2k_legacy, _k2y_legacy, _forward_legacy, _adjoint_legacy, _selfadjoint_legacy, _plan_legacy, _offload_legacy, _solve_legacy
from ._nufft_class_methods_cpu import _k2y_legacy_host, _y2k_legacy_host, _selfadjoint_legacy_host, _forward_legacy_host, _adjoint_legacy_host, _solve_legacy_host
def __init__(self, device_indx=None, legacy=None):
"""
Constructor.
:param None:
:type None: Python NoneType
:return: NUFFT: the pynufft.NUFFT instance
:rtype: NUFFT: the pynufft.NUFFT class
:Example:
>>> from pynufft import NUFFT
>>> NufftObj = NUFFT()
or
>>> from pynufft import NUFFT, helper
>>> device = helper.device_list()[0]
>>> NufftObj = NUFFT(device) # for first acceleration device in the system
"""
if device_indx is None:
self._init__cpu()
self.processor = 'cpu'
else:
if legacy is True:
self._init__device(device_indx)
self.processor = 'hsa_legacy'
else:
self._init__device(device_indx)
self.processor = 'hsa'
def __del__(self):
if self.processor is 'hsa' or 'hsa_legacy':
self.release()
else:
pass
def plan(self, *args, **kwargs):
"""
Plan the NUFFT object with the geometry provided.
:param om: The M off-grid locates in the frequency domain,
which is normalized between [-pi, pi]
:param Nd: The matrix size of the equispaced image.
Example: Nd=(256,256) for a 2D image;
Nd = (128,128,128) for a 3D image
:param Kd: The matrix size of the oversampled frequency grid.
Example: Kd=(512,512) for 2D image;
Kd = (256,256,256) for a 3D image
:param Jd: The interpolator size.
Example: Jd=(6,6) for 2D image;
Jd = (6,6,6) for a 3D image
:param ft_axes: (Optional) The axes for Fourier transform.
The default is all axes if 'None' is given.
:type om: numpy.float array, matrix size = M * ndims
:type Nd: tuple, ndims integer elements.
:type Kd: tuple, ndims integer elements.
:type Jd: tuple, ndims integer elements.
:type ft_axes: None, or tuple with optional integer elements.
:returns: 0
:rtype: int, float
:ivar Nd: initial value: Nd
:ivar Kd: initial value: Kd
:ivar Jd: initial value: Jd
:ivar ft_axes: initial value: None
:Example:
>>> from pynufft import NUFFT
>>> NufftObj = NUFFT()
>>> NufftObj.plan(om, Nd, Kd, Jd)
or
>>> NufftObj.plan(om, Nd, Kd, Jd, ft_axes)
"""
func = {'cpu': self._plan_cpu,
'hsa': self._plan_device,
'hsa_legacy': self._plan_legacy}
return func.get(self.processor)(*args, **kwargs)
def forward(self, *args, **kwargs):
"""
Forward NUFFT (host code)
:param x: The input numpy array, with the size of Nd
:type: numpy array with the dtype of numpy.complex64
:return: y: The output numpy array, with the size of (M,)
:rtype: numpy array with the dtype of numpy.complex64
"""
func = {'cpu': self._forward_cpu,
'hsa': self._forward_host,
'hsa_legacy': self._forward_legacy_host}
return func.get(self.processor)(*args, **kwargs)
def adjoint(self, *args, **kwargs):
"""
Adjoint NUFFT (host code)
:param y: The input numpy array, with the size of (M,)
:type: numpy array with the dtype of numpy.complex64
:return: x: The output numpy array,
with the size of Nd or Nd
:rtype: numpy array with the dtype of numpy.complex64
"""
func = {'cpu': self._adjoint_cpu,
'hsa': self._adjoint_host,
'hsa_legacy': self._adjoint_legacy_host}
return func.get(self.processor)(*args, **kwargs)
def selfadjoint(self, *args, **kwargs):
"""
selfadjoint NUFFT (host code)
:param x: The input numpy array, with size=Nd
:type: numpy array with dtype =numpy.complex64
:return: x: The output numpy array, with size=Nd
:rtype: numpy array with dtype =numpy.complex64
"""
func = {'cpu': self._selfadjoint_cpu,
'hsa': self._selfadjoint_host,
'hsa_legacy': self._selfadjoint_legacy_host}
return func.get(self.processor)(*args, **kwargs)
def solve(self, *args, **kwargs):
"""
Solve NUFFT (host code)
:param y: data, numpy.complex64. The shape = (M,)
:param solver: 'cg', 'L1TVOLS', 'lsmr', 'lsqr', 'dc', 'bicg',
'bicgstab', 'cg', 'gmres','lgmres'
:param maxiter: the number of iterations
:type y: numpy array, dtype = numpy.complex64
:type solver: string
:type maxiter: int
:return: numpy array with size Nd.
"""
func = {'cpu': self._solve_cpu,
'hsa': self._solve_host,
'hsa_legacy': self._solve_legacy_host}
return func.get(self.processor)(*args, **kwargs)
def xx2k(self, *args, **kwargs):
func = {'cpu': self._xx2k_cpu,
'hsa': self._xx2k_host,
'hsa_legacy': self._xx2k_host}
return func.get(self.processor)(*args, **kwargs)
def k2xx(self, *args, **kwargs):
func = {'cpu': self._k2xx_cpu,
'hsa': self._k2xx_host,
'hsa_legacy': self._k2xx_host}
return func.get(self.processor)(*args, **kwargs)
def x2xx(self, *args, **kwargs):
func = {'cpu': self._x2xx_cpu,
'hsa': self._x2xx_host,
'hsa_legacy': self._x2xx_host}
return func.get(self.processor)(*args, **kwargs)
def xx2x(self, *args, **kwargs):
func = {'cpu': self._xx2x_cpu,
'hsa': self._xx2x_host,
'hsa_legacy': self._xx2x_host}
return func.get(self.processor)(*args, **kwargs)
def k2y(self, *args, **kwargs):
func = {'cpu': self._k2y_cpu,
'hsa': self._k2y_host,
'hsa_legacy': self._k2y_legacy_host}
return func.get(self.processor)(*args, **kwargs)
def y2k(self, *args, **kwargs):
func = {'cpu': self._y2k_cpu,
'hsa': self._y2k_host,
'hsa_legacy': self._y2k_legacy_host}
return func.get(self.processor)(*args, **kwargs)
def k2yk2(self, *args, **kwargs):
func = {'cpu': self._k2yk2_cpu,
'hsa': self._k2yk2_host,
'hsa_legacy': self._k2yk2_host}
return func.get(self.processor)(*args, **kwargs)
# def adjoint_many2one(self, *args, **kwargs):
# func = {'cpu': self._adjoint_many2one_cpu,
# 'hsa': self._adjoint_many2one_host,
# 'hsa_legacy': self._adjoint_many2one_legacy_host}
# return func.get(self.processor)(*args, **kwargs)
#
# def forward_one2many(self, *args, **kwargs):
# func = {'cpu': self._forward_one2many_cpu,
# 'hsa': self._forward_one2many_host,
# 'hsa_legacy': self._forward_one2many_legacy_host}
# return func.get(self.processor)(*args, **kwargs)
# def selfadjoint_one2many2one(self, *args, **kwargs):
# func = {'cpu': self._selfadjoint_one2many2one_cpu,
# 'hsa': self._selfadjoint_one2many2one_host,
# 'hsa_legacy': self._selfadjoint_one2many2one_legacy_host}
# return func.get(self.processor)(*args, **kwargs)
def k2xx_one2one(self, *args, **kwargs):
func = {'cpu':self._k2xx_one2one_cpu}
return func.get(self.processor)(*args, **kwargs)
def xx2k_one2one(self, *args, **kwargs):
func = {'cpu':self._xx2k_one2one_cpu}
return func.get(self.processor)(*args, **kwargs)
def k2y2k(self, *args, **kwargs):
func = {'cpu': self._k2y2k_cpu}
return func.get(self.processor)(*args, **kwargs)
# def set_sense(self, *args, **kwargs):
# func = {'cpu': self._set_sense_cpu,
# 'hsa': self._set_sense_host,
# 'hsa_legacy': self._set_sense_host}
# return func.get(self.processor)(*args, **kwargs)
# def reset_sense(self, *args, **kwargs):
# func = {'cpu': self._reset_sense_cpu,
# 'hsa': self._reset_sense_host,
# 'hsa_legacy': self._reset_sense_host}
# return func.get(self.processor)(*args, **kwargs)
