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

Martin Reinecke's avatar
Martin Reinecke committed
30
31
auto c64 = py::dtype("complex64");
auto c128 = py::dtype("complex128");
32
auto c256 = py::dtype("complex256");
Martin Reinecke's avatar
Martin Reinecke committed
33
34
auto f32 = py::dtype("float32");
auto f64 = py::dtype("float64");
35
auto f128 = py::dtype("float128");
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
36
37

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

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

Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
53
shape_t makeaxes(const py::array &in, py::object axes)
Martin Reinecke's avatar
more    
Martin Reinecke committed
54
  {
Martin Reinecke's avatar
Martin Reinecke committed
55
  if (axes.is(py::none()))
Martin Reinecke's avatar
more    
Martin Reinecke committed
56
    {
57
    shape_t res(size_t(in.ndim()));
Martin Reinecke's avatar
Martin Reinecke committed
58
59
60
    for (size_t i=0; i<res.size(); ++i)
      res[i]=i;
    return res;
Martin Reinecke's avatar
more    
Martin Reinecke committed
61
    }
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
62
  auto tmp=axes.cast<shape_t>();
Martin Reinecke's avatar
fix    
Martin Reinecke committed
63
  if ((tmp.size()>size_t(in.ndim())) || (tmp.size()==0))
Martin Reinecke's avatar
more    
Martin Reinecke committed
64
    throw runtime_error("bad axes argument");
Martin Reinecke's avatar
Martin Reinecke committed
65
66
  for (auto sz: tmp)
    if (sz>=size_t(in.ndim()))
Martin Reinecke's avatar
more    
Martin Reinecke committed
67
      throw runtime_error("invalid axis number");
Martin Reinecke's avatar
Martin Reinecke committed
68
  return tmp;
Martin Reinecke's avatar
more    
Martin Reinecke committed
69
70
  }

71
72
73
74
#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; \
75
  if (dtype.is(T3)) return func<long double> args; \
76
77
  throw runtime_error("unsupported data type");

Martin Reinecke's avatar
Martin Reinecke committed
78
template<typename T> py::array xfftn_internal(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
79
  const shape_t &axes, long double fct, bool inplace, bool fwd, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
80
  {
Martin Reinecke's avatar
Martin Reinecke committed
81
  auto dims(copy_shape(in));
82
  py::array res = inplace ? in : py::array_t<complex<T>>(dims);
Martin Reinecke's avatar
Martin Reinecke committed
83
84
85
86
87
88
89
90
  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;
  c2c(dims, s_in, s_out, axes, fwd, d_in, d_out, T(fct), nthreads);
  }
Martin Reinecke's avatar
Martin Reinecke committed
91
92
  return res;
  }
Martin Reinecke's avatar
Martin Reinecke committed
93

94
py::array xfftn(const py::array &a, py::object axes, double fct, bool inplace,
Martin Reinecke's avatar
Martin Reinecke committed
95
  bool fwd, size_t nthreads)
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
96
  {
97
  DISPATCH(a, c128, c64, c256, xfftn_internal, (a, makeaxes(a, axes), fct,
Martin Reinecke's avatar
Martin Reinecke committed
98
           inplace, fwd, nthreads))
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
99
  }
100

