pypocketfft.cc 23.1 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
10
 *  Copyright (C) 2019 Peter Bell
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
11
 *  \author Martin Reinecke
12
 *  \author Peter Bell
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
13 14
 */

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

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

21
namespace {
22

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

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

29 30 31 32
// 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;

33 34 35 36 37 38
using c64 = std::complex<float>;
using c128 = std::complex<double>;
using clong = std::complex<ldbl_t>;
using f32 = float;
using f64 = double;
using flong = ldbl_t;
Martin Reinecke's avatar
Martin Reinecke committed
39
auto None = py::none();
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
Martin Reinecke committed
57
shape_t makeaxes(const py::array &in, const py::object &axes)
Martin Reinecke's avatar
more  
Martin Reinecke committed
58
  {
59
  if (axes.is_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
    }
66 67 68
  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
69
    throw runtime_error("bad axes argument");
70 71 72 73 74 75 76 77
  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
78 79
  }

80
#define DISPATCH(arr, T1, T2, T3, func, args) \
81
  { \
82 83 84
  if (py::isinstance<py::array_t<T1>>(arr)) return func<double> args; \
  if (py::isinstance<py::array_t<T2>>(arr)) return func<float> args;  \
  if (py::isinstance<py::array_t<T3>>(arr)) return func<ldbl_t> args; \
85 86
  throw runtime_error("unsupported data type"); \
  }
87

88 89 90 91 92 93 94 95 96
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,
Martin Reinecke's avatar
Martin Reinecke committed
97
  const shape_t &axes, size_t fct=1, int delta=0)
98 99 100 101
  {
  if (inorm==0) return T(1);
  size_t N(1);
  for (auto a: axes)
102
    N *= fct * size_t(int64_t(shape[a])+delta);
103 104 105
  return norm_fct<T>(inorm, N);
  }

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

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

Martin Reinecke's avatar
Martin Reinecke committed
164 165
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
166
  {
167
  if (a.dtype().kind() == 'c')
Martin Reinecke's avatar
Martin Reinecke committed
168 169
    DISPATCH(a, c128, c64, clong, c2c_internal, (a, axes_, forward,
             inorm, out_, nthreads))
170

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

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

Martin Reinecke's avatar
Martin Reinecke committed
196 197
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
198
  {
Martin Reinecke's avatar
Martin Reinecke committed
199
  DISPATCH(in, f64, f32, flong, r2c_internal, (in, axes_, forward, inorm, out_,
Martin Reinecke's avatar
Martin Reinecke committed
200
    nthreads))
Martin Reinecke's avatar
Martin Reinecke committed
201
  }
202

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

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

231
template<typename T> py::array dct_internal(const py::array &in,
232
  const py::object &axes_, int type, int inorm, py::object &out_,
233 234 235 236 237 238 239 240 241 242 243
  size_t nthreads)
  {
  auto axes = makeaxes(in, axes_);
  auto dims(copy_shape(in));
  py::array res = prepare_output<T>(out_, dims);
  auto s_in=copy_strides(in);
  auto s_out=copy_strides(res);
  auto d_in=reinterpret_cast<const T *>(in.data());
  auto d_out=reinterpret_cast<T *>(res.mutable_data());
  {
  py::gil_scoped_release release;
244 245
  T fct = (type==1) ? norm_fct<T>(inorm, dims, axes, 2, -1)
                    : norm_fct<T>(inorm, dims, axes, 2);
246
  bool ortho = inorm == 1;
Martin Reinecke's avatar
Martin Reinecke committed
247
  pocketfft::dct(dims, s_in, s_out, axes, type, d_in, d_out, fct, ortho,
248
    nthreads);
249 250 251 252
  }
  return res;
  }

253
py::array dct(const py::array &in, int type, const py::object &axes_,
254
  int inorm, py::object &out_, size_t nthreads)
255
  {
256
  if ((type<1) || (type>4)) throw invalid_argument("invalid DCT type");
257 258
  DISPATCH(in, f64, f32, flong, dct_internal, (in, axes_, type, inorm, out_,
    nthreads))
259 260
  }

261
template<typename T> py::array dst_internal(const py::array &in,
262
  const py::object &axes_, int type, int inorm, py::object &out_,
263 264 265 266 267 268 269 270 271 272 273
  size_t nthreads)
  {
  auto axes = makeaxes(in, axes_);
  auto dims(copy_shape(in));
  py::array res = prepare_output<T>(out_, dims);
  auto s_in=copy_strides(in);
  auto s_out=copy_strides(res);
  auto d_in=reinterpret_cast<const T *>(in.data());
  auto d_out=reinterpret_cast<T *>(res.mutable_data());
  {
  py::gil_scoped_release release;
274 275
  T fct = (type==1) ? norm_fct<T>(inorm, dims, axes, 2, 1)
                    : norm_fct<T>(inorm, dims, axes, 2);
276
  bool ortho = inorm == 1;
Martin Reinecke's avatar
Martin Reinecke committed
277
  pocketfft::dst(dims, s_in, s_out, axes, type, d_in, d_out, fct, ortho,
278
    nthreads);
279 280 281 282
  }
  return res;
  }

283
py::array dst(const py::array &in, int type, const py::object &axes_,
284
  int inorm, py::object &out_, size_t nthreads)
285
  {
286
  if ((type<1) || (type>4)) throw invalid_argument("invalid DST type");
287
  DISPATCH(in, f64, f32, flong, dst_internal, (in, axes_, type, inorm,
Martin Reinecke's avatar
Martin Reinecke committed
288
    out_, nthreads))
289 290
  }

Martin Reinecke's avatar
Martin Reinecke committed
291
template<typename T> py::array c2r_internal(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
292 293
  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
294
  {
Martin Reinecke's avatar
Martin Reinecke committed
295
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
more  
Martin Reinecke committed
296
  size_t axis = axes.back();
Martin Reinecke's avatar
Martin Reinecke committed
297 298 299
  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])
300
    throw invalid_argument("bad lastsize");
Martin Reinecke's avatar
more  
Martin Reinecke committed
301
  dims_out[axis] = lastsize;
Martin Reinecke's avatar
Martin Reinecke committed
302
  py::array res = prepare_output<T>(out_, dims_out);
Martin Reinecke's avatar
Martin Reinecke committed
303 304 305 306 307 308
  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;
309
  T fct = norm_fct<T>(inorm, dims_out, axes);
310 311
  pocketfft::c2r(dims_out, s_in, s_out, axes, forward, d_in, d_out, fct,
    nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
312
  }
Martin Reinecke's avatar
more  
Martin Reinecke committed
313 314
  return res;
  }
