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

19
namespace {
20

21
using namespace std;
Martin Reinecke's avatar
Martin Reinecke committed
22
23
using pocketfft::shape_t;
using pocketfft::stride_t;
Martin Reinecke's avatar
Martin Reinecke committed
24

Martin Reinecke's avatar
Martin Reinecke committed
25
26
namespace py = pybind11;

27
28
29
30
// 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
31
32
auto c64 = py::dtype("complex64");
auto c128 = py::dtype("complex128");
33
auto clong = py::dtype("longcomplex");
Martin Reinecke's avatar
Martin Reinecke committed
34
35
auto f32 = py::dtype("float32");
auto f64 = py::dtype("float64");
36
auto flong = py::dtype("longfloat");
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
37
38

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

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

Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
54
shape_t makeaxes(const py::array &in, py::object axes)
Martin Reinecke's avatar
more    
Martin Reinecke committed
55
  {
Martin Reinecke's avatar
Martin Reinecke committed
56
  if (axes.is(py::none()))
Martin Reinecke's avatar
more    
Martin Reinecke committed
57
    {
58
    shape_t res(size_t(in.ndim()));
Martin Reinecke's avatar
Martin Reinecke committed
59
60
61
    for (size_t i=0; i<res.size(); ++i)
      res[i]=i;
    return res;
Martin Reinecke's avatar
more    
Martin Reinecke committed
62
    }
63
64
65
  auto tmp=axes.cast<std::vector<ptrdiff_t>>();
  auto ndim = in.ndim();
  if ((tmp.size()>size_t(ndim)) || (tmp.size()==0))
Martin Reinecke's avatar
more    
Martin Reinecke committed
66
    throw runtime_error("bad axes argument");
67
68
69
70
71
72
73
74
  for (auto& sz: tmp)
    {
    if (sz<0)
      sz += ndim;
    if ((sz>=ndim) || (sz<0))
      throw invalid_argument("axes exceeds dimensionality of output");
    }
  return shape_t(tmp.begin(), tmp.end());
Martin Reinecke's avatar
more    
Martin Reinecke committed
75
76
  }

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

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

Martin Reinecke's avatar
Martin Reinecke committed
121
122
template<typename T> py::array c2c_sym_internal(const py::array &in,
  const shape_t &axes, bool forward, int inorm, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
123
124
125
126
127
128
129
130
131
132
  {
  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);
Martin Reinecke's avatar
Martin Reinecke committed
133
  pocketfft::r2c(dims, s_in, s_out, axes, forward, d_in, d_out, fct, nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
134
135
136
  // now fill in second half
  using namespace pocketfft::detail;
  ndarr<complex<T>> ares(res.mutable_data(), dims, s_out);
137
  rev_iter iter(ares, axes);
Martin Reinecke's avatar
Martin Reinecke committed
138
139
  while(iter.remaining()>0)
    {
140
141
    auto v = ares[iter.ofs()];
    ares[iter.rev_ofs()] = conj(v);
Martin Reinecke's avatar
Martin Reinecke committed
142
143
144
145
146
147
    iter.advance();
    }
  }
  return res;
  }

Martin Reinecke's avatar
Martin Reinecke committed
148
149
py::array c2c(const py::array &a, py::object axes, bool forward, int inorm,
  bool inplace, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
150
  {
151
  if (a.dtype().kind() == 'c')
Martin Reinecke's avatar
Martin Reinecke committed
152
153
    DISPATCH(a, c128, c64, clong, c2c_internal, (a, makeaxes(a, axes), forward,
             inorm, inplace, nthreads))
154

155
  if (inplace) throw runtime_error("cannot do this operation in-place");
Martin Reinecke's avatar
Martin Reinecke committed
156
157
  DISPATCH(a, f64, f32, flong, c2c_sym_internal, (a, makeaxes(a, axes), forward,
           inorm, nthreads))
Martin Reinecke's avatar
Martin Reinecke committed
158
  }
159

Martin Reinecke's avatar
Martin Reinecke committed
160
161
template<typename T> py::array r2c_internal(const py::array &in,
  py::object axes_, bool forward, int inorm, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
162
  {
Martin Reinecke's avatar
Martin Reinecke committed
163
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
Martin Reinecke committed
164
  auto dims_in(copy_shape(in)), dims_out(dims_in);
Martin Reinecke's avatar
more    
Martin Reinecke committed
165
  dims_out[axes.back()] = (dims_out[axes.back()]>>1)+1;
Martin Reinecke's avatar
Martin Reinecke committed
166
  py::array res = py::array_t<complex<T>>(dims_out);
Martin Reinecke's avatar
Martin Reinecke committed
167
168
169
170
171
172
  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;
173
  T fct = norm_fct<T>(inorm, dims_in, axes);
174
175
  pocketfft::r2c(dims_in, s_in, s_out, axes, forward, d_in, d_out, fct,
    nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
176
  }
Martin Reinecke's avatar
more    
Martin Reinecke committed
177
178
  return res;
  }
179

Martin Reinecke's avatar
Martin Reinecke committed
180
py::array r2c(const py::array &in, py::object axes_, bool forward, int inorm,
Martin Reinecke's avatar
Martin Reinecke committed
181
  size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
182
  {
Martin Reinecke's avatar
Martin Reinecke committed
183
184
  DISPATCH(in, f64, f32, flong, r2c_internal, (in, axes_, forward, inorm,
    nthreads))
Martin Reinecke's avatar
Martin Reinecke committed
185
  }
186

Martin Reinecke's avatar
Martin Reinecke committed
187
188
189
template<typename T> py::array r2r_fftpack_internal(const py::array &in,
  py::object axes_, bool input_halfcomplex, bool forward, int inorm,
  bool inplace, size_t nthreads)
190
  {
Martin Reinecke's avatar
Martin Reinecke committed
191
  auto axes = makeaxes(in, axes_);
192
193
  auto dims(copy_shape(in));
  py::array res = inplace ? in : py::array_t<T>(dims);
Martin Reinecke's avatar
Martin Reinecke committed
194
195
196
197
198
199
  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;
Martin Reinecke's avatar
Martin Reinecke committed
200
  T fct = norm_fct<T>(inorm, dims, axes);
201
202
  pocketfft::r2r_fftpack(dims, s_in, s_out, axes, !input_halfcomplex, forward,
    d_in, d_out, fct, nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
203
  }
204
205
  return res;
  }
206

Martin Reinecke's avatar
Martin Reinecke committed
207
208
209
py::array r2r_fftpack(const py::array &in, py::object axes_,
  bool input_halfcomplex, bool forward, int inorm, bool inplace,
  size_t nthreads)
210
  {
Martin Reinecke's avatar
Martin Reinecke committed
211
212
  DISPATCH(in, f64, f32, flong, r2r_fftpack_internal, (in, axes_,
    input_halfcomplex, forward, inorm, inplace, nthreads))
213
  }
214

Martin Reinecke's avatar
Martin Reinecke committed
215
216
template<typename T> py::array c2r_internal(const py::array &in,
  py::object axes_, size_t lastsize, bool forward, int inorm, size_t nthreads)
Martin Reinecke's avatar
more    
Martin Reinecke committed
217
  {
Martin Reinecke's avatar
Martin Reinecke committed
218
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
more    
Martin Reinecke committed
219
  size_t axis = axes.back();
Martin Reinecke's avatar
Martin Reinecke committed
220
221
222
  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
223
224
    throw runtime_error("bad lastsize");
  dims_out[axis] = lastsize;
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
225
  py::array res = py::array_t<T>(dims_out);
Martin Reinecke's avatar
Martin Reinecke committed
226
227
228
229
230
231
  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;
232
  T fct = norm_fct<T>(inorm, dims_out, axes);
233
234
  pocketfft::c2r(dims_out, s_in, s_out, axes, forward, d_in, d_out, fct,
    nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
235
  }
Martin Reinecke's avatar
more    
Martin Reinecke committed
236
237
  return res;
  }
238

Martin Reinecke's avatar
Martin Reinecke committed
239
240
py::array c2r(const py::array &in, py::object axes_, size_t lastsize,
  bool forward, int inorm, size_t nthreads)
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
241
  {
Martin Reinecke's avatar
Martin Reinecke committed
242
243
  DISPATCH(in, c128, c64, clong, c2r_internal, (in, axes_, lastsize, forward,
    inorm, nthreads))
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
244
  }
Martin Reinecke's avatar
Martin Reinecke committed
245

246
template<typename T> py::array separable_hartley_internal(const py::array &in,
247
  py::object axes_, int inorm, bool inplace, size_t nthreads)
Martin Reinecke's avatar
more    
Martin Reinecke committed
248
  {
Martin Reinecke's avatar
Martin Reinecke committed
249
  auto dims(copy_shape(in));
250
  py::array res = inplace ? in : py::array_t<T>(dims);
Martin Reinecke's avatar
Martin Reinecke committed
251
252
253
254
255
256
257
  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;
258
  T fct = norm_fct<T>(inorm, dims, axes);
Martin Reinecke's avatar
Martin Reinecke committed
259
  pocketfft::r2r_hartley(dims, s_in, s_out, axes, d_in, d_out, fct, nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
260
  }
Martin Reinecke's avatar
Martin Reinecke committed
261
262
  return res;
  }
263

264
py::array separable_hartley(const py::array &in, py::object axes_, int inorm,
Martin Reinecke's avatar
Martin Reinecke committed
265
  bool inplace, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
266
  {
267
268
  DISPATCH(in, f64, f32, flong, separable_hartley_internal, (in, axes_, inorm,
    inplace, nthreads))
Martin Reinecke's avatar
Martin Reinecke committed
269
  }
270

Martin Reinecke's avatar
fixes    
Martin Reinecke committed
271
template<typename T>py::array complex2hartley(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
272
  const py::array &tmp, py::object axes_, bool inplace)
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
273
  {
Martin Reinecke's avatar
Martin Reinecke committed
274
  using namespace pocketfft::detail;
Martin Reinecke's avatar
Martin Reinecke committed
275
  auto dims_out(copy_shape(in));
Martin Reinecke's avatar
Martin Reinecke committed
276
  py::array out = inplace ? in : py::array_t<T>(dims_out);
277
  cndarr<cmplx<T>> atmp(tmp.data(), copy_shape(tmp), copy_strides(tmp));
Martin Reinecke's avatar
Martin Reinecke committed
278
  ndarr<T> aout(out.mutable_data(), copy_shape(out), copy_strides(out));
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
279
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
Martin Reinecke committed
280
281
  {
  py::gil_scoped_release release;
282
283
  simple_iter iin(atmp);
  rev_iter iout(aout, axes);
284
285
  if (iin.remaining()!=iout.remaining()) throw runtime_error("oops");
  while(iin.remaining()>0)
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
286
    {
287
288
289
    auto v = atmp[iin.ofs()];
    aout[iout.ofs()] = v.r+v.i;
    aout[iout.rev_ofs()] = v.r-v.i;
290
    iin.advance(); iout.advance();
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
291
    }
Martin Reinecke's avatar
Martin Reinecke committed
292
  }
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
293
294
  return out;
  }
295

296
py::array genuine_hartley(const py::array &in, py::object axes_, int inorm,
Martin Reinecke's avatar
Martin Reinecke committed
297
  bool inplace, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
298
  {
299
300
  auto tmp = r2c(in, axes_, true, inorm, nthreads);
  DISPATCH(in, f64, f32, flong, complex2hartley, (in, tmp, axes_, inplace))
Martin Reinecke's avatar
Martin Reinecke committed
301
  }
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
302

Martin Reinecke's avatar
Martin Reinecke committed
303
const char *pypocketfft_DS = R"""(Fast Fourier and Hartley transforms.
Martin Reinecke's avatar
Martin Reinecke committed
304
305

This module supports
306
- single, double, and long double precision
Martin Reinecke's avatar
Martin Reinecke committed
307
308
309
310
311
312
- 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
313
)""";
Martin Reinecke's avatar
Martin Reinecke committed
314

Martin Reinecke's avatar
Martin Reinecke committed
315
const char *c2c_DS = R"""(Performs a complex FFT.
316
317
318

Parameters
----------
319
320
321
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.
322
axes : list of integers
323
    The axes along which the FFT is carried out.
324
    If not set, all axes will be transformed.
Martin Reinecke's avatar
Martin Reinecke committed
325
326
forward : bool
    If `True`, a negative sign is used in the exponent, else a positive one.
327
328
329
330
331
332
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.
333
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
334
    if False, returns the result in a new array and leaves the input unchanged.
335
    if True, stores the result in the input array and returns a handle to it.
336
337
    NOTE: if `a` is real-valued and `inplace` is `True`, an exception will be
    raised!
Martin Reinecke's avatar
Martin Reinecke committed
338
339
340
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
341
342
343

Returns
-------
344
np.ndarray (same shape as `a`, complex type with same accuracy as `a`)
345
    The transformed data.
Martin Reinecke's avatar
Martin Reinecke committed
346
)""";
347

Martin Reinecke's avatar
Martin Reinecke committed
348
const char *r2c_DS = R"""(Performs an FFT whose input is strictly real.
Martin Reinecke's avatar
Martin Reinecke committed
349
350
351

Parameters
----------
352
a : numpy.ndarray (any real type)
Martin Reinecke's avatar
Martin Reinecke committed
353
354
    The input data
axes : list of integers
355
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
356
    If not set, all axes will be transformed in ascending order.
Martin Reinecke's avatar
Martin Reinecke committed
357
358
forward : bool
    If `True`, a negative sign is used in the exponent, else a positive one.
359
360
361
362
363
364
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
365
366
367
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
368
369
370

Returns
-------
371
np.ndarray (complex type with same accuracy as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
372
    The transformed data. The shape is identical to that of the input array,
Martin Reinecke's avatar
Martin Reinecke committed
373
374
    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
375
)""";
Martin Reinecke's avatar
Martin Reinecke committed
376

Martin Reinecke's avatar
Martin Reinecke committed
377
const char *c2r_DS = R"""(Performs an FFT whose output is strictly real.
Martin Reinecke's avatar
Martin Reinecke committed
378
379
380

Parameters
----------
381
a : numpy.ndarray (any complex type)
Martin Reinecke's avatar
Martin Reinecke committed
382
383
    The input data
axes : list of integers
384
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
385
386
387
    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.
Martin Reinecke's avatar
Martin Reinecke committed
388
389
forward : bool
    If `True`, a negative sign is used in the exponent, else a positive one.
390
391
392
393
394
395
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
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

Returns
-------
402
np.ndarray (real type with same accuracy as `a`)
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 *r2r_fftpack_DS = R"""(Performs a real-valued FFT using the FFTPACK storage scheme.
409
410
411

Parameters
----------
412
a : numpy.ndarray (any real type)
Martin Reinecke's avatar
Martin Reinecke committed
413
414
415
416
417
418
419
420
421
    The input data
axes : list of integers
    The axes along which the FFT is carried out.
    If not set, all axes will be transformed.
input_halfcomplex : bool
    if True, the transform will go from halfcomplex to real format, else in
    the other direction
forward : bool
    If `True`, a negative sign is used in the exponent, else a positive one.
422
423
424
425
426
427
inorm : int
    Normalization type
      0 : no normalization
      1 : divide by sqrt(N)
      2 : divide by N
    where N is the length of `axis`.
428
429
430
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
431
432
433
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
434
435
436

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

441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
const char *separable_hartley_DS = R"""(Performs a separable Hartley transform.
For every requested axis, a 1D forward Fourier transform is carried out, and
the real and imaginary parts of the result are added before the next axis is
processed.

Parameters
----------
a : numpy.ndarray (np.float32 or np.float64)
    The input data
axes : list of integers
    The axes along which the transform is carried out.
    If not set, all axes will be transformed.
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.
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.
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).

Returns
-------
np.ndarray (same shape and data type as `a`)
    The transformed data
)""";

const char *genuine_hartley_DS = R"""(Performs a full Hartley transform.
A full Fourier transform is carried out over the requested axes, and the
sum of real and imaginary parts of the result is stored in the output
array. For a single transformed axis, this is identical to `separable_hartley`,
but when transforming multiple axes, the results are different.
477
478
479
480
481
482

Parameters
----------
a : numpy.ndarray (np.float32 or np.float64)
    The input data
axes : list of integers
483
    The axes along which the transform is carried out.
484
    If not set, all axes will be transformed.
485
486
487
488
489
490
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.
491
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
492
    if False, returns the result in a new array and leaves the input unchanged.
493
    if True, stores the result in the input array and returns a handle to it.
Martin Reinecke's avatar
Martin Reinecke committed
494
495
496
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
497
498
499

Returns
-------
500
np.ndarray (same shape and data type as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
501
    The transformed data
Martin Reinecke's avatar
Martin Reinecke committed
502
)""";
503

Martin Reinecke's avatar
Martin Reinecke committed
504
505
506
507
} // unnamed namespace

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

Martin Reinecke's avatar
Martin Reinecke committed
510
  m.doc() = pypocketfft_DS;
Martin Reinecke's avatar
Martin Reinecke committed
511
  m.def("c2c",&c2c, c2c_DS, "a"_a, "axes"_a=py::none(), "forward"_a=true,
512
    "inorm"_a=0, "inplace"_a=false, "nthreads"_a=1);
Martin Reinecke's avatar
Martin Reinecke committed
513
514
515
516
517
518
519
  m.def("r2c",&r2c, r2c_DS, "a"_a, "axes"_a=py::none(), "forward"_a=true,
    "inorm"_a=0, "nthreads"_a=1);
  m.def("c2r",&c2r, c2r_DS, "a"_a, "axes"_a=py::none(), "lastsize"_a=0,
    "forward"_a=true, "inorm"_a=0, "nthreads"_a=1);
  m.def("r2r_fftpack",&r2r_fftpack, r2r_fftpack_DS, "a"_a, "axes"_a,
    "input_halfcomplex"_a, "forward"_a, "inorm"_a=0, "inplace"_a=false,
    "nthreads"_a=1);
520
521
522
523
  m.def("separable_hartley",&separable_hartley, separable_hartley_DS, "a"_a,
    "axes"_a=py::none(), "inorm"_a=0, "inplace"_a=false, "nthreads"_a=1);
  m.def("genuine_hartley",&genuine_hartley, genuine_hartley_DS, "a"_a,
    "axes"_a=py::none(), "inorm"_a=0, "inplace"_a=false, "nthreads"_a=1);
Martin Reinecke's avatar
Martin Reinecke committed
524
  }