pypocketfft.cc 16 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;
Martin Reinecke's avatar
Martin Reinecke committed
246
  auto dims_out(copy_shape(in));
Martin Reinecke's avatar
Martin Reinecke committed
247
248
249
  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
250
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
Martin Reinecke committed
251
252
  {
  py::gil_scoped_release release;
253
254
255
256
  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
257
    {
258
259
260
261
    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
262
    }
Martin Reinecke's avatar
Martin Reinecke committed
263
  }
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
264
265
  return out;
  }
266

267
py::array mycomplex2hartley(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
268
  const py::array &tmp, py::object axes_, bool inplace)
269
  {
270
  DISPATCH(in, f64, f32, flong, complex2hartley, (in, tmp, axes_, inplace))
271
  }
272

273
py::array hartley2(const py::array &in, py::object axes_, int inorm,
Martin Reinecke's avatar
Martin Reinecke committed
274
  bool inplace, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
275
  {
276
  return mycomplex2hartley(in, rfftn(in, axes_, inorm, nthreads), axes_,
Martin Reinecke's avatar
Martin Reinecke committed
277
278
    inplace);
  }
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
279

Martin Reinecke's avatar
Martin Reinecke committed
280
const char *pypocketfft_DS = R"""(Fast Fourier and Hartley transforms.
Martin Reinecke's avatar
Martin Reinecke committed
281
282
283
284
285
286
287
288
289

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

Martin Reinecke's avatar
Martin Reinecke committed
292
const char *fftn_DS = R"""(
293
294
295
296
297
298
299
Performs a forward complex FFT.

Parameters
----------
a : numpy.ndarray (np.complex64 or np.complex128)
    The input data
axes : list of integers
300
    The axes along which the FFT is carried out.
301
    If not set, all axes will be transformed.
302
303
304
305
306
307
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.
308
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
309
    if False, returns the result in a new array and leaves the input unchanged.
310
    if True, stores the result in the input array and returns a handle to it.
Martin Reinecke's avatar
Martin Reinecke committed
311
312
313
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
314
315
316
317

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

Martin Reinecke's avatar
Martin Reinecke committed
321
const char *ifftn_DS = R"""(Performs a backward complex FFT.
322
323
324
325
326
327

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

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

Martin Reinecke's avatar
Martin Reinecke committed
349
const char *rfftn_DS = R"""(Performs a forward real-valued FFT.
Martin Reinecke's avatar
Martin Reinecke committed
350
351
352
353
354
355

Parameters
----------
a : numpy.ndarray (np.float32 or np.float64)
    The input data
axes : list of integers
356
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
357
    If not set, all axes will be transformed in ascending order.
358
359
360
361
362
363
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
364
365
366
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
367
368
369
370

Returns
-------
np.ndarray (np.complex64 or np.complex128)
Martin Reinecke's avatar
Martin Reinecke committed
371
    The transformed data. The shape is identical to that of the input array,
Martin Reinecke's avatar
Martin Reinecke committed
372
373
    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
374
)""";
Martin Reinecke's avatar
Martin Reinecke committed
375

Martin Reinecke's avatar
Martin Reinecke committed
376
const char *rfft_scipy_DS = R"""(Performs a forward real-valued FFT.
377
378
379
380
381
382
383

Parameters
----------
a : numpy.ndarray (np.float32 or np.float64)
    The input data
axis : int
    The axis along which the FFT is carried out.
384
385
386
387
388
389
inorm : int
    Normalization type
      0 : no normalization
      1 : divide by sqrt(N)
      2 : divide by N
    where N is the length of `axis`.
390
391
392
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
393
394
395
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
396
397
398
399
400
401
402

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
403
)""";
404

Martin Reinecke's avatar
Martin Reinecke committed
405
const char *irfftn_DS = R"""(Performs a backward real-valued FFT.
Martin Reinecke's avatar
Martin Reinecke committed
406
407
408
409
410
411

Parameters
----------
a : numpy.ndarray (np.complex64 or np.complex128)
    The input data
axes : list of integers
412
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
413
414
415
    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.
416
417
418
419
420
421
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
422
423
424
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
425
426
427
428

Returns
-------
np.ndarray (np.float32 or np.float64)
Martin Reinecke's avatar
Martin Reinecke committed
429
430
    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
431
    entries.
Martin Reinecke's avatar
Martin Reinecke committed
432
)""";
Martin Reinecke's avatar
Martin Reinecke committed
433

Martin Reinecke's avatar
Martin Reinecke committed
434
const char *irfft_scipy_DS = R"""(Performs a backward real-valued FFT.
435
436
437
438
439
440
441
442

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.
443
444
445
446
447
448
inorm : int
    Normalization type
      0 : no normalization
      1 : divide by sqrt(N)
      2 : divide by N
    where N is the length of `axis`.
449
450
451
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
452
453
454
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
455
456
457
458
459

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
460
)""";
461

Martin Reinecke's avatar
Martin Reinecke committed
462
const char *hartley_DS = R"""(Performs a Hartley transform.
463
464
465
466
467
468
469
470
471
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
472
    The axes along which the transform is carried out.
473
    If not set, all axes will be transformed.
474
475
476
477
478
479
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.
480
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
481
    if False, returns the result in a new array and leaves the input unchanged.
482
    if True, stores the result in the input array and returns a handle to it.
Martin Reinecke's avatar
Martin Reinecke committed
483
484
485
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
486
487
488
489

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

Martin Reinecke's avatar
Martin Reinecke committed
493
494
495
496
} // unnamed namespace

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

Martin Reinecke's avatar
Martin Reinecke committed
499
  m.doc() = pypocketfft_DS;
500
  m.def("fftn",&fftn, fftn_DS, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
501
    "inplace"_a=false, "nthreads"_a=1);
502
  m.def("ifftn",&ifftn, ifftn_DS, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
503
    "inplace"_a=false, "nthreads"_a=1);
504
  m.def("rfftn",&rfftn, rfftn_DS, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
505
    "nthreads"_a=1);
506
  m.def("rfft_scipy",&rfft_scipy, rfft_scipy_DS, "a"_a, "axis"_a, "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
507
    "inplace"_a=false, "nthreads"_a=1);
508
  m.def("irfftn",&irfftn, irfftn_DS, "a"_a, "axes"_a=py::none(), "lastsize"_a=0,
509
510
511
512
    "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
513
    "inplace"_a=false, "nthreads"_a=1);
514
  m.def("hartley2",&hartley2, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
515
    "inplace"_a=false, "nthreads"_a=1);
Martin Reinecke's avatar
Martin Reinecke committed
516
517
  m.def("complex2hartley",&mycomplex2hartley, "in"_a, "tmp"_a, "axes"_a,
    "inplace"_a=false);
Martin Reinecke's avatar
Martin Reinecke committed
518
  }