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
    }
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
#define DISPATCH(arr, T1, T2, T3, func, args) \
76
  { \
77
78
79
  auto dtype = arr.dtype(); \
  if (dtype.is(T1)) return func<double> args; \
  if (dtype.is(T2)) return func<float> args; \
80
  if (dtype.is(T3)) return func<ldbl_t> args; \
81
82
  throw runtime_error("unsupported data type"); \
  }
83

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

Martin Reinecke's avatar
Martin Reinecke committed
119
template<typename T> py::array sym_rfftn_internal(const py::array &in,
120
  const shape_t &axes, int inorm, bool fwd, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
121
122
123
124
125
126
127
128
129
130
131
132
133
134
  {
  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);
135
  rev_iter iter(ares, axes);
Martin Reinecke's avatar
Martin Reinecke committed
136
137
  while(iter.remaining()>0)
    {
138
139
    auto v = ares[iter.ofs()];
    ares[iter.rev_ofs()] = conj(v);
Martin Reinecke's avatar
Martin Reinecke committed
140
141
    iter.advance();
    }
142
143
144
145
146
147
148
149
150
151
  if (!fwd)
    {
    simple_iter iter(ares);
    while(iter.remaining()>0)
      {
      auto v = ares[iter.ofs()];
      ares[iter.ofs()] = conj(v);
      iter.advance();
      }
    }
Martin Reinecke's avatar
Martin Reinecke committed
152
153
154
155
  }
  return res;
  }

156
157
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
158
  {
159
160
161
  if (a.dtype().kind() == 'c')
    DISPATCH(a, c128, c64, clong, xfftn_internal, (a, makeaxes(a, axes), inorm,
             inplace, fwd, nthreads))
162

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

168
169
170
171
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); }

172
173
174
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
175

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

