pypocketfft.cc 17.6 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;

30
31
32
33
// Only instantiate long double transforms if they offer more precision
using ldbl_t = typename std::conditional<
  sizeof(long double)==sizeof(double), double, long double>::type;

Martin Reinecke's avatar
Martin Reinecke committed
34
35
auto c64 = py::dtype("complex64");
auto c128 = py::dtype("complex128");
36
auto clong = py::dtype("longcomplex");
Martin Reinecke's avatar
Martin Reinecke committed
37
38
auto f32 = py::dtype("float32");
auto f64 = py::dtype("float64");
39
auto flong = py::dtype("longfloat");
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
40
41

shape_t copy_shape(const py::array &arr)
Martin Reinecke's avatar
Martin Reinecke committed
42
  {
43
  shape_t res(size_t(arr.ndim()));
Martin Reinecke's avatar
Martin Reinecke committed
44
  for (size_t i=0; i<res.size(); ++i)
45
    res[i] = size_t(arr.shape(int(i)));
Martin Reinecke's avatar
Martin Reinecke committed
46
47
  return res;
  }
Martin Reinecke's avatar
Martin Reinecke committed
48

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

Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
57
shape_t makeaxes(const py::array &in, py::object axes)
Martin Reinecke's avatar
more    
Martin Reinecke committed
58
  {
Martin Reinecke's avatar
Martin Reinecke committed
59
  if (axes.is(py::none()))
Martin Reinecke's avatar
more    
Martin Reinecke committed
60
    {
61
    shape_t res(size_t(in.ndim()));
Martin Reinecke's avatar
Martin Reinecke committed
62
63
64
    for (size_t i=0; i<res.size(); ++i)
      res[i]=i;
    return res;
Martin Reinecke's avatar
more    
Martin Reinecke committed
65
    }
66
67
68
  auto tmp=axes.cast<std::vector<ptrdiff_t>>();
  auto ndim = in.ndim();
  if ((tmp.size()>size_t(ndim)) || (tmp.size()==0))
Martin Reinecke's avatar
more    
Martin Reinecke committed
69
    throw runtime_error("bad axes argument");
70
71
72
73
74
75
76
77
  for (auto& sz: tmp)
    {
    if (sz<0)
      sz += ndim;
    if ((sz>=ndim) || (sz<0))
      throw invalid_argument("axes exceeds dimensionality of output");
    }
  return shape_t(tmp.begin(), tmp.end());
Martin Reinecke's avatar
more    
Martin Reinecke committed
78
79
  }

80
#define DISPATCH(arr, T1, T2, T3, func, args) \
81
  { \
82
83
84
  auto dtype = arr.dtype(); \
  if (dtype.is(T1)) return func<double> args; \
  if (dtype.is(T2)) return func<float> args; \
85
  if (dtype.is(T3)) return func<ldbl_t> args; \
86
87
  throw runtime_error("unsupported data type"); \
  }
88

89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
template<typename T> T norm_fct(int inorm, size_t N)
  {
  if (inorm==0) return T(1);
  if (inorm==2) return T(1/ldbl_t(N));
  if (inorm==1) return T(1/sqrt(ldbl_t(N)));
  throw invalid_argument("invalid value for inorm (must be 0, 1, or 2)");
  }

template<typename T> T norm_fct(int inorm, const shape_t &shape,
  const shape_t &axes)
  {
  if (inorm==0) return T(1);
  size_t N(1);
  for (auto a: axes)
    N *= shape[a];
  return norm_fct<T>(inorm, N);
  }