315

Martin Reinecke's avatar
Martin Reinecke committed
316 317
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
318
  {
Martin Reinecke's avatar
Martin Reinecke committed
319
  DISPATCH(in, c128, c64, clong, c2r_internal, (in, axes_, lastsize, forward,
Martin Reinecke's avatar
Martin Reinecke committed
320
    inorm, out_, nthreads))
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
321
  }
Martin Reinecke's avatar
Martin Reinecke committed
322

323
template<typename T> py::array separable_hartley_internal(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
324
  const py::object &axes_, int inorm, py::object &out_, size_t nthreads)
Martin Reinecke's avatar
more  
Martin Reinecke committed
325
  {
Martin Reinecke's avatar
Martin Reinecke committed
326
  auto dims(copy_shape(in));
Martin Reinecke's avatar
Martin Reinecke committed
327
  py::array res = prepare_output<T>(out_, dims);
Martin Reinecke's avatar
Martin Reinecke committed
328 329 330 331 332 333 334
  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;
335
  T fct = norm_fct<T>(inorm, dims, axes);
336 337
  pocketfft::r2r_separable_hartley(dims, s_in, s_out, axes, d_in, d_out, fct,
    nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
338
  }
Martin Reinecke's avatar
Martin Reinecke committed
339 340
  return res;
  }
341

Martin Reinecke's avatar
Martin Reinecke committed
342 343
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
344
  {
345
  DISPATCH(in, f64, f32, flong, separable_hartley_internal, (in, axes_, inorm,
Martin Reinecke's avatar
Martin Reinecke committed
346
    out_, nthreads))
Martin Reinecke's avatar
Martin Reinecke committed
347
  }
348

