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

Martin Reinecke's avatar
Martin Reinecke committed
19
//
20
// Python interface
Martin Reinecke's avatar
Martin Reinecke committed
21
22
//

23
namespace {
24

25
using namespace std;
Martin Reinecke's avatar
Martin Reinecke committed
26
27
using pocketfft::shape_t;
using pocketfft::stride_t;
Martin Reinecke's avatar
Martin Reinecke committed
28

Martin Reinecke's avatar
Martin Reinecke committed
29
30
namespace py = pybind11;

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

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

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

Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
58
shape_t makeaxes(const py::array &in, py::object axes)
Martin Reinecke's avatar
more  
Martin Reinecke committed
59
  {
Martin Reinecke's avatar
Martin Reinecke committed
60
  if (axes.is(py::none()))
Martin Reinecke's avatar
more  
Martin Reinecke committed
61
    {
62
    shape_t res(size_t(in.ndim()));
Martin Reinecke's avatar
Martin Reinecke committed
63
64
65
    for (size_t i=0; i<res.size(); ++i)
      res[i]=i;
    return res;
Martin Reinecke's avatar
more  
Martin Reinecke committed
66
    }
67
68
69
  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
70
    throw runtime_error("bad axes argument");
71
72
73
74
75
76
77
78
  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
79
80
  }

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

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

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

Martin Reinecke's avatar
Martin Reinecke committed
152
153
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
154
  {
155
  if (a.dtype().kind() == 'c')
Martin Reinecke's avatar
Martin Reinecke committed
156
157
    DISPATCH(a, c128, c64, clong, c2c_internal, (a, makeaxes(a, axes), forward,
             inorm, inplace, nthreads))
158

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

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

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

Martin Reinecke's avatar
Martin Reinecke committed
190
191
192
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)
193
  {
Martin Reinecke's avatar
Martin Reinecke committed
194
  auto axes = makeaxes(in, axes_);
195
196
  auto dims(copy_shape(in));
  py::array res = inplace ? in : py::array_t<T>(dims);
Martin Reinecke's avatar
Martin Reinecke committed
197
198
199
200
201
202
  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
203
204
205
  T fct = norm_fct<T>(inorm, dims, axes);
  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
206
  }
207
208
  return res;
  }
209

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

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

Martin Reinecke's avatar
Martin Reinecke committed
241
242
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
243
  {
Martin Reinecke's avatar
Martin Reinecke committed
244
245
  DISPATCH(in, c128, c64, clong, c2r_internal, (in, axes_, lastsize, forward,
    inorm, nthreads))
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
246
  }
Martin Reinecke's avatar
Martin Reinecke committed
247
248

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

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

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

298
py::array mycomplex2hartley(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
299
  const py::array &tmp, py::object axes_, bool inplace)
300
  {
301
  DISPATCH(in, f64, f32, flong, complex2hartley, (in, tmp, axes_, inplace))
302
  }
303

