pypocketfft.cc 16 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

124 125 126
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); }
127

128 129 130
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
131

Martin Reinecke's avatar
Martin Reinecke committed
132
template<typename T> py::array rfftn_internal(const py::array &in,
133
  py::object axes_, int inorm, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
134
  {
Martin Reinecke's avatar
Martin Reinecke committed
135
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
Martin Reinecke committed
136
  auto dims_in(copy_shape(in)), dims_out(dims_in);
Martin Reinecke's avatar
more  
Martin Reinecke committed
137
  dims_out[axes.back()] = (dims_out[axes.back()]>>1)+1;
Martin Reinecke's avatar
Martin Reinecke committed
138
  py::array res = py::array_t<complex<T>>(dims_out);
Martin Reinecke's avatar
Martin Reinecke committed
139 140 141 142 143 144
  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;
145 146
  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
147
  }
Martin Reinecke's avatar
more  
Martin Reinecke committed
148 149
  return res;
  }
150

151
py::array rfftn(const py::array &in, py::object axes_, int inorm,
Martin Reinecke's avatar
Martin Reinecke committed
152
  size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
153
  {
154
  DISPATCH(in, f64, f32, flong, rfftn_internal, (in, axes_, inorm, nthreads))
Martin Reinecke's avatar
Martin Reinecke committed
155
  }
156

157
template<typename T> py::array xrfft_scipy(const py::array &in,
158
  size_t axis, int inorm, bool inplace, bool fwd, size_t nthreads)
159 160 161
  {
  auto dims(copy_shape(in));
  py::array res = inplace ? in : py::array_t<T>(dims);
Martin Reinecke's avatar
Martin Reinecke committed
162 163 164 165 166 167
  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;
168 169
  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
170
  }
171 172
  return res;
  }
173

174 175
py::array rfft_scipy(const py::array &in, size_t axis, int inorm,
  bool inplace, size_t nthreads)
176
  {
177
  DISPATCH(in, f64, f32, flong, xrfft_scipy, (in, axis, inorm, inplace, true,
Martin Reinecke's avatar
Martin Reinecke committed
178
    nthreads))
179
  }
180

181
py::array irfft_scipy(const py::array &in, size_t axis, int inorm,
Martin Reinecke's avatar
Martin Reinecke committed
182
  bool inplace, size_t nthreads)
183
  {
184
  DISPATCH(in, f64, f32, flong, xrfft_scipy, (in, axis, inorm, inplace, false,
Martin Reinecke's avatar
Martin Reinecke committed
185
    nthreads))
186
  }
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
187
template<typename T> py::array irfftn_internal(const py::array &in,
188
  py::object axes_, size_t lastsize, int inorm, size_t nthreads)
Martin Reinecke's avatar
more  
Martin Reinecke committed
189
  {
Martin Reinecke's avatar
Martin Reinecke committed
190
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
more  
Martin Reinecke committed
191
  size_t axis = axes.back();
Martin Reinecke's avatar
Martin Reinecke committed
192 193 194
  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
195 196
    throw runtime_error("bad lastsize");
  dims_out[axis] = lastsize;
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
197
  py::array res = py::array_t<T>(dims_out);
Martin Reinecke's avatar
Martin Reinecke committed
198 199 200 201 202 203
  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;
204 205
  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
206
  }
Martin Reinecke's avatar
more  
Martin Reinecke committed
207 208
  return res;
  }
209

Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
210
py::array irfftn(const py::array &in, py::object axes_, size_t lastsize,
211
  int inorm, size_t nthreads)
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
212
  {
213
  DISPATCH(in, c128, c64, clong, irfftn_internal, (in, axes_, lastsize, inorm,
Martin Reinecke's avatar
Martin Reinecke committed
214
    nthreads))
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
215
  }
Martin Reinecke's avatar
Martin Reinecke committed
216 217

template<typename T> py::array hartley_internal(const py::array &in,
218
  py::object axes_, int inorm, bool inplace, size_t nthreads)
Martin Reinecke's avatar
more  
Martin Reinecke committed
219
  {
Martin Reinecke's avatar
Martin Reinecke committed
220
  auto dims(copy_shape(in));
221
  py::array res = inplace ? in : py::array_t<T>(dims);
Martin Reinecke's avatar
Martin Reinecke committed
222 223 224 225 226 227 228
  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;
229 230
  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
231
  }
Martin Reinecke's avatar
Martin Reinecke committed
232 233
  return res;
  }
234

