pypocketfft.cc 13.1 KB
Newer Older
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
1 2 3 4 5 6
/*
 * This file is part of pocketfft.
 * Licensed under a 3-clause BSD style license - see LICENSE.md
 */

/*
7
 *  Python interface.
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
8
 *
9
 *  Copyright (C) 2019 Max-Planck-Society
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
10 11 12
 *  \author Martin Reinecke
 */

Martin Reinecke's avatar
Martin Reinecke committed
13 14
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
Martin Reinecke's avatar
more  
Martin Reinecke committed
15
#include <pybind11/stl.h>
Martin Reinecke's avatar
Martin Reinecke committed
16

17
#include "pocketfft_hdronly.h"
Martin Reinecke's avatar
Martin Reinecke committed
18

Martin Reinecke's avatar
Martin Reinecke committed
19
//
20
// Python interface
Martin Reinecke's avatar
Martin Reinecke committed
21 22
//

23
namespace {
24

25 26
using namespace std;
using namespace pocketfft;
Martin Reinecke's avatar
Martin Reinecke committed
27

Martin Reinecke's avatar
Martin Reinecke committed
28 29
namespace py = pybind11;

Martin Reinecke's avatar
Martin Reinecke committed
30 31
auto c64 = py::dtype("complex64");
auto c128 = py::dtype("complex128");
32
//auto c256 = py::dtype("complex256");
Martin Reinecke's avatar
Martin Reinecke committed
33 34
auto f32 = py::dtype("float32");
auto f64 = py::dtype("float64");
35
//auto f128 = py::dtype("float128");
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
36 37

shape_t copy_shape(const py::array &arr)
Martin Reinecke's avatar
Martin Reinecke committed
38
  {
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
39
  shape_t res(arr.ndim());
Martin Reinecke's avatar
Martin Reinecke committed
40 41 42 43
  for (size_t i=0; i<res.size(); ++i)
    res[i] = arr.shape(i);
  return res;
  }
Martin Reinecke's avatar
Martin Reinecke committed
44

Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
45
stride_t copy_strides(const py::array &arr)
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
46
  {
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
47
  stride_t res(arr.ndim());
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
48
  for (size_t i=0; i<res.size(); ++i)
Martin Reinecke's avatar
Martin Reinecke committed
49
    res[i] = arr.strides(i);
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
50 51
  return res;
  }
Martin Reinecke's avatar
Martin Reinecke committed
52

Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
53
shape_t makeaxes(const py::array &in, py::object axes)
Martin Reinecke's avatar
more  
Martin Reinecke committed
54
  {
Martin Reinecke's avatar
Martin Reinecke committed
55
  if (axes.is(py::none()))
Martin Reinecke's avatar
more  
Martin Reinecke committed
56
    {
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
57
    shape_t res(in.ndim());
Martin Reinecke's avatar
Martin Reinecke committed
58 59 60
    for (size_t i=0; i<res.size(); ++i)
      res[i]=i;
    return res;
Martin Reinecke's avatar
more  
Martin Reinecke committed
61
    }
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
62
  auto tmp=axes.cast<shape_t>();
Martin Reinecke's avatar
fix  
Martin Reinecke committed
63
  if ((tmp.size()>size_t(in.ndim())) || (tmp.size()==0))
Martin Reinecke's avatar
more  
Martin Reinecke committed
64
    throw runtime_error("bad axes argument");
Martin Reinecke's avatar
Martin Reinecke committed
65 66
  for (auto sz: tmp)
    if (sz>=size_t(in.ndim()))
Martin Reinecke's avatar
more  
Martin Reinecke committed
67
      throw runtime_error("invalid axis number");
Martin Reinecke's avatar
Martin Reinecke committed
68
  return tmp;
Martin Reinecke's avatar
more  
Martin Reinecke committed
69 70
  }

71 72 73 74 75 76 77
#define DISPATCH(arr, T1, T2, T3, func, args) \
  auto dtype = arr.dtype(); \
  if (dtype.is(T1)) return func<double> args; \
  if (dtype.is(T2)) return func<float> args; \
/*  if (dtype.is(T3)) return func<long double> args; */ \
  throw runtime_error("unsupported data type");

Martin Reinecke's avatar
Martin Reinecke committed
78
template<typename T> py::array xfftn_internal(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
79
  const shape_t &axes, double fct, bool inplace, bool fwd, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
80
  {
Martin Reinecke's avatar
Martin Reinecke committed
81
  auto dims(copy_shape(in));
82
  py::array res = inplace ? in : py::array_t<complex<T>>(dims);
Martin Reinecke's avatar
sync  
Martin Reinecke committed
83 84
  c2c(dims, copy_strides(in), copy_strides(res), axes, fwd,
    reinterpret_cast<const complex<T> *>(in.data()),
Martin Reinecke's avatar
Martin Reinecke committed
85
    reinterpret_cast<complex<T> *>(res.mutable_data()), T(fct), nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
86 87
  return res;
  }
Martin Reinecke's avatar
Martin Reinecke committed
88

89
py::array xfftn(const py::array &a, py::object axes, double fct, bool inplace,
Martin Reinecke's avatar
Martin Reinecke committed
90
  bool fwd, size_t nthreads)
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
91
  {
92
  DISPATCH(a, c128, c64, c256, xfftn_internal, (a, makeaxes(a, axes), fct,
Martin Reinecke's avatar
Martin Reinecke committed
93
           inplace, fwd, nthreads))
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
94
  }
95

Martin Reinecke's avatar
Martin Reinecke committed
96 97
py::array fftn(const py::array &a, py::object axes, double fct, bool inplace, size_t nthreads)
  { return xfftn(a, axes, fct, inplace, true, nthreads); }
98

Martin Reinecke's avatar
Martin Reinecke committed
99 100
py::array ifftn(const py::array &a, py::object axes, double fct, bool inplace, size_t nthreads)
  { return xfftn(a, axes, fct, inplace, false, nthreads); }
Martin Reinecke's avatar
Martin Reinecke committed
101

Martin Reinecke's avatar
Martin Reinecke committed
102
template<typename T> py::array rfftn_internal(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
103
  py::object axes_, T fct, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
104
  {
Martin Reinecke's avatar
Martin Reinecke committed
105
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
Martin Reinecke committed
106
  auto dims_in(copy_shape(in)), dims_out(dims_in);
Martin Reinecke's avatar
more  
Martin Reinecke committed
107
  dims_out[axes.back()] = (dims_out[axes.back()]>>1)+1;
Martin Reinecke's avatar
Martin Reinecke committed
108
  py::array res = py::array_t<complex<T>>(dims_out);
Martin Reinecke's avatar
sync  
Martin Reinecke committed
109 110
  r2c(dims_in, copy_strides(in), copy_strides(res), axes,
    reinterpret_cast<const T *>(in.data()),
Martin Reinecke's avatar
Martin Reinecke committed
111
    reinterpret_cast<complex<T> *>(res.mutable_data()), T(fct), nthreads);
Martin Reinecke's avatar
more  
Martin Reinecke committed
112 113
  return res;
  }
114

Martin Reinecke's avatar
Martin Reinecke committed
115
py::array rfftn(const py::array &in, py::object axes_, double fct, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
116
  {
Martin Reinecke's avatar
Martin Reinecke committed
117
  DISPATCH(in, f64, f32, f128, rfftn_internal, (in, axes_, fct, nthreads))
Martin Reinecke's avatar
Martin Reinecke committed
118
  }
119

120
template<typename T> py::array xrfft_scipy(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
121
  size_t axis, double fct, bool inplace, bool fwd, size_t nthreads)
122 123 124
  {
  auto dims(copy_shape(in));
  py::array res = inplace ? in : py::array_t<T>(dims);
Martin Reinecke's avatar
Martin Reinecke committed
125
  r2r_fftpack(dims, copy_strides(in), copy_strides(res), axis, fwd,
Martin Reinecke's avatar
sync  
Martin Reinecke committed
126
    reinterpret_cast<const T *>(in.data()),
Martin Reinecke's avatar
Martin Reinecke committed
127
    reinterpret_cast<T *>(res.mutable_data()), T(fct), nthreads);
128 129
  return res;
  }
130

Martin Reinecke's avatar
Martin Reinecke committed
131
py::array rfft_scipy(const py::array &in, size_t axis, double fct, bool inplace, size_t nthreads)
132
  {
Martin Reinecke's avatar
Martin Reinecke committed
133
  DISPATCH(in, f64, f32, f128, xrfft_scipy, (in, axis, fct, inplace, true, nthreads))
134
  }
135

Martin Reinecke's avatar
Martin Reinecke committed
136
py::array irfft_scipy(const py::array &in, size_t axis, double fct,
Martin Reinecke's avatar
Martin Reinecke committed
137
  bool inplace, size_t nthreads)
138
  {
Martin Reinecke's avatar
Martin Reinecke committed
139
  DISPATCH(in, f64, f32, f128, xrfft_scipy, (in, axis, fct, inplace, false, nthreads))
140
  }
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
141
template<typename T> py::array irfftn_internal(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
142
  py::object axes_, size_t lastsize, T fct, size_t nthreads)
Martin Reinecke's avatar
more  
Martin Reinecke committed
143
  {
Martin Reinecke's avatar
Martin Reinecke committed
144
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
more  
Martin Reinecke committed
145
  size_t axis = axes.back();
Martin Reinecke's avatar
Martin Reinecke committed
146 147 148
  shape_t dims_in(copy_shape(in)), dims_out=dims_in;
  if (lastsize==0) lastsize=2*dims_in[axis]-1;
  if ((lastsize/2) + 1 != dims_in[axis])
Martin Reinecke's avatar
more  
Martin Reinecke committed
149 150
    throw runtime_error("bad lastsize");
  dims_out[axis] = lastsize;
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
151
  py::array res = py::array_t<T>(dims_out);
Martin Reinecke's avatar
bug fix  
Martin Reinecke committed
152
  c2r(dims_out, copy_strides(in), copy_strides(res), axes,
Martin Reinecke's avatar
sync  
Martin Reinecke committed
153
    reinterpret_cast<const complex<T> *>(in.data()),
Martin Reinecke's avatar
Martin Reinecke committed
154
    reinterpret_cast<T *>(res.mutable_data()), T(fct), nthreads);
Martin Reinecke's avatar
more  
Martin Reinecke committed
155 156
  return res;
  }
157

Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
158
py::array irfftn(const py::array &in, py::object axes_, size_t lastsize,
Martin Reinecke's avatar
Martin Reinecke committed
159
  double fct, size_t nthreads)
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
160
  {
Martin Reinecke's avatar
Martin Reinecke committed
161
  DISPATCH(in, c128, c64, c256, irfftn_internal, (in, axes_, lastsize, fct, nthreads))
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
162
  }
Martin Reinecke's avatar
Martin Reinecke committed
163 164

template<typename T> py::array hartley_internal(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
165
  py::object axes_, double fct, bool inplace, size_t nthreads)
Martin Reinecke's avatar
more  
Martin Reinecke committed
166
  {
Martin Reinecke's avatar
Martin Reinecke committed
167
  auto dims(copy_shape(in));
168
  py::array res = inplace ? in : py::array_t<T>(dims);
Martin Reinecke's avatar
Martin Reinecke committed
169
  r2r_hartley(dims, copy_strides(in), copy_strides(res), makeaxes(in, axes_),
Martin Reinecke's avatar
sync  
Martin Reinecke committed
170
    reinterpret_cast<const T *>(in.data()),
Martin Reinecke's avatar
Martin Reinecke committed
171
    reinterpret_cast<T *>(res.mutable_data()), T(fct), nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
172 173
  return res;
  }
174

175
py::array hartley(const py::array &in, py::object axes_, double fct,
Martin Reinecke's avatar
Martin Reinecke committed
176
  bool inplace, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
177
  {
Martin Reinecke's avatar
Martin Reinecke committed
178
  DISPATCH(in, f64, f32, f128, hartley_internal, (in, axes_, fct, inplace, nthreads))
Martin Reinecke's avatar
Martin Reinecke committed
179
  }
180

Martin Reinecke's avatar
fixes  
Martin Reinecke committed
181
template<typename T>py::array complex2hartley(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
182
  const py::array &tmp, py::object axes_, bool inplace)
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
183
  {
Martin Reinecke's avatar
Martin Reinecke committed
184
  using namespace pocketfft::detail;
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
185
  int ndim = in.ndim();
Martin Reinecke's avatar
Martin Reinecke committed
186
  auto dims_out(copy_shape(in));
Martin Reinecke's avatar
Martin Reinecke committed
187 188 189
  py::array out = inplace ? in : py::array_t<T>(dims_out);
  ndarr<cmplx<T>> atmp(tmp.data(), copy_shape(tmp), copy_strides(tmp));
  ndarr<T> aout(out.mutable_data(), copy_shape(out), copy_strides(out));
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
190
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
Martin Reinecke committed
191
  size_t axis = axes.back();
Martin Reinecke's avatar
Martin Reinecke committed
192 193
  multi_iter<1,cmplx<T>,T> it(atmp, aout, axis);
  vector<bool> swp(ndim,false);
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
194 195
  for (auto i: axes)
    if (i!=axis)
Martin Reinecke's avatar
Martin Reinecke committed
196
      swp[i] = true;
Martin Reinecke's avatar
Martin Reinecke committed
197
  while(it.remaining()>0)
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
198
    {
Martin Reinecke's avatar
Martin Reinecke committed
199
    ptrdiff_t rofs = 0;
Martin Reinecke's avatar
Martin Reinecke committed
200
    for (size_t i=0; i<it.pos.size(); ++i)
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
201
      {
Martin Reinecke's avatar
Martin Reinecke committed
202
      if (i==axis) continue;
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
203
      if (!swp[i])
Martin Reinecke's avatar
Martin Reinecke committed
204
        rofs += it.pos[i]*it.oarr.stride(i);
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
205 206
      else
        {
Martin Reinecke's avatar
Martin Reinecke committed
207 208
        auto x = (it.pos[i]==0) ? 0 : it.iarr.shape(i)-it.pos[i];
        rofs += x*it.oarr.stride(i);
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
209 210
        }
      }
Martin Reinecke's avatar
Martin Reinecke committed
211
    it.advance(1);
Martin Reinecke's avatar
Martin Reinecke committed
212
    for (size_t i=0; i<it.length_in(); ++i)
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
213
      {
Martin Reinecke's avatar
Martin Reinecke committed
214 215
      auto re = it.in(i).r;
      auto im = it.in(i).i;
Martin Reinecke's avatar
Martin Reinecke committed
216
      auto rev_i = (i==0) ? 0 : it.length_out()-i;
Martin Reinecke's avatar
Martin Reinecke committed
217 218
      it.out(i) = re+im;
      aout[rofs + rev_i*it.stride_out()] = re-im;
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
219 220 221 222
      }
    }
  return out;
  }
223

224
py::array mycomplex2hartley(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
225
  const py::array &tmp, py::object axes_, bool inplace)
226
  {
227
  DISPATCH(in, f64, f32, f128, complex2hartley, (in, tmp, axes_, inplace))
228
  }
229

Martin Reinecke's avatar
Martin Reinecke committed
230
py::array hartley2(const py::array &in, py::object axes_, double fct,
Martin Reinecke's avatar
Martin Reinecke committed
231 232
  bool inplace, size_t nthreads)
  { return mycomplex2hartley(in, rfftn(in, axes_, fct, nthreads), axes_, inplace); }
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
233

Martin Reinecke's avatar
Martin Reinecke committed
234
const char *pypocketfft_DS = R"DELIM(Fast Fourier and Hartley transforms.
Martin Reinecke's avatar
Martin Reinecke committed
235 236 237 238 239 240 241 242 243 244 245

This module supports
- single and double precision
- complex and real-valued transforms
- multi-dimensional transforms

For two- and higher-dimensional transforms the code will use SSE2 and AVX
vector instructions for faster execution if these are supported by the CPU and
were enabled during compilation.
)DELIM";

246 247 248 249 250 251 252 253
const char *fftn_DS = R"DELIM(
Performs a forward complex FFT.

Parameters
----------
a : numpy.ndarray (np.complex64 or np.complex128)
    The input data
axes : list of integers
254
    The axes along which the FFT is carried out.
255
    If not set, all axes will be transformed.
Martin Reinecke's avatar
Martin Reinecke committed
256
fct : float
Martin Reinecke's avatar
Martin Reinecke committed
257
    Normalization factor
258
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
259
    if False, returns the result in a new array and leaves the input unchanged.
260
    if True, stores the result in the input array and returns a handle to it.
261 262 263 264

Returns
-------
np.ndarray (same shape and data type as a)
265
    The transformed data.
266 267
)DELIM";

Martin Reinecke's avatar
Martin Reinecke committed
268
const char *ifftn_DS = R"DELIM(Performs a backward complex FFT.
269 270 271 272 273 274

Parameters
----------
a : numpy.ndarray (np.complex64 or np.complex128)
    The input data
axes : list of integers
275
    The axes along which the FFT is carried out.
276
    If not set, all axes will be transformed.
Martin Reinecke's avatar
Martin Reinecke committed
277
fct : float
Martin Reinecke's avatar
Martin Reinecke committed
278
    Normalization factor
279
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
280
    if False, returns the result in a new array and leaves the input unchanged.
281
    if True, stores the result in the input array and returns a handle to it.
282 283 284 285

Returns
-------
np.ndarray (same shape and data type as a)
Martin Reinecke's avatar
Martin Reinecke committed
286
    The transformed data
287 288
)DELIM";

Martin Reinecke's avatar
Martin Reinecke committed
289
const char *rfftn_DS = R"DELIM(Performs a forward real-valued FFT.
Martin Reinecke's avatar
Martin Reinecke committed
290 291 292 293 294 295

Parameters
----------
a : numpy.ndarray (np.float32 or np.float64)
    The input data
axes : list of integers
296
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
297 298
    If not set, all axes will be transformed in ascending order.
fct : float
Martin Reinecke's avatar
Martin Reinecke committed
299
    Normalization factor
Martin Reinecke's avatar
Martin Reinecke committed
300 301 302 303

Returns
-------
np.ndarray (np.complex64 or np.complex128)
Martin Reinecke's avatar
Martin Reinecke committed
304
    The transformed data. The shape is identical to that of the input array,
Martin Reinecke's avatar
Martin Reinecke committed
305 306 307 308
    except for the axis that was transformed last. If the length of that axis
    was n on input, it is n//2+1 on output.
)DELIM";

309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330
const char *rfft_scipy_DS = R"DELIM(Performs a forward real-valued FFT.

Parameters
----------
a : numpy.ndarray (np.float32 or np.float64)
    The input data
axis : int
    The axis along which the FFT is carried out.
fct : float
    Normalization factor
inplace : bool
    if False, returns the result in a new array and leaves the input unchanged.
    if True, stores the result in the input array and returns a handle to it.

Returns
-------
np.ndarray (np.float32 or np.float64)
    The transformed data. The shape is identical to that of the input array.
    Along the transformed axis, values are arranged in
    FFTPACK half-complex order, i.e. `a[0].re, a[1].re, a[1].im, a[2].re ...`.
)DELIM";

Martin Reinecke's avatar
Martin Reinecke committed
331
const char *irfftn_DS = R"DELIM(Performs a backward real-valued FFT.
Martin Reinecke's avatar
Martin Reinecke committed
332 333 334 335 336 337

Parameters
----------
a : numpy.ndarray (np.complex64 or np.complex128)
    The input data
axes : list of integers
338
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
339 340 341 342
    If not set, all axes will be transformed in ascending order.
lastsize : the output size of the last axis to be transformed.
    If the corresponding input axis has size n, this can be 2*n-2 or 2*n-1.
fct : float
Martin Reinecke's avatar
Martin Reinecke committed
343
    Normalization factor
Martin Reinecke's avatar
Martin Reinecke committed
344 345 346 347

Returns
-------
np.ndarray (np.float32 or np.float64)
Martin Reinecke's avatar
Martin Reinecke committed
348 349
    The transformed data. The shape is identical to that of the input array,
    except for the axis that was transformed last, which has now `lastsize`
Martin Reinecke's avatar
Martin Reinecke committed
350 351 352
    entries.
)DELIM";

353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373
const char *irfft_scipy_DS = R"DELIM(Performs a backward real-valued FFT.

Parameters
----------
a : numpy.ndarray (np.float32 or np.float64)
    The input data. Along the transformed axis, values are expected in
    FFTPACK half-complex order, i.e. `a[0].re, a[1].re, a[1].im, a[2].re ...`.
axis : int
    The axis along which the FFT is carried out.
fct : float
    Normalization factor
inplace : bool
    if False, returns the result in a new array and leaves the input unchanged.
    if True, stores the result in the input array and returns a handle to it.

Returns
-------
np.ndarray (np.float32 or np.float64)
    The transformed data. The shape is identical to that of the input array.
)DELIM";