Martin Reinecke's avatar
Martin Reinecke committed
101
102
py::array fftn(const py::array &a, py::object axes, double fct, bool inplace,
  size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
103
  { return xfftn(a, axes, fct, inplace, true, nthreads); }
104

Martin Reinecke's avatar
Martin Reinecke committed
105
106
py::array ifftn(const py::array &a, py::object axes, double fct, bool inplace,
  size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
107
  { return xfftn(a, axes, fct, inplace, false, nthreads); }
Martin Reinecke's avatar
Martin Reinecke committed
108

Martin Reinecke's avatar
Martin Reinecke committed
109
template<typename T> py::array rfftn_internal(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
110
  py::object axes_, long double fct, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
111
  {
Martin Reinecke's avatar
Martin Reinecke committed
112
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
Martin Reinecke committed
113
  auto dims_in(copy_shape(in)), dims_out(dims_in);
Martin Reinecke's avatar
more    
Martin Reinecke committed
114
  dims_out[axes.back()] = (dims_out[axes.back()]>>1)+1;
Martin Reinecke's avatar
Martin Reinecke committed
115
  py::array res = py::array_t<complex<T>>(dims_out);
Martin Reinecke's avatar
Martin Reinecke committed
116
117
118
119
120
121
122
123
  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;
  r2c(dims_in, s_in, s_out, axes, d_in, d_out, T(fct), nthreads);
  }
Martin Reinecke's avatar
more    
Martin Reinecke committed
124
125
  return res;
  }
126

Martin Reinecke's avatar
Martin Reinecke committed
127
128
py::array rfftn(const py::array &in, py::object axes_, double fct,
  size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
129
  {
Martin Reinecke's avatar
Martin Reinecke committed
130
  DISPATCH(in, f64, f32, f128, rfftn_internal, (in, axes_, fct, nthreads))
Martin Reinecke's avatar
Martin Reinecke committed
131
  }
132

133
template<typename T> py::array xrfft_scipy(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
134
  size_t axis, long double fct, bool inplace, bool fwd, size_t nthreads)
135
136
137
  {
  auto dims(copy_shape(in));
  py::array res = inplace ? in : py::array_t<T>(dims);
Martin Reinecke's avatar
Martin Reinecke committed
138
139
140
141
142
143
144
145
  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;
  r2r_fftpack(dims, s_in, s_out, axis, fwd, d_in, d_out, T(fct), nthreads);
  }
146
147
  return res;
  }
148

Martin Reinecke's avatar
Martin Reinecke committed
149
150
py::array rfft_scipy(const py::array &in, size_t axis, double fct, bool inplace,
  size_t nthreads)
151
  {
Martin Reinecke's avatar
Martin Reinecke committed
152
153
  DISPATCH(in, f64, f32, f128, xrfft_scipy, (in, axis, fct, inplace, true,
    nthreads))
154
  }
155

Martin Reinecke's avatar
Martin Reinecke committed
156
py::array irfft_scipy(const py::array &in, size_t axis, double fct,
Martin Reinecke's avatar
Martin Reinecke committed
157
  bool inplace, size_t nthreads)
158
  {
Martin Reinecke's avatar
Martin Reinecke committed
159
160
  DISPATCH(in, f64, f32, f128, xrfft_scipy, (in, axis, fct, inplace, false,
    nthreads))
161
  }
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
162
template<typename T> py::array irfftn_internal(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
163
  py::object axes_, size_t lastsize, long double fct, size_t nthreads)
Martin Reinecke's avatar
more    
Martin Reinecke committed
164
  {
Martin Reinecke's avatar
Martin Reinecke committed
165
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
more    
Martin Reinecke committed
166
  size_t axis = axes.back();
Martin Reinecke's avatar
Martin Reinecke committed
167
168
169
  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
170
171
    throw runtime_error("bad lastsize");
  dims_out[axis] = lastsize;
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
172
  py::array res = py::array_t<T>(dims_out);
Martin Reinecke's avatar
Martin Reinecke committed
173
174
175
176
177
178
179
180
  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;
  c2r(dims_out, s_in, s_out, axes, d_in, d_out, T(fct), nthreads);
  }
Martin Reinecke's avatar
more    
Martin Reinecke committed
181
182
  return res;
  }
183

Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
184
py::array irfftn(const py::array &in, py::object axes_, size_t lastsize,
Martin Reinecke's avatar
Martin Reinecke committed
185
  double fct, size_t nthreads)
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
186
  {
Martin Reinecke's avatar
Martin Reinecke committed
187
188
  DISPATCH(in, c128, c64, c256, irfftn_internal, (in, axes_, lastsize, fct,
    nthreads))
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
189
  }
Martin Reinecke's avatar
Martin Reinecke committed
190
191

template<typename T> py::array hartley_internal(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
192
  py::object axes_, long double fct, bool inplace, size_t nthreads)
Martin Reinecke's avatar
more    
Martin Reinecke committed
193
  {
Martin Reinecke's avatar
Martin Reinecke committed
194
  auto dims(copy_shape(in));
195
  py::array res = inplace ? in : py::array_t<T>(dims);
Martin Reinecke's avatar
Martin Reinecke committed
196
197
198
199
200
201
202
203
204
  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;
  r2r_hartley(dims, s_in, s_out, axes, d_in, d_out, T(fct), nthreads);
  }
Martin Reinecke's avatar
Martin Reinecke committed
205
206
  return res;
  }
