pypocketfft.cc 16.5 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
    }
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
66
  auto tmp=axes.cast<shape_t>();
Martin Reinecke's avatar
fix  
Martin Reinecke committed
67
  if ((tmp.size()>size_t(in.ndim())) || (tmp.size()==0))
Martin Reinecke's avatar
more  
Martin Reinecke committed
68
    throw runtime_error("bad axes argument");
Martin Reinecke's avatar
Martin Reinecke committed
69
70
  for (auto sz: tmp)
    if (sz>=size_t(in.ndim()))
Martin Reinecke's avatar
more  
Martin Reinecke committed
71
      throw runtime_error("invalid axis number");
Martin Reinecke's avatar
Martin Reinecke committed
72
  return tmp;
Martin Reinecke's avatar
more  
Martin Reinecke committed
73
74
  }

75
76
77
78
#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; \
79
  if (dtype.is(T3)) return func<ldbl_t> args; \
80
81
  throw runtime_error("unsupported data type");

82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
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
100
template<typename T> py::array xfftn_internal(const py::array &in,
101
  const shape_t &axes, int inorm, bool inplace, bool fwd, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
102
  {
Martin Reinecke's avatar
Martin Reinecke committed
103
  auto dims(copy_shape(in));
104
  py::array res = inplace ? in : py::array_t<complex<T>>(dims);
Martin Reinecke's avatar
Martin Reinecke committed
105
106
107
108
109
110
  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;
111
112
  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
113
  }
Martin Reinecke's avatar
Martin Reinecke committed
114
115
  return res;
  }
Martin Reinecke's avatar
Martin Reinecke committed
116

117
118
py::array xfftn(const py::array &a, py::object axes, int inorm,
  bool inplace, bool fwd, size_t nthreads)
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
119
  {
120
  DISPATCH(a, c128, c64, clong, xfftn_internal, (a, makeaxes(a, axes), inorm,
Martin Reinecke's avatar
Martin Reinecke committed
121
           inplace, fwd, nthreads))
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
122
  }
123

124
125
126
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); }
127

128
129
130
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
131