Martin Reinecke's avatar
fixes  
Martin Reinecke committed
349
template<typename T>py::array complex2hartley(const py::array &in,
Martin Reinecke's avatar
Martin Reinecke committed
350
  const py::array &tmp, const py::object &axes_, py::object &out_)
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
351
  {
Martin Reinecke's avatar
Martin Reinecke committed
352
  using namespace pocketfft::detail;
Martin Reinecke's avatar
Martin Reinecke committed
353
  auto dims_out(copy_shape(in));
Martin Reinecke's avatar
Martin Reinecke committed
354
  py::array out = prepare_output<T>(out_, dims_out);
355
  cndarr<cmplx<T>> atmp(tmp.data(), copy_shape(tmp), copy_strides(tmp));
Martin Reinecke's avatar
Martin Reinecke committed
356
  ndarr<T> aout(out.mutable_data(), copy_shape(out), copy_strides(out));
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
357
  auto axes = makeaxes(in, axes_);
Martin Reinecke's avatar
Martin Reinecke committed
358 359
  {
  py::gil_scoped_release release;
360 361
  simple_iter iin(atmp);
  rev_iter iout(aout, axes);
362
  while(iin.remaining()>0)
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
363
    {
364 365 366
    auto v = atmp[iin.ofs()];
    aout[iout.ofs()] = v.r+v.i;
    aout[iout.rev_ofs()] = v.r-v.i;
367
    iin.advance(); iout.advance();
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
368
    }
Martin Reinecke's avatar
Martin Reinecke committed
369
  }
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
370 371
  return out;
  }
372

Martin Reinecke's avatar
Martin Reinecke committed
373 374
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
375
  {
Martin Reinecke's avatar
Martin Reinecke committed
376 377
  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
378
  }
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
379

380 381 382 383 384 385
size_t good_size(size_t n, bool real)
  {
  using namespace pocketfft::detail;
  return real ? util::good_size_real(n) : util::good_size_cmplx(n);
  }

Martin Reinecke's avatar
Martin Reinecke committed
386
const char *pypocketfft_DS = R"""(Fast Fourier and Hartley transforms.
Martin Reinecke's avatar
Martin Reinecke committed
387 388

This module supports
389
- single, double, and long double precision
Martin Reinecke's avatar
Martin Reinecke committed
390 391 392 393 394 395
- 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
396
)""";
Martin Reinecke's avatar
Martin Reinecke committed
397

Martin Reinecke's avatar
Martin Reinecke committed
398
const char *c2c_DS = R"""(Performs a complex FFT.
399 400 401

Parameters
----------
402 403 404
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.
405
axes : list of integers
406
    The axes along which the FFT is carried out.
407
    If not set, all axes will be transformed.
Martin Reinecke's avatar
Martin Reinecke committed
408 409
forward : bool
    If `True`, a negative sign is used in the exponent, else a positive one.
410 411 412 413 414 415
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
416
out : numpy.ndarray (same shape as `a`, complex type with same accuracy as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
417
    May be identical to `a`, but if it isn't, it 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).
422 423 424

Returns
-------
Martin Reinecke's avatar
Martin Reinecke committed
425
numpy.ndarray (same shape as `a`, complex type with same accuracy as `a`)
426
    The transformed data.
Martin Reinecke's avatar
Martin Reinecke committed
427
)""";
428

Martin Reinecke's avatar
Martin Reinecke committed
429
const char *r2c_DS = R"""(Performs an FFT whose input is strictly real.
Martin Reinecke's avatar
Martin Reinecke committed
430 431 432

Parameters
----------
433
a : numpy.ndarray (any real type)
Martin Reinecke's avatar
Martin Reinecke committed
434 435
    The input data
axes : list of integers
436
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
437
    If not set, all axes will be transformed in ascending order.
Martin Reinecke's avatar
Martin Reinecke committed
438 439
forward : bool
    If `True`, a negative sign is used in the exponent, else a positive one.
440 441 442 443 444 445
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
446 447
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
448
    Must not overlap with `a`.
Martin Reinecke's avatar
Martin Reinecke committed
449
    If None, a new array is allocated to store the output.
Martin Reinecke's avatar
Martin Reinecke committed
450 451 452
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
453 454 455

Returns
-------
Martin Reinecke's avatar
Martin Reinecke committed
456
numpy.ndarray (complex type with same accuracy as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
457
    The transformed data. The shape is identical to that of the input array,
Martin Reinecke's avatar
Martin Reinecke committed
458 459
    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
460
)""";
Martin Reinecke's avatar
Martin Reinecke committed
461

