Add numpy-based implementation of spectral ops (PaddlePaddle#33)

* add numpy reference implementation of spectral ops

paddle/fluid/operators/spectral_op.cc

@@ -578,7 +578,7 @@ void exec_fft(const DeviceContext& ctx, const Tensor* x, Tensor* out,
   TransCompute<platform::CPUDeviceContext, To>(ndim, ctx, transposed_output,
                                                out, reverse_dim_permute);
 }
-} // namespace anonymous
+}  // anonymous namespace
 
 template <typename Ti, typename To>
 struct FFTC2CFunctor<platform::CPUDeviceContext, Ti, To> {
@@ -640,7 +640,7 @@ T compute_factor(int64_t size, FFTNormMode normalization) {
   PADDLE_THROW(
       platform::errors::InvalidArgument("Unsupported normalization type"));
 }
-} //namespace anonymous
+}  // anonymous namespace
 
 template <typename Ti, typename To>
 struct FFTC2CFunctor<platform::CPUDeviceContext, Ti, To> {

paddle/fluid/operators/spectral_op.cu

@@ -520,9 +520,10 @@ static inline PlanLRUCache& cufft_get_plan_cache(int64_t device_index) {
   return *plan_caches[device_index];
 }
 
-// Execute a general unnormalized fft operation (can be c2c, onesided r2c or onesided c2r)
+// Execute a general unnormalized fft operation (can be c2c, onesided r2c or
+// onesided c2r)
 template <typename DeviceContext, typename Ti, typename To>
-void exec_fft(const DeviceContext& ctx, const Tensor* X, Tensor* out, 
+void exec_fft(const DeviceContext& ctx, const Tensor* X, Tensor* out,
               const std::vector<int64_t>& dim, bool forward) {
   const auto x_dims = framework::vectorize(X->dims());
   const auto out_dims = framework::vectorize(out->dims());
@@ -753,8 +754,8 @@ struct FFTC2CFunctor<platform::CUDADeviceContext, Ti, To> {
           std::min(static_cast<size_t>(kMaxCUFFTNdim), working_axes.size());
       first_dims.assign(working_axes.end() - max_dims, working_axes.end());
 
-      exec_fft<platform::CUDADeviceContext, Ti, To>(
-          ctx, p_working_tensor, p_out, first_dims, forward);
+      exec_fft<platform::CUDADeviceContext, Ti, To>(ctx, p_working_tensor,
+                                                    p_out, first_dims, forward);
       working_axes.resize(working_axes.size() - max_dims);
       first_dims.clear();
 
@@ -781,8 +782,8 @@ struct FFTC2RFunctor<platform::CUDADeviceContext, Ti, To> {
       framework::Tensor x_copy(X->type());
       x_copy.mutable_data<Ti>(X->dims(), ctx.GetPlace());
       framework::TensorCopy(*X, ctx.GetPlace(), &x_copy);
-      exec_fft<platform::CUDADeviceContext, Ti, To>(ctx, &x_copy, out,
-                                                    axes, forward);
+      exec_fft<platform::CUDADeviceContext, Ti, To>(ctx, &x_copy, out, axes,
+                                                    forward);
     } else {
       framework::Tensor temp_tensor;
       temp_tensor.mutable_data<Ti>(X->dims(), ctx.GetPlace());
@@ -791,8 +792,8 @@ struct FFTC2RFunctor<platform::CUDADeviceContext, Ti, To> {
       FFTC2CFunctor<platform::CUDADeviceContext, Ti, Ti> c2c_functor;
       c2c_functor(ctx, X, &temp_tensor, dims, FFTNormMode::none, forward);
 
-      exec_fft<platform::CUDADeviceContext, Ti, To>(
-          ctx, &temp_tensor, out, {axes.back()}, forward);
+      exec_fft<platform::CUDADeviceContext, Ti, To>(ctx, &temp_tensor, out,
+                                                    {axes.back()}, forward);
     }
     exec_normalization<platform::CUDADeviceContext, To>(
         ctx, out, out, normalization, out_dims, axes);
@@ -809,8 +810,8 @@ struct FFTR2CFunctor<platform::CUDADeviceContext, Ti, To> {
     framework::Tensor* r2c_out = out;
     const std::vector<int64_t> last_dim{axes.back()};
     std::vector<int64_t> out_dims = framework::vectorize(out->dims());
-    exec_fft<platform::CUDADeviceContext, Ti, To>(ctx, X, r2c_out,
-                                                  last_dim, forward);
+    exec_fft<platform::CUDADeviceContext, Ti, To>(ctx, X, r2c_out, last_dim,
+                                                  forward);
 
     // Step2: C2C transform on the remaining dimension
     framework::Tensor c2c_out;

python/paddle/fluid/tests/unittests/fft/spectral_op_np.py

@@ -0,0 +1,101 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+from functools import partial
+from numpy import asarray
+from numpy.fft._pocketfft import _raw_fft, _raw_fftnd, _get_forward_norm, _get_backward_norm, _cook_nd_args
+
+
+def _fftc2c(a, n=None, axis=-1, norm=None, forward=None):
+    a = asarray(a)
+    if n is None:
+        n = a.shape[axis]
+    if forward:
+        inv_norm = _get_forward_norm(n, norm)
+    else:
+        inv_norm = _get_backward_norm(n, norm)
+    output = _raw_fft(a, n, axis, False, forward, inv_norm)
+    return output
+
+
+def _fftr2c(a, n=None, axis=-1, norm=None, forward=None):
+    a = asarray(a)
+    if n is None:
+        n = a.shape[axis]
+    if forward:
+        inv_norm = _get_forward_norm(n, norm)
+    else:
+        inv_norm = _get_backward_norm(n, norm)
+    output = _raw_fft(a, n, axis, True, True, inv_norm)
+    if not forward:
+        output = output.conj()
+    return output
+
+
+def _fftc2r(a, n=None, axis=-1, norm=None, forward=None):
+    a = asarray(a)
+    if n is None:
+        n = (a.shape[axis] - 1) * 2
+    if forward:
+        inv_norm = _get_forward_norm(n, norm)
+    else:
+        inv_norm = _get_backward_norm(n, norm)
+    output = _raw_fft(a.conj()
+                      if forward else a, n, axis, True, False, inv_norm)
+    return output
+
+
+def fft_c2c(x, axes, normalization, forward):
+    f = partial(_fftc2c, forward=forward)
+    y = _raw_fftnd(x, s=None, axes=axes, function=f, norm=normalization)
+    return y
+
+
+def fft_c2c_backward(dy, axes, normalization, forward):
+    f = partial(_fftc2c, forward=forward)
+    dx = _raw_fftnd(dy, s=None, axes=axes, function=f, norm=normalization)
+    return dx
+
+
+def fft_r2c(x, axes, normalization, forward, onesided):
+    a = asarray(x)
+    s, axes = _cook_nd_args(a, axes=axes)
+    if onesided:
+        a = _fftr2c(a, s[-1], axes[-1], normalization, forward)
+        for ii in range(len(axes) - 1):
+            a = _fftc2c(a, s[ii], axes[ii], normalization, forward)
+    else:
+        a = fft_c2c(x, axes, normalization, forward)
+    return a
+
+
+def fft_r2c_backward(dy, x, axes, normalization, forward, onesided):
+    a = dy
+    if not onesided:
+        a = fft_c2c_backward(a, axes, normalization, forward).real
+    else:
+        pad_widths = [(0, 0)] * a.ndim
+        last_axis = axes[-1]
+        if last_axis < 0:
+            last_axis += a.ndim
+        last_dim_size = a.shape[last_axis]
+        pad_widths[last_axis] = (0, x.shape[last_axis] - last_dim_size)
+        a = np.pad(a, pad_width=pad_widths)
+        a = fft_c2c_backward(a, axes, normalization, forward).real
+    return a
+
+
+def fft_c2r(x, axes, normalization, forward, last_dim_size):
+    pass