pypocketfft.cc 13.3 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
#pragma GCC visibility push(hidden)
Martin Reinecke's avatar
Martin Reinecke committed
18

19
#include "pocketfft_hdronly.h"
Martin Reinecke's avatar
Martin Reinecke committed
20

Martin Reinecke's avatar
Martin Reinecke committed
21
//
22
// Python interface
Martin Reinecke's avatar
Martin Reinecke committed
23 24
//

25
namespace {
26

27 28 29
using namespace std;
using namespace pocketfft;
using namespace pocketfft::detail;
Martin Reinecke's avatar
Martin Reinecke committed
30

Martin Reinecke's avatar
Martin Reinecke committed
31 32
namespace py = pybind11;

Martin Reinecke's avatar
Martin Reinecke committed
33 34
auto c64 = py::dtype("complex64");
auto c128 = py::dtype("complex128");
Martin Reinecke's avatar
Martin Reinecke committed
35 36
auto f32 = py::dtype("float32");
auto f64 = py::dtype("float64");
Martin Reinecke's avatar
Martin Reinecke committed
37

Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
38 39 40 41 42 43 44 45 46 47
bool tcheck(const py::array &arr, const py::object &t1, const py::object &t2)
  {
  if (arr.dtype().is(t1))
    return true;
  if (arr.dtype().is(t2))
    return false;
  throw runtime_error("unsupported data type");
  }

shape_t copy_shape(const py::array &arr)
Martin Reinecke's avatar
Martin Reinecke committed
48
  {
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
49
  shape_t res(arr.ndim());
Martin Reinecke's avatar
Martin Reinecke committed
50 51 52 53
  for (size_t i=0; i<res.size(); ++i)
    res[i] = arr.shape(i);
  return res;
  }
Martin Reinecke's avatar
Martin Reinecke committed
54

Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
55
stride_t copy_strides(const py::array &arr)
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
56
  {
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
57
  stride_t res(arr.ndim());
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
58
  for (size_t i=0; i<res.size(); ++i)
Martin Reinecke's avatar
Martin Reinecke committed
59
    res[i] = arr.strides(i);
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
60 61
  return res;
  }
Martin Reinecke's avatar
Martin Reinecke committed
62

Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
63
shape_t makeaxes(const py::array &in, py::object axes)
Martin Reinecke's avatar
more  
Martin Reinecke committed
64
  {
Martin Reinecke's avatar
Martin Reinecke committed
65
  if (axes.is(py::none()))
Martin Reinecke's avatar
more  
Martin Reinecke committed
66
    {
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
67
    shape_t res(in.ndim());
Martin Reinecke's avatar
Martin Reinecke committed
68 69 70
    for (size_t i=0; i<res.size(); ++i)
      res[i]=i;
    return res;
Martin Reinecke's avatar
more  
Martin Reinecke committed
71
    }
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
72
  auto tmp=axes.cast<shape_t>();
Martin Reinecke's avatar
fix  
Martin Reinecke committed
73
  if ((tmp.size()>size_t(in.ndim())) || (tmp.size()==0))
Martin Reinecke's avatar
more  
Martin Reinecke committed
74
    throw runtime_error("bad axes argument");
Martin Reinecke's avatar
Martin Reinecke committed
75 76
  for (auto sz: tmp)
    if (sz>=size_t(in.ndim()))
Martin Reinecke's avatar
more  
Martin Reinecke committed
77
      throw runtime_error("invalid axis number");
Martin Reinecke's avatar
Martin Reinecke committed
78
  return tmp;
Martin Reinecke's avatar
more  
Martin Reinecke committed
79 80
  }

Martin Reinecke's avatar
Martin Reinecke committed
81
template<typename T> py::array xfftn_internal(const py::array &in,
82
  const shape_t &axes, double fct, bool inplace, bool fwd)
Martin Reinecke's avatar
Martin Reinecke committed
83
  {
Martin Reinecke's avatar
Martin Reinecke committed
84
  auto dims(copy_shape(in));
85
  py::array res = inplace ? in : py::array_t<complex<T>>(dims);
Martin Reinecke's avatar
Martin Reinecke committed
86 87
  ndarr<cmplx<T>> ain(in.data(), dims, copy_strides(in));
  ndarr<cmplx<T>> aout(res.mutable_data(), dims, copy_strides(res));
Martin Reinecke's avatar
Martin Reinecke committed
88
  general_c<T>(ain, aout, axes, fwd, fct);
Martin Reinecke's avatar
Martin Reinecke committed
89 90
  return res;
  }
Martin Reinecke's avatar
Martin Reinecke committed
91

92 93
py::array xfftn(const py::array &a, py::object axes, double fct, bool inplace,
  bool fwd)
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
94
  {
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
95
  return tcheck(a, c128, c64) ?
96 97
    xfftn_internal<double>(a, makeaxes(a, axes), fct, inplace, fwd) :
    xfftn_internal<float> (a, makeaxes(a, axes), fct, inplace, fwd);
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
98
  }
99 100 101 102
py::array fftn(const py::array &a, py::object axes, double fct, bool inplace)
  { return xfftn(a, axes, fct, inplace, true); }
py::array ifftn(const py::array &a, py::object axes, double fct, bool inplace)
  { return xfftn(a, axes, fct, inplace, false); }
Martin Reinecke's avatar
Martin Reinecke committed
103

Martin Reinecke's avatar
Martin Reinecke committed
104 105
template<typename T> py::array rfftn_internal(const py::array &in,
  py::object axes_, T fct)
Martin Reinecke's avatar
Martin Reinecke committed
106
  {
Martin Reinecke's avatar
Martin Reinecke committed
107
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
Martin Reinecke committed
108
  auto dims_in(copy_shape(in)), dims_out(dims_in);
Martin Reinecke's avatar
more  
Martin Reinecke committed
109
  dims_out[axes.back()] = (dims_out[axes.back()]>>1)+1;
Martin Reinecke's avatar
Martin Reinecke committed
110
  py::array res = py::array_t<complex<T>>(dims_out);
Martin Reinecke's avatar
Martin Reinecke committed
111 112
  ndarr<T> ain(in.data(), dims_in, copy_strides(in));
  ndarr<cmplx<T>> aout(res.mutable_data(), dims_out, copy_strides(res));
Martin Reinecke's avatar
Martin Reinecke committed
113
  general_r2c<T>(ain, aout, axes.back(), fct);
114
  if (axes.size()==1) return res;
Martin Reinecke's avatar
Martin Reinecke committed
115
  shape_t axes2(axes.begin(), --axes.end());
Martin Reinecke's avatar
Martin Reinecke committed
116
  general_c<T>(aout, aout, axes2, true, 1.);
Martin Reinecke's avatar
more  
Martin Reinecke committed
117 118
  return res;
  }
Martin Reinecke's avatar
Martin Reinecke committed
119 120
py::array rfftn(const py::array &in, py::object axes_, double fct)
  {
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
121 122
  return tcheck(in, f64, f32) ? rfftn_internal<double>(in, axes_, fct)
                              : rfftn_internal<float> (in, axes_, fct);
Martin Reinecke's avatar
Martin Reinecke committed
123
  }
124
template<typename T> py::array xrfft_scipy(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
125
  size_t axis, double fct, bool inplace, bool fwd)
126 127 128 129 130
  {
  auto dims(copy_shape(in));
  py::array res = inplace ? in : py::array_t<T>(dims);
  ndarr<T> ain(in.data(), dims, copy_strides(in));
  ndarr<T> aout(res.mutable_data(), dims, copy_strides(res));
Martin Reinecke's avatar
Martin Reinecke committed
131
  general_r<T>(ain, aout, axis, fwd, fct);
132 133
  return res;
  }
Martin Reinecke's avatar
Martin Reinecke committed
134
py::array rfft_scipy(const py::array &in, size_t axis, double fct, bool inplace)
135 136 137 138 139
  {
  return tcheck(in, f64, f32) ?
    xrfft_scipy<double>(in, axis, fct, inplace, true) :
    xrfft_scipy<float> (in, axis, fct, inplace, true);
  }
Martin Reinecke's avatar
Martin Reinecke committed
140 141
py::array irfft_scipy(const py::array &in, size_t axis, double fct,
  bool inplace)
142 143 144 145 146
  {
  return tcheck(in, f64, f32) ?
    xrfft_scipy<double>(in, axis, fct, inplace, false) :
    xrfft_scipy<float> (in, axis, fct, inplace, false);
  }
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
147 148
template<typename T> py::array irfftn_internal(const py::array &in,
  py::object axes_, size_t lastsize, T fct)
Martin Reinecke's avatar
more  
Martin Reinecke committed
149
  {
Martin Reinecke's avatar
Martin Reinecke committed
150
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
Martin Reinecke committed
151 152
  py::array inter = (axes.size()==1) ? in :
    xfftn_internal<T>(in, shape_t(axes.begin(), --axes.end()), 1., false, false);
Martin Reinecke's avatar
more  
Martin Reinecke committed
153 154

  size_t axis = axes.back();
Martin Reinecke's avatar
Martin Reinecke committed
155
  if (lastsize==0) lastsize=2*inter.shape(axis)-1;
Martin Reinecke's avatar
Martin Reinecke committed
156
  if (ptrdiff_t(lastsize/2) + 1 != inter.shape(axis))
Martin Reinecke's avatar
more  
Martin Reinecke committed
157
    throw runtime_error("bad lastsize");
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
158
  auto dims_out(copy_shape(inter));
Martin Reinecke's avatar
more  
Martin Reinecke committed
159
  dims_out[axis] = lastsize;
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
160
  py::array res = py::array_t<T>(dims_out);
Martin Reinecke's avatar
Martin Reinecke committed
161 162
  ndarr<cmplx<T>> ain(inter.data(), copy_shape(inter), copy_strides(inter));
  ndarr<T> aout(res.mutable_data(), dims_out, copy_strides(res));
Martin Reinecke's avatar
Martin Reinecke committed
163
  general_c2r<T>(ain, aout, axis, fct);
Martin Reinecke's avatar
more  
Martin Reinecke committed
164 165
  return res;
  }
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
166 167 168
py::array irfftn(const py::array &in, py::object axes_, size_t lastsize,
  double fct)
  {
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
169 170 171
  return tcheck(in, c128, c64) ?
    irfftn_internal<double>(in, axes_, lastsize, fct) :
    irfftn_internal<float> (in, axes_, lastsize, fct);
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
172
  }
Martin Reinecke's avatar
Martin Reinecke committed
173 174

template<typename T> py::array hartley_internal(const py::array &in,
175
  py::object axes_, double fct, bool inplace)
Martin Reinecke's avatar
more  
Martin Reinecke committed
176
  {
Martin Reinecke's avatar
Martin Reinecke committed
177
  auto dims(copy_shape(in));
178
  py::array res = inplace ? in : py::array_t<T>(dims);
Martin Reinecke's avatar
Martin Reinecke committed
179 180
  ndarr<T> ain(in.data(), copy_shape(in), copy_strides(in));
  ndarr<T> aout(res.mutable_data(), copy_shape(res), copy_strides(res));
Martin Reinecke's avatar
Martin Reinecke committed
181
  general_hartley<T>(ain, aout, makeaxes(in, axes_), fct);
Martin Reinecke's avatar
Martin Reinecke committed
182 183
  return res;
  }
184 185
py::array hartley(const py::array &in, py::object axes_, double fct,
  bool inplace)
Martin Reinecke's avatar
Martin Reinecke committed
186
  {
187 188 189
  return tcheck(in, f64, f32) ?
    hartley_internal<double>(in, axes_, fct, inplace) :
    hartley_internal<float> (in, axes_, fct, inplace);
Martin Reinecke's avatar
Martin Reinecke committed
190
  }
191

Martin Reinecke's avatar
fixes  
Martin Reinecke committed
192
template<typename T>py::array complex2hartley(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
193
  const py::array &tmp, py::object axes_, bool inplace)
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
194 195
  {
  int ndim = in.ndim();
Martin Reinecke's avatar
Martin Reinecke committed
196
  auto dims_out(copy_shape(in));
Martin Reinecke's avatar
Martin Reinecke committed
197 198 199
  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
200
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
Martin Reinecke committed
201
  size_t axis = axes.back();
Martin Reinecke's avatar
Martin Reinecke committed
202 203
  multi_iter<1,cmplx<T>,T> it(atmp, aout, axis);
  vector<bool> swp(ndim,false);
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
204 205
  for (auto i: axes)
    if (i!=axis)
Martin Reinecke's avatar
Martin Reinecke committed
206
      swp[i] = true;
Martin Reinecke's avatar
Martin Reinecke committed
207
  while(it.remaining()>0)
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
208
    {
Martin Reinecke's avatar
Martin Reinecke committed
209
    ptrdiff_t rofs = 0;
Martin Reinecke's avatar
Martin Reinecke committed
210
    for (size_t i=0; i<it.pos.size(); ++i)
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
211
      {
Martin Reinecke's avatar
Martin Reinecke committed
212
      if (i==axis) continue;
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
213
      if (!swp[i])
Martin Reinecke's avatar
Martin Reinecke committed
214
        rofs += it.pos[i]*it.oarr.stride(i);
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
215 216
      else
        {
Martin Reinecke's avatar
Martin Reinecke committed
217 218
        auto x = (it.pos[i]==0) ? 0 : it.iarr.shape(i)-it.pos[i];
        rofs += x*it.oarr.stride(i);
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
219 220
        }
      }
Martin Reinecke's avatar
Martin Reinecke committed
221
    it.advance(1);
Martin Reinecke's avatar
Martin Reinecke committed
222
    for (size_t i=0; i<it.length_in(); ++i)
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
223
      {
Martin Reinecke's avatar
Martin Reinecke committed
224 225
      auto re = it.in(i).r;
      auto im = it.in(i).i;
Martin Reinecke's avatar
Martin Reinecke committed
226
      auto rev_i = (i==0) ? 0 : it.length_out()-i;
Martin Reinecke's avatar
Martin Reinecke committed
227 228
      it.out(i) = re+im;
      aout[rofs + rev_i*it.stride_out()] = re-im;
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
229 230 231 232
      }
    }
  return out;
  }
233
py::array mycomplex2hartley(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
234
  const py::array &tmp, py::object axes_, bool inplace)
235
  {
Martin Reinecke's avatar
Martin Reinecke committed
236 237
  return tcheck(in, f64, f32) ? complex2hartley<double>(in, tmp, axes_, inplace)
                              : complex2hartley<float> (in, tmp, axes_, inplace);
238
  }
Martin Reinecke's avatar
Martin Reinecke committed
239 240 241
py::array hartley2(const py::array &in, py::object axes_, double fct,
  bool inplace)
  { return mycomplex2hartley(in, rfftn(in, axes_, fct), axes_, inplace); }
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
242

Martin Reinecke's avatar
Martin Reinecke committed
243
const char *pypocketfft_DS = R"DELIM(Fast Fourier and Hartley transforms.
Martin Reinecke's avatar
Martin Reinecke committed
244 245 246 247 248 249 250 251 252 253 254

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.
)DELIM";

