pypocketfft.cc 17.2 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
    }
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
66
  auto tmp=axes.cast<shape_t>();
Martin Reinecke's avatar
fix  
Martin Reinecke committed
67
  if ((tmp.size()>size_t(in.ndim())) || (tmp.size()==0))
Martin Reinecke's avatar
more  
Martin Reinecke committed
68
    throw runtime_error("bad axes argument");
Martin Reinecke's avatar
Martin Reinecke committed
69 70
  for (auto sz: tmp)
    if (sz>=size_t(in.ndim()))
Martin Reinecke's avatar
more  
Martin Reinecke committed
71
      throw runtime_error("invalid axis number");
Martin Reinecke's avatar
Martin Reinecke committed
72
  return tmp;
Martin Reinecke's avatar
more  
Martin Reinecke committed
73 74
  }

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

82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99
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
100
template<typename T> py::array xfftn_internal(const py::array &in,
101
  const shape_t &axes, int inorm, bool inplace, bool fwd, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
102
  {
Martin Reinecke's avatar
Martin Reinecke committed
103
  auto dims(copy_shape(in));
104
  py::array res = inplace ? in : py::array_t<complex<T>>(dims);
Martin Reinecke's avatar
Martin Reinecke committed
105 106 107 108 109 110
  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;
111 112
  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
113
  }
Martin Reinecke's avatar
Martin Reinecke committed
114 115
  return res;
  }
Martin Reinecke's avatar
Martin Reinecke committed
116

117 118
py::array xfftn(const py::array &a, py::object axes, int inorm,
  bool inplace, bool fwd, size_t nthreads)
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
119
  {
120
  DISPATCH(a, c128, c64, clong, xfftn_internal, (a, makeaxes(a, axes), inorm,
Martin Reinecke's avatar
Martin Reinecke committed
121
           inplace, fwd, nthreads))
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
122
  }
123

Martin Reinecke's avatar
Martin Reinecke committed
124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140
template<typename T> py::array sym_rfftn_internal(const py::array &in,
  py::object axes_, int inorm, size_t nthreads)
  {
  auto axes = makeaxes(in, axes_);
  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);
  r2c(dims, s_in, s_out, axes, d_in, d_out, fct, nthreads);
  // 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 152 153 154
    iter.advance();
    }
  }
  return res;
  }

