pypocketfft.cc 17.4 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
Martin Reinecke committed
37
auto None = py::none();
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
38
39

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

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

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

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

87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
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
105
106
107
108
109
110
111
112
113
114
template<typename T> py::array_t<T> prepare_output(py::object &out_,
  shape_t &dims)
  {
  if (out_.is(None)) return py::array_t<T>(dims);
  auto tmp = out_.cast<py::array_t<T>>();
  if (!tmp.is(out_)) // a new object was created during casting
    throw runtime_error("unexpected data type for output array");
  return tmp;
  }

Martin Reinecke's avatar
Martin Reinecke committed
115
template<typename T> py::array c2c_internal(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
116
117
  const py::object &axes_, bool forward, int inorm, py::object &out_,
  size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
118
  {
Martin Reinecke's avatar
Martin Reinecke committed
119
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
Martin Reinecke committed
120
  auto dims(copy_shape(in));
Martin Reinecke's avatar
Martin Reinecke committed
121
  auto res = prepare_output<complex<T>>(out_, dims);
Martin Reinecke's avatar
Martin Reinecke committed
122
123
124
125
126
127
  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;
128
  T fct = norm_fct<T>(inorm, dims, axes);
Martin Reinecke's avatar
Martin Reinecke committed
129
  pocketfft::c2c(dims, s_in, s_out, axes, forward, d_in, d_out, fct, nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
130
  }
Martin Reinecke's avatar
Martin Reinecke committed
131
132
  return res;
  }
Martin Reinecke's avatar
Martin Reinecke committed
133

Martin Reinecke's avatar
Martin Reinecke committed
134
template<typename T> py::array c2c_sym_internal(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
135
136
  const py::object &axes_, bool forward, int inorm, py::object &out_,
  size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
137
  {
Martin Reinecke's avatar
Martin Reinecke committed
138
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
Martin Reinecke committed
139
  auto dims(copy_shape(in));
Martin Reinecke's avatar
Martin Reinecke committed
140
  auto res = prepare_output<complex<T>>(out_, dims);
Martin Reinecke's avatar
Martin Reinecke committed
141
142
143
144
145
146
147
  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
148
  pocketfft::r2c(dims, s_in, s_out, axes, forward, d_in, d_out, fct, nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
149
150
151
  // now fill in second half
  using namespace pocketfft::detail;
  ndarr<complex<T>> ares(res.mutable_data(), dims, s_out);
152
  rev_iter iter(ares, axes);
Martin Reinecke's avatar
Martin Reinecke committed
153
154
  while(iter.remaining()>0)
    {
155
156
    auto v = ares[iter.ofs()];
    ares[iter.rev_ofs()] = conj(v);
Martin Reinecke's avatar
Martin Reinecke committed
157
158
159
160
161
162
    iter.advance();
    }
  }
  return res;
  }

Martin Reinecke's avatar
Martin Reinecke committed
163
164
py::array c2c(const py::array &a, const py::object &axes_, bool forward,
  int inorm, py::object &out_, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
165
  {
166
  if (a.dtype().kind() == 'c')
Martin Reinecke's avatar
Martin Reinecke committed
167
168
    DISPATCH(a, c128, c64, clong, c2c_internal, (a, axes_, forward,
             inorm, out_, nthreads))
169

Martin Reinecke's avatar
Martin Reinecke committed
170
171
  DISPATCH(a, f64, f32, flong, c2c_sym_internal, (a, axes_, forward,
           inorm, out_, nthreads))
Martin Reinecke's avatar
Martin Reinecke committed
172
  }
173

Martin Reinecke's avatar
Martin Reinecke committed
174
template<typename T> py::array r2c_internal(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
175
176
  const py::object &axes_, bool forward, int inorm, py::object &out_,
  size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
177
  {
Martin Reinecke's avatar
Martin Reinecke committed
178
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
Martin Reinecke committed
179
  auto dims_in(copy_shape(in)), dims_out(dims_in);
Martin Reinecke's avatar
more  
Martin Reinecke committed
180
  dims_out[axes.back()] = (dims_out[axes.back()]>>1)+1;
Martin Reinecke's avatar
Martin Reinecke committed
181
  py::array res = prepare_output<complex<T>>(out_, dims_out);
Martin Reinecke's avatar
Martin Reinecke committed
182
183
184
185
186
187
  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;
188
  T fct = norm_fct<T>(inorm, dims_in, axes);
189
190
  pocketfft::r2c(dims_in, s_in, s_out, axes, forward, 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

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

Martin Reinecke's avatar
Martin Reinecke committed
202
template<typename T> py::array r2r_fftpack_internal(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
203
204
  const py::object &axes_, bool input_halfcomplex, bool forward, int inorm,
  py::object &out_, size_t nthreads)
205
  {
Martin Reinecke's avatar
Martin Reinecke committed
206
  auto axes = makeaxes(in, axes_);
207
  auto dims(copy_shape(in));
Martin Reinecke's avatar
Martin Reinecke committed
208
  py::array res = prepare_output<T>(out_, dims);
Martin Reinecke's avatar
Martin Reinecke committed
209
210
211
212
213
214
  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
215
  T fct = norm_fct<T>(inorm, dims, axes);
216
217
  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
218
  }
219
220
  return res;
  }
221

Martin Reinecke's avatar
Martin Reinecke committed
222
223
py::array r2r_fftpack(const py::array &in, const py::object &axes_,
  bool input_halfcomplex, bool forward, int inorm, py::object &out_,
Martin Reinecke's avatar
Martin Reinecke committed
224
  size_t nthreads)
225
  {
Martin Reinecke's avatar
Martin Reinecke committed
226
  DISPATCH(in, f64, f32, flong, r2r_fftpack_internal, (in, axes_,
Martin Reinecke's avatar
Martin Reinecke committed
227
    input_halfcomplex, forward, inorm, out_, nthreads))
228
  }
229

Martin Reinecke's avatar
Martin Reinecke committed
230
template<typename T> py::array c2r_internal(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
231
232
  const py::object &axes_, size_t lastsize, bool forward, int inorm,
  py::object &out_, 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
Martin Reinecke committed
241
  py::array res = prepare_output<T>(out_, 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
  T fct = norm_fct<T>(inorm, dims_out, axes);
249
250
  pocketfft::c2r(dims_out, s_in, s_out, axes, forward, d_in, d_out, fct,
    nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
251
  }
Martin Reinecke's avatar
more  
Martin Reinecke committed
252
253
  return res;
  }
254

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

262
template<typename T> py::array separable_hartley_internal(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
263
  const py::object &axes_, int inorm, py::object &out_, size_t nthreads)
Martin Reinecke's avatar
more  
Martin Reinecke committed
264
  {
Martin Reinecke's avatar
Martin Reinecke committed
265
  auto dims(copy_shape(in));
Martin Reinecke's avatar
Martin Reinecke committed
266
  py::array res = prepare_output<T>(out_, dims);
Martin Reinecke's avatar
Martin Reinecke committed
267
268
269
270
271
272
273
  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;
274
  T fct = norm_fct<T>(inorm, dims, axes);
Martin Reinecke's avatar
Martin Reinecke committed
275
  pocketfft::r2r_hartley(dims, s_in, s_out, axes, d_in, d_out, fct, nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
276
  }
Martin Reinecke's avatar
Martin Reinecke committed
277
278
  return res;
  }
279

Martin Reinecke's avatar
Martin Reinecke committed
280
281
py::array separable_hartley(const py::array &in, const py::object &axes_,
  int inorm, py::object &out_, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
282
  {
283
  DISPATCH(in, f64, f32, flong, separable_hartley_internal, (in, axes_, inorm,
Martin Reinecke's avatar
Martin Reinecke committed
284
    out_, nthreads))
Martin Reinecke's avatar
Martin Reinecke committed
285
  }
286

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

Martin Reinecke's avatar
Martin Reinecke committed
312
313
py::array genuine_hartley(const py::array &in, const py::object &axes_,
  int inorm, py::object &out_, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
314
  {
Martin Reinecke's avatar
Martin Reinecke committed
315
316
  auto tmp = r2c(in, axes_, true, inorm, None, nthreads);
  DISPATCH(in, f64, f32, flong, complex2hartley, (in, tmp, axes_, out_))
Martin Reinecke's avatar
Martin Reinecke committed
317
  }
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
318

Martin Reinecke's avatar
Martin Reinecke committed
319
const char *pypocketfft_DS = R"""(Fast Fourier and Hartley transforms.
Martin Reinecke's avatar
Martin Reinecke committed
320
321

This module supports
322
- single, double, and long double precision
Martin Reinecke's avatar
Martin Reinecke committed
323
324
325
326
327
328
- 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
329
)""";
Martin Reinecke's avatar
Martin Reinecke committed
330

Martin Reinecke's avatar
Martin Reinecke committed
331
const char *c2c_DS = R"""(Performs a complex FFT.
332
333
334

Parameters
----------
335
336
337
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.
338
axes : list of integers
339
    The axes along which the FFT is carried out.
340
    If not set, all axes will be transformed.
Martin Reinecke's avatar
Martin Reinecke committed
341
342
forward : bool
    If `True`, a negative sign is used in the exponent, else a positive one.
343
344
345
346
347
348
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.
Martin Reinecke's avatar
Martin Reinecke committed
349
350
351
out : numpy.ndarray (same shape as `a`, complex type with same accuracy as `a`)
    May be identical to `a`.
    If None, a new array is allocated to store the output.
Martin Reinecke's avatar
Martin Reinecke committed
352
353
354
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
355
356
357

Returns
-------
Martin Reinecke's avatar
Martin Reinecke committed
358
numpy.ndarray (same shape as `a`, complex type with same accuracy as `a`)
359
    The transformed data.
Martin Reinecke's avatar
Martin Reinecke committed
360
)""";
361

Martin Reinecke's avatar
Martin Reinecke committed
362
const char *r2c_DS = R"""(Performs an FFT whose input is strictly real.
Martin Reinecke's avatar
Martin Reinecke committed
363
364
365

Parameters
----------
366
a : numpy.ndarray (any real type)
Martin Reinecke's avatar
Martin Reinecke committed
367
368
    The input data
axes : list of integers
369
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
370
    If not set, all axes will be transformed in ascending order.
Martin Reinecke's avatar
Martin Reinecke committed
371
372
forward : bool
    If `True`, a negative sign is used in the exponent, else a positive one.
373
374
375
376
377
378
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
379
380
381
out : numpy.ndarray (complex type with same accuracy as `a`)
    For the required shape, see the `Returns` section.
    If None, a new array is allocated to store the output.
Martin Reinecke's avatar
Martin Reinecke committed
382
383
384
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
385
386
387

Returns
-------
Martin Reinecke's avatar
Martin Reinecke committed
388
numpy.ndarray (complex type with same accuracy as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
389
    The transformed data. The shape is identical to that of the input array,
Martin Reinecke's avatar
Martin Reinecke committed
390
391
    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
392
)""";
Martin Reinecke's avatar
Martin Reinecke committed
393

Martin Reinecke's avatar
Martin Reinecke committed
394
const char *c2r_DS = R"""(Performs an FFT whose output is strictly real.
Martin Reinecke's avatar
Martin Reinecke committed
395
396
397

Parameters
----------
398
a : numpy.ndarray (any complex type)
Martin Reinecke's avatar
Martin Reinecke committed
399
400
    The input data
axes : list of integers
401
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
402
403
404
    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
405
406
forward : bool
    If `True`, a negative sign is used in the exponent, else a positive one.
407
408
409
410
411
412
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
413
414
415
out : numpy.ndarray (real type with same accuracy as `a`)
    For the required shape, see the `Returns` section.
    If None, a new array is allocated to store the output.
Martin Reinecke's avatar
Martin Reinecke committed
416
417
418
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
419
420
421

Returns
-------
Martin Reinecke's avatar
Martin Reinecke committed
422
numpy.ndarray (real type with same accuracy as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
423
424
    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
425
    entries.
Martin Reinecke's avatar
Martin Reinecke committed
426
)""";
Martin Reinecke's avatar
Martin Reinecke committed
427

Martin Reinecke's avatar
Martin Reinecke committed
428
const char *r2r_fftpack_DS = R"""(Performs a real-valued FFT using the FFTPACK storage scheme.
429
430
431

Parameters
----------
432
a : numpy.ndarray (any real type)
Martin Reinecke's avatar
Martin Reinecke committed
433
434
435
436
437
438
439
440
441
    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.
442
443
444
445
446
447
inorm : int
    Normalization type
      0 : no normalization
      1 : divide by sqrt(N)
      2 : divide by N
    where N is the length of `axis`.
Martin Reinecke's avatar
Martin Reinecke committed
448
449
450
out : numpy.ndarray (same shape and data type as `a`)
    May be identical to `a`.
    If None, a new array is allocated to store the output.
Martin Reinecke's avatar
Martin Reinecke committed
451
452
453
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
454
455
456

Returns
-------
Martin Reinecke's avatar
Martin Reinecke committed
457
numpy.ndarray (same shape and data type as `a`)
458
    The transformed data. The shape is identical to that of the input array.
Martin Reinecke's avatar
Martin Reinecke committed
459
)""";
460

461
462
463
464
465
466
467
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
----------
Martin Reinecke's avatar
Martin Reinecke committed
468
a : numpy.ndarray (any real type)
469
470
471
472
473
474
475
476
477
478
    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.
Martin Reinecke's avatar
Martin Reinecke committed
479
480
481
out : numpy.ndarray (same shape and data type as `a`)
    May be identical to `a`.
    If None, a new array is allocated to store the output.
482
483
484
485
486
487
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).

Returns
-------
Martin Reinecke's avatar
Martin Reinecke committed
488
numpy.ndarray (same shape and data type as `a`)
489
490
491
492
493
494
495
496
    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.
497
498
499

Parameters
----------
Martin Reinecke's avatar
Martin Reinecke committed
500
a : numpy.ndarray (any real type)
501
502
    The input data
axes : list of integers
503
    The axes along which the transform is carried out.
504
    If not set, all axes will be transformed.
505
506
507
508
509
510
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.
Martin Reinecke's avatar
Martin Reinecke committed
511
512
513
out : numpy.ndarray (same shape and data type as `a`)
    May be identical to `a`.
    If None, a new array is allocated to store the output.
Martin Reinecke's avatar
Martin Reinecke committed
514
515
516
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
517
518
519

Returns
-------
Martin Reinecke's avatar
Martin Reinecke committed
520
numpy.ndarray (same shape and data type as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
521
    The transformed data
Martin Reinecke's avatar
Martin Reinecke committed
522
)""";
523

Martin Reinecke's avatar
Martin Reinecke committed
524
525
526
527
} // unnamed namespace

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

Martin Reinecke's avatar
Martin Reinecke committed
530
  m.doc() = pypocketfft_DS;
Martin Reinecke's avatar
Martin Reinecke committed
531
532
533
534
535
536
  m.def("c2c",&c2c, c2c_DS, "a"_a, "axes"_a=None, "forward"_a=true,
    "inorm"_a=0, "out"_a=None, "nthreads"_a=1);
  m.def("r2c",&r2c, r2c_DS, "a"_a, "axes"_a=None, "forward"_a=true,
    "inorm"_a=0, "out"_a=None, "nthreads"_a=1);
  m.def("c2r",&c2r, c2r_DS, "a"_a, "axes"_a=None, "lastsize"_a=0,
    "forward"_a=true, "inorm"_a=0, "out"_a=None, "nthreads"_a=1);
Martin Reinecke's avatar
Martin Reinecke committed
537
  m.def("r2r_fftpack",&r2r_fftpack, r2r_fftpack_DS, "a"_a, "axes"_a,
Martin Reinecke's avatar
Martin Reinecke committed
538
    "input_halfcomplex"_a, "forward"_a, "inorm"_a=0, "out"_a=None,
Martin Reinecke's avatar
Martin Reinecke committed
539
    "nthreads"_a=1);
540
  m.def("separable_hartley",&separable_hartley, separable_hartley_DS, "a"_a,
Martin Reinecke's avatar
Martin Reinecke committed
541
    "axes"_a=None, "inorm"_a=0, "out"_a=None, "nthreads"_a=1);
542
  m.def("genuine_hartley",&genuine_hartley, genuine_hartley_DS, "a"_a,
Martin Reinecke's avatar
Martin Reinecke committed
543
    "axes"_a=None, "inorm"_a=0, "out"_a=None, "nthreads"_a=1);
Martin Reinecke's avatar
Martin Reinecke committed
544
  }