255 256 257 258 259 260 261 262
const char *fftn_DS = R"DELIM(
Performs a forward complex FFT.

Parameters
----------
a : numpy.ndarray (np.complex64 or np.complex128)
    The input data
axes : list of integers
263
    The axes along which the FFT is carried out.
264
    If not set, all axes will be transformed.
Martin Reinecke's avatar
Martin Reinecke committed
265
fct : float
Martin Reinecke's avatar
Martin Reinecke committed
266
    Normalization factor
267
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
268
    if False, returns the result in a new array and leaves the input unchanged.
269
    if True, stores the result in the input array and returns a handle to it.
270 271 272 273

Returns
-------
np.ndarray (same shape and data type as a)
274
    The transformed data.
275 276
)DELIM";

Martin Reinecke's avatar
Martin Reinecke committed
277
const char *ifftn_DS = R"DELIM(Performs a backward complex FFT.
278 279 280 281 282 283

Parameters
----------
a : numpy.ndarray (np.complex64 or np.complex128)
    The input data
axes : list of integers
284
    The axes along which the FFT is carried out.
285
    If not set, all axes will be transformed.
Martin Reinecke's avatar
Martin Reinecke committed
286
fct : float
Martin Reinecke's avatar
Martin Reinecke committed
287
    Normalization factor