py::array fftn(const py::array &a, py::object axes, int inorm, bool inplace,
  size_t nthreads)
  {
155 156 157
    if (a.dtype().kind() == 'c')
      return xfftn(a, axes, inorm, inplace, true, nthreads);

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

162 163 164
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
165

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

185
py::array rfftn(const py::array &in, py::object axes_, int inorm,
Martin Reinecke's avatar
Martin Reinecke committed
186
  size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
187
  {
188
  DISPATCH(in, f64, f32, flong, rfftn_internal, (in, axes_, inorm, nthreads))
Martin Reinecke's avatar
Martin Reinecke committed
189
  }
190

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

208 209
py::array rfft_scipy(const py::array &in, size_t axis, int inorm,
  bool inplace, size_t nthreads)
210
  {
211
  DISPATCH(in, f64, f32, flong, xrfft_scipy, (in, axis, inorm, inplace, true,
Martin Reinecke's avatar
Martin Reinecke committed
212
    nthreads))
213
  }
214

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

Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
244
py::array irfftn(const py::array &in, py::object axes_, size_t lastsize,
245
  int inorm, size_t nthreads)
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
246
  {
247
  DISPATCH(in, c128, c64, clong, irfftn_internal, (in, axes_, lastsize, inorm,
Martin Reinecke's avatar
Martin Reinecke committed
248
    nthreads))
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
249
  }
Martin Reinecke's avatar
Martin Reinecke committed
250 251

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

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

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

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

307
py::array hartley2(const py::array &in, py::object axes_, int inorm,
Martin Reinecke's avatar
Martin Reinecke committed
308
  bool inplace, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
309
  {
310
  return mycomplex2hartley(in, rfftn(in, axes_, inorm, nthreads), axes_,
Martin Reinecke's avatar
Martin Reinecke committed
311 312
    inplace);
  }
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
313

Martin Reinecke's avatar
Martin Reinecke committed
314
const char *pypocketfft_DS = R"""(Fast Fourier and Hartley transforms.
Martin Reinecke's avatar
Martin Reinecke committed
315 316 317 318 319 320 321 322 323

This module supports
- single and double precision
- complex and real-valued transforms
- multi-dimensional transforms

For two- and higher-dimensional transforms the code will use SSE2 and AVX
vector instructions for faster execution if these are supported by the CPU and
were enabled during compilation.
Martin Reinecke's avatar
Martin Reinecke committed
324
)""";
Martin Reinecke's avatar
Martin Reinecke committed
325

Martin Reinecke's avatar
Martin Reinecke committed
326
const char *fftn_DS = R"""(
327 328 329 330
Performs a forward complex FFT.

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

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

Martin Reinecke's avatar
Martin Reinecke committed
358
const char *ifftn_DS = R"""(Performs a backward complex FFT.
359 360 361

Parameters
----------
362
a : numpy.ndarray (any complex type)
363 364
    The input data
axes : list of integers
365
    The axes along which the FFT is carried out.
366
    If not set, all axes will be transformed.
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 axes.
373
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
374
    if False, returns the result in a new array and leaves the input unchanged.
375
    if True, stores the result in the input array and returns a handle to it.
Martin Reinecke's avatar
Martin Reinecke committed
376 377 378
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
379 380 381 382

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

Martin Reinecke's avatar
Martin Reinecke committed
386
const char *rfftn_DS = R"""(Performs a forward real-valued FFT.
Martin Reinecke's avatar
Martin Reinecke committed
387 388 389

Parameters
----------
390
a : numpy.ndarray (any real type)
Martin Reinecke's avatar
Martin Reinecke committed
391 392
    The input data
axes : list of integers
393
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
394
    If not set, all axes will be transformed in ascending order.
395 396 397 398 399 400
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
401 402 403
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
404 405 406

Returns
-------
407
np.ndarray (complex type with same accuracy as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
408
    The transformed data. The shape is identical to that of the input array,
Martin Reinecke's avatar
Martin Reinecke committed
409 410
    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
411
)""";
Martin Reinecke's avatar
Martin Reinecke committed
412

Martin Reinecke's avatar
Martin Reinecke committed
413
const char *rfft_scipy_DS = R"""(Performs a forward real-valued FFT.
414 415 416

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

Returns
-------
436
np.ndarray (same shape and data type as `a`)
437 438 439
    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
440
)""";
441

Martin Reinecke's avatar
Martin Reinecke committed
442
const char *irfftn_DS = R"""(Performs a backward real-valued FFT.
Martin Reinecke's avatar
Martin Reinecke committed
443 444 445

Parameters
----------
446
a : numpy.ndarray (any complex type)
Martin Reinecke's avatar
Martin Reinecke committed
447 448
    The input data
axes : list of integers
449
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
450 451 452
    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.
453 454 455 456 457 458
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
459 460 461
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
462 463 464

Returns
-------
465
np.ndarray (real type with same accuracy as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
466 467
    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
468
    entries.
Martin Reinecke's avatar
Martin Reinecke committed
469
)""";
Martin Reinecke's avatar
Martin Reinecke committed
470

Martin Reinecke's avatar
Martin Reinecke committed
471
const char *irfft_scipy_DS = R"""(Performs a backward real-valued FFT.
472 473 474

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

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

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

Returns
-------
526
np.ndarray (same shape and data type as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
527
    The transformed data
Martin Reinecke's avatar
Martin Reinecke committed
528
)""";
529

Martin Reinecke's avatar
Martin Reinecke committed
530 531 532 533
} // unnamed namespace

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

Martin Reinecke's avatar
Martin Reinecke committed
536
  m.doc() = pypocketfft_DS;
537
  m.def("fftn",&fftn, fftn_DS, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
538
    "inplace"_a=false, "nthreads"_a=1);
539
  m.def("ifftn",&ifftn, ifftn_DS, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
540
    "inplace"_a=false, "nthreads"_a=1);
541
  m.def("rfftn",&rfftn, rfftn_DS, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
542
    "nthreads"_a=1);
543
  m.def("rfft_scipy",&rfft_scipy, rfft_scipy_DS, "a"_a, "axis"_a, "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
544
    "inplace"_a=false, "nthreads"_a=1);
545
  m.def("irfftn",&irfftn, irfftn_DS, "a"_a, "axes"_a=py::none(), "lastsize"_a=0,
546 547 548 549
    "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
550
    "inplace"_a=false, "nthreads"_a=1);
551
  m.def("hartley2",&hartley2, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
552
    "inplace"_a=false, "nthreads"_a=1);
Martin Reinecke's avatar
Martin Reinecke committed
553 554
  m.def("complex2hartley",&mycomplex2hartley, "in"_a, "tmp"_a, "axes"_a,
    "inplace"_a=false);
Martin Reinecke's avatar
Martin Reinecke committed
555
  }