207

208
py::array hartley(const py::array &in, py::object axes_, double fct,
Martin Reinecke's avatar
Martin Reinecke committed
209
  bool inplace, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
210
  {
Martin Reinecke's avatar
Martin Reinecke committed
211
212
  DISPATCH(in, f64, f32, f128, hartley_internal, (in, axes_, fct, inplace,
    nthreads))
Martin Reinecke's avatar
Martin Reinecke committed
213
  }
214

Martin Reinecke's avatar
fixes    
Martin Reinecke committed
215
template<typename T>py::array complex2hartley(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
216
  const py::array &tmp, py::object axes_, bool inplace)
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
217
  {
Martin Reinecke's avatar
Martin Reinecke committed
218
  using namespace pocketfft::detail;
219
  size_t ndim = size_t(in.ndim());
Martin Reinecke's avatar
Martin Reinecke committed
220
  auto dims_out(copy_shape(in));
Martin Reinecke's avatar
Martin Reinecke committed
221
222
223
  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
224
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
Martin Reinecke committed
225
226
  {
  py::gil_scoped_release release;
Martin Reinecke's avatar
Martin Reinecke committed
227
  size_t axis = axes.back();
Martin Reinecke's avatar
Martin Reinecke committed
228
229
  multi_iter<1,cmplx<T>,T> it(atmp, aout, axis);
  vector<bool> swp(ndim,false);
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
230
231
  for (auto i: axes)
    if (i!=axis)
Martin Reinecke's avatar
Martin Reinecke committed
232
      swp[i] = true;
Martin Reinecke's avatar
Martin Reinecke committed
233
  while(it.remaining()>0)
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
234
    {
Martin Reinecke's avatar
Martin Reinecke committed
235
    ptrdiff_t rofs = 0;
Martin Reinecke's avatar
Martin Reinecke committed
236
    for (size_t i=0; i<it.pos.size(); ++i)
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
237
      {
Martin Reinecke's avatar
Martin Reinecke committed
238
      if (i==axis) continue;
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
239
      if (!swp[i])
240
        rofs += ptrdiff_t(it.pos[i])*it.oarr.stride(i);
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
241
242
      else
        {
243
        auto x = ptrdiff_t((it.pos[i]==0) ? 0 : it.iarr.shape(i)-it.pos[i]);
Martin Reinecke's avatar
Martin Reinecke committed
244
        rofs += x*it.oarr.stride(i);
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
245
246
        }
      }
Martin Reinecke's avatar
Martin Reinecke committed
247
    it.advance(1);
Martin Reinecke's avatar
Martin Reinecke committed
248
    for (size_t i=0; i<it.length_in(); ++i)
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
249
      {
Martin Reinecke's avatar
Martin Reinecke committed
250
251
      auto re = it.in(i).r;
      auto im = it.in(i).i;
252
      auto rev_i = ptrdiff_t((i==0) ? 0 : it.length_out()-i);
Martin Reinecke's avatar
Martin Reinecke committed
253
254
      it.out(i) = re+im;
      aout[rofs + rev_i*it.stride_out()] = re-im;
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
255
256
      }
    }
Martin Reinecke's avatar
Martin Reinecke committed
257
  }
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
258
259
  return out;
  }
260

261
py::array mycomplex2hartley(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
262
  const py::array &tmp, py::object axes_, bool inplace)
263
  {
264
  DISPATCH(in, f64, f32, f128, complex2hartley, (in, tmp, axes_, inplace))
265
  }
266

Martin Reinecke's avatar
Martin Reinecke committed
267
py::array hartley2(const py::array &in, py::object axes_, double fct,
Martin Reinecke's avatar
Martin Reinecke committed
268
  bool inplace, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
269
270
271
272
  {
  return mycomplex2hartley(in, rfftn(in, axes_, fct, nthreads), axes_,
    inplace);
  }
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
273

Martin Reinecke's avatar
Martin Reinecke committed
274
const char *pypocketfft_DS = R"""(Fast Fourier and Hartley transforms.
Martin Reinecke's avatar
Martin Reinecke committed
275
276
277
278
279
280
281
282
283

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

Martin Reinecke's avatar
Martin Reinecke committed
286
const char *fftn_DS = R"""(
287
288
289
290
291
292
293
Performs a forward complex FFT.

Parameters
----------
a : numpy.ndarray (np.complex64 or np.complex128)
    The input data
axes : list of integers
294
    The axes along which the FFT is carried out.
295
    If not set, all axes will be transformed.
Martin Reinecke's avatar
Martin Reinecke committed
296
fct : float
Martin Reinecke's avatar
Martin Reinecke committed
297
    Normalization factor
298
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
299
    if False, returns the result in a new array and leaves the input unchanged.
300
    if True, stores the result in the input array and returns a handle to it.
Martin Reinecke's avatar
Martin Reinecke committed
301
302
303
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
304
305
306
307

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

Martin Reinecke's avatar
Martin Reinecke committed
311
const char *ifftn_DS = R"""(Performs a backward complex FFT.
312
313
314
315
316
317

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

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

Martin Reinecke's avatar
Martin Reinecke committed
335
const char *rfftn_DS = R"""(Performs a forward real-valued FFT.
Martin Reinecke's avatar
Martin Reinecke committed
336
337
338
339
340
341

Parameters
----------
a : numpy.ndarray (np.float32 or np.float64)
    The input data
axes : list of integers
342
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
343
344
    If not set, all axes will be transformed in ascending order.
fct : float
Martin Reinecke's avatar
Martin Reinecke committed
345
    Normalization factor
Martin Reinecke's avatar
Martin Reinecke committed
346
347
348
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
349
350
351
352

Returns
-------
np.ndarray (np.complex64 or np.complex128)
Martin Reinecke's avatar
Martin Reinecke committed
353
    The transformed data. The shape is identical to that of the input array,
Martin Reinecke's avatar
Martin Reinecke committed
354
355
    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
356
)""";
Martin Reinecke's avatar
Martin Reinecke committed
357

Martin Reinecke's avatar
Martin Reinecke committed
358
const char *rfft_scipy_DS = R"""(Performs a forward real-valued FFT.
359
360
361
362
363
364
365
366
367
368
369
370

Parameters
----------
a : numpy.ndarray (np.float32 or np.float64)
    The input data
axis : int
    The axis along which the FFT is carried out.
fct : float
    Normalization factor
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
371
372
373
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
374
375
376
377
378
379
380

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
381
)""";
382

Martin Reinecke's avatar
Martin Reinecke committed
383
const char *irfftn_DS = R"""(Performs a backward real-valued FFT.
Martin Reinecke's avatar
Martin Reinecke committed
384
385
386
387
388
389

Parameters
----------
a : numpy.ndarray (np.complex64 or np.complex128)
    The input data
axes : list of integers
390
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
391
392
393
394
    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.
fct : float
Martin Reinecke's avatar
Martin Reinecke committed
395
    Normalization factor
Martin Reinecke's avatar
Martin Reinecke committed
396
397
398
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
399
400
401
402

Returns
-------
np.ndarray (np.float32 or np.float64)
Martin Reinecke's avatar
Martin Reinecke committed
403
404
    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
405
    entries.
Martin Reinecke's avatar
Martin Reinecke committed
406
)""";
Martin Reinecke's avatar
Martin Reinecke committed
407

Martin Reinecke's avatar
Martin Reinecke committed
408
const char *irfft_scipy_DS = R"""(Performs a backward real-valued FFT.
409
410
411
412
413
414
415
416
417
418
419
420
421

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.
fct : float
    Normalization factor
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
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).
425
426
427
428
429

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
430
)""";
431

Martin Reinecke's avatar
Martin Reinecke committed
432
const char *hartley_DS = R"""(Performs a Hartley transform.
433
434
435
436
437
438
439
440
441
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
442
    The axes along which the transform is carried out.
443
    If not set, all axes will be transformed.
Martin Reinecke's avatar
Martin Reinecke committed
444
fct : float
Martin Reinecke's avatar
Martin Reinecke committed
445
    Normalization factor
446
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
447
    if False, returns the result in a new array and leaves the input unchanged.
448
    if True, stores the result in the input array and returns a handle to it.
Martin Reinecke's avatar
Martin Reinecke committed
449
450
451
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
452
453
454
455

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

Martin Reinecke's avatar
Martin Reinecke committed
459
460
461
462
} // unnamed namespace

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

Martin Reinecke's avatar
Martin Reinecke committed
465
  m.doc() = pypocketfft_DS;
466
  m.def("fftn",&fftn, fftn_DS, "a"_a, "axes"_a=py::none(), "fct"_a=1.,
Martin Reinecke's avatar
Martin Reinecke committed
467
    "inplace"_a=false, "nthreads"_a=1);
468
  m.def("ifftn",&ifftn, ifftn_DS, "a"_a, "axes"_a=py::none(), "fct"_a=1.,
Martin Reinecke's avatar
Martin Reinecke committed
469
    "inplace"_a=false, "nthreads"_a=1);
Martin Reinecke's avatar
Martin Reinecke committed
470
471
  m.def("rfftn",&rfftn, rfftn_DS, "a"_a, "axes"_a=py::none(), "fct"_a=1.,
    "nthreads"_a=1);
472
  m.def("rfft_scipy",&rfft_scipy, rfft_scipy_DS, "a"_a, "axis"_a, "fct"_a=1.,
Martin Reinecke's avatar
Martin Reinecke committed
473
    "inplace"_a=false, "nthreads"_a=1);
474
  m.def("irfftn",&irfftn, irfftn_DS, "a"_a, "axes"_a=py::none(), "lastsize"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
475
    "fct"_a=1., "nthreads"_a=1);
476
  m.def("irfft_scipy",&irfft_scipy, irfft_scipy_DS, "a"_a, "axis"_a, "fct"_a=1.,
Martin Reinecke's avatar
Martin Reinecke committed
477
    "inplace"_a=false, "nthreads"_a=1);
478
  m.def("hartley",&hartley, hartley_DS, "a"_a, "axes"_a=py::none(), "fct"_a=1.,
Martin Reinecke's avatar
Martin Reinecke committed
479
    "inplace"_a=false, "nthreads"_a=1);
Martin Reinecke's avatar
Martin Reinecke committed
480
  m.def("hartley2",&hartley2, "a"_a, "axes"_a=py::none(), "fct"_a=1.,
Martin Reinecke's avatar
Martin Reinecke committed
481
    "inplace"_a=false, "nthreads"_a=1);
Martin Reinecke's avatar
Martin Reinecke committed
482
483
  m.def("complex2hartley",&mycomplex2hartley, "in"_a, "tmp"_a, "axes"_a,
    "inplace"_a=false);
Martin Reinecke's avatar
Martin Reinecke committed
484
  }