Martin Reinecke's avatar
Martin Reinecke committed
462
const char *c2r_DS = R"""(Performs an FFT whose output is strictly real.
Martin Reinecke's avatar
Martin Reinecke committed
463 464 465

Parameters
----------
466
a : numpy.ndarray (any complex type)
Martin Reinecke's avatar
Martin Reinecke committed
467 468
    The input data
axes : list of integers
469
    The axes along which the FFT is carried out.
Martin Reinecke's avatar
Martin Reinecke committed
470 471 472
    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
473 474
forward : bool
    If `True`, a negative sign is used in the exponent, else a positive one.
475 476 477 478 479 480
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
481 482
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
483
    Must not overlap with `a`.
Martin Reinecke's avatar
Martin Reinecke committed
484
    If None, a new array is allocated to store the output.
Martin Reinecke's avatar
Martin Reinecke committed
485 486 487
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
488 489 490

Returns
-------
Martin Reinecke's avatar
Martin Reinecke committed
491
numpy.ndarray (real type with same accuracy as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
492 493
    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
494
    entries.
Martin Reinecke's avatar
Martin Reinecke committed
495
)""";
Martin Reinecke's avatar
Martin Reinecke committed
496

Martin Reinecke's avatar
Martin Reinecke committed
497
const char *r2r_fftpack_DS = R"""(Performs a real-valued FFT using the FFTPACK storage scheme.
498 499 500

Parameters
----------
501
a : numpy.ndarray (any real type)
Martin Reinecke's avatar
Martin Reinecke committed
502 503 504 505
    The input data
axes : list of integers
    The axes along which the FFT is carried out.
    If not set, all axes will be transformed.
506 507 508 509
real2hermitian : bool
    if True, the input is purely real and the output will have Hermitian
    symmetry and be stored in FFTPACK's halfcomplex ordering, otherwise the
    opposite.
Martin Reinecke's avatar
Martin Reinecke committed
510 511
forward : bool
    If `True`, a negative sign is used in the exponent, else a positive one.
512 513 514 515 516 517
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
518
out : numpy.ndarray (same shape and data type as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
519
    May be identical to `a`, but if it isn't, it must not overlap with `a`.
Martin Reinecke's avatar
Martin Reinecke committed
520
    If None, a new array is allocated to store the output.
Martin Reinecke's avatar
Martin Reinecke committed
521 522 523
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
524 525 526

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

531 532 533 534 535 536 537
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
538
a : numpy.ndarray (any real type)
539 540 541 542 543 544 545 546 547 548
    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
549
out : numpy.ndarray (same shape and data type as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
550
    May be identical to `a`, but if it isn't, it must not overlap with `a`.
Martin Reinecke's avatar
Martin Reinecke committed
551
    If None, a new array is allocated to store the output.
552 553 554 555 556 557
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
558
numpy.ndarray (same shape and data type as `a`)
559 560 561 562 563 564 565 566
    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.
567 568 569

Parameters
----------
Martin Reinecke's avatar
Martin Reinecke committed
570
a : numpy.ndarray (any real type)
571 572
    The input data
axes : list of integers
573
    The axes along which the transform is carried out.
574
    If not set, all axes will be transformed.
575 576 577 578 579 580
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
581
out : numpy.ndarray (same shape and data type as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
582
    May be identical to `a`, but if it isn't, it must not overlap with `a`.
Martin Reinecke's avatar
Martin Reinecke committed
583
    If None, a new array is allocated to store the output.
Martin Reinecke's avatar
Martin Reinecke committed
584 585 586
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).
587 588 589

Returns
-------
Martin Reinecke's avatar
Martin Reinecke committed
590
numpy.ndarray (same shape and data type as `a`)
Martin Reinecke's avatar
Martin Reinecke committed
591
    The transformed data
Martin Reinecke's avatar
Martin Reinecke committed
592
)""";
593

Martin Reinecke's avatar
Martin Reinecke committed
594 595 596 597 598 599 600 601 602 603 604 605 606 607
const char *dct_DS = R"""(Performs a discrete cosine transform.

Parameters
----------
a : numpy.ndarray (any real type)
    The input data
type : integer
    the type of DCT. Must be in [1; 4].
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
608
      1 : make transform orthogonal and divide by sqrt(N)
Martin Reinecke's avatar
Martin Reinecke committed
609 610 611 612
      2 : divide by N
    where N is the product of n_i for every transformed axis i.
    n_i is 2*(<axis_length>-1 for type 1 and 2*<axis length>
    for types 2, 3, 4.
613 614 615 616 617 618 619
    Making the transform orthogonal involves the following additional steps
    for every 1D sub-transform:
      Type 1 : multiply first and last input value by sqrt(2)
               divide first and last output value by sqrt(2)
      Type 2 : divide first output value by sqrt(2)
      Type 3 : multiply first input value by sqrt(2)
      Type 4 : nothing
Martin Reinecke's avatar
Martin Reinecke committed
620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646
out : numpy.ndarray (same shape and data type as `a`)
    May be identical to `a`, but if it isn't, it must not overlap with `a`.
    If None, a new array is allocated to store the output.
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).

Returns
-------
numpy.ndarray (same shape and data type as `a`)
    The transformed data
)""";

