added normalization options to curfft following numpy

theano/gpuarray/fft.py

@@ -128,6 +128,7 @@ class CuRFFTOp(Op):
         thunk.lazy = False
 
         return thunk
+curfft_op = CuRFFTOp()
 
 
 class CuIRFFTOp(Op):
@@ -234,27 +235,40 @@ class CuIRFFTOp(Op):
         thunk.lazy = False
 
         return thunk
+cuirfft_op = CuIRFFTOp()
 
 
-def curfft(inputs):
+def curfft(inputs, norm=None):
     """
-    Performs the real unitary fast Fourier Transform normalized
-    by :math:`\sqrt n`.
+    Performs the real-valued input fast Fourier Transform using the
+    gpuarray backend.
 
     Parameters
     ----------
     inputs
         Array of real-valued float32 of size (m, n), containing m inputs of
         length n.
+    norm : {None, 'ortho', 'no_norm'}
+        Normalization of transform. Following numpy, default *None* normalizes
+        only the inverse transform by n, 'ortho' yields the unitary transform
+        (:math:`1/\sqrt n` forward and back). In addition, 'no_norm' leaves
+        the transform unnormalized.
     """
-    fft_op = CuRFFTOp()
-    return fft_op(inputs) / T.sqrt(((inputs.shape[1:]).prod()).astype('float32'))
 
+    cond_norm = _unitary(norm)
+    if cond_norm is None:
+        return curfft_op(inputs)
+    elif cond_norm == "ortho":
+        return curfft_op(inputs) / T.sqrt(((inputs.shape[1:]).prod())
+                                          .astype('float32'))
+    elif cond_norm == "no_norm":
+        return curfft_op(inputs)
 
-def cuirfft(inputs):
+
+def cuirfft(inputs, norm=None):
     """
-    Performs the real unitary fast inverse Fourier Transform normalized
-    by :math:`\sqrt n`.
+    Performs the real-valued output inverse Fourier Transform using the
+    gpuarray backend.
 
     Parameters
     ----------
@@ -262,7 +276,26 @@ def cuirfft(inputs):
         Array of float32 of size (m, n/2+1, 2), containing m inputs with n/2+1
         non-trivial elements and real and imaginary parts stored as separate
         arrays.
+    norm : {None, 'ortho', 'no_norm'}
+        Normalization of transform. Following numpy, default *None* normalizes
+        only the inverse transform by n, 'ortho' yields the unitary transform
+        (:math:`1/\sqrt n` forward and back). In addition, 'no_norm' leaves
+        the transform unnormalized.
     """
-    ifft_op = CuIRFFTOp()
-    return ifft_op(inputs) / T.sqrt((((inputs.shape[1:-1] - 1) * 2).prod())
-                                    .astype('float32'))
+
+    cond_norm = _unitary(norm)
+    if cond_norm is None:
+        return cuirfft_op(inputs) / (((inputs.shape[1:-1] - 1) * 2).prod()
+                                     .astype('float32'))
+    if cond_norm == "ortho":
+        return cuirfft_op(inputs) / T.sqrt((((inputs.shape[1:-1] - 1) * 2).prod())
+                                           .astype('float32'))
+    if cond_norm == "no_norm":
+        return cuirfft_op(inputs)
+
+
+def _unitary(norm):
+    if norm not in (None, "ortho", "no_norm"):
+        raise ValueError("Invalid value %s for norm, must be None, 'ortho' or "
+                         "'no norm'" % norm)
+    return norm

theano/gpuarray/tests/test_fft.py

@@ -22,11 +22,13 @@ if not scikits_cuda_available:  # noqa
 if not pycuda_available:  # noqa
     raise SkipTest('Optional package pycuda not available')
 
+# Transform sizes
+N = 64
+
 
 class TestFFT(unittest.TestCase):
 
     def test_rfft(self):
-        N = 64
         inputs_val = np.random.random((1, N)).astype('float32')
         inputs = theano.shared(inputs_val)
 
@@ -36,12 +38,11 @@ class TestFFT(unittest.TestCase):
         res_rfft_comp = (np.asarray(res_rfft[:, :, 0]) +
                          1j * np.asarray(res_rfft[:, :, 1]))
 
-        rfft_ref = numpy.fft.rfft(inputs_val, N, 1, norm='ortho')
+        rfft_ref = numpy.fft.rfft(inputs_val, N, 1)
 
         utt.assert_allclose(rfft_ref, res_rfft_comp)
 
     def test_irfft(self):
-        N = 64
         inputs_val = np.random.random((1, N)).astype('float32')
         inputs = theano.shared(inputs_val)
 
@@ -57,7 +58,6 @@ class TestFFT(unittest.TestCase):
         utt.assert_allclose(inputs_val, np.asarray(res_ifft))
 
     def test_type(self):
-        N = 64
         inputs_val = np.random.random((1, N)).astype('float64')
         inputs = theano.shared(inputs_val)
 
@@ -65,3 +65,48 @@ class TestFFT(unittest.TestCase):
             theano.gpuarray.fft.curfft(inputs)
         with self.assertRaises(AssertionError):
             theano.gpuarray.fft.cuirfft(inputs)
+
+    def test_norm(self):
+        inputs_val = np.random.random((1, N)).astype('float32')
+        inputs = theano.shared(inputs_val)
+
+        # Unitary normalization
+        rfft = theano.gpuarray.fft.curfft(inputs, norm='ortho')
+        f_rfft = theano.function([], rfft, mode=mode_with_gpu)
+        res_rfft = f_rfft()
+        res_rfft_comp = (np.asarray(res_rfft[:, :, 0]) +
+                         1j * np.asarray(res_rfft[:, :, 1]))
+
+        rfft_ref_ortho = numpy.fft.rfft(inputs_val, N, 1, norm='ortho')
+
+        utt.assert_allclose(rfft_ref_ortho, res_rfft_comp)
+
+        # No normalization
+        rfft = theano.gpuarray.fft.curfft(inputs, norm='no_norm')
+        f_rfft = theano.function([], rfft, mode=mode_with_gpu)
+        res_rfft = f_rfft()
+        res_rfft_comp = (np.asarray(res_rfft[:, :, 0]) +
+                         1j * np.asarray(res_rfft[:, :, 1]))
+
+        utt.assert_allclose(rfft_ref_ortho * np.sqrt(N), res_rfft_comp)
+
+        # Inverse FFT inputs
+        inputs_val = np.random.random((1, N // 2 + 1, 2)).astype('float32')
+        inputs = theano.shared(inputs_val)
+        inputs_ref = inputs_val[:, :, 0] + 1j * inputs_val[:, :, 1]
+
+        # Unitary normalization inverse FFT
+        irfft = theano.gpuarray.fft.cuirfft(inputs, norm='ortho')
+        f_irfft = theano.function([], irfft, mode=mode_with_gpu)
+        res_irfft = f_irfft()
+
+        irfft_ref_ortho = numpy.fft.irfft(inputs_ref, norm='ortho')
+
+        utt.assert_allclose(irfft_ref_ortho, res_irfft)
+
+        # No normalization inverse FFT
+        irfft = theano.gpuarray.fft.cuirfft(inputs, norm='no_norm')
+        f_irfft = theano.function([], irfft, mode=mode_with_gpu)
+        res_irfft = f_irfft()
+
+        utt.assert_allclose(irfft_ref_ortho * np.sqrt(N), res_irfft)