ENH: Major reorganization and added 3D interface.

Author: marchdf

README.rst

@@ -1,6 +1,6 @@
 Python bindings to a subset of the `NUFFT algorithm
-<http://www.cims.nyu.edu/cmcl/nufft/nufft.html>`_. Only the 1D and 2D cases are
-implemented.
+<http://www.cims.nyu.edu/cmcl/nufft/nufft.html>`_. 1D, 2D, and 3D
+cases are implemented.
 
 Usage
 -----
@@ -21,7 +21,8 @@ You can specify your required precision using ``eps=1e-15``. The default is
 Authors and License
 -------------------
 
-Python bindings by Dan Foreman-Mackey and Thomas Arildsen but the code that actually does the work
-is from the Greengard lab at NYU (see `the website
-<http://www.cims.nyu.edu/cmcl/nufft/nufft.html>`_). The Fortran code is BSD
-licensed and the Python bindings are MIT licensed.
+Python bindings by Dan Foreman-Mackey, Thomas Arildsen, and
+Marc T. Henry de Frahan but the code that actually does the work is
+from the Greengard lab at NYU (see `the website
+<http://www.cims.nyu.edu/cmcl/nufft/nufft.html>`_). The Fortran code
+is BSD licensed and the Python bindings are MIT licensed.

nufft/__init__.py

@@ -8,6 +8,14 @@
     __NUFFT_SETUP__ = False
 
 if not __NUFFT_SETUP__:
-    __all__ = ["nufft1d1freqs", "nufft1d1", "nufft1d2", "nufft1d3",
-               "nufft2d1", "nufft2d2", "nufft2d3"]
-    from .nufft import nufft1d1freqs, nufft1d1, nufft1d2, nufft1d3, nufft2d1, nufft2d2, nufft2d3
+    __all__ = ["nufft1d1freqs",
+               "nufft1d1",
+               "nufft1d2",
+               "nufft1d3",
+               "nufft2d1",
+               "nufft2d2",
+               "nufft2d3",
+               "nufft3d1",
+               "nufft3d2",
+               "nufft3d3"]
+    from .nufft import nufft1d1freqs, nufft1d1, nufft1d2, nufft1d3, nufft2d1, nufft2d2, nufft2d3, nufft3d1, nufft3d2, nufft3d3

nufft/nufft.py

@@ -2,8 +2,16 @@
 
 from __future__ import division, print_function
 
-__all__ = ["nufft1d1freqs", "nufft1d1", "nufft1d2", "nufft1d3",
-           "nufft2d1", "nufft2d2", "nufft2d3"]
+__all__ = ["nufft1d1freqs",
+           "nufft1d1",
+           "nufft1d2",
+           "nufft1d3",
+           "nufft2d1",
+           "nufft2d2",
+           "nufft2d3",
+           "nufft3d1",
+           "nufft3d2",
+           "nufft3d3"]
 
 import numpy as np
 from ._nufft import (
@@ -13,6 +21,9 @@
     dirft2d1, nufft2d1f90,
     dirft2d2, nufft2d2f90,
     dirft2d3, nufft2d3f90,
+    dirft3d1, nufft3d1f90,
+    dirft3d2, nufft3d2f90,
+    dirft3d3, nufft3d3f90,
 )
 
 
@@ -151,8 +162,10 @@ def nufft2d1(x, y, z, ms, mt, df=1.0, eps=1e-15, iflag=1, direct=False):
     :type y: array
     :param z: non-equispaced function values
     :type z: array
-    :param ms: number of frequencies
+    :param ms: number of frequencies (x-direction)
     :type ms: int
+    :param mt: number of frequencies (y-direction)
+    :type mt: int
     :param df: frequency spacing
     :type df: double
     :param eps: tolerance for NUFFT
@@ -231,8 +244,10 @@ def nufft2d3(x, y, z, f, g, eps=1e-15, iflag=1, direct=False):
     :type y: array
     :param z: non-equispaced function values
     :type z: array
-    :param f: non-equispaced frequencies
+    :param f: non-equispaced frequencies (x-direction)
     :type f: array
+    :param g: non-equispaced frequencies (y-direction)
+    :type g: array
     :param eps: tolerance for NUFFT
     :type eps: double
     :param iflag: sign for the exponential (0 means :math:`-i`, greater than 0 means :math:`+i`)
@@ -264,3 +279,147 @@ def nufft2d3(x, y, z, f, g, eps=1e-15, iflag=1, direct=False):
         if flag:
             raise RuntimeError("nufft2d3 failed with code {0}".format(flag))
     return p / len(x)
