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

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

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

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

19
namespace {
20

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

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

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

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

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

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

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

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

86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103
template<typename T> T norm_fct(int inorm, size_t N)
  {
  if (inorm==0) return T(1);
  if (inorm==2) return T(1/ldbl_t(N));
  if (inorm==1) return T(1/sqrt(ldbl_t(N)));
  throw invalid_argument("invalid value for inorm (must be 0, 1, or 2)");
  }

template<typename T> T norm_fct(int inorm, const shape_t &shape,
  const shape_t &axes)
  {
  if (inorm==0) return T(1);
  size_t N(1);
  for (auto a: axes)
    N *= shape[a];
  return norm_fct<T>(inorm, N);
  }

Martin Reinecke's avatar
Martin Reinecke committed
104 105
template<typename T> py::array c2c_internal(const py::array &in,
  const shape_t &axes, bool forward, int inorm, bool inplace, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
106
  {
Martin Reinecke's avatar
Martin Reinecke committed
107
  auto dims(copy_shape(in));
108
  py::array res = inplace ? in : py::array_t<complex<T>>(dims);
Martin Reinecke's avatar
Martin Reinecke committed
109 110 111 112 113 114
  auto s_in=copy_strides(in);
  auto s_out=copy_strides(res);
  auto d_in=reinterpret_cast<const complex<T> *>(in.data());
  auto d_out=reinterpret_cast<complex<T> *>(res.mutable_data());
  {
  py::gil_scoped_release release;
115
  T fct = norm_fct<T>(inorm, dims, axes);
Martin Reinecke's avatar
Martin Reinecke committed
116
  pocketfft::c2c(dims, s_in, s_out, axes, forward, d_in, d_out, fct, nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
117
  }
Martin Reinecke's avatar
Martin Reinecke committed
118 119
  return res;
  }
Martin Reinecke's avatar
Martin Reinecke committed
120

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

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

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

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

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

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

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

Martin Reinecke's avatar
Martin Reinecke committed
215 216
template<typename T> py::array c2r_internal(const py::array &in,
  py::object axes_, size_t lastsize, bool forward, int inorm, size_t nthreads)
Martin Reinecke's avatar
more  
Martin Reinecke committed
217
  {
Martin Reinecke's avatar
Martin Reinecke committed
218
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
more  
Martin Reinecke committed
219
  size_t axis = axes.back();
Martin Reinecke's avatar
Martin Reinecke committed
220 221 222
  shape_t dims_in(copy_shape(in)), dims_out=dims_in;
  if (lastsize==0) lastsize=2*dims_in[axis]-1;
  if ((lastsize/2) + 1 != dims_in[axis])
Martin Reinecke's avatar
more  
Martin Reinecke committed
223 224
    throw runtime_error("bad lastsize");
  dims_out[axis] = lastsize;
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
225
  py::array res = py::array_t<T>(dims_out);
Martin Reinecke's avatar
Martin Reinecke committed
226 227 228 229 230 231
  auto s_in=copy_strides(in);
  auto s_out=copy_strides(res);
  auto d_in=reinterpret_cast<const complex<T> *>(in.data());
  auto d_out=reinterpret_cast<T *>(res.mutable_data());
  {
  py::gil_scoped_release release;
232
  T fct = norm_fct<T>(inorm, dims_out, axes);
233 234
  pocketfft::c2r(dims_out, s_in, s_out, axes, forward, d_in, d_out, fct,
    nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
235
  }
Martin Reinecke's avatar
more  
Martin Reinecke committed
236 237
  return res;
  }
238

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

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

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

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

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

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

This module supports
306
- single, double, and long double precision
Martin Reinecke's avatar
Martin Reinecke committed
307 308 309 310 311 312
- complex and real-valued transforms
- multi-dimensional transforms

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

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

Parameters
----------
319 320 321
a : numpy.ndarray (any complex or real type)
    The input data. If its type is real, a more efficient real-to-complex
    transform will be used.
322
axes : list of integers
323
    The axes along which the FFT is carried out.
324
    If not set, all axes will be transformed.
Martin Reinecke's avatar
Martin Reinecke committed
325 326
forward : bool
    If `True`, a negative sign is used in the exponent, else a positive one.
327 328 329 330 331 332
inorm : int
    Normalization type
      0 : no normalization
      1 : divide by sqrt(N)
      2 : divide by N
    where N is the product of the lengths of the transformed axes.
333
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
334
    if False, returns the result in a new array and leaves the input unchanged.
335
    if True, stores the result in the input array and returns a handle to it.
336 337
    NOTE: if `a` is real-valued and `inplace` is `True`, an exception will be
    raised!
Martin Reinecke's avatar
Martin Reinecke committed
338 339 340
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
341 342 343

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

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

Parameters
----------
352
a : numpy.ndarray (any real type)
Martin Reinecke's avatar
Martin Reinecke committed
353 354
    The input data
axes : list of integers
355
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
356
    If not set, all axes will be transformed in ascending order.
Martin Reinecke's avatar
Martin Reinecke committed
357 358
forward : bool
    If `True`, a negative sign is used in the exponent, else a positive one.
359 360 361 362 363 364
inorm : int
    Normalization type
      0 : no normalization
      1 : divide by sqrt(N)
      2 : divide by N
    where N is the product of the lengths of the transformed input axes.
Martin Reinecke's avatar
Martin Reinecke committed
365 366 367
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
Martin Reinecke's avatar
Martin Reinecke committed
368 369 370

Returns
-------
371
np.ndarray (complex type with same accuracy as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
372
    The transformed data. The shape is identical to that of the input array,
Martin Reinecke's avatar
Martin Reinecke committed
373 374
    except for the axis that was transformed last. If the length of that axis
    was n on input, it is n//2+1 on output.
Martin Reinecke's avatar
Martin Reinecke committed
375
)""";
Martin Reinecke's avatar
Martin Reinecke committed
376

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

Parameters
----------
381
a : numpy.ndarray (any complex type)
Martin Reinecke's avatar
Martin Reinecke committed
382 383
    The input data
axes : list of integers
384
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
385 386 387
    If not set, all axes will be transformed in ascending order.
lastsize : the output size of the last axis to be transformed.
    If the corresponding input axis has size n, this can be 2*n-2 or 2*n-1.
Martin Reinecke's avatar
Martin Reinecke committed
388 389
forward : bool
    If `True`, a negative sign is used in the exponent, else a positive one.
390 391 392 393 394 395
inorm : int
    Normalization type
      0 : no normalization
      1 : divide by sqrt(N)
      2 : divide by N
    where N is the product of the lengths of the transformed output axes.
Martin Reinecke's avatar
Martin Reinecke committed
396 397 398
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
Martin Reinecke's avatar
Martin Reinecke committed
399 400 401

Returns
-------
402
np.ndarray (real type with same accuracy as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
403 404
    The transformed data. The shape is identical to that of the input array,
    except for the axis that was transformed last, which has now `lastsize`
Martin Reinecke's avatar
Martin Reinecke committed
405
    entries.
Martin Reinecke's avatar
Martin Reinecke committed
406
)""";
Martin Reinecke's avatar
Martin Reinecke committed
407

Martin Reinecke's avatar
Martin Reinecke committed
408
const char *r2r_fftpack_DS = R"""(Performs a real-valued FFT using the FFTPACK storage scheme.
409 410 411

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

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

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

Parameters
----------
a : numpy.ndarray (np.float32 or np.float64)
    The input data
axes : list of integers
    The axes along which the transform is carried out.
    If not set, all axes will be transformed.
inorm : int
    Normalization type
      0 : no normalization
      1 : divide by sqrt(N)
      2 : divide by N
    where N is the product of the lengths of the transformed axes.
inplace : bool
    if False, returns the result in a new array and leaves the input unchanged.
    if True, stores the result in the input array and returns a handle to it.
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).

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

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

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

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

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

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

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