195
py::array rfftn(const py::array &in, py::object axes_, int inorm,
Martin Reinecke's avatar
Martin Reinecke committed
196
  size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
197
  {
198
  DISPATCH(in, f64, f32, flong, rfftn_internal, (in, axes_, inorm, nthreads))
Martin Reinecke's avatar
Martin Reinecke committed
199
  }
200

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

218
219
py::array rfft_scipy(const py::array &in, size_t axis, int inorm,
  bool inplace, size_t nthreads)
220
  {
221
  DISPATCH(in, f64, f32, flong, xrfft_scipy, (in, axis, inorm, inplace, true,
Martin Reinecke's avatar
Martin Reinecke committed
222
    nthreads))
223
  }
224

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

Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
254
py::array irfftn(const py::array &in, py::object axes_, size_t lastsize,
255
  int inorm, size_t nthreads)
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
256
  {
257
  DISPATCH(in, c128, c64, clong, irfftn_internal, (in, axes_, lastsize, inorm,
Martin Reinecke's avatar
Martin Reinecke committed
258
    nthreads))
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
259
  }
Martin Reinecke's avatar
Martin Reinecke committed
260
261

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

279
py::array hartley(const py::array &in, py::object axes_, int inorm,
Martin Reinecke's avatar
Martin Reinecke committed
280
  bool inplace, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
281
  {
282
  DISPATCH(in, f64, f32, flong, hartley_internal, (in, axes_, inorm, inplace,
Martin Reinecke's avatar
Martin Reinecke committed
283
    nthreads))
Martin Reinecke's avatar
Martin Reinecke committed
284
  }
285

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

311
py::array mycomplex2hartley(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
312
  const py::array &tmp, py::object axes_, bool inplace)
313
  {
314
  DISPATCH(in, f64, f32, flong, complex2hartley, (in, tmp, axes_, inplace))
315
  }
316

317
py::array hartley2(const py::array &in, py::object axes_, int inorm,
Martin Reinecke's avatar
Martin Reinecke committed
318
  bool inplace, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
319
  {
320
  return mycomplex2hartley(in, rfftn(in, axes_, inorm, nthreads), axes_,
Martin Reinecke's avatar
Martin Reinecke committed
321
322
    inplace);
  }
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
323

Martin Reinecke's avatar
Martin Reinecke committed
324
const char *pypocketfft_DS = R"""(Fast Fourier and Hartley transforms.
Martin Reinecke's avatar
Martin Reinecke committed
325
326
327
328
329
330
331
332
333

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

Martin Reinecke's avatar
Martin Reinecke committed
336
const char *fftn_DS = R"""(
337
338
339
340
Performs a forward complex FFT.

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

Returns
-------
364
np.ndarray (same shape as `a`, complex type with same accuracy as `a`)
365
    The transformed data.
Martin Reinecke's avatar
Martin Reinecke committed
366
)""";
367

Martin Reinecke's avatar
Martin Reinecke committed
368
const char *ifftn_DS = R"""(Performs a backward complex FFT.
369
370
371

Parameters
----------
372
373
374
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.
375
axes : list of integers
376
    The axes along which the FFT is carried out.
377
    If not set, all axes will be transformed.
378
379
380
381
382
383
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.
384
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
385
    if False, returns the result in a new array and leaves the input unchanged.
386
    if True, stores the result in the input array and returns a handle to it.
387
388
    NOTE: if `a` is real-valued and `inplace` is `True`, an exception will be
    raised!
Martin Reinecke's avatar
Martin Reinecke committed
389
390
391
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
392
393
394

Returns
-------
395
np.ndarray (same shape as `a`, complex type with same accuracy as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
396
    The transformed data
Martin Reinecke's avatar
Martin Reinecke committed
397
)""";
398

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

Parameters
----------
403
a : numpy.ndarray (any real type)
Martin Reinecke's avatar
Martin Reinecke committed
404
405
    The input data
axes : list of integers
406
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
407
    If not set, all axes will be transformed in ascending order.
408
409
410
411
412
413
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
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).
Martin Reinecke's avatar
Martin Reinecke committed
417
418
419

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

Martin Reinecke's avatar
Martin Reinecke committed
426
const char *rfft_scipy_DS = R"""(Performs a forward real-valued FFT.
427
428
429

Parameters
----------
430
a : numpy.ndarray (any real type)
431
432
433
    The input data
axis : int
    The axis along which the FFT is carried out.
434
435
436
437
438
439
inorm : int
    Normalization type
      0 : no normalization
      1 : divide by sqrt(N)
      2 : divide by N
    where N is the length of `axis`.
440
441
442
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
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).
446
447
448

Returns
-------
449
np.ndarray (same shape and data type as `a`)
450
451
452
    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
453
)""";
454

Martin Reinecke's avatar
Martin Reinecke committed
455
const char *irfftn_DS = R"""(Performs a backward real-valued FFT.
Martin Reinecke's avatar
Martin Reinecke committed
456
457
458

Parameters
----------
459
a : numpy.ndarray (any complex type)
Martin Reinecke's avatar
Martin Reinecke committed
460
461
    The input data
axes : list of integers
462
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
463
464
465
    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.
466
467
468
469
470
471
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
472
473
474
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
475
476
477

Returns
-------
478
np.ndarray (real type with same accuracy as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
479
480
    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
481
    entries.
Martin Reinecke's avatar
Martin Reinecke committed
482
)""";
Martin Reinecke's avatar
Martin Reinecke committed
483

Martin Reinecke's avatar
Martin Reinecke committed
484
const char *irfft_scipy_DS = R"""(Performs a backward real-valued FFT.
485
486
487

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

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

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

Returns
-------
539
np.ndarray (same shape and data type as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
540
    The transformed data
Martin Reinecke's avatar
Martin Reinecke committed
541
)""";
542

Martin Reinecke's avatar
Martin Reinecke committed
543
544
545
546
} // unnamed namespace

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

Martin Reinecke's avatar
Martin Reinecke committed
549
  m.doc() = pypocketfft_DS;
550
  m.def("fftn",&fftn, fftn_DS, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
551
    "inplace"_a=false, "nthreads"_a=1);
552
  m.def("ifftn",&ifftn, ifftn_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("rfftn",&rfftn, rfftn_DS, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
555
    "nthreads"_a=1);
556
  m.def("rfft_scipy",&rfft_scipy, rfft_scipy_DS, "a"_a, "axis"_a, "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
557
    "inplace"_a=false, "nthreads"_a=1);
558
  m.def("irfftn",&irfftn, irfftn_DS, "a"_a, "axes"_a=py::none(), "lastsize"_a=0,
559
560
561
562
    "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
563
    "inplace"_a=false, "nthreads"_a=1);
564
  m.def("hartley2",&hartley2, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
565
    "inplace"_a=false, "nthreads"_a=1);
Martin Reinecke's avatar
Martin Reinecke committed
566
567
  m.def("complex2hartley",&mycomplex2hartley, "in"_a, "tmp"_a, "axes"_a,
    "inplace"_a=false);
Martin Reinecke's avatar
Martin Reinecke committed
568
  }