288
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
289
    if False, returns the result in a new array and leaves the input unchanged.
290
    if True, stores the result in the input array and returns a handle to it.
291 292 293 294

Returns
-------
np.ndarray (same shape and data type as a)
Martin Reinecke's avatar
Martin Reinecke committed
295
    The transformed data
296 297
)DELIM";

Martin Reinecke's avatar
Martin Reinecke committed
298
const char *rfftn_DS = R"DELIM(Performs a forward real-valued FFT.
Martin Reinecke's avatar
Martin Reinecke committed
299 300 301 302 303 304

Parameters
----------
a : numpy.ndarray (np.float32 or np.float64)
    The input data
axes : list of integers
305
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
306 307
    If not set, all axes will be transformed in ascending order.
fct : float
Martin Reinecke's avatar
Martin Reinecke committed
308
    Normalization factor
Martin Reinecke's avatar
Martin Reinecke committed
309 310 311 312

Returns
-------
np.ndarray (np.complex64 or np.complex128)
Martin Reinecke's avatar
Martin Reinecke committed
313
    The transformed data. The shape is identical to that of the input array,
Martin Reinecke's avatar
Martin Reinecke committed
314 315 316 317
    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.
)DELIM";

318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339
const char *rfft_scipy_DS = R"DELIM(Performs a forward real-valued FFT.

