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

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

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

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

19
namespace {
20

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

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

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

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

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

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

Martin Reinecke's avatar
Martin Reinecke committed
55
shape_t makeaxes(const py::array &in, const py::object &axes)
Martin Reinecke's avatar
more  
Martin Reinecke committed
56
  {
Martin Reinecke's avatar
Martin Reinecke committed
57
  if (axes.is(None))
Martin Reinecke's avatar
more  
Martin Reinecke committed
58
    {
59
    shape_t res(size_t(in.ndim()));
Martin Reinecke's avatar
Martin Reinecke committed
60 61 62
    for (size_t i=0; i<res.size(); ++i)
      res[i]=i;
    return res;
Martin Reinecke's avatar
more  
Martin Reinecke committed
63
    }
64 65 66
  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
67
    throw runtime_error("bad axes argument");
68 69 70 71 72 73 74 75
  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
76 77
  }

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

87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104
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
105 106 107 108 109 110 111 112 113 114
template<typename T> py::array_t<T> prepare_output(py::object &out_,
  shape_t &dims)
  {
  if (out_.is(None)) return py::array_t<T>(dims);
  auto tmp = out_.cast<py::array_t<T>>();
  if (!tmp.is(out_)) // a new object was created during casting
    throw runtime_error("unexpected data type for output array");
  return tmp;
  }

Martin Reinecke's avatar
Martin Reinecke committed
115
template<typename T> py::array c2c_internal(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
116 117
  const py::object &axes_, bool forward, int inorm, py::object &out_,
  size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
118
  {
Martin Reinecke's avatar
Martin Reinecke committed
119
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
Martin Reinecke committed
120
  auto dims(copy_shape(in));
Martin Reinecke's avatar
Martin Reinecke committed
121
  auto res = prepare_output<complex<T>>(out_, dims);
Martin Reinecke's avatar
Martin Reinecke committed
122 123 124 125 126 127
  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;
128
  T fct = norm_fct<T>(inorm, dims, axes);
Martin Reinecke's avatar
Martin Reinecke committed
129
  pocketfft::c2c(dims, s_in, s_out, axes, forward, d_in, d_out, fct, nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
130
  }
Martin Reinecke's avatar
Martin Reinecke committed
131 132
  return res;
  }
Martin Reinecke's avatar
Martin Reinecke committed
133

Martin Reinecke's avatar
Martin Reinecke committed
134
template<typename T> py::array c2c_sym_internal(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
135 136
  const py::object &axes_, bool forward, int inorm, py::object &out_,
  size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
137
  {
Martin Reinecke's avatar
Martin Reinecke committed
138
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
Martin Reinecke committed
139
  auto dims(copy_shape(in));
Martin Reinecke's avatar
Martin Reinecke committed
140
  auto res = prepare_output<complex<T>>(out_, dims);
Martin Reinecke's avatar
Martin Reinecke committed
141 142 143 144 145 146 147
  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
148
  pocketfft::r2c(dims, s_in, s_out, axes, forward, d_in, d_out, fct, nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
149 150 151
  // now fill in second half
  using namespace pocketfft::detail;
  ndarr<complex<T>> ares(res.mutable_data(), dims, s_out);
152
  rev_iter iter(ares, axes);
Martin Reinecke's avatar
Martin Reinecke committed
153 154
  while(iter.remaining()>0)
    {
155 156
    auto v = ares[iter.ofs()];
    ares[iter.rev_ofs()] = conj(v);
Martin Reinecke's avatar
Martin Reinecke committed
157 158 159 160 161 162
    iter.advance();
    }
  }
  return res;
  }

Martin Reinecke's avatar
Martin Reinecke committed
163 164
py::array c2c(const py::array &a, const py::object &axes_, bool forward,
  int inorm, py::object &out_, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
165
  {
166
  if (a.dtype().kind() == 'c')
Martin Reinecke's avatar
Martin Reinecke committed
167 168
    DISPATCH(a, c128, c64, clong, c2c_internal, (a, axes_, forward,
             inorm, out_, nthreads))
169

Martin Reinecke's avatar
Martin Reinecke committed
170 171
  DISPATCH(a, f64, f32, flong, c2c_sym_internal, (a, axes_, forward,
           inorm, out_, nthreads))
Martin Reinecke's avatar
Martin Reinecke committed
172
  }
173

Martin Reinecke's avatar
Martin Reinecke committed
174
template<typename T> py::array r2c_internal(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
175 176
  const py::object &axes_, bool forward, int inorm, py::object &out_,
  size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
177
  {
Martin Reinecke's avatar
Martin Reinecke committed
178
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
Martin Reinecke committed
179
  auto dims_in(copy_shape(in)), dims_out(dims_in);
Martin Reinecke's avatar
more  
Martin Reinecke committed
180
  dims_out[axes.back()] = (dims_out[axes.back()]>>1)+1;
Martin Reinecke's avatar
Martin Reinecke committed
181
  py::array res = prepare_output<complex<T>>(out_, dims_out);
Martin Reinecke's avatar
Martin Reinecke committed
182 183 184 185 186 187
  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;
188
  T fct = norm_fct<T>(inorm, dims_in, axes);
189 190
  pocketfft::r2c(dims_in, s_in, s_out, axes, forward, d_in, d_out, fct,
    nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
191
  }
Martin Reinecke's avatar
more  
Martin Reinecke committed
192 193
  return res;
  }
194

Martin Reinecke's avatar
Martin Reinecke committed
195 196
py::array r2c(const py::array &in, const py::object &axes_, bool forward,
  int inorm, py::object &out_, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
197
  {
Martin Reinecke's avatar
Martin Reinecke committed
198
  DISPATCH(in, f64, f32, flong, r2c_internal, (in, axes_, forward, inorm, out_,
Martin Reinecke's avatar
Martin Reinecke committed
199
    nthreads))
Martin Reinecke's avatar
Martin Reinecke committed
200
  }
201

Martin Reinecke's avatar
Martin Reinecke committed
202
template<typename T> py::array r2r_fftpack_internal(const py::array &in,
203
  const py::object &axes_, bool r2hc, bool forward, int inorm,
Martin Reinecke's avatar
Martin Reinecke committed
204
  py::object &out_, size_t nthreads)
205
  {
Martin Reinecke's avatar
Martin Reinecke committed
206
  auto axes = makeaxes(in, axes_);
207
  auto dims(copy_shape(in));
Martin Reinecke's avatar
Martin Reinecke committed
208
  py::array res = prepare_output<T>(out_, dims);
Martin Reinecke's avatar
Martin Reinecke committed
209 210 211 212 213 214
  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
215
  T fct = norm_fct<T>(inorm, dims, axes);
216
  pocketfft::r2r_fftpack(dims, s_in, s_out, axes, r2hc, forward,
217
    d_in, d_out, fct, nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
218
  }
219 220
  return res;
  }
221

Martin Reinecke's avatar
Martin Reinecke committed
222
py::array r2r_fftpack(const py::array &in, const py::object &axes_,
223
  bool r2hc, bool forward, int inorm, py::object &out_,
Martin Reinecke's avatar
Martin Reinecke committed
224
  size_t nthreads)
225
  {
226 227
  DISPATCH(in, f64, f32, flong, r2r_fftpack_internal, (in, axes_, r2hc,
    forward, inorm, out_, nthreads))
228
  }
229

Martin Reinecke's avatar
Martin Reinecke committed
230
template<typename T> py::array c2r_internal(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
231 232
  const py::object &axes_, size_t lastsize, bool forward, int inorm,
  py::object &out_, size_t nthreads)
Martin Reinecke's avatar
more  
Martin Reinecke committed
233
  {
Martin Reinecke's avatar
Martin Reinecke committed
234
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
more  
Martin Reinecke committed
235
  size_t axis = axes.back();
Martin Reinecke's avatar
Martin Reinecke committed
236 237 238
  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
239 240
    throw runtime_error("bad lastsize");
  dims_out[axis] = lastsize;
Martin Reinecke's avatar
Martin Reinecke committed
241
  py::array res = prepare_output<T>(out_, dims_out);
Martin Reinecke's avatar
Martin Reinecke committed
242 243 244 245 246 247
  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;
248
  T fct = norm_fct<T>(inorm, dims_out, axes);
249 250
  pocketfft::c2r(dims_out, s_in, s_out, axes, forward, d_in, d_out, fct,
    nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
251
  }
Martin Reinecke's avatar
more  
Martin Reinecke committed
252 253
  return res;
  }
254

Martin Reinecke's avatar
Martin Reinecke committed
255 256
py::array c2r(const py::array &in, const py::object &axes_, size_t lastsize,
  bool forward, int inorm, py::object &out_, size_t nthreads)
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
257
  {
Martin Reinecke's avatar
Martin Reinecke committed
258
  DISPATCH(in, c128, c64, clong, c2r_internal, (in, axes_, lastsize, forward,
Martin Reinecke's avatar
Martin Reinecke committed
259
    inorm, out_, nthreads))
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
260
  }
Martin Reinecke's avatar
Martin Reinecke committed
261

262
template<typename T> py::array separable_hartley_internal(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
263
  const py::object &axes_, int inorm, py::object &out_, size_t nthreads)
Martin Reinecke's avatar
more  
Martin Reinecke committed
264
  {
Martin Reinecke's avatar
Martin Reinecke committed
265
  auto dims(copy_shape(in));
Martin Reinecke's avatar
Martin Reinecke committed
266
  py::array res = prepare_output<T>(out_, dims);
Martin Reinecke's avatar
Martin Reinecke committed
267 268 269 270 271 272 273
  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;
274
  T fct = norm_fct<T>(inorm, dims, axes);
275 276
  pocketfft::r2r_separable_hartley(dims, s_in, s_out, axes, d_in, d_out, fct,
    nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
277
  }
Martin Reinecke's avatar
Martin Reinecke committed
278 279
  return res;
  }
280

Martin Reinecke's avatar
Martin Reinecke committed
281 282
py::array separable_hartley(const py::array &in, const py::object &axes_,
  int inorm, py::object &out_, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
283
  {
284
  DISPATCH(in, f64, f32, flong, separable_hartley_internal, (in, axes_, inorm,
Martin Reinecke's avatar
Martin Reinecke committed
285
    out_, nthreads))
Martin Reinecke's avatar
Martin Reinecke committed
286
  }
287

Martin Reinecke's avatar
fixes  
Martin Reinecke committed
288
template<typename T>py::array complex2hartley(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
289
  const py::array &tmp, const py::object &axes_, py::object &out_)
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
290
  {
Martin Reinecke's avatar
Martin Reinecke committed
291
  using namespace pocketfft::detail;
Martin Reinecke's avatar
Martin Reinecke committed
292
  auto dims_out(copy_shape(in));
Martin Reinecke's avatar
Martin Reinecke committed
293
  py::array out = prepare_output<T>(out_, dims_out);
294
  cndarr<cmplx<T>> atmp(tmp.data(), copy_shape(tmp), copy_strides(tmp));
Martin Reinecke's avatar
Martin Reinecke committed
295
  ndarr<T> aout(out.mutable_data(), copy_shape(out), copy_strides(out));
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
296
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
Martin Reinecke committed
297 298
  {
  py::gil_scoped_release release;
299 300
  simple_iter iin(atmp);
  rev_iter iout(aout, axes);
301 302
  if (iin.remaining()!=iout.remaining()) throw runtime_error("oops");
  while(iin.remaining()>0)
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
303
    {
304 305 306
    auto v = atmp[iin.ofs()];
    aout[iout.ofs()] = v.r+v.i;
    aout[iout.rev_ofs()] = v.r-v.i;
307
    iin.advance(); iout.advance();
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
308
    }
Martin Reinecke's avatar
Martin Reinecke committed
309
  }
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
310 311
  return out;
  }
312

Martin Reinecke's avatar
Martin Reinecke committed
313 314
py::array genuine_hartley(const py::array &in, const py::object &axes_,
  int inorm, py::object &out_, size_t nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
315
  {
Martin Reinecke's avatar
Martin Reinecke committed
316 317
  auto tmp = r2c(in, axes_, true, inorm, None, nthreads);
  DISPATCH(in, f64, f32, flong, complex2hartley, (in, tmp, axes_, out_))
Martin Reinecke's avatar
Martin Reinecke committed
318
  }
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
319

Martin Reinecke's avatar
Martin Reinecke committed
320
const char *pypocketfft_DS = R"""(Fast Fourier and Hartley transforms.
Martin Reinecke's avatar
Martin Reinecke committed
321 322

This module supports
323
- single, double, and long double precision
Martin Reinecke's avatar
Martin Reinecke committed
324 325 326 327 328 329
- 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
330
)""";
Martin Reinecke's avatar
Martin Reinecke committed
331

Martin Reinecke's avatar
Martin Reinecke committed
332
const char *c2c_DS = R"""(Performs a complex FFT.
333 334 335

Parameters
----------
336 337 338
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.
339
axes : list of integers
340
    The axes along which the FFT is carried out.
341
    If not set, all axes will be transformed.
Martin Reinecke's avatar
Martin Reinecke committed
342 343
forward : bool
    If `True`, a negative sign is used in the exponent, else a positive one.
344 345 346 347 348 349
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.
Martin Reinecke's avatar
Martin Reinecke committed
350
out : numpy.ndarray (same shape as `a`, complex type with same accuracy as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
351
    May be identical to `a`, but if it isn't, it must not overlap with `a`.
Martin Reinecke's avatar
Martin Reinecke committed
352
    If None, a new array is allocated to store the output.
Martin Reinecke's avatar
Martin Reinecke committed
353 354 355
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
356 357 358

Returns
-------
Martin Reinecke's avatar
Martin Reinecke committed
359
numpy.ndarray (same shape as `a`, complex type with same accuracy as `a`)
360
    The transformed data.
Martin Reinecke's avatar
Martin Reinecke committed
361
)""";
362

Martin Reinecke's avatar
Martin Reinecke committed
363
const char *r2c_DS = R"""(Performs an FFT whose input is strictly real.
Martin Reinecke's avatar
Martin Reinecke committed
364 365 366

Parameters
----------
367
a : numpy.ndarray (any real type)
Martin Reinecke's avatar
Martin Reinecke committed
368 369
    The input data
axes : list of integers
370
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
371
    If not set, all axes will be transformed in ascending order.
Martin Reinecke's avatar
Martin Reinecke committed
372 373
forward : bool
    If `True`, a negative sign is used in the exponent, else a positive one.
374 375 376 377 378 379
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
380 381
out : numpy.ndarray (complex type with same accuracy as `a`)
    For the required shape, see the `Returns` section.
Martin Reinecke's avatar
Martin Reinecke committed
382
    Must not overlap with `a`.
Martin Reinecke's avatar
Martin Reinecke committed
383
    If None, a new array is allocated to store the output.
Martin Reinecke's avatar
Martin Reinecke committed
384 385 386
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
387 388 389

Returns
-------
Martin Reinecke's avatar
Martin Reinecke committed
390
numpy.ndarray (complex type with same accuracy as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
391
    The transformed data. The shape is identical to that of the input array,
Martin Reinecke's avatar
Martin Reinecke committed
392 393
    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
394
)""";
Martin Reinecke's avatar
Martin Reinecke committed
395

Martin Reinecke's avatar
Martin Reinecke committed
396
const char *c2r_DS = R"""(Performs an FFT whose output is strictly real.
Martin Reinecke's avatar
Martin Reinecke committed
397 398 399

Parameters
----------
400
a : numpy.ndarray (any complex type)
Martin Reinecke's avatar
Martin Reinecke committed
401 402
    The input data
axes : list of integers
403
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
404 405 406
    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
407 408
forward : bool
    If `True`, a negative sign is used in the exponent, else a positive one.
409 410 411 412 413 414
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
415 416
out : numpy.ndarray (real type with same accuracy as `a`)
    For the required shape, see the `Returns` section.
Martin Reinecke's avatar
Martin Reinecke committed
417
    Must not overlap with `a`.
Martin Reinecke's avatar
Martin Reinecke committed
418
    If None, a new array is allocated to store the output.
Martin Reinecke's avatar
Martin Reinecke committed
419 420 421
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
422 423 424

Returns
-------
Martin Reinecke's avatar
Martin Reinecke committed
425
numpy.ndarray (real type with same accuracy as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
426 427
    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
428
    entries.
Martin Reinecke's avatar
Martin Reinecke committed
429
)""";
Martin Reinecke's avatar
Martin Reinecke committed
430

Martin Reinecke's avatar
Martin Reinecke committed
431
const char *r2r_fftpack_DS = R"""(Performs a real-valued FFT using the FFTPACK storage scheme.
432 433 434

Parameters
----------
435
a : numpy.ndarray (any real type)
Martin Reinecke's avatar
Martin Reinecke committed
436 437 438 439
    The input data
axes : list of integers
    The axes along which the FFT is carried out.
    If not set, all axes will be transformed.
440
r2hc : bool
Martin Reinecke's avatar
typo  
Martin Reinecke committed
441
    if True, the input is purely real and the output will be halfcomplex,
442
    otherwise the opposite
Martin Reinecke's avatar
Martin Reinecke committed
443 444
forward : bool
    If `True`, a negative sign is used in the exponent, else a positive one.
445 446 447 448 449 450
inorm : int
    Normalization type
      0 : no normalization
      1 : divide by sqrt(N)
      2 : divide by N
    where N is the length of `axis`.
Martin Reinecke's avatar
Martin Reinecke committed
451
out : numpy.ndarray (same shape and data type as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
452
    May be identical to `a`, but if it isn't, it must not overlap with `a`.
Martin Reinecke's avatar
Martin Reinecke committed
453
    If None, a new array is allocated to store the output.
Martin Reinecke's avatar
Martin Reinecke committed
454 455 456
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
457 458 459

Returns
-------
Martin Reinecke's avatar
Martin Reinecke committed
460
numpy.ndarray (same shape and data type as `a`)
461
    The transformed data. The shape is identical to that of the input array.
Martin Reinecke's avatar
Martin Reinecke committed
462
)""";
463

464 465 466 467 468 469 470
const char *separable_hartley_DS = R"""(Performs a separable Hartley transform.
For every requested axis, a 1D forward Fourier transform is carried out, and
the real and imaginary parts of the result are added before the next axis is
processed.

Parameters
----------
Martin Reinecke's avatar
Martin Reinecke committed
471
a : numpy.ndarray (any real type)
472 473 474 475 476 477 478 479 480 481
    The input data
axes : list of integers
    The axes along which the transform is carried out.
    If not set, all axes will be transformed.
inorm : int
    Normalization type
      0 : no normalization
      1 : divide by sqrt(N)
      2 : divide by N
    where N is the product of the lengths of the transformed axes.
Martin Reinecke's avatar
Martin Reinecke committed
482
out : numpy.ndarray (same shape and data type as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
483
    May be identical to `a`, but if it isn't, it must not overlap with `a`.
Martin Reinecke's avatar
Martin Reinecke committed
484
    If None, a new array is allocated to store the output.
485 486 487 488 489 490
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).

Returns
-------
Martin Reinecke's avatar
Martin Reinecke committed
491
numpy.ndarray (same shape and data type as `a`)
492 493 494 495 496 497 498 499
    The transformed data
)""";

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

Parameters
----------
Martin Reinecke's avatar
Martin Reinecke committed
503
a : numpy.ndarray (any real type)
504 505
    The input data
axes : list of integers
506
    The axes along which the transform is carried out.
507
    If not set, all axes will be transformed.
508 509 510 511 512 513
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.
Martin Reinecke's avatar
Martin Reinecke committed
514
out : numpy.ndarray (same shape and data type as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
515
    May be identical to `a`, but if it isn't, it must not overlap with `a`.
Martin Reinecke's avatar
Martin Reinecke committed
516
    If None, a new array is allocated to store the output.
Martin Reinecke's avatar
Martin Reinecke committed
517 518 519
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
520 521 522

Returns
-------
Martin Reinecke's avatar
Martin Reinecke committed
523
numpy.ndarray (same shape and data type as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
524
    The transformed data
Martin Reinecke's avatar
Martin Reinecke committed
525
)""";
526

Martin Reinecke's avatar
Martin Reinecke committed
527 528 529 530
} // unnamed namespace

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

Martin Reinecke's avatar
Martin Reinecke committed
533
  m.doc() = pypocketfft_DS;
534
  m.def("c2c", c2c, c2c_DS, "a"_a, "axes"_a=None, "forward"_a=true,
Martin Reinecke's avatar
Martin Reinecke committed
535
    "inorm"_a=0, "out"_a=None, "nthreads"_a=1);
536
  m.def("r2c", r2c, r2c_DS, "a"_a, "axes"_a=None, "forward"_a=true,
Martin Reinecke's avatar
Martin Reinecke committed
537
    "inorm"_a=0, "out"_a=None, "nthreads"_a=1);
538
  m.def("c2r", c2r, c2r_DS, "a"_a, "axes"_a=None, "lastsize"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
539
    "forward"_a=true, "inorm"_a=0, "out"_a=None, "nthreads"_a=1);
540 541 542
  m.def("r2r_fftpack", r2r_fftpack, r2r_fftpack_DS, "a"_a, "axes"_a,
    "r2hc"_a, "forward"_a, "inorm"_a=0, "out"_a=None, "nthreads"_a=1);
  m.def("separable_hartley", separable_hartley, separable_hartley_DS, "a"_a,
Martin Reinecke's avatar
Martin Reinecke committed
543
    "axes"_a=None, "inorm"_a=0, "out"_a=None, "nthreads"_a=1);
544
  m.def("genuine_hartley", genuine_hartley, genuine_hartley_DS, "a"_a,
Martin Reinecke's avatar
Martin Reinecke committed
545
    "axes"_a=None, "inorm"_a=0, "out"_a=None, "nthreads"_a=1);
Martin Reinecke's avatar
Martin Reinecke committed
546
  }