Martin Reinecke's avatar
Martin Reinecke committed
132
template<typename T> py::array rfftn_internal(const py::array &in,
133
  py::object axes_, int inorm, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
134
  {
Martin Reinecke's avatar
Martin Reinecke committed
135
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
Martin Reinecke committed
136
  auto dims_in(copy_shape(in)), dims_out(dims_in);
Martin Reinecke's avatar
more  
Martin Reinecke committed
137
  dims_out[axes.back()] = (dims_out[axes.back()]>>1)+1;
Martin Reinecke's avatar
Martin Reinecke committed
138
  py::array res = py::array_t<complex<T>>(dims_out);
Martin Reinecke's avatar
Martin Reinecke committed
139
140
141
142
143
144
  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;
145
146
  T fct = norm_fct<T>(inorm, dims_in, axes);
  r2c(dims_in, s_in, s_out, axes, d_in, d_out, fct, nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
147
  }
Martin Reinecke's avatar
more  
Martin Reinecke committed
148
149
  return res;
  }
150

151
py::array rfftn(const py::array &in, py::object axes_, int inorm,
Martin Reinecke's avatar
Martin Reinecke committed
152
  size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
153
  {
154
  DISPATCH(in, f64, f32, flong, rfftn_internal, (in, axes_, inorm, nthreads))
Martin Reinecke's avatar
Martin Reinecke committed
155
  }
156

157
template<typename T> py::array xrfft_scipy(const py::array &in,
158
  size_t axis, int inorm, bool inplace, bool fwd, size_t nthreads)
159
160
161
  {
  auto dims(copy_shape(in));
  py::array res = inplace ? in : py::array_t<T>(dims);
Martin Reinecke's avatar
Martin Reinecke committed
162
163
164
165
166
167
  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;
168
169
  T fct = norm_fct<T>(inorm, dims[axis]);
  r2r_fftpack(dims, s_in, s_out, axis, fwd, d_in, d_out, fct, nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
170
  }
171
172
  return res;
  }
173

174
175
py::array rfft_scipy(const py::array &in, size_t axis, int inorm,
  bool inplace, size_t nthreads)
176
  {
177
  DISPATCH(in, f64, f32, flong, xrfft_scipy, (in, axis, inorm, inplace, true,
Martin Reinecke's avatar
Martin Reinecke committed
178
    nthreads))
179
  }
180

181
py::array irfft_scipy(const py::array &in, size_t axis, int inorm,
Martin Reinecke's avatar
Martin Reinecke committed
182
  bool inplace, size_t nthreads)
183
  {
184
  DISPATCH(in, f64, f32, flong, xrfft_scipy, (in, axis, inorm, inplace, false,
Martin Reinecke's avatar
Martin Reinecke committed
185
    nthreads))
186
  }
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
187
template<typename T> py::array irfftn_internal(const py::array &in,
188
  py::object axes_, size_t lastsize, int inorm, size_t nthreads)
Martin Reinecke's avatar
more  
Martin Reinecke committed
189
  {
Martin Reinecke's avatar
Martin Reinecke committed
190
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
more  
Martin Reinecke committed
191
  size_t axis = axes.back();
Martin Reinecke's avatar
Martin Reinecke committed
192
193
194
  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
195
196
    throw runtime_error("bad lastsize");
  dims_out[axis] = lastsize;
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
197
  py::array res = py::array_t<T>(dims_out);
Martin Reinecke's avatar
Martin Reinecke committed
198
199
200
201
202
203
  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;
204
205
  T fct = norm_fct<T>(inorm, dims_out, axes);
  c2r(dims_out, s_in, s_out, axes, d_in, d_out, fct, nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
206
  }
Martin Reinecke's avatar
more  
Martin Reinecke committed
207
208
  return res;
  }
209

Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
210
py::array irfftn(const py::array &in, py::object axes_, size_t lastsize,
211
  int inorm, size_t nthreads)
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
212
  {
213
  DISPATCH(in, c128, c64, clong, irfftn_internal, (in, axes_, lastsize, inorm,
Martin Reinecke's avatar
Martin Reinecke committed
214
    nthreads))
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
215
  }
Martin Reinecke's avatar
Martin Reinecke committed
216
217

template<typename T> py::array hartley_internal(const py::array &in,
218
  py::object axes_, int inorm, bool inplace, size_t nthreads)
Martin Reinecke's avatar
more  
Martin Reinecke committed
219
  {
Martin Reinecke's avatar
Martin Reinecke committed
220
  auto dims(copy_shape(in));
221
  py::array res = inplace ? in : py::array_t<T>(dims);
Martin Reinecke's avatar
Martin Reinecke committed
222
223
224
225
226
227
228
  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;
229
230
  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
231
  }
Martin Reinecke's avatar
Martin Reinecke committed
232
233
  return res;
  }
234

235
py::array hartley(const py::array &in, py::object axes_, int inorm,
Martin Reinecke's avatar
Martin Reinecke committed
236
  bool inplace, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
237
  {
238
  DISPATCH(in, f64, f32, flong, hartley_internal, (in, axes_, inorm, inplace,
Martin Reinecke's avatar
Martin Reinecke committed
239
    nthreads))
Martin Reinecke's avatar
Martin Reinecke committed
240
  }
241

Martin Reinecke's avatar
fixes  
Martin Reinecke committed
242
template<typename T>py::array complex2hartley(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
243
  const py::array &tmp, py::object axes_, bool inplace)
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
244
  {
Martin Reinecke's avatar
Martin Reinecke committed
245
  using namespace pocketfft::detail;
246
  size_t ndim = size_t(in.ndim());
Martin Reinecke's avatar
Martin Reinecke committed
247
  auto dims_out(copy_shape(in));
Martin Reinecke's avatar
Martin Reinecke committed
248
249
250
  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
251
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
Martin Reinecke committed
252
253
  {
  py::gil_scoped_release release;
Martin Reinecke's avatar
Martin Reinecke committed
254
  size_t axis = axes.back();
Martin Reinecke's avatar
Martin Reinecke committed
255
256
  multi_iter<1,cmplx<T>,T> it(atmp, aout, axis);
  vector<bool> swp(ndim,false);
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
257
258
  for (auto i: axes)
    if (i!=axis)
Martin Reinecke's avatar
Martin Reinecke committed
259
      swp[i] = true;
Martin Reinecke's avatar
Martin Reinecke committed
260
  while(it.remaining()>0)
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
261
    {
Martin Reinecke's avatar
Martin Reinecke committed
262
    ptrdiff_t rofs = 0;
Martin Reinecke's avatar
Martin Reinecke committed
263
    for (size_t i=0; i<it.pos.size(); ++i)
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
264
      {
Martin Reinecke's avatar
Martin Reinecke committed
265
      if (i==axis) continue;
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
266
      if (!swp[i])
267
        rofs += ptrdiff_t(it.pos[i])*it.oarr.stride(i);
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
268
269
      else
        {
270
        auto x = ptrdiff_t((it.pos[i]==0) ? 0 : it.iarr.shape(i)-it.pos[i]);
Martin Reinecke's avatar
Martin Reinecke committed
271
        rofs += x*it.oarr.stride(i);
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
272
273
        }
      }
Martin Reinecke's avatar
Martin Reinecke committed
274
    it.advance(1);
Martin Reinecke's avatar
Martin Reinecke committed
275
    for (size_t i=0; i<it.length_in(); ++i)
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
276
      {
Martin Reinecke's avatar
Martin Reinecke committed
277
278
      auto re = it.in(i).r;
      auto im = it.in(i).i;
279
      auto rev_i = ptrdiff_t((i==0) ? 0 : it.length_out()-i);
Martin Reinecke's avatar
Martin Reinecke committed
280
281
      it.out(i) = re+im;
      aout[rofs + rev_i*it.stride_out()] = re-im;
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
282
283
      }
    }
Martin Reinecke's avatar
Martin Reinecke committed
284
  }
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
285
286
  return out;
  }
287

288
py::array mycomplex2hartley(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
289
  const py::array &tmp, py::object axes_, bool inplace)
290
  {
291
  DISPATCH(in, f64, f32, flong, complex2hartley, (in, tmp, axes_, inplace))
292
  }
293

294
py::array hartley2(const py::array &in, py::object axes_, int inorm,
Martin Reinecke's avatar
Martin Reinecke committed
295
  bool inplace, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
296
  {
297
  return mycomplex2hartley(in, rfftn(in, axes_, inorm, nthreads), axes_,
Martin Reinecke's avatar
Martin Reinecke committed
298
299
    inplace);
  }
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
300

Martin Reinecke's avatar
Martin Reinecke committed
301
const char *pypocketfft_DS = R"""(Fast Fourier and Hartley transforms.
Martin Reinecke's avatar
Martin Reinecke committed
302
303
304
305
306
307
308
309
310

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.
Martin Reinecke's avatar
Martin Reinecke committed
311
)""";
Martin Reinecke's avatar
Martin Reinecke committed
312

Martin Reinecke's avatar
Martin Reinecke committed
313
const char *fftn_DS = R"""(
314
315
316
317
318
319
320
Performs a forward complex FFT.

Parameters
----------
a : numpy.ndarray (np.complex64 or np.complex128)
    The input data
axes : list of integers
321
    The axes along which the FFT is carried out.
322
    If not set, all axes will be transformed.
323
324
325
326
327
328
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.
329
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
330
    if False, returns the result in a new array and leaves the input unchanged.
331
    if True, stores the result in the input array and returns a handle to it.
Martin Reinecke's avatar
Martin Reinecke committed
332
333
334
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
335
336
337
338

Returns
-------
np.ndarray (same shape and data type as a)
339
    The transformed data.
Martin Reinecke's avatar
Martin Reinecke committed
340
)""";
341

Martin Reinecke's avatar
Martin Reinecke committed
342
const char *ifftn_DS = R"""(Performs a backward complex FFT.
343
344
345
346
347
348

Parameters
----------
a : numpy.ndarray (np.complex64 or np.complex128)
    The input data
axes : list of integers
349
    The axes along which the FFT is carried out.
350
    If not set, all axes will be transformed.
351
352
353
354
355
356
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.
357
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
358
    if False, returns the result in a new array and leaves the input unchanged.
359
    if True, stores the result in the input array and returns a handle to it.
Martin Reinecke's avatar
Martin Reinecke committed
360
361
362
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
363
364
365
366

Returns
-------
np.ndarray (same shape and data type as a)
Martin Reinecke's avatar
Martin Reinecke committed
367
    The transformed data
Martin Reinecke's avatar
Martin Reinecke committed
368
)""";
369

Martin Reinecke's avatar
Martin Reinecke committed
370
const char *rfftn_DS = R"""(Performs a forward real-valued FFT.
Martin Reinecke's avatar
Martin Reinecke committed
371
372
373
374
375
376

Parameters
----------
a : numpy.ndarray (np.float32 or np.float64)
    The input data
axes : list of integers
377
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
378
    If not set, all axes will be transformed in ascending order.
379
380
381
382
383
384
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
385
386
387
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
388
389
390
391

Returns
-------
np.ndarray (np.complex64 or np.complex128)
Martin Reinecke's avatar
Martin Reinecke committed
392
    The transformed data. The shape is identical to that of the input array,
Martin Reinecke's avatar
Martin Reinecke committed
393
394
    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
395
)""";
Martin Reinecke's avatar
Martin Reinecke committed
396

Martin Reinecke's avatar
Martin Reinecke committed
397
const char *rfft_scipy_DS = R"""(Performs a forward real-valued FFT.
398
399
400
401
402
403
404

Parameters
----------
a : numpy.ndarray (np.float32 or np.float64)
    The input data
axis : int
    The axis along which the FFT is carried out.
405
406
407
408
409
410
inorm : int
    Normalization type
      0 : no normalization
      1 : divide by sqrt(N)
      2 : divide by N
    where N is the length of `axis`.
411
412
413
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
414
415
416
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
417
418
419
420
421
422
423

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 ...`.
Martin Reinecke's avatar
Martin Reinecke committed
424
)""";
425

Martin Reinecke's avatar
Martin Reinecke committed
426
const char *irfftn_DS = R"""(Performs a backward real-valued FFT.
Martin Reinecke's avatar
Martin Reinecke committed
427
428
429
430
431
432

Parameters
----------
a : numpy.ndarray (np.complex64 or np.complex128)
    The input data
axes : list of integers
433
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
434
435
436
    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.
437
438
439
440
441
442
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
443
444
445
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
446
447
448
449

Returns
-------
np.ndarray (np.float32 or np.float64)
Martin Reinecke's avatar
Martin Reinecke committed
450
451
    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
452
    entries.
Martin Reinecke's avatar
Martin Reinecke committed
453
)""";
Martin Reinecke's avatar
Martin Reinecke committed
454

Martin Reinecke's avatar
Martin Reinecke committed
455
const char *irfft_scipy_DS = R"""(Performs a backward real-valued FFT.
456
457
458
459
460
461
462
463

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.
464
465
466
467
468
469
inorm : int
    Normalization type
      0 : no normalization
      1 : divide by sqrt(N)
      2 : divide by N
    where N is the length of `axis`.
470
471
472
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
473
474
475
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
476
477
478
479
480

Returns
-------
np.ndarray (np.float32 or np.float64)
    The transformed data. The shape is identical to that of the input array.
Martin Reinecke's avatar
Martin Reinecke committed
481
)""";
482

Martin Reinecke's avatar
Martin Reinecke committed
483
const char *hartley_DS = R"""(Performs a Hartley transform.
484
485
486
487
488
489
490
491
492
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
493
    The axes along which the transform is carried out.
494
    If not set, all axes will be transformed.
495
496
497
498
499
500
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.
501
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
502
    if False, returns the result in a new array and leaves the input unchanged.
503
    if True, stores the result in the input array and returns a handle to it.
Martin Reinecke's avatar
Martin Reinecke committed
504
505
506
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
507
508
509
510

Returns
-------
np.ndarray (same shape and data type as a)
Martin Reinecke's avatar
Martin Reinecke committed
511
    The transformed data
Martin Reinecke's avatar
Martin Reinecke committed
512
)""";
513

Martin Reinecke's avatar
Martin Reinecke committed
514
515
516
517
} // unnamed namespace

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

Martin Reinecke's avatar
Martin Reinecke committed
520
  m.doc() = pypocketfft_DS;
521
  m.def("fftn",&fftn, fftn_DS, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
522
    "inplace"_a=false, "nthreads"_a=1);
523
  m.def("ifftn",&ifftn, ifftn_DS, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
524
    "inplace"_a=false, "nthreads"_a=1);
525
  m.def("rfftn",&rfftn, rfftn_DS, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
526
    "nthreads"_a=1);
527
  m.def("rfft_scipy",&rfft_scipy, rfft_scipy_DS, "a"_a, "axis"_a, "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
528
    "inplace"_a=false, "nthreads"_a=1);
529
  m.def("irfftn",&irfftn, irfftn_DS, "a"_a, "axes"_a=py::none(), "lastsize"_a=0,
530
531
532
533
    "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
534
    "inplace"_a=false, "nthreads"_a=1);
535
  m.def("hartley2",&hartley2, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
536
    "inplace"_a=false, "nthreads"_a=1);
Martin Reinecke's avatar
Martin Reinecke committed
537
538
  m.def("complex2hartley",&mycomplex2hartley, "in"_a, "tmp"_a, "axes"_a,
    "inplace"_a=false);
Martin Reinecke's avatar
Martin Reinecke committed
539
  }