pypocketfft.cc 17.3 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

Martin Reinecke's avatar
Martin Reinecke committed
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
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))
    }
  }
164

165
166
167
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
168

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

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

194
template<typename T> py::array xrfft_scipy(const py::array &in,
195
  size_t axis, int inorm, bool inplace, bool fwd, size_t nthreads)
196
197
198
  {
  auto dims(copy_shape(in));
  py::array res = inplace ? in : py::array_t<T>(dims);
Martin Reinecke's avatar
Martin Reinecke committed
199
200
201
202
203
204
  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;
205
206
  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
207
  }
208
209
  return res;
  }
210

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

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

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

template<typename T> py::array hartley_internal(const py::array &in,
255
  py::object axes_, int inorm, bool inplace, size_t nthreads)
Martin Reinecke's avatar
more  
Martin Reinecke committed
256
  {
Martin Reinecke's avatar
Martin Reinecke committed
257
  auto dims(copy_shape(in));
258
  py::array res = inplace ? in : py::array_t<T>(dims);
Martin Reinecke's avatar
Martin Reinecke committed
259
260
261
262
263
264
265
  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;
266
267
  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
268
  }
Martin Reinecke's avatar
Martin Reinecke committed
269
270
  return res;
  }
271

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

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

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

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

Martin Reinecke's avatar
Martin Reinecke committed
317
const char *pypocketfft_DS = R"""(Fast Fourier and Hartley transforms.
Martin Reinecke's avatar
Martin Reinecke committed
318
319
320
321
322
323
324
325
326

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
327
)""";
Martin Reinecke's avatar
Martin Reinecke committed
328

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

Parameters
----------
334
335
336
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.
337
axes : list of integers
338
    The axes along which the FFT is carried out.
339
    If not set, all axes will be transformed.
340
341
342
343
344
345
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.
346
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
347
    if False, returns the result in a new array and leaves the input unchanged.
348
    if True, stores the result in the input array and returns a handle to it.
349
350
    NOTE: if `a` is real-valued and `inplace` is `True`, an exception will be
    raised!
Martin Reinecke's avatar
Martin Reinecke committed
351
352
353
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
354
355
356

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

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

Parameters
----------
365
a : numpy.ndarray (any complex type)
366
367
    The input data
axes : list of integers
368
    The axes along which the FFT is carried out.
369
    If not set, all axes will be transformed.
370
371
372
373
374
375
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.
376
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
377
    if False, returns the result in a new array and leaves the input unchanged.
378
    if True, stores the result in the input array and returns a handle to it.
Martin Reinecke's avatar
Martin Reinecke committed
379
380
381
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
382
383
384
385

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

Martin Reinecke's avatar
Martin Reinecke committed
389
const char *rfftn_DS = R"""(Performs a forward real-valued FFT.
Martin Reinecke's avatar
Martin Reinecke committed
390
391
392

Parameters
----------
393
a : numpy.ndarray (any real type)
Martin Reinecke's avatar
Martin Reinecke committed
394
395
    The input data
axes : list of integers
396
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
397
    If not set, all axes will be transformed in ascending order.
398
399
400
401
402
403
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
404
405
406
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
407
408
409

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

Martin Reinecke's avatar
Martin Reinecke committed
416
const char *rfft_scipy_DS = R"""(Performs a forward real-valued FFT.
417
418
419

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

Returns
-------
439
np.ndarray (same shape and data type as `a`)
440
441
442
    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
443
)""";
444

Martin Reinecke's avatar
Martin Reinecke committed
445
const char *irfftn_DS = R"""(Performs a backward real-valued FFT.
Martin Reinecke's avatar
Martin Reinecke committed
446
447
448

Parameters
----------
449
a : numpy.ndarray (any complex type)
Martin Reinecke's avatar
Martin Reinecke committed
450
451
    The input data
axes : list of integers
452
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
453
454
455
    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.
456
457
458
459
460
461
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
462
463
464
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
465
466
467

Returns
-------
468
np.ndarray (real type with same accuracy as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
469
470
    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
471
    entries.
Martin Reinecke's avatar
Martin Reinecke committed
472
)""";
Martin Reinecke's avatar
Martin Reinecke committed
473

Martin Reinecke's avatar
Martin Reinecke committed
474
const char *irfft_scipy_DS = R"""(Performs a backward real-valued FFT.
475
476
477

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

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

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

Returns
-------
529
np.ndarray (same shape and data type as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
530
    The transformed data
Martin Reinecke's avatar
Martin Reinecke committed
531
)""";
532

Martin Reinecke's avatar
Martin Reinecke committed
533
534
535
536
} // unnamed namespace

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

Martin Reinecke's avatar
Martin Reinecke committed
539
  m.doc() = pypocketfft_DS;
540
  m.def("fftn",&fftn, fftn_DS, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
541
    "inplace"_a=false, "nthreads"_a=1);
542
  m.def("ifftn",&ifftn, ifftn_DS, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
543
    "inplace"_a=false, "nthreads"_a=1);
544
  m.def("rfftn",&rfftn, rfftn_DS, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
545
    "nthreads"_a=1);
546
  m.def("rfft_scipy",&rfft_scipy, rfft_scipy_DS, "a"_a, "axis"_a, "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
547
    "inplace"_a=false, "nthreads"_a=1);
548
  m.def("irfftn",&irfftn, irfftn_DS, "a"_a, "axes"_a=py::none(), "lastsize"_a=0,
549
550
551
552
    "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
553
    "inplace"_a=false, "nthreads"_a=1);
554
  m.def("hartley2",&hartley2, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
555
    "inplace"_a=false, "nthreads"_a=1);
Martin Reinecke's avatar
Martin Reinecke committed
556
557
  m.def("complex2hartley",&mycomplex2hartley, "in"_a, "tmp"_a, "axes"_a,
    "inplace"_a=false);
Martin Reinecke's avatar
Martin Reinecke committed
558
  }