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

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

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

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

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

23
namespace {
24

25 26
using namespace std;
using namespace pocketfft;
Martin Reinecke's avatar
Martin Reinecke committed
27

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

Martin Reinecke's avatar
Martin Reinecke committed
30 31
auto c64 = py::dtype("complex64");
auto c128 = py::dtype("complex128");
32
auto c256 = py::dtype("complex256");
Martin Reinecke's avatar
Martin Reinecke committed
33 34
auto f32 = py::dtype("float32");
auto f64 = py::dtype("float64");
35
auto f128 = py::dtype("float128");
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
36 37

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

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

Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
53
shape_t makeaxes(const py::array &in, py::object axes)
Martin Reinecke's avatar
more  
Martin Reinecke committed
54
  {
Martin Reinecke's avatar
Martin Reinecke committed
55
  if (axes.is(py::none()))
Martin Reinecke's avatar
more  
Martin Reinecke committed
56
    {
57
    shape_t res(size_t(in.ndim()));
Martin Reinecke's avatar
Martin Reinecke committed
58 59 60
    for (size_t i=0; i<res.size(); ++i)
      res[i]=i;
    return res;
Martin Reinecke's avatar
more  
Martin Reinecke committed
61
    }
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
62
  auto tmp=axes.cast<shape_t>();
Martin Reinecke's avatar
fix  
Martin Reinecke committed
63
  if ((tmp.size()>size_t(in.ndim())) || (tmp.size()==0))
Martin Reinecke's avatar
more  
Martin Reinecke committed
64
    throw runtime_error("bad axes argument");
Martin Reinecke's avatar
Martin Reinecke committed
65 66
  for (auto sz: tmp)
    if (sz>=size_t(in.ndim()))
Martin Reinecke's avatar
more  
Martin Reinecke committed
67
      throw runtime_error("invalid axis number");
Martin Reinecke's avatar
Martin Reinecke committed
68
  return tmp;
Martin Reinecke's avatar
more  
Martin Reinecke committed
69 70
  }

71 72 73 74
#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; \
75
  if (dtype.is(T3)) return func<long double> args; \
76 77
  throw runtime_error("unsupported data type");

Martin Reinecke's avatar
Martin Reinecke committed
78
template<typename T> py::array xfftn_internal(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
79
  const shape_t &axes, long double fct, bool inplace, bool fwd, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
80
  {
Martin Reinecke's avatar
Martin Reinecke committed
81
  auto dims(copy_shape(in));
82
  py::array res = inplace ? in : py::array_t<complex<T>>(dims);
Martin Reinecke's avatar
sync  
Martin Reinecke committed
83 84
  c2c(dims, copy_strides(in), copy_strides(res), axes, fwd,
    reinterpret_cast<const complex<T> *>(in.data()),
Martin Reinecke's avatar
Martin Reinecke committed
85
    reinterpret_cast<complex<T> *>(res.mutable_data()), T(fct), nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
86 87
  return res;
  }
Martin Reinecke's avatar
Martin Reinecke committed
88

89
py::array xfftn(const py::array &a, py::object axes, double fct, bool inplace,
Martin Reinecke's avatar
Martin Reinecke committed
90
  bool fwd, size_t nthreads)
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
91
  {
92
  DISPATCH(a, c128, c64, c256, xfftn_internal, (a, makeaxes(a, axes), fct,
Martin Reinecke's avatar
Martin Reinecke committed
93
           inplace, fwd, nthreads))
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
94
  }
95

Martin Reinecke's avatar
Martin Reinecke committed
96 97
py::array fftn(const py::array &a, py::object axes, double fct, bool inplace,
  size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
98
  { return xfftn(a, axes, fct, inplace, true, nthreads); }
99

Martin Reinecke's avatar
Martin Reinecke committed
100 101
py::array ifftn(const py::array &a, py::object axes, double fct, bool inplace,
  size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
102
  { return xfftn(a, axes, fct, inplace, false, nthreads); }
Martin Reinecke's avatar
Martin Reinecke committed
103

Martin Reinecke's avatar
Martin Reinecke committed
104
template<typename T> py::array rfftn_internal(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
105
  py::object axes_, long double fct, size_t nthreads)
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
sync  
Martin Reinecke committed
111 112
  r2c(dims_in, copy_strides(in), copy_strides(res), axes,
    reinterpret_cast<const T *>(in.data()),
Martin Reinecke's avatar
Martin Reinecke committed
113
    reinterpret_cast<complex<T> *>(res.mutable_data()), T(fct), nthreads);
Martin Reinecke's avatar
more  
Martin Reinecke committed
114 115
  return res;
  }
116

Martin Reinecke's avatar
Martin Reinecke committed
117 118
py::array rfftn(const py::array &in, py::object axes_, double fct,
  size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
119
  {
Martin Reinecke's avatar
Martin Reinecke committed
120
  DISPATCH(in, f64, f32, f128, rfftn_internal, (in, axes_, fct, nthreads))
Martin Reinecke's avatar
Martin Reinecke committed
121
  }
122

123
template<typename T> py::array xrfft_scipy(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
124
  size_t axis, long double fct, bool inplace, bool fwd, size_t nthreads)
125 126 127
  {
  auto dims(copy_shape(in));
  py::array res = inplace ? in : py::array_t<T>(dims);
Martin Reinecke's avatar
Martin Reinecke committed
128
  r2r_fftpack(dims, copy_strides(in), copy_strides(res), axis, fwd,
Martin Reinecke's avatar
sync  
Martin Reinecke committed
129
    reinterpret_cast<const T *>(in.data()),
Martin Reinecke's avatar
Martin Reinecke committed
130
    reinterpret_cast<T *>(res.mutable_data()), T(fct), nthreads);
131 132
  return res;
  }
133

Martin Reinecke's avatar
Martin Reinecke committed
134 135
py::array rfft_scipy(const py::array &in, size_t axis, double fct, bool inplace,
  size_t nthreads)
136
  {
Martin Reinecke's avatar
Martin Reinecke committed
137 138
  DISPATCH(in, f64, f32, f128, xrfft_scipy, (in, axis, fct, inplace, true,
    nthreads))
139
  }
140

Martin Reinecke's avatar
Martin Reinecke committed
141
py::array irfft_scipy(const py::array &in, size_t axis, double fct,
Martin Reinecke's avatar
Martin Reinecke committed
142
  bool inplace, size_t nthreads)
143
  {
Martin Reinecke's avatar
Martin Reinecke committed
144 145
  DISPATCH(in, f64, f32, f128, xrfft_scipy, (in, axis, fct, inplace, false,
    nthreads))
146
  }
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
147
template<typename T> py::array irfftn_internal(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
148
  py::object axes_, size_t lastsize, long double fct, size_t nthreads)
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
more  
Martin Reinecke committed
151
  size_t axis = axes.back();
Martin Reinecke's avatar
Martin Reinecke committed
152 153 154
  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
155 156
    throw runtime_error("bad lastsize");
  dims_out[axis] = lastsize;
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
157
  py::array res = py::array_t<T>(dims_out);
Martin Reinecke's avatar
bug fix  
Martin Reinecke committed
158
  c2r(dims_out, copy_strides(in), copy_strides(res), axes,
Martin Reinecke's avatar
sync  
Martin Reinecke committed
159
    reinterpret_cast<const complex<T> *>(in.data()),
Martin Reinecke's avatar
Martin Reinecke committed
160
    reinterpret_cast<T *>(res.mutable_data()), T(fct), nthreads);
Martin Reinecke's avatar
more  
Martin Reinecke committed
161 162
  return res;
  }
163

Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
164
py::array irfftn(const py::array &in, py::object axes_, size_t lastsize,
Martin Reinecke's avatar
Martin Reinecke committed
165
  double fct, size_t nthreads)
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
166
  {
Martin Reinecke's avatar
Martin Reinecke committed
167 168
  DISPATCH(in, c128, c64, c256, irfftn_internal, (in, axes_, lastsize, fct,
    nthreads))
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
169
  }
Martin Reinecke's avatar
Martin Reinecke committed
170 171

template<typename T> py::array hartley_internal(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
172
  py::object axes_, long double fct, bool inplace, size_t nthreads)
Martin Reinecke's avatar
more  
Martin Reinecke committed
173
  {
Martin Reinecke's avatar
Martin Reinecke committed
174
  auto dims(copy_shape(in));
175
  py::array res = inplace ? in : py::array_t<T>(dims);
Martin Reinecke's avatar
Martin Reinecke committed
176
  r2r_hartley(dims, copy_strides(in), copy_strides(res), makeaxes(in, axes_),
Martin Reinecke's avatar
sync  
Martin Reinecke committed
177
    reinterpret_cast<const T *>(in.data()),
Martin Reinecke's avatar
Martin Reinecke committed
178
    reinterpret_cast<T *>(res.mutable_data()), T(fct), nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
179 180
  return res;
  }
181

182
py::array hartley(const py::array &in, py::object axes_, double fct,
Martin Reinecke's avatar
Martin Reinecke committed
183
  bool inplace, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
184
  {
Martin Reinecke's avatar
Martin Reinecke committed
185 186
  DISPATCH(in, f64, f32, f128, hartley_internal, (in, axes_, fct, inplace,
    nthreads))
Martin Reinecke's avatar
Martin Reinecke committed
187
  }
188

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

232
py::array mycomplex2hartley(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
233
  const py::array &tmp, py::object axes_, bool inplace)
234
  {
235
  DISPATCH(in, f64, f32, f128, complex2hartley, (in, tmp, axes_, inplace))
236
  }
237

Martin Reinecke's avatar
Martin Reinecke committed
238
py::array hartley2(const py::array &in, py::object axes_, double fct,
Martin Reinecke's avatar
Martin Reinecke committed
239
  bool inplace, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
240 241 242 243
  {
  return mycomplex2hartley(in, rfftn(in, axes_, fct, nthreads), axes_,
    inplace);
  }
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
244

Martin Reinecke's avatar
Martin Reinecke committed
245
const char *pypocketfft_DS = R"""(Fast Fourier and Hartley transforms.
Martin Reinecke's avatar
Martin Reinecke committed
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.
Martin Reinecke's avatar
Martin Reinecke committed
255
)""";
Martin Reinecke's avatar
Martin Reinecke committed
256

Martin Reinecke's avatar
Martin Reinecke committed
257
const char *fftn_DS = R"""(
258 259 260 261 262 263 264
Performs a forward complex FFT.

Parameters
----------
a : numpy.ndarray (np.complex64 or np.complex128)
    The input data
axes : list of integers
265
    The axes along which the FFT is carried out.
266
    If not set, all axes will be transformed.
Martin Reinecke's avatar
Martin Reinecke committed
267
fct : float
Martin Reinecke's avatar
Martin Reinecke committed
268
    Normalization factor
269
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
270
    if False, returns the result in a new array and leaves the input unchanged.
271
    if True, stores the result in the input array and returns a handle to it.
Martin Reinecke's avatar
Martin Reinecke committed
272 273 274
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
275 276 277 278

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

Martin Reinecke's avatar
Martin Reinecke committed
282
const char *ifftn_DS = R"""(Performs a backward complex FFT.
283 284 285 286 287 288

Parameters
----------
a : numpy.ndarray (np.complex64 or np.complex128)
    The input data
axes : list of integers
289
    The axes along which the FFT is carried out.
290
    If not set, all axes will be transformed.
Martin Reinecke's avatar
Martin Reinecke committed
291
fct : float
Martin Reinecke's avatar
Martin Reinecke committed
292
    Normalization factor
293
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
294
    if False, returns the result in a new array and leaves the input unchanged.
295
    if True, stores the result in the input array and returns a handle to it.
Martin Reinecke's avatar
Martin Reinecke committed
296 297 298
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
299 300 301 302

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

Martin Reinecke's avatar
Martin Reinecke committed
306
const char *rfftn_DS = R"""(Performs a forward real-valued FFT.
Martin Reinecke's avatar
Martin Reinecke committed
307 308 309 310 311 312

Parameters
----------
a : numpy.ndarray (np.float32 or np.float64)
    The input data
axes : list of integers
313
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
314 315
    If not set, all axes will be transformed in ascending order.
fct : float
Martin Reinecke's avatar
Martin Reinecke committed
316
    Normalization factor
Martin Reinecke's avatar
Martin Reinecke committed
317 318 319
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
320 321 322 323

Returns
-------
np.ndarray (np.complex64 or np.complex128)
Martin Reinecke's avatar
Martin Reinecke committed
324
    The transformed data. The shape is identical to that of the input array,
Martin Reinecke's avatar
Martin Reinecke committed
325 326
    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
327
)""";
Martin Reinecke's avatar
Martin Reinecke committed
328

Martin Reinecke's avatar
Martin Reinecke committed
329
const char *rfft_scipy_DS = R"""(Performs a forward real-valued FFT.
330 331 332 333 334 335 336 337 338 339 340 341

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.
Martin Reinecke's avatar
Martin Reinecke committed
342 343 344
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
345 346 347 348 349 350 351

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
352
)""";
353

Martin Reinecke's avatar
Martin Reinecke committed
354
const char *irfftn_DS = R"""(Performs a backward real-valued FFT.
Martin Reinecke's avatar
Martin Reinecke committed
355 356 357 358 359 360

Parameters
----------
a : numpy.ndarray (np.complex64 or np.complex128)
    The input data
axes : list of integers
361
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
362 363 364 365
    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
366
    Normalization factor
Martin Reinecke's avatar
Martin Reinecke committed
367 368 369
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
370 371 372 373

Returns
-------
np.ndarray (np.float32 or np.float64)
Martin Reinecke's avatar
Martin Reinecke committed
374 375
    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
376
    entries.
Martin Reinecke's avatar
Martin Reinecke committed
377
)""";
Martin Reinecke's avatar
Martin Reinecke committed
378

Martin Reinecke's avatar
Martin Reinecke committed
379
const char *irfft_scipy_DS = R"""(Performs a backward real-valued FFT.
380 381 382 383 384 385 386 387 388 389 390 391 392

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.
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

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
401
)""";
402

Martin Reinecke's avatar
Martin Reinecke committed
403
const char *hartley_DS = R"""(Performs a Hartley transform.
404 405 406 407 408 409 410 411 412
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
413
    The axes along which the transform is carried out.
414
    If not set, all axes will be transformed.
Martin Reinecke's avatar
Martin Reinecke committed
415
fct : float
Martin Reinecke's avatar
Martin Reinecke committed
416
    Normalization factor
417
inplace : bool
Martin Reinecke's avatar
Martin Reinecke committed
418
    if False, returns the result in a new array and leaves the input unchanged.
419
    if True, stores the result in the input array and returns a handle to it.
Martin Reinecke's avatar
Martin Reinecke committed
420 421 422
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
423 424 425 426

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

Martin Reinecke's avatar
Martin Reinecke committed
430 431 432 433
} // unnamed namespace

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

Martin Reinecke's avatar
Martin Reinecke committed
436
  m.doc() = pypocketfft_DS;
437
  m.def("fftn",&fftn, fftn_DS, "a"_a, "axes"_a=py::none(), "fct"_a=1.,
Martin Reinecke's avatar
Martin Reinecke committed
438
    "inplace"_a=false, "nthreads"_a=1);
439
  m.def("ifftn",&ifftn, ifftn_DS, "a"_a, "axes"_a=py::none(), "fct"_a=1.,
Martin Reinecke's avatar
Martin Reinecke committed
440
    "inplace"_a=false, "nthreads"_a=1);
Martin Reinecke's avatar
Martin Reinecke committed
441 442
  m.def("rfftn",&rfftn, rfftn_DS, "a"_a, "axes"_a=py::none(), "fct"_a=1.,
    "nthreads"_a=1);
443
  m.def("rfft_scipy",&rfft_scipy, rfft_scipy_DS, "a"_a, "axis"_a, "fct"_a=1.,
Martin Reinecke's avatar
Martin Reinecke committed
444
    "inplace"_a=false, "nthreads"_a=1);
445
  m.def("irfftn",&irfftn, irfftn_DS, "a"_a, "axes"_a=py::none(), "lastsize"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
446
    "fct"_a=1., "nthreads"_a=1);
447
  m.def("irfft_scipy",&irfft_scipy, irfft_scipy_DS, "a"_a, "axis"_a, "fct"_a=1.,
Martin Reinecke's avatar
Martin Reinecke committed
448
    "inplace"_a=false, "nthreads"_a=1);
449
  m.def("hartley",&hartley, hartley_DS, "a"_a, "axes"_a=py::none(), "fct"_a=1.,
Martin Reinecke's avatar
Martin Reinecke committed
450
    "inplace"_a=false, "nthreads"_a=1);
Martin Reinecke's avatar
Martin Reinecke committed
451
  m.def("hartley2",&hartley2, "a"_a, "axes"_a=py::none(), "fct"_a=1.,
Martin Reinecke's avatar
Martin Reinecke committed
452
    "inplace"_a=false, "nthreads"_a=1);
Martin Reinecke's avatar
Martin Reinecke committed
453 454
  m.def("complex2hartley",&mycomplex2hartley, "in"_a, "tmp"_a, "axes"_a,
    "inplace"_a=false);
Martin Reinecke's avatar
Martin Reinecke committed
455
  }