+
+
+def nufft3d1(x, y, z, c, ms, mt, mu, df=1.0, eps=1e-15, iflag=1, direct=False):
+    """
+    NUFFT type 1 in three dimensions
+
+    :param x: non-equispaced locations
+    :type x: array
+    :param y: non-equispaced locations
+    :type y: array
+    :param z: non-equispaced locations
+    :type z: array
+    :param c: non-equispaced function values
+    :type c: array
+    :param ms: number of frequencies (x-direction)
+    :type ms: int
+    :param mt: number of frequencies (y-direction)
+    :type mt: int
+    :param mu: number of frequencies (z-direction)
+    :type mu: int
+    :param df: frequency spacing
+    :type df: double
+    :param eps: tolerance for NUFFT
+    :type eps: double
+    :param iflag: sign for the exponential (0 means :math:`-i`, greater than 0 means :math:`+i`)
+    :type iflag: int
+    :param direct: use direct NUFFT methods
+    :type direct: bool
+    :return: function in integer frequency space
+    :rtype: array
+    """
+    # Make sure that the data are properly formatted.
+    x = np.ascontiguousarray(x, dtype=np.float64)
+    y = np.ascontiguousarray(y, dtype=np.float64)
+    z = np.ascontiguousarray(z, dtype=np.float64)
+    c = np.ascontiguousarray(c, dtype=np.complex128)
+    if len(x) != len(y) or len(y) != len(z) or len(z) != len(c):
+        raise ValueError("Dimension mismatch")
+
+    # Run the Fortran code.
+    if direct:
+        p = dirft3d1(x * df, y * df, z * df, c, iflag, ms, mt, mu)
+    else:
+        p, flag = nufft3d1f90(x * df, y * df, z * df,
+                              c, iflag, eps, ms, mt, mu)
+        # Check the output and return.
+        if flag:
+            raise RuntimeError("nufft3d1 failed with code {0}".format(flag))
+    return p
+
+
+def nufft3d2(x, y, z, p, df=1.0, eps=1e-15, iflag=1, direct=False):
+    """
+    NUFFT type 2 in three dimensions
+
+    :param x: non-equispaced locations
+    :type x: array
+    :param y: non-equispaced locations
+    :type y: array
+    :param z: non-equispaced locations
+    :type z: array
+    :param p: function in integer frequency space
+    :type p: array
+    :param df: frequency spacing
+    :type df: double
+    :param eps: tolerance for NUFFT
+    :type eps: double
+    :param iflag: sign for the exponential (0 means :math:`-i`, greater than 0 means :math:`+i`)
+    :type iflag: int
+    :param direct: use direct NUFFT methods
+    :type direct: bool
+    :return: function evaluated at non-equispaced locations
+    :rtype: array
+    """
+    # Make sure that the data are properly formatted.
+    x = np.ascontiguousarray(x, dtype=np.float64)
+    y = np.ascontiguousarray(y, dtype=np.float64)
+    z = np.ascontiguousarray(z, dtype=np.float64)
+    p = np.ascontiguousarray(p, dtype=np.complex128)
+    if len(x) != len(y) or len(y) != len(z):
+        raise ValueError("Dimension mismatch")
+
+    # Run the Fortran code.
+    if direct:
+        c = dirft3d2(x * df, y * df, z * df, iflag, p)
+    else:
+        c, flag = nufft3d2f90(x * df, y * df, z * df, iflag, eps, p)
+        # Check the output and return.
+        if flag:
+            raise RuntimeError("nufft3d2 failed with code {0}".format(flag))
+    return c
+
+
+def nufft3d3(x, y, z, c, f, g, h, eps=1e-15, iflag=1, direct=False):
+    """
+    NUFFT type 3 in three dimensions
+
+    :param x: non-equispaced locations
+    :type x: array
+    :param y: non-equispaced locations
+    :type y: array
+    :param z: non-equispaced locations
+    :type z: array
+    :param c: non-equispaced function values
+    :type c: array
+    :param f: non-equispaced frequencies (x-direction)
+    :type f: array
+    :param g: non-equispaced frequencies (y-direction)
+    :type g: array
+    :param h: non-equispaced frequencies (z-direction)
+    :type h: array
+    :param eps: tolerance for NUFFT
+    :type eps: double
+    :param iflag: sign for the exponential (0 means :math:`-i`, greater than 0 means :math:`+i`)
+    :type iflag: int
+    :param direct: use direct NUFFT methods
+    :type direct: bool
+    :return: function in non-equispaced frequency space
+    :rtype: array
+    """
+    # Make sure that the data are properly formatted.
+    x = np.ascontiguousarray(x, dtype=np.float64)
+    y = np.ascontiguousarray(y, dtype=np.float64)
+    z = np.ascontiguousarray(z, dtype=np.float64)
+    c = np.ascontiguousarray(c, dtype=np.complex128)
+    if len(x) != len(y) or len(y) != len(z) or len(z) != len(c):
+        raise ValueError("Dimension mismatch")
+
+    # Make sure that the frequencies are of the right type.
+    f = np.ascontiguousarray(f, dtype=np.float64)
+    g = np.ascontiguousarray(g, dtype=np.float64)
+    h = np.ascontiguousarray(h, dtype=np.float64)
+    if len(f) != len(g) or len(g) != len(h):
+        raise ValueError("Dimension mismatch")
+
+    # Run the Fortran code.
+    if direct:
+        p = dirft3d3(x, y, z, c, iflag, f, g, h)
+    else:
+        p, flag = nufft3d3f90(x, y, z, c, iflag, eps, f, g, h)
+        # Check the output and return.
+        if flag:
+            raise RuntimeError("nufft3d3 failed with code {0}".format(flag))
+    return p / len(x)