304
py::array hartley2(const py::array &in, py::object axes_, int inorm,
Martin Reinecke's avatar
Martin Reinecke committed
305
  bool inplace, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
306
  {
Martin Reinecke's avatar
Martin Reinecke committed
307
  return mycomplex2hartley(in, r2c(in, axes_, true, inorm, nthreads), axes_,
Martin Reinecke's avatar
Martin Reinecke committed
308
309
    inplace);
  }
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
310

Martin Reinecke's avatar
Martin Reinecke committed
311
const char *pypocketfft_DS = R"""(Fast Fourier and Hartley transforms.
Martin Reinecke's avatar
Martin Reinecke committed
312
313

This module supports
314
- single, double, and long double precision
Martin Reinecke's avatar
Martin Reinecke committed
315
316
317
318
319
320
- 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
321
)""";
Martin Reinecke's avatar
Martin Reinecke committed
322

Martin Reinecke's avatar
Martin Reinecke committed
323
const char *c2c_DS = R"""(Performs a complex FFT.
324
325
326

Parameters
----------
327
328
329
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.
330
axes : list of integers
331
    The axes along which the FFT is carried out.
332
    If not set, all axes will be transformed.
Martin Reinecke's avatar
Martin Reinecke committed
333
334
forward : bool
    If `True`, a negative sign is used in the exponent, else a positive one.
335
336
337
338
339
340
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.
341
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
342
    if False, returns the result in a new array and leaves the input unchanged.
343
    if True, stores the result in the input array and returns a handle to it.
344
345
    NOTE: if `a` is real-valued and `inplace` is `True`, an exception will be
    raised!
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).
349
350
351

Returns
-------
352
np.ndarray (same shape as `a`, complex type with same accuracy as `a`)
353
    The transformed data.
Martin Reinecke's avatar
Martin Reinecke committed
354
)""";
355

Martin Reinecke's avatar
Martin Reinecke committed
356
const char *r2c_DS = R"""(Performs an FFT whose input is strictly real.
Martin Reinecke's avatar
Martin Reinecke committed
357
358
359

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

Returns
-------
379
np.ndarray (complex type with same accuracy as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
380
    The transformed data. The shape is identical to that of the input array,
Martin Reinecke's avatar
Martin Reinecke committed
381
382
    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
383
)""";
Martin Reinecke's avatar
Martin Reinecke committed
384

Martin Reinecke's avatar
Martin Reinecke committed
385
const char *c2r_DS = R"""(Performs an FFT whose output is strictly real.
Martin Reinecke's avatar
Martin Reinecke committed
386
387
388

Parameters
----------
389
a : numpy.ndarray (any complex type)
Martin Reinecke's avatar
Martin Reinecke committed
390
391
    The input data
axes : list of integers
392
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
393
394
395
    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
396
397
forward : bool
    If `True`, a negative sign is used in the exponent, else a positive one.
398
399
400
401
402
403
inorm : int
    Normalization type
      0 : no normalization
      1 : divide by sqrt(N)
      2 : divide by N
    where N is the product of the lengths of the transformed output axes.
Martin Reinecke's avatar
Martin Reinecke committed
404
405
406
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
Martin Reinecke's avatar
Martin Reinecke committed
407
408
409

Returns
-------
410
np.ndarray (real type with same accuracy as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
411
412
    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
413
    entries.
Martin Reinecke's avatar
Martin Reinecke committed
414
)""";
Martin Reinecke's avatar
Martin Reinecke committed
415

Martin Reinecke's avatar
Martin Reinecke committed
416
const char *r2r_fftpack_DS = R"""(Performs a real-valued FFT using the FFTPACK storage scheme.
417
418
419

Parameters
----------
420
a : numpy.ndarray (any real type)
Martin Reinecke's avatar
Martin Reinecke committed
421
422
423
424
425
426
427
428
429
    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.
430
431
432
433
434
435
inorm : int
    Normalization type
      0 : no normalization
      1 : divide by sqrt(N)
      2 : divide by N
    where N is the length of `axis`.
436
437
438
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
439
440
441
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
442
443
444

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

Martin Reinecke's avatar
Martin Reinecke committed
449
const char *hartley_DS = R"""(Performs a Hartley transform.
450
451
452
453
454
455
456
457
458
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
459
    The axes along which the transform is carried out.
460
    If not set, all axes will be transformed.
461
462
463
464
465
466
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.
467
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
468
    if False, returns the result in a new array and leaves the input unchanged.
469
    if True, stores the result in the input array and returns a handle to it.
Martin Reinecke's avatar
Martin Reinecke committed
470
471
472
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
473
474
475

Returns
-------
476
np.ndarray (same shape and data type as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
477
    The transformed data
Martin Reinecke's avatar
Martin Reinecke committed
478
)""";
479

Martin Reinecke's avatar
Martin Reinecke committed
480
481
482
483
} // unnamed namespace

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

Martin Reinecke's avatar
Martin Reinecke committed
486
  m.doc() = pypocketfft_DS;
Martin Reinecke's avatar
Martin Reinecke committed
487
  m.def("c2c",&c2c, c2c_DS, "a"_a, "axes"_a=py::none(), "forward"_a=true,
488
    "inorm"_a=0, "inplace"_a=false, "nthreads"_a=1);
Martin Reinecke's avatar
Martin Reinecke committed
489
490
491
492
493
494
495
  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);
496
  m.def("hartley",&hartley, hartley_DS, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
497
    "inplace"_a=false, "nthreads"_a=1);
498
  m.def("hartley2",&hartley2, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
499
    "inplace"_a=false, "nthreads"_a=1);
Martin Reinecke's avatar
Martin Reinecke committed
500
501
  m.def("complex2hartley",&mycomplex2hartley, "in"_a, "tmp"_a, "axes"_a,
    "inplace"_a=false);
Martin Reinecke's avatar
Martin Reinecke committed
502
  }