const char *dst_DS = R"""(Performs a discrete sine transform.

Parameters
----------
a : numpy.ndarray (any real type)
    The input data
type : integer
    the type of DST. Must be in [1; 4].
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
647
      1 : make transform orthogonal and divide by sqrt(N)
Martin Reinecke's avatar
Martin Reinecke committed
648 649 650 651
      2 : divide by N
    where N is the product of n_i for every transformed axis i.
    n_i is 2*(<axis_length>+1 for type 1 and 2*<axis length>
    for types 2, 3, 4.
652 653 654 655 656 657
    Making the transform orthogonal involves the following additional steps
    for every 1D sub-transform:
      Type 1 : nothing
      Type 2 : divide first output value by sqrt(2)
      Type 3 : multiply first input value by sqrt(2)
      Type 4 : nothing
Martin Reinecke's avatar
Martin Reinecke committed
658 659 660 661 662 663 664 665 666 667 668 669 670
out : numpy.ndarray (same shape and data type as `a`)
    May be identical to `a`, but if it isn't, it must not overlap with `a`.
    If None, a new array is allocated to store the output.
nthreads : int
    Number of threads to use. If 0, use the system default (typically governed
    by the `OMP_NUM_THREADS` environment variable).

Returns
-------
numpy.ndarray (same shape and data type as `a`)
    The transformed data
)""";

671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686
const char * good_size_DS = R"""(Returns a good length to pad an FFT to.

Parameters
----------
n : int
    Minimum transform length
real : bool, optional
    True if either input or output of FFT should be fully real.

Returns
-------
out : int
    The smallest fast size >= n

)""";

Martin Reinecke's avatar
Martin Reinecke committed
687 688 689 690
} // unnamed namespace

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

Martin Reinecke's avatar
Martin Reinecke committed
693
  m.doc() = pypocketfft_DS;
694
  m.def("c2c", c2c, c2c_DS, "a"_a, "axes"_a=None, "forward"_a=true,
Martin Reinecke's avatar
Martin Reinecke committed
695
    "inorm"_a=0, "out"_a=None, "nthreads"_a=1);
696
  m.def("r2c", r2c, r2c_DS, "a"_a, "axes"_a=None, "forward"_a=true,
Martin Reinecke's avatar
Martin Reinecke committed
697
    "inorm"_a=0, "out"_a=None, "nthreads"_a=1);
698
  m.def("c2r", c2r, c2r_DS, "a"_a, "axes"_a=None, "lastsize"_a=0,
Martin Reinecke's avatar
Martin Reinecke committed
699
    "forward"_a=true, "inorm"_a=0, "out"_a=None, "nthreads"_a=1);
700
  m.def("r2r_fftpack", r2r_fftpack, r2r_fftpack_DS, "a"_a, "axes"_a,
701
    "real2hermitian"_a, "forward"_a, "inorm"_a=0, "out"_a=None, "nthreads"_a=1);
702
  m.def("separable_hartley", separable_hartley, separable_hartley_DS, "a"_a,
Martin Reinecke's avatar
Martin Reinecke committed
703
    "axes"_a=None, "inorm"_a=0, "out"_a=None, "nthreads"_a=1);
704
  m.def("genuine_hartley", genuine_hartley, genuine_hartley_DS, "a"_a,
Martin Reinecke's avatar
Martin Reinecke committed
705
    "axes"_a=None, "inorm"_a=0, "out"_a=None, "nthreads"_a=1);
Martin Reinecke's avatar
Martin Reinecke committed
706
  m.def("dct", dct, dct_DS, "a"_a, "type"_a, "axes"_a=None, "inorm"_a=0,
707
    "out"_a=None, "nthreads"_a=1);
Martin Reinecke's avatar
Martin Reinecke committed
708
  m.def("dst", dst, dst_DS, "a"_a, "type"_a, "axes"_a=None, "inorm"_a=0,
709
    "out"_a=None, "nthreads"_a=1);
710
  m.def("good_size", good_size, good_size_DS, "n"_a, "real"_a=false);
Martin Reinecke's avatar
Martin Reinecke committed
711
  }