Martin Reinecke's avatar
Martin Reinecke committed
107
template<typename T> py::array xfftn_internal(const py::array &in,
108
  const shape_t &axes, int inorm, bool inplace, bool fwd, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
109
  {
Martin Reinecke's avatar
Martin Reinecke committed
110
  auto dims(copy_shape(in));
111
  py::array res = inplace ? in : py::array_t<complex<T>>(dims);
Martin Reinecke's avatar
Martin Reinecke committed
112
113
114
115
116
117
  auto s_in=copy_strides(in);
  auto s_out=copy_strides(res);
  auto d_in=reinterpret_cast<const complex<T> *>(in.data());
  auto d_out=reinterpret_cast<complex<T> *>(res.mutable_data());
  {
  py::gil_scoped_release release;
118
119
  T fct = norm_fct<T>(inorm, dims, axes);
  c2c(dims, s_in, s_out, axes, fwd, d_in, d_out, fct, nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
120
  }
Martin Reinecke's avatar
Martin Reinecke committed
121
122
  return res;
  }
Martin Reinecke's avatar
Martin Reinecke committed
123

Martin Reinecke's avatar
Martin Reinecke committed
124
template<typename T> py::array sym_rfftn_internal(const py::array &in,
125
  const shape_t &axes, int inorm, bool fwd, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
126
127
128
129
130
131
132
133
134
135
  {
  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);
136
  r2c(dims, s_in, s_out, axes, fwd, d_in, d_out, fct, nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
137
138
139
  // now fill in second half
  using namespace pocketfft::detail;
  ndarr<complex<T>> ares(res.mutable_data(), dims, s_out);
140
  rev_iter iter(ares, axes);
Martin Reinecke's avatar
Martin Reinecke committed
141
142
  while(iter.remaining()>0)
    {
143
144
    auto v = ares[iter.ofs()];
    ares[iter.rev_ofs()] = conj(v);
Martin Reinecke's avatar
Martin Reinecke committed
145
146
147
148
149
150
    iter.advance();
    }
  }
  return res;
  }

151
152
py::array xfftn(const py::array &a, py::object axes, int inorm,
  bool inplace, bool fwd, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
153
  {
154
155
156
  if (a.dtype().kind() == 'c')
    DISPATCH(a, c128, c64, clong, xfftn_internal, (a, makeaxes(a, axes), inorm,
             inplace, fwd, nthreads))
157

158
159
160
  if (inplace) throw runtime_error("cannot do this operation in-place");
  DISPATCH(a, f64, f32, flong, sym_rfftn_internal, (a, makeaxes(a, axes), inorm,
           fwd, nthreads))
Martin Reinecke's avatar
Martin Reinecke committed
161
  }
162

163
164
165
166
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); }

167
168
169
py::array ifftn(const py::array &a, py::object axes, int inorm,
  bool inplace, size_t nthreads)
  { return xfftn(a, axes, inorm, inplace, false, nthreads); }
Martin Reinecke's avatar
Martin Reinecke committed
170

Martin Reinecke's avatar
Martin Reinecke committed
171
template<typename T> py::array rfftn_internal(const py::array &in,
172
  py::object axes_, int inorm, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
173
  {
Martin Reinecke's avatar
Martin Reinecke committed
174
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
Martin Reinecke committed
175
  auto dims_in(copy_shape(in)), dims_out(dims_in);
Martin Reinecke's avatar
more    
Martin Reinecke committed
176
  dims_out[axes.back()] = (dims_out[axes.back()]>>1)+1;
Martin Reinecke's avatar
Martin Reinecke committed
177
  py::array res = py::array_t<complex<T>>(dims_out);
Martin Reinecke's avatar
Martin Reinecke committed
178
179
180
181
182
183
  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;
184
  T fct = norm_fct<T>(inorm, dims_in, axes);
185
  r2c(dims_in, s_in, s_out, axes, true, d_in, d_out, fct, nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
186
  }
Martin Reinecke's avatar
more    
Martin Reinecke committed
187
188
  return res;
  }
189

190
py::array rfftn(const py::array &in, py::object axes_, int inorm,
Martin Reinecke's avatar
Martin Reinecke committed
191
  size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
192
  {
193
  DISPATCH(in, f64, f32, flong, rfftn_internal, (in, axes_, inorm, nthreads))
Martin Reinecke's avatar
Martin Reinecke committed
194
  }
195

196
template<typename T> py::array xrfft_scipy(const py::array &in,
197
  size_t axis, int inorm, bool inplace, bool fwd, size_t nthreads)
198
199
200
  {
  auto dims(copy_shape(in));
  py::array res = inplace ? in : py::array_t<T>(dims);
Martin Reinecke's avatar
Martin Reinecke committed
201
202
203
204
205
206
  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<T *>(res.mutable_data());
  {
  py::gil_scoped_release release;
207
  T fct = norm_fct<T>(inorm, dims[axis]);
208
  r2r_fftpack(dims, s_in, s_out, {axis}, fwd, fwd, d_in, d_out, fct, nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
209
  }
210
211
  return res;
  }
212

213
214
py::array rfft_scipy(const py::array &in, size_t axis, int inorm,
  bool inplace, size_t nthreads)
215
  {
216
  DISPATCH(in, f64, f32, flong, xrfft_scipy, (in, axis, inorm, inplace, true,
Martin Reinecke's avatar
Martin Reinecke committed
217
    nthreads))
218
  }
219

220
py::array irfft_scipy(const py::array &in, size_t axis, int inorm,
Martin Reinecke's avatar
Martin Reinecke committed
221
  bool inplace, size_t nthreads)
222
  {
223
  DISPATCH(in, f64, f32, flong, xrfft_scipy, (in, axis, inorm, inplace, false,
Martin Reinecke's avatar
Martin Reinecke committed
224
    nthreads))
225
  }
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
226
template<typename T> py::array irfftn_internal(const py::array &in,
227
  py::object axes_, size_t lastsize, int inorm, size_t nthreads)
Martin Reinecke's avatar
more    
Martin Reinecke committed
228
  {
Martin Reinecke's avatar
Martin Reinecke committed
229
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
more    
Martin Reinecke committed
230
  size_t axis = axes.back();
Martin Reinecke's avatar
Martin Reinecke committed
231
232
233
  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
234
235
    throw runtime_error("bad lastsize");
  dims_out[axis] = lastsize;
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
236
  py::array res = py::array_t<T>(dims_out);
Martin Reinecke's avatar
Martin Reinecke committed
237
238
239
240
241
242
  auto s_in=copy_strides(in);
  auto s_out=copy_strides(res);
  auto d_in=reinterpret_cast<const complex<T> *>(in.data());
  auto d_out=reinterpret_cast<T *>(res.mutable_data());
  {
  py::gil_scoped_release release;
243
  T fct = norm_fct<T>(inorm, dims_out, axes);
244
  c2r(dims_out, s_in, s_out, axes, false, d_in, d_out, fct, nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
245
  }
Martin Reinecke's avatar
more    
Martin Reinecke committed
246
247
  return res;
  }
248

Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
249
py::array irfftn(const py::array &in, py::object axes_, size_t lastsize,
250
  int inorm, size_t nthreads)
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
251
  {
252
  DISPATCH(in, c128, c64, clong, irfftn_internal, (in, axes_, lastsize, inorm,
Martin Reinecke's avatar
Martin Reinecke committed
253
    nthreads))
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
254
  }
Martin Reinecke's avatar
Martin Reinecke committed
255
256

template<typename T> py::array hartley_internal(const py::array &in,
257
  py::object axes_, int inorm, bool inplace, size_t nthreads)
Martin Reinecke's avatar
more    
Martin Reinecke committed
258
  {
Martin Reinecke's avatar
Martin Reinecke committed
259
  auto dims(copy_shape(in));
260
  py::array res = inplace ? in : py::array_t<T>(dims);
Martin Reinecke's avatar
Martin Reinecke committed
261
262
263
264
265
266
267
  auto axes = makeaxes(in, axes_);
  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<T *>(res.mutable_data());
  {
  py::gil_scoped_release release;
268
269
  T fct = norm_fct<T>(inorm, dims, axes);
  r2r_hartley(dims, s_in, s_out, axes, d_in, d_out, fct, nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
270
  }
Martin Reinecke's avatar
Martin Reinecke committed
271
272
  return res;
  }
273

274
py::array hartley(const py::array &in, py::object axes_, int inorm,
Martin Reinecke's avatar
Martin Reinecke committed
275
  bool inplace, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
276
  {
277
  DISPATCH(in, f64, f32, flong, hartley_internal, (in, axes_, inorm, inplace,
Martin Reinecke's avatar
Martin Reinecke committed
278
    nthreads))
Martin Reinecke's avatar
Martin Reinecke committed
279
  }
280

Martin Reinecke's avatar
fixes    
Martin Reinecke committed
281
template<typename T>py::array complex2hartley(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
282
  const py::array &tmp, py::object axes_, bool inplace)
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
283
  {
Martin Reinecke's avatar
Martin Reinecke committed
284
  using namespace pocketfft::detail;
Martin Reinecke's avatar
Martin Reinecke committed
285
  auto dims_out(copy_shape(in));
Martin Reinecke's avatar
Martin Reinecke committed
286
  py::array out = inplace ? in : py::array_t<T>(dims_out);
287
  cndarr<cmplx<T>> atmp(tmp.data(), copy_shape(tmp), copy_strides(tmp));
Martin Reinecke's avatar
Martin Reinecke committed
288
  ndarr<T> aout(out.mutable_data(), copy_shape(out), copy_strides(out));
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
289
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
Martin Reinecke committed
290
291
  {
  py::gil_scoped_release release;
292
293
  simple_iter iin(atmp);
  rev_iter iout(aout, axes);
294
295
  if (iin.remaining()!=iout.remaining()) throw runtime_error("oops");
  while(iin.remaining()>0)
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
296
    {
297
298
299
    auto v = atmp[iin.ofs()];
    aout[iout.ofs()] = v.r+v.i;
    aout[iout.rev_ofs()] = v.r-v.i;
300
    iin.advance(); iout.advance();
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
301
    }
Martin Reinecke's avatar
Martin Reinecke committed
302
  }
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
303
304
  return out;
  }
305

306
py::array mycomplex2hartley(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
307
  const py::array &tmp, py::object axes_, bool inplace)
308
  {
309
  DISPATCH(in, f64, f32, flong, complex2hartley, (in, tmp, axes_, inplace))
310
  }
311

312
py::array hartley2(const py::array &in, py::object axes_, int inorm,
Martin Reinecke's avatar
Martin Reinecke committed
313
  bool inplace, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
314
  {
315
  return mycomplex2hartley(in, rfftn(in, axes_, inorm, nthreads), axes_,
Martin Reinecke's avatar
Martin Reinecke committed
316
317
    inplace);
  }
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
318

Martin Reinecke's avatar
Martin Reinecke committed
319
const char *pypocketfft_DS = R"""(Fast Fourier and Hartley transforms.
Martin Reinecke's avatar
Martin Reinecke committed
320
321

This module supports
322
- single, double, and long double precision
Martin Reinecke's avatar
Martin Reinecke committed
323
324
325
326
327
328
- 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.
Martin Reinecke's avatar
Martin Reinecke committed
329
)""";
Martin Reinecke's avatar
Martin Reinecke committed
330

Martin Reinecke's avatar
Martin Reinecke committed
331
const char *fftn_DS = R"""(
332
333
334
335
Performs a forward complex FFT.

Parameters
----------
336
337
338
a : numpy.ndarray (any complex or real type)
    The input data. If its type is real, a more efficient real-to-complex
    transform will be used.
339
axes : list of integers
340
    The axes along which the FFT is carried out.
341
    If not set, all axes will be transformed.
342
343
344
345
346
347
inorm : int
    Normalization type
      0 : no normalization
      1 : divide by sqrt(N)
      2 : divide by N
    where N is the product of the lengths of the transformed axes.
348
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
349
    if False, returns the result in a new array and leaves the input unchanged.
350
    if True, stores the result in the input array and returns a handle to it.
351
352
    NOTE: if `a` is real-valued and `inplace` is `True`, an exception will be
    raised!
Martin Reinecke's avatar
Martin Reinecke committed
353
354
355
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
356
357
358

Returns
-------
359
np.ndarray (same shape as `a`, complex type with same accuracy as `a`)
360
    The transformed data.
Martin Reinecke's avatar
Martin Reinecke committed
361
)""";
362

Martin Reinecke's avatar
Martin Reinecke committed
363
const char *ifftn_DS = R"""(Performs a backward complex FFT.
364
365
366

Parameters
----------
367
368
369
a : numpy.ndarray (any complex or real type)
    The input data. If its type is real, a more efficient real-to-complex
    transform will be used.
370
axes : list of integers
371
    The axes along which the FFT is carried out.
372
    If not set, all axes will be transformed.
373
374
375
376
377
378
inorm : int
    Normalization type
      0 : no normalization
      1 : divide by sqrt(N)
      2 : divide by N
    where N is the product of the lengths of the transformed axes.
379
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
380
    if False, returns the result in a new array and leaves the input unchanged.
381
    if True, stores the result in the input array and returns a handle to it.
382
383
    NOTE: if `a` is real-valued and `inplace` is `True`, an exception will be
    raised!
Martin Reinecke's avatar
Martin Reinecke committed
384
385
386
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
387
388
389

Returns
-------
390
np.ndarray (same shape as `a`, complex type with same accuracy as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
391
    The transformed data
Martin Reinecke's avatar
Martin Reinecke committed
392
)""";
393

Martin Reinecke's avatar
Martin Reinecke committed
394
const char *rfftn_DS = R"""(Performs a forward real-valued FFT.
Martin Reinecke's avatar
Martin Reinecke committed
395
396
397

Parameters
----------
398
a : numpy.ndarray (any real type)
Martin Reinecke's avatar
Martin Reinecke committed
399
400
    The input data
axes : list of integers
401
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
402
    If not set, all axes will be transformed in ascending order.
403
404
405
406
407
408
inorm : int
    Normalization type
      0 : no normalization
      1 : divide by sqrt(N)
      2 : divide by N
    where N is the product of the lengths of the transformed input axes.
Martin Reinecke's avatar
Martin Reinecke committed
409
410
411
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
Martin Reinecke's avatar
Martin Reinecke committed
412
413
414

Returns
-------
415
np.ndarray (complex type with same accuracy as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
416
    The transformed data. The shape is identical to that of the input array,
Martin Reinecke's avatar
Martin Reinecke committed
417
418
    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.
Martin Reinecke's avatar
Martin Reinecke committed
419
)""";
Martin Reinecke's avatar
Martin Reinecke committed
420

Martin Reinecke's avatar
Martin Reinecke committed
421
const char *rfft_scipy_DS = R"""(Performs a forward real-valued FFT.
422
423
424

Parameters
----------
425
a : numpy.ndarray (any real type)
426
427
428
    The input data
axis : int
    The axis along which the FFT is carried out.
429
430
431
432
433
434
inorm : int
    Normalization type
      0 : no normalization
      1 : divide by sqrt(N)
      2 : divide by N
    where N is the length of `axis`.
435
436
437
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.
Martin Reinecke's avatar
Martin Reinecke committed
438
439
440
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
441
442
443

Returns
-------
444
np.ndarray (same shape and data type as `a`)
445
446
447
    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 ...`.
Martin Reinecke's avatar
Martin Reinecke committed
448
)""";
449

Martin Reinecke's avatar
Martin Reinecke committed
450
const char *irfftn_DS = R"""(Performs a backward real-valued FFT.
Martin Reinecke's avatar
Martin Reinecke committed
451
452
453

Parameters
----------
454
a : numpy.ndarray (any complex type)
Martin Reinecke's avatar
Martin Reinecke committed
455
456
    The input data
axes : list of integers
457
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
458
459
460
    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.
461
462
463
464
465
466
inorm : int
    Normalization type
      0 : no normalization
      1 : divide by sqrt(N)
      2 : divide by N
    where N is the product of the lengths of the transformed output axes.
Martin Reinecke's avatar
Martin Reinecke committed
467
468
469
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
Martin Reinecke's avatar
Martin Reinecke committed
470
471
472

Returns
-------
473
np.ndarray (real type with same accuracy as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
474
475
    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
476
    entries.
Martin Reinecke's avatar
Martin Reinecke committed
477
)""";
Martin Reinecke's avatar
Martin Reinecke committed
478

Martin Reinecke's avatar
Martin Reinecke committed
479
const char *irfft_scipy_DS = R"""(Performs a backward real-valued FFT.
480
481
482

Parameters
----------
483
a : numpy.ndarray (any real type)
484
485
486
487
    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.
488
489
490
491
492
493
inorm : int
    Normalization type
      0 : no normalization
      1 : divide by sqrt(N)
      2 : divide by N
    where N is the length of `axis`.
494
495
496
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.
Martin Reinecke's avatar
Martin Reinecke committed
497
498
499
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
500
501
502

Returns
-------
503
np.ndarray (same shape and data type as `a`)
504
    The transformed data. The shape is identical to that of the input array.
Martin Reinecke's avatar
Martin Reinecke committed
505
)""";
506

Martin Reinecke's avatar
Martin Reinecke committed
507
const char *hartley_DS = R"""(Performs a Hartley transform.
508
509
510
511
512
513
514
515
516
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
517
    The axes along which the transform is carried out.
518
    If not set, all axes will be transformed.
519
520
521
522
523
524
inorm : int
    Normalization type
      0 : no normalization
      1 : divide by sqrt(N)
      2 : divide by N
    where N is the product of the lengths of the transformed axes.
525
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
526
    if False, returns the result in a new array and leaves the input unchanged.
527
    if True, stores the result in the input array and returns a handle to it.
Martin Reinecke's avatar
Martin Reinecke committed
528
529
530
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
531
532
533

Returns
-------
534
np.ndarray (same shape and data type as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
535
    The transformed data
Martin Reinecke's avatar
Martin Reinecke committed
536
)""";
537

Martin Reinecke's avatar
Martin Reinecke committed
538
539
540
541
} // unnamed namespace

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

Martin Reinecke's avatar
Martin Reinecke committed
544
  m.doc() = pypocketfft_DS;
545
  m.def("fftn",&fftn, fftn_DS, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
546
    "inplace"_a=false, "nthreads"_a=1);
547
  m.def("ifftn",&ifftn, ifftn_DS, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
548
    "inplace"_a=false, "nthreads"_a=1);
549
  m.def("rfftn",&rfftn, rfftn_DS, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
550
    "nthreads"_a=1);
551
  m.def("rfft_scipy",&rfft_scipy, rfft_scipy_DS, "a"_a, "axis"_a, "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
552
    "inplace"_a=false, "nthreads"_a=1);
553
  m.def("irfftn",&irfftn, irfftn_DS, "a"_a, "axes"_a=py::none(), "lastsize"_a=0,
554
555
556
557
    "inorm"_a=0, "nthreads"_a=1);
  m.def("irfft_scipy",&irfft_scipy, irfft_scipy_DS, "a"_a, "axis"_a,
    "inorm"_a=0, "inplace"_a=false, "nthreads"_a=1);
  m.def("hartley",&hartley, hartley_DS, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
558
    "inplace"_a=false, "nthreads"_a=1);
559
  m.def("hartley2",&hartley2, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
560
    "inplace"_a=false, "nthreads"_a=1);
Martin Reinecke's avatar
Martin Reinecke committed
561
562
  m.def("complex2hartley",&mycomplex2hartley, "in"_a, "tmp"_a, "axes"_a,
    "inplace"_a=false);
Martin Reinecke's avatar
Martin Reinecke committed
563
  }