374
const char *hartley_DS = R"DELIM(Performs a Hartley transform.
375 376 377 378 379 380 381 382 383
For every requested axis, a 1D forward Fourier transform is carried out,
and the sum of real and imaginary parts of the result is stored in the output
array.

Parameters
----------
a : numpy.ndarray (np.float32 or np.float64)
    The input data
axes : list of integers
384
    The axes along which the transform is carried out.
385
    If not set, all axes will be transformed.
Martin Reinecke's avatar
Martin Reinecke committed
386
fct : float
Martin Reinecke's avatar
Martin Reinecke committed
387
    Normalization factor
388
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
389
    if False, returns the result in a new array and leaves the input unchanged.
390
    if True, stores the result in the input array and returns a handle to it.
391 392 393 394

Returns
-------
np.ndarray (same shape and data type as a)
Martin Reinecke's avatar
Martin Reinecke committed
395
    The transformed data
396 397
)DELIM";

Martin Reinecke's avatar
Martin Reinecke committed
398 399 400 401
} // unnamed namespace

PYBIND11_MODULE(pypocketfft, m)
  {
Martin Reinecke's avatar
more  
Martin Reinecke committed
402 403
  using namespace pybind11::literals;

Martin Reinecke's avatar
Martin Reinecke committed
404
  m.doc() = pypocketfft_DS;
405
  m.def("fftn",&fftn, fftn_DS, "a"_a, "axes"_a=py::none(), "fct"_a=1.,
Martin Reinecke's avatar
Martin Reinecke committed
406
    "inplace"_a=false, "nthreads"_a=1);
407
  m.def("ifftn",&ifftn, ifftn_DS, "a"_a, "axes"_a=py::none(), "fct"_a=1.,
Martin Reinecke's avatar
Martin Reinecke committed
408 409
    "inplace"_a=false, "nthreads"_a=1);
  m.def("rfftn",&rfftn, rfftn_DS, "a"_a, "axes"_a=py::none(), "fct"_a=1., "nthreads"_a=1);
410
  m.def("rfft_scipy",&rfft_scipy, rfft_scipy_DS, "a"_a, "axis"_a, "fct"_a=1.,
Martin Reinecke's avatar
Martin Reinecke committed
411
    "inplace"_a=false, "nthreads"_a=1);
412
  m.def("irfftn",&irfftn, irfftn_DS, "a"_a, "axes"_a=py::none(), "lastsize"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
413
    "fct"_a=1., "nthreads"_a=1);
414
  m.def("irfft_scipy",&irfft_scipy, irfft_scipy_DS, "a"_a, "axis"_a, "fct"_a=1.,
Martin Reinecke's avatar
Martin Reinecke committed
415
    "inplace"_a=false, "nthreads"_a=1);
416
  m.def("hartley",&hartley, hartley_DS, "a"_a, "axes"_a=py::none(), "fct"_a=1.,
Martin Reinecke's avatar
Martin Reinecke committed
417
    "inplace"_a=false, "nthreads"_a=1);
Martin Reinecke's avatar
Martin Reinecke committed
418
  m.def("hartley2",&hartley2, "a"_a, "axes"_a=py::none(), "fct"_a=1.,
Martin Reinecke's avatar
Martin Reinecke committed
419
    "inplace"_a=false, "nthreads"_a=1);
Martin Reinecke's avatar
Martin Reinecke committed
420 421
  m.def("complex2hartley",&mycomplex2hartley, "in"_a, "tmp"_a, "axes"_a,
    "inplace"_a=false);
Martin Reinecke's avatar
Martin Reinecke committed
422
  }