merge master

pypocketfft.cc

@@ -121,9 +121,46 @@ py::array xfftn(const py::array &a, py::object axes, int inorm,
            inplace, fwd, nthreads))
   }
 
-py::array fftn(const py::array &a, py::object axes, int inorm,
-  bool inplace, size_t nthreads)
-  { return xfftn(a, axes, inorm, inplace, true, nthreads); }
+template<typename T> py::array sym_rfftn_internal(const py::array &in,
+  py::object axes_, int inorm, size_t nthreads)
+  {
+  auto axes = makeaxes(in, axes_);
+  auto dims(copy_shape(in));
+  py::array res = py::array_t<complex<T>>(dims);
+  auto s_in=copy_strides(in);
+  auto s_out=copy_strides(res);
+  auto d_in=reinterpret_cast<const T *>(in.data());
+  auto d_out=reinterpret_cast<complex<T> *>(res.mutable_data());
+  {
+  py::gil_scoped_release release;
+  T fct = norm_fct<T>(inorm, dims, axes);
+  r2c(dims, s_in, s_out, axes, d_in, d_out, fct, nthreads);
+  // now fill in second half
+  using namespace pocketfft::detail;
+  ndarr<complex<T>> ares(res.mutable_data(), dims, s_out);
+  rev_iter<complex<T>> iter(ares, axes);
+  while(iter.remaining()>0)
+    {
+    auto v = ares.get(iter.ofs());
+    ares.set(iter.rev_ofs(), conj(v));
+    iter.advance();
+    }
+  }
+  return res;
+  }
+
+py::array fftn(const py::array &a, py::object axes, int inorm, bool inplace,
+  size_t nthreads)
+  {
+  try {
+    return xfftn(a, axes, inorm, inplace, true, nthreads);
+    }
+  catch (runtime_error &)
+    {
+    if (inplace) throw runtime_error("cannot do this operation in-place");
+    DISPATCH(a, f64, f32, flong, sym_rfftn_internal, (a, axes, inorm, nthreads))
+    }
+  }
 
 py::array ifftn(const py::array &a, py::object axes, int inorm,
   bool inplace, size_t nthreads)