235
py::array hartley(const py::array &in, py::object axes_, int inorm,
Martin Reinecke's avatar
Martin Reinecke committed
236
  bool inplace, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
237
  {
238
  DISPATCH(in, f64, f32, flong, hartley_internal, (in, axes_, inorm, inplace,
Martin Reinecke's avatar
Martin Reinecke committed
239
    nthreads))
Martin Reinecke's avatar
Martin Reinecke committed
240
  }
241

Martin Reinecke's avatar
fixes  
Martin Reinecke committed
242
template<typename T>py::array complex2hartley(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
243
  const py::array &tmp, py::object axes_, bool inplace)
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
244
  {
Martin Reinecke's avatar
Martin Reinecke committed
245
  using namespace pocketfft::detail;
Martin Reinecke's avatar
Martin Reinecke committed
246
  auto dims_out(copy_shape(in));
Martin Reinecke's avatar
Martin Reinecke committed
247 248 249
  py::array out = inplace ? in : py::array_t<T>(dims_out);
  ndarr<cmplx<T>> atmp(tmp.data(), copy_shape(tmp), copy_strides(tmp));
  ndarr<T> aout(out.mutable_data(), copy_shape(out), copy_strides(out));
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
250
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
Martin Reinecke committed
251 252
  {
  py::gil_scoped_release release;
253 254 255 256
  simple_iter<cmplx<T>> iin(atmp);
  rev_iter<T> iout(aout, axes);
  if (iin.remaining()!=iout.remaining()) throw runtime_error("oops");
  while(iin.remaining()>0)
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
257
    {
258 259 260 261
    auto v = atmp.get(iin.ofs());
    aout.set(iout.ofs(), v.r+v.i);
    aout.set(iout.rev_ofs(), v.r-v.i);
    iin.advance(); iout.advance();
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
262
    }
Martin Reinecke's avatar
Martin Reinecke committed
263
  }
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
264 265
  return out;
  }
266

267
py::array mycomplex2hartley(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
268
  const py::array &tmp, py::object axes_, bool inplace)
269
  {
270
  DISPATCH(in, f64, f32, flong, complex2hartley, (in, tmp, axes_, inplace))
271
  }
272

273
py::array hartley2(const py::array &in, py::object axes_, int inorm,
Martin Reinecke's avatar
Martin Reinecke committed
274
  bool inplace, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
275
  {
276
  return mycomplex2hartley(in, rfftn(in, axes_, inorm, nthreads), axes_,
Martin Reinecke's avatar
Martin Reinecke committed
277 278
    inplace);
  }
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
279

Martin Reinecke's avatar
Martin Reinecke committed
280
const char *pypocketfft_DS = R"""(Fast Fourier and Hartley transforms.
Martin Reinecke's avatar
Martin Reinecke committed
281 282 283 284 285 286 287 288 289

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
290
)""";
Martin Reinecke's avatar
Martin Reinecke committed
291

Martin Reinecke's avatar
Martin Reinecke committed
292
const char *fftn_DS = R"""(
293 294 295 296 297 298 299
Performs a forward complex FFT.

Parameters
----------
a : numpy.ndarray (np.complex64 or np.complex128)
    The input data
axes : list of integers
300
    The axes along which the FFT is carried out.
301
    If not set, all axes will be transformed.
302 303 304 305 306 307
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.
308
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
309
    if False, returns the result in a new array and leaves the input unchanged.
310
    if True, stores the result in the input array and returns a handle to it.
Martin Reinecke's avatar
Martin Reinecke committed
311 312 313
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
314 315 316 317

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

Martin Reinecke's avatar
Martin Reinecke committed
321
const char *ifftn_DS = R"""(Performs a backward complex FFT.
322 323 324 325 326 327

Parameters
----------
a : numpy.ndarray (np.complex64 or np.complex128)
    The input data
axes : list of integers
328
    The axes along which the FFT is carried out.
329
    If not set, all axes will be transformed.
330 331 332 333 334 335
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.
336
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
337
    if False, returns the result in a new array and leaves the input unchanged.
338
    if True, stores the result in the input array and returns a handle to it.
Martin Reinecke's avatar
Martin Reinecke committed
339 340 341
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
342 343 344 345

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

Martin Reinecke's avatar
Martin Reinecke committed
349
const char *rfftn_DS = R"""(Performs a forward real-valued FFT.
Martin Reinecke's avatar
Martin Reinecke committed
350 351 352 353 354 355

Parameters
----------
a : numpy.ndarray (np.float32 or np.float64)
    The input data
axes : list of integers
356
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
357
    If not set, all axes will be transformed in ascending order.
358 359 360 361 362 363
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
364 365 366
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
367 368 369 370

Returns
-------
np.ndarray (np.complex64 or np.complex128)
Martin Reinecke's avatar
Martin Reinecke committed
371
    The transformed data. The shape is identical to that of the input array,
Martin Reinecke's avatar
Martin Reinecke committed
372 373
    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
374
)""";
Martin Reinecke's avatar
Martin Reinecke committed
375

Martin Reinecke's avatar
Martin Reinecke committed
376
const char *rfft_scipy_DS = R"""(Performs a forward real-valued FFT.
377 378 379 380 381 382 383

Parameters
----------
a : numpy.ndarray (np.float32 or np.float64)
    The input data
axis : int
    The axis along which the FFT is carried out.
384 385 386 387 388 389
inorm : int
    Normalization type
      0 : no normalization
      1 : divide by sqrt(N)
      2 : divide by N
    where N is the length of `axis`.
390 391 392
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
393 394 395
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
396 397 398 399 400 401 402

Returns
-------
np.ndarray (np.float32 or np.float64)
    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
403
)""";
404

Martin Reinecke's avatar
Martin Reinecke committed
405
const char *irfftn_DS = R"""(Performs a backward real-valued FFT.
Martin Reinecke's avatar
Martin Reinecke committed
406 407 408 409 410 411

Parameters
----------
a : numpy.ndarray (np.complex64 or np.complex128)
    The input data
axes : list of integers
412
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
413 414 415
    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.
416 417 418 419 420 421
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
422 423 424
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
425 426 427 428

Returns
-------
np.ndarray (np.float32 or np.float64)
Martin Reinecke's avatar
Martin Reinecke committed
429 430
    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
431
    entries.
Martin Reinecke's avatar
Martin Reinecke committed
432
)""";
Martin Reinecke's avatar
Martin Reinecke committed
433

Martin Reinecke's avatar
Martin Reinecke committed
434
const char *irfft_scipy_DS = R"""(Performs a backward real-valued FFT.
435 436 437 438 439 440 441 442

Parameters
----------
a : numpy.ndarray (np.float32 or np.float64)
    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.
443 444 445 446 447 448
inorm : int
    Normalization type
      0 : no normalization
      1 : divide by sqrt(N)
      2 : divide by N
    where N is the length of `axis`.
449 450 451
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
452 453 454
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
455 456 457 458 459

Returns
-------
np.ndarray (np.float32 or np.float64)
    The transformed data. The shape is identical to that of the input array.
Martin Reinecke's avatar
Martin Reinecke committed
460
)""";
461

Martin Reinecke's avatar
Martin Reinecke committed
462
const char *hartley_DS = R"""(Performs a Hartley transform.
463 464 465 466 467 468 469 470 471
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
472
    The axes along which the transform is carried out.
473
    If not set, all axes will be transformed.
474 475 476 477 478 479
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.
480
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
481
    if False, returns the result in a new array and leaves the input unchanged.
482
    if True, stores the result in the input array and returns a handle to it.
Martin Reinecke's avatar
Martin Reinecke committed
483 484 485
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
486 487 488 489

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

Martin Reinecke's avatar
Martin Reinecke committed
493 494 495 496
} // unnamed namespace

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

Martin Reinecke's avatar
Martin Reinecke committed
499
  m.doc() = pypocketfft_DS;
500
  m.def("fftn",&fftn, fftn_DS, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
501
    "inplace"_a=false, "nthreads"_a=1);
502
  m.def("ifftn",&ifftn, ifftn_DS, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
503
    "inplace"_a=false, "nthreads"_a=1);
504
  m.def("rfftn",&rfftn, rfftn_DS, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
505
    "nthreads"_a=1);
506
  m.def("rfft_scipy",&rfft_scipy, rfft_scipy_DS, "a"_a, "axis"_a, "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
507
    "inplace"_a=false, "nthreads"_a=1);
508
  m.def("irfftn",&irfftn, irfftn_DS, "a"_a, "axes"_a=py::none(), "lastsize"_a=0,
509 510 511 512
    "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
513
    "inplace"_a=false, "nthreads"_a=1);
514
  m.def("hartley2",&hartley2, "a"_a, "axes"_a=py::none(), "inorm"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
515
    "inplace"_a=false, "nthreads"_a=1);
Martin Reinecke's avatar
Martin Reinecke committed
516 517
  m.def("complex2hartley",&mycomplex2hartley, "in"_a, "tmp"_a, "axes"_a,
    "inplace"_a=false);
Martin Reinecke's avatar
Martin Reinecke committed
518
  }