Parameters
----------
a : numpy.ndarray (np.float32 or np.float64)
    The input data
axis : int
    The axis along which the FFT is carried out.
fct : float
    Normalization factor
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.

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 ...`.
)DELIM";

Martin Reinecke's avatar
Martin Reinecke committed
340
const char *irfftn_DS = R"DELIM(Performs a backward real-valued FFT.
Martin Reinecke's avatar
Martin Reinecke committed
341 342 343 344 345 346

Parameters
----------
a : numpy.ndarray (np.complex64 or np.complex128)
    The input data
axes : list of integers
347
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
348 349 350 351
    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.
fct : float
Martin Reinecke's avatar
Martin Reinecke committed
352
    Normalization factor
Martin Reinecke's avatar
Martin Reinecke committed
353 354 355 356

Returns
-------
np.ndarray (np.float32 or np.float64)
Martin Reinecke's avatar
Martin Reinecke committed
357 358
    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
359 360 361
    entries.
)DELIM";

362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382
const char *irfft_scipy_DS = R"DELIM(Performs a backward real-valued FFT.

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.
fct : float
    Normalization factor
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.

Returns
-------
np.ndarray (np.float32 or np.float64)
    The transformed data. The shape is identical to that of the input array.
)DELIM";

383
const char *hartley_DS = R"DELIM(Performs a Hartley transform.
384 385 386 387 388 389 390 391 392
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
393
    The axes along which the transform is carried out.
394
    If not set, all axes will be transformed.
Martin Reinecke's avatar
Martin Reinecke committed
395
fct : float
Martin Reinecke's avatar
Martin Reinecke committed
396
    Normalization factor
397
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
398
    if False, returns the result in a new array and leaves the input unchanged.
399
    if True, stores the result in the input array and returns a handle to it.
400 401 402 403

Returns
-------
np.ndarray (same shape and data type as a)
Martin Reinecke's avatar
Martin Reinecke committed
404
    The transformed data
405 406
)DELIM";

Martin Reinecke's avatar
Martin Reinecke committed
407 408
} // unnamed namespace

409 410
#pragma GCC visibility pop

Martin Reinecke's avatar
Martin Reinecke committed
411 412
PYBIND11_MODULE(pypocketfft, m)
  {
Martin Reinecke's avatar
more  
Martin Reinecke committed
413 414
  using namespace pybind11::literals;

Martin Reinecke's avatar
Martin Reinecke committed
415
  m.doc() = pypocketfft_DS;
416 417 418 419
  m.def("fftn",&fftn, fftn_DS, "a"_a, "axes"_a=py::none(), "fct"_a=1.,
    "inplace"_a=false);
  m.def("ifftn",&ifftn, ifftn_DS, "a"_a, "axes"_a=py::none(), "fct"_a=1.,
    "inplace"_a=false);
Martin Reinecke's avatar
Martin Reinecke committed
420
  m.def("rfftn",&rfftn, rfftn_DS, "a"_a, "axes"_a=py::none(), "fct"_a=1.);
421
  m.def("rfft_scipy",&rfft_scipy, rfft_scipy_DS, "a"_a, "axis"_a, "fct"_a=1.,
422
    "inplace"_a=false);
423 424
  m.def("irfftn",&irfftn, irfftn_DS, "a"_a, "axes"_a=py::none(), "lastsize"_a=0,
    "fct"_a=1.);
425 426
  m.def("irfft_scipy",&irfft_scipy, irfft_scipy_DS, "a"_a, "axis"_a, "fct"_a=1.,
    "inplace"_a=false);
427 428
  m.def("hartley",&hartley, hartley_DS, "a"_a, "axes"_a=py::none(), "fct"_a=1.,
    "inplace"_a=false);
Martin Reinecke's avatar
Martin Reinecke committed
429 430 431 432
  m.def("hartley2",&hartley2, "a"_a, "axes"_a=py::none(), "fct"_a=1.,
    "inplace"_a=false);
  m.def("complex2hartley",&mycomplex2hartley, "in"_a, "tmp"_a, "axes"_a,
    "inplace"_a=false);
Martin Reinecke's avatar
Martin Reinecke committed
433
  }