pypocketfft.cc 17.6 KB
Newer Older
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
1 2 3 4 5 6
/*
 * This file is part of pocketfft.
 * Licensed under a 3-clause BSD style license - see LICENSE.md
 */

/*
7
 *  Python interface.
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
8
 *
9
 *  Copyright (C) 2019 Max-Planck-Society
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
10 11 12
 *  \author Martin Reinecke
 */

Martin Reinecke's avatar
Martin Reinecke committed
13 14
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
Martin Reinecke's avatar
more  
Martin Reinecke committed
15
#include <pybind11/stl.h>
Martin Reinecke's avatar
Martin Reinecke committed
16

17
#include "pocketfft_hdronly.h"
Martin Reinecke's avatar
Martin Reinecke committed
18

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

23
namespace {
24

25 26
using namespace std;
using namespace pocketfft;
Martin Reinecke's avatar
Martin Reinecke committed
27

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

30 31 32 33
// Only instantiate long double transforms if they offer more precision
using ldbl_t = typename std::conditional<
  sizeof(long double)==sizeof(double), double, long double>::type;

Martin Reinecke's avatar
Martin Reinecke committed
34 35
auto c64 = py::dtype("complex64");
auto c128 = py::dtype("complex128");
36
auto clong = py::dtype("longcomplex");
Martin Reinecke's avatar
Martin Reinecke committed
37 38
auto f32 = py::dtype("float32");
auto f64 = py::dtype("float64");
39
auto flong = py::dtype("longfloat");
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
40 41

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

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

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

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

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

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

151 152
py::array xfftn(const py::array &a, py::object axes, int inorm,
  bool inplace, bool fwd, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
153
  {
154 155 156
  if (a.dtype().kind() == 'c')
    DISPATCH(a, c128, c64, clong, xfftn_internal, (a, makeaxes(a, axes), inorm,
             inplace, fwd, nthreads))
157

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

163 164 165 166
py::array fftn(const py::array &a, py::object axes, int inorm, bool inplace,
  size_t nthreads)
  { return xfftn(a, axes, inorm, inplace, true, nthreads); }

167 168 169
py::array ifftn(const py::array &a, py::object axes, int inorm,
  bool inplace, size_t nthreads)
  { return xfftn(a, axes, inorm, inplace, false, nthreads); }
Martin Reinecke's avatar
Martin Reinecke committed
170

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

190
py::array rfftn(const py::array &in, py::object axes_, int inorm,
Martin Reinecke's avatar
Martin Reinecke committed
191
  size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
192
  {
193
  DISPATCH(in, f64, f32, flong, rfftn_internal, (in, axes_, inorm, nthreads))
Martin Reinecke's avatar
Martin Reinecke committed
194
  }
195

196
template<typename T> py::array xrfft_scipy(const py::array &in,
197
  size_t axis, int inorm, bool inplace, bool fwd, size_t nthreads)
198 199 200
  {
  auto dims(copy_shape(in));
  py::array res = inplace ? in : py::array_t<T>(dims);
Martin Reinecke's avatar
Martin Reinecke committed
201 202 203 204 205 206
  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;
207
  T fct = norm_fct<T>(inorm, dims[axis]);
208
  r2r_fftpack(dims, s_in, s_out, {axis}, fwd, fwd, d_in, d_out, fct, nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
209
  }
210 211
  return res;
  }
212

213 214
py::array rfft_scipy(const py::array &in, size_t axis, int inorm,
  bool inplace, size_t nthreads)
215
  {
216
  DISPATCH(in, f64, f32, flong, xrfft_scipy, (in, axis, inorm, inplace, true,
Martin Reinecke's avatar
Martin Reinecke committed
217
    nthreads))
218
  }
219

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

Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
249
py::array irfftn(const py::array &in, py::object axes_, size_t lastsize,
250
  int inorm, size_t nthreads)
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
251
  {
252
  DISPATCH(in, c128, c64, clong, irfftn_internal, (in, axes_, lastsize, inorm,
Martin Reinecke's avatar
Martin Reinecke committed
253
    nthreads))
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
254
  }
Martin Reinecke's avatar
Martin Reinecke committed
255 256

template<typename T> py::array hartley_internal(const py::array &in,
257
  py::object axes_, int inorm, bool inplace, size_t nthreads)
Martin Reinecke's avatar
more  
Martin Reinecke committed
258
  {
Martin Reinecke's avatar
Martin Reinecke committed
259
  auto dims(copy_shape(in));
260
  py::array res = inplace ? in : py::array_t<T>(dims);
Martin Reinecke's avatar
Martin Reinecke committed
261 262 263 264 265 266 267
  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;
268 269
  T fct = norm_fct<T>(inorm, dims, axes);
  r2r_hartley(dims, s_in, s_out, axes, d_in, d_out, fct, nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
270
  }
Martin Reinecke's avatar
Martin Reinecke committed
271 272
  return res;
  }
273

274
py::array hartley(const py::array &in, py::object axes_, int inorm,
Martin Reinecke's avatar
Martin Reinecke committed
275
  bool inplace, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
276
  {
277
  DISPATCH(in, f64, f32, flong, hartley_internal, (in, axes_, inorm, inplace,
Martin Reinecke's avatar
Martin Reinecke committed
278
    nthreads))
Martin Reinecke's avatar
Martin Reinecke committed
279
  }
280

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

306
py::array mycomplex2hartley(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
307
  const py::array &tmp, py::object axes_, bool inplace)
308
  {
309
  DISPATCH(in, f64, f32, flong, complex2hartley, (in, tmp, axes_, inplace))
310
  }
311

312
py::array hartley2(const py::array &in, py::object axes_, int inorm,
Martin Reinecke's avatar
Martin Reinecke committed
313
  bool inplace, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
314
  {
315
  return mycomplex2hartley(in, rfftn(in, axes_, inorm, nthreads), axes_,
Martin Reinecke's avatar
Martin Reinecke committed
316 317
    inplace);
  }
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 *fftn_DS = R"""(
332 333 334 335
Performs a forward complex FFT.

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

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

Martin Reinecke's avatar
Martin Reinecke committed
363
const char *ifftn_DS = R"""(Performs a backward complex FFT.
364 365 366

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

Returns
-------
390
np.ndarray (same shape as `a`, complex type with same accuracy as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
391
    The transformed data
Martin Reinecke's avatar
Martin Reinecke committed
392
)""";
393

Martin Reinecke's avatar
Martin Reinecke committed
394
const char *rfftn_DS = R"""(Performs a forward real-valued FFT.
Martin Reinecke's avatar
Martin Reinecke committed
395 396 397

Parameters
----------
398
a : numpy.ndarray (any real 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
    If not set, all axes will be transformed in ascending order.
403 404 405 406 407 408
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
409 410 411
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
412 413 414

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

Martin Reinecke's avatar
Martin Reinecke committed
421
const char *rfft_scipy_DS = R"""(Performs a forward real-valued FFT.
422 423 424

Parameters
----------
425
a : numpy.ndarray (any real type)
426 427 428
    The input data
axis : int
    The axis along which the FFT is carried out.
429 430 431 432 433 434
inorm : int
    Normalization type
      0 : no normalization
      1 : divide by sqrt(N)
      2 : divide by N
    where N is the length of `axis`.
435 436 437
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
438 439 440
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
441 442 443

Returns
-------
444
np.ndarray (same shape and data type as `a`)
445 446 447
    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
448
)""";
449

Martin Reinecke's avatar
Martin Reinecke committed
450
const char *irfftn_DS = R"""(Performs a backward real-valued FFT.
Martin Reinecke's avatar
Martin Reinecke committed
451 452 453

Parameters
----------
454
a : numpy.ndarray (any complex type)
Martin Reinecke's avatar
Martin Reinecke committed
455 456
    The input data
axes : list of integers
457
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
458 459 460
    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.
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 output axes.
Martin Reinecke's avatar
Martin Reinecke committed
467 468 469
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
470 471 472

Returns
-------
473
np.ndarray (real type with same accuracy as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
474 475
    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
476
    entries.
Martin Reinecke's avatar
Martin Reinecke committed
477
)""";
Martin Reinecke's avatar
Martin Reinecke committed
478

Martin Reinecke's avatar
Martin Reinecke committed
479
const char *irfft_scipy_DS = R"""(Performs a backward real-valued FFT.
480 481 482

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

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

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

Returns
-------
534
np.ndarray (same shape and data type as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
535
    The transformed data
Martin Reinecke's avatar
Martin Reinecke committed
536
)""";
537

Martin Reinecke's avatar
Martin Reinecke committed
538 539 540 541
} // unnamed namespace

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

Martin Reinecke's avatar
Martin Reinecke committed
544
  m.doc() = pypocketfft_DS;
545
  m.def("fftn",&fftn, fftn_DS, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
546
    "inplace"_a=false, "nthreads"_a=1);
547
  m.def("ifftn",&ifftn, ifftn_DS, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
548
    "inplace"_a=false, "nthreads"_a=1);
549
  m.def("rfftn",&rfftn, rfftn_DS, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
550
    "nthreads"_a=1);
551
  m.def("rfft_scipy",&rfft_scipy, rfft_scipy_DS, "a"_a, "axis"_a, "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
552
    "inplace"_a=false, "nthreads"_a=1);
553
  m.def("irfftn",&irfftn, irfftn_DS, "a"_a, "axes"_a=py::none(), "lastsize"_a=0,
554 555 556 557
    "inorm"_a=0, "nthreads"_a=1);
  m.def("irfft_scipy",&irfft_scipy, irfft_scipy_DS, "a"_a, "axis"_a,
    "inorm"_a=0, "inplace"_a=false, "nthreads"_a=1);
  m.def("hartley",&hartley, hartley_DS, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
558
    "inplace"_a=false, "nthreads"_a=1);
559
  m.def("hartley2",&hartley2, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
560
    "inplace"_a=false, "nthreads"_a=1);
Martin Reinecke's avatar
Martin Reinecke committed
561 562
  m.def("complex2hartley",&mycomplex2hartley, "in"_a, "tmp"_a, "axes"_a,
    "inplace"_a=false);
Martin Reinecke's avatar
Martin Reinecke committed
563
  }