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

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

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

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

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

23
namespace {
24

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

This module supports
314
- single, double, and long double precision
Martin Reinecke's avatar
Martin Reinecke committed
315 316 317 318 319 320
- complex and real-valued transforms
- multi-dimensional transforms

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

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

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

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

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

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

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

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

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

Returns
-------
410
np.ndarray (real type with same accuracy as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
411 412
    The transformed data. The shape is identical to that of the input array,
    except for the axis that was transformed last, which has now `lastsize`
Martin Reinecke's avatar
Martin Reinecke committed
413
    entries.
Martin Reinecke's avatar
Martin Reinecke committed
414
)""";
Martin Reinecke's avatar
Martin Reinecke committed
415

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

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

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

Martin Reinecke's avatar
Martin Reinecke committed
449
const char *hartley_DS = R"""(Performs a Hartley transform.
450 451 452 453 454 455 456 457 458
For every requested axis, a 1D forward Fourier transform is carried out,
and the sum of real and imaginary parts of the result is stored in the output
array.

Parameters
----------
a : numpy.ndarray (np.float32 or np.float64)
    The input data
axes : list of integers
459
    The axes along which the transform is carried out.
460
    If not set, all axes will be transformed.
461 462 463 464 465 466
inorm : int
    Normalization type
      0 : no normalization
      1 : divide by sqrt(N)
      2 : divide by N
    where N is the product of the lengths of the transformed axes.
467
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
468
    if False, returns the result in a new array and leaves the input unchanged.
469
    if True, stores the result in the input array and returns a handle to it.
Martin Reinecke's avatar
Martin Reinecke committed
470 471 472
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
473 474 475

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

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

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

Martin Reinecke's avatar
Martin Reinecke committed
486
  m.doc() = pypocketfft_DS;
Martin Reinecke's avatar
Martin Reinecke committed
487
  m.def("c2c",&c2c, c2c_DS, "a"_a, "axes"_a=py::none(), "forward"_a=true,
488
    "inorm"_a=0, "inplace"_a=false, "nthreads"_a=1);
Martin Reinecke's avatar
Martin Reinecke committed
489 490 491 492 493 494 495
  m.def("r2c",&r2c, r2c_DS, "a"_a, "axes"_a=py::none(), "forward"_a=true,
    "inorm"_a=0, "nthreads"_a=1);
  m.def("c2r",&c2r, c2r_DS, "a"_a, "axes"_a=py::none(), "lastsize"_a=0,
    "forward"_a=true, "inorm"_a=0, "nthreads"_a=1);
  m.def("r2r_fftpack",&r2r_fftpack, r2r_fftpack_DS, "a"_a, "axes"_a,
    "input_halfcomplex"_a, "forward"_a, "inorm"_a=0, "inplace"_a=false,
    "nthreads"_a=1);
496
  m.def("hartley",&hartley, hartley_DS, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
497
    "inplace"_a=false, "nthreads"_a=1);
498
  m.def("hartley2",&hartley2, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
499
    "inplace"_a=false, "nthreads"_a=1);
Martin Reinecke's avatar
Martin Reinecke committed
500 501
  m.def("complex2hartley",&mycomplex2hartley, "in"_a, "tmp"_a, "axes"_a,
    "inplace"_a=false);
Martin Reinecke's avatar
Martin Reinecke committed
502
  }