FFT design considerations

wds15 · January 26, 2017, 7:06pm

Hi!

Having an FFT in Stan would be really cool!

One of the very cool applications of a FFT transform (and its inverse) is the possibility to calculate the convolution as a simple product of two functions in the transformed space. This is for example handy if one wants to get the difference distribution of two quantities; one simply FFT transforms the function, multiplies them and then takes the inverse of that product.

Long story short: If this is possible with what is planned then by all means it should be taken in consideration for the design. That would be quite useful I think.

Best,
Sebastian

thel · January 30, 2017, 6:46pm

It’s very inconvenient to do this without complex numbers. I think the options here are

Implement a complex type and vectorized complex arithmetic.
Establish a convention for treating 2-by-n matrices as complex vectors, providing complex arithmetic.
Not expose any complex arithmetic, but give users a convolution function.

I think 1 has the character of tilting at windmills for pretty thin benefit, it’s a truly huge project. 2 seems like a design disaster (picturing nightmares if our handling of matrices evolves or eigen’s does). We can’t prevent users from doing this in an ad hoc way but I don’t want to bless it. 3 is a bit unsatisfying but has the benefit of letting us add the pieces one at a time. If after a few months we have a zoo of signal processing functions, we may want to re-evaluate this.

Thel

Bob_Carpenter · January 30, 2017, 6:56pm

Complex type isn’t enough of a priority now. And I
think the derivatives are pretty tricky to do in general.

Is the second option what we were talking about before?
If so, I don’t see a problem with that.

I don’t understand the third option. Is the convolution
why people want these operations? Or is it more things
which is why you think libraries may grow?

Bob

betanalpha · January 30, 2017, 7:05pm

Forget convolution for now – there are more subtle issues involved due to the discretization assumed in the DFT. Right now let’s just start with

real[] fft_amps(real[] x);
real[] fft_phases(real[] x);
real[,] fft_amps_phases(real[] x);

thel · January 30, 2017, 10:39pm

With regards to 2 I’m just saying that I don’t want to give people a bunch of functions like:

real[2,] pointwise_multiply_as_if_complex_vector(real[2,-] z, real[2,-] w)

as a brittle workaround for the lack of a complex type. People may slice the thing we give them and do all kinds of manipulation, I’m just kind of wary of creating a pseudo-type. So I’m saying that we shouldn’t really worry about convolution for now, but later if there’s a need we can specific functions for convolution.

Michael’s signatures are exactly right, though I’m tempted to add one more: real[,] fft_real_imag(real[] x);

Bob_Carpenter · January 30, 2017, 10:50pm

I’d suggest starting with just one function, getting
feedback and merging, then going on to the rest.

Bob

thel · January 30, 2017, 10:57pm

That’s absolutely the plan. Just figured Sebastian deserved a response on convolutions.

wds15 · January 31, 2017, 8:38am

Hi!

Ok; let’s keep things simple. The Fourier transform is super handy for a convolution, but I am not an expert on the discrete FFT nor is this an easy request as this entails to transform and back-transform. Still, if we think that convolutions are feasible at some point, then we should probably design the internal code with this application in mind? But as Bob says, getting started with the simple function is a great start.

Sebastian

Bob_Carpenter · January 31, 2017, 5:59pm

The actual FFT code is sitting inside of Eigen, so it’s
out of our hands unless we want to rewrite it.

Bob

wds15 · January 31, 2017, 7:16pm

Oh, did not know that. This makes it much easier…just expose what they did sounds as the way to go then.

Bob_Carpenter · January 31, 2017, 7:18pm

That’s the issue—how to do that. Their API is based
on complex numbers, which we don’t have in Stan.

Bob

nikunj410 · November 14, 2021, 7:01pm

Hi Bob,

I am working on using convolution operators to model spatial population dynamics. Right now my go-to solution is to use nested loops in the function block to implement the convolution. Ideally, the simulations can be made much faster by solving convolution by FFT.

I saw that the Stan team just released support for complex numbers but I could not find a FFT (iFFT) function. Is the stan team working on providing FFT related functions to solve convolutions? thank you in advance

Bob_Carpenter · February 16, 2022, 10:52pm

Yes, and if we’re fast, they should be in the next release (in about 3 months or so). I already have them implemented in the C++ tests for complex numbers and now @WardBrian is finishing off the vector and matrix wrappers for complex numbers.

WardBrian · February 16, 2022, 10:56pm

I think at this point it wouldn’t be premature to start the work on adding the FFT functions to the math library. It would provide a nice stress test for my work on the compiler if nothing else

Bob_Carpenter · March 9, 2022, 9:58pm

I’ve got both 1D and 2D FFT and inverse FFT working on the math branch feature/20-fft (this has been a feature request since before we moved the repos). All the reverse-mode autodiff seems to be working, but some of the forward mode is failing. And I’m only getting results matching NumPy to two decimal places on the 2D case. I don’t know if this is my poor implementation or NumPy’s!

syclik · March 10, 2022, 1:18am

Awesome!!!

ahartikainen · March 10, 2022, 8:27am

Sounds great, I added some links to numpy source.

This is from numpy docs

https://numpy.org/doc/stable/reference/generated/numpy.fft.fft2.html

fft2 is just fftn with a different default for axes .

https://numpy.org/doc/stable/reference/generated/numpy.fft.fftn.html#numpy.fft.fftn

And the main function which iterates over the axes

github.com

numpy/numpy/blob/4adc87dff15a247e417d50f10cc4def8e1c17a03/numpy/fft/_pocketfft.py#L701


      
              s = list(s)
              if axes is None:
                  axes = list(range(-len(s), 0))
              if len(s) != len(axes):
                  raise ValueError("Shape and axes have different lengths.")
              if invreal and shapeless:
                  s[-1] = (a.shape[axes[-1]] - 1) * 2
              return s, axes
          
          
          def _raw_fftnd(a, s=None, axes=None, function=fft, norm=None):
              a = asarray(a)
              s, axes = _cook_nd_args(a, s, axes)
              itl = list(range(len(axes)))
              itl.reverse()
              for ii in itl:
                  a = function(a, n=s[ii], axis=axes[ii], norm=norm)
              return a
          
          
          def _fftn_dispatcher(a, s=None, axes=None, norm=None):

And here is the fft function

github.com

numpy/numpy/blob/4adc87dff15a247e417d50f10cc4def8e1c17a03/numpy/fft/_pocketfft.py#L123


      
              except KeyError:
                  raise ValueError(f'Invalid norm value {norm}; should be "backward", '
                                   '"ortho" or "forward".') from None
          
          
          def _fft_dispatcher(a, n=None, axis=None, norm=None):
              return (a,)
          
          
          @array_function_dispatch(_fft_dispatcher)
          def fft(a, n=None, axis=-1, norm=None):
              """
              Compute the one-dimensional discrete Fourier Transform.
          
              This function computes the one-dimensional *n*-point discrete Fourier
              Transform (DFT) with the efficient Fast Fourier Transform (FFT)
              algorithm [CT].
          
              Parameters
              ----------
              a : array_like

and its raw function

github.com

numpy/numpy/blob/4adc87dff15a247e417d50f10cc4def8e1c17a03/numpy/fft/_pocketfft.py#L49


      
          
          
          array_function_dispatch = functools.partial(
              overrides.array_function_dispatch, module='numpy.fft')
          
          
          # `inv_norm` is a float by which the result of the transform needs to be
          # divided. This replaces the original, more intuitive 'fct` parameter to avoid
          # divisions by zero (or alternatively additional checks) in the case of
          # zero-length axes during its computation.
          def _raw_fft(a, n, axis, is_real, is_forward, inv_norm):
              axis = normalize_axis_index(axis, a.ndim)
              if n is None:
                  n = a.shape[axis]
          
              fct = 1/inv_norm
          
              if a.shape[axis] != n:
                  s = list(a.shape)
                  index = [slice(None)]*len(s)
                  if s[axis] > n:

And here is a .c file for real and complex cases

github.com

numpy/numpy/blob/4adc87dff15a247e417d50f10cc4def8e1c17a03/numpy/fft/_pocketfft.c#L2195


      
          
          WARN_UNUSED_RESULT static int rfft_forward(rfft_plan plan, double c[], double fct)
            {
            if (plan->packplan)
              return rfftp_forward(plan->packplan,c,fct);
            else // if (plan->blueplan)
              return rfftblue_forward(plan->blueplan,c,fct);
            }
          
          static PyObject *
          execute_complex(PyObject *a1, int is_forward, double fct)
          {
              PyArrayObject *data = (PyArrayObject *)PyArray_FromAny(a1,
                      PyArray_DescrFromType(NPY_CDOUBLE), 1, 0,
                      NPY_ARRAY_ENSURECOPY | NPY_ARRAY_DEFAULT |
                      NPY_ARRAY_ENSUREARRAY | NPY_ARRAY_FORCECAST,
                      NULL);
              if (!data) return NULL;
          
              int npts = PyArray_DIM(data, PyArray_NDIM(data) - 1);
              cfft_plan plan=NULL;

github.com

numpy/numpy/blob/4adc87dff15a247e417d50f10cc4def8e1c17a03/numpy/fft/_pocketfft.c#L2230


      
              if (plan) destroy_cfft_plan(plan);
              Py_END_ALLOW_THREADS;
              if (fail) {
                Py_XDECREF(data);
                return PyErr_NoMemory();
              }
              return (PyObject *)data;
          }
          
          static PyObject *
          execute_real_forward(PyObject *a1, double fct)
          {
              rfft_plan plan=NULL;
              int fail = 0;
              PyArrayObject *data = (PyArrayObject *)PyArray_FromAny(a1,
                      PyArray_DescrFromType(NPY_DOUBLE), 1, 0,
                      NPY_ARRAY_DEFAULT | NPY_ARRAY_ENSUREARRAY | NPY_ARRAY_FORCECAST,
                      NULL);
              if (!data) return NULL;
          
              int ndim = PyArray_NDIM(data);

github.com

numpy/numpy/blob/4adc87dff15a247e417d50f10cc4def8e1c17a03/numpy/fft/_pocketfft.c#L2285


      
              }
              if (fail) {
                Py_XDECREF(data);
                Py_XDECREF(ret);
                return PyErr_NoMemory();
              }
              Py_DECREF(data);
              return (PyObject *)ret;
          }
          static PyObject *
          execute_real_backward(PyObject *a1, double fct)
          {
              rfft_plan plan=NULL;
              PyArrayObject *data = (PyArrayObject *)PyArray_FromAny(a1,
                      PyArray_DescrFromType(NPY_CDOUBLE), 1, 0,
                      NPY_ARRAY_DEFAULT | NPY_ARRAY_ENSUREARRAY | NPY_ARRAY_FORCECAST,
                      NULL);
              if (!data) return NULL;
              int npts = PyArray_DIM(data, PyArray_NDIM(data) - 1);
              PyArrayObject *ret = (PyArrayObject *)PyArray_Empty(PyArray_NDIM(data),
                      PyArray_DIMS(data), PyArray_DescrFromType(NPY_DOUBLE), 0);

Bob_Carpenter · March 18, 2022, 2:29pm

Was this a request to go beyond fft2? The rest of the fftn are easy, but I thought I’d start simple with the main ones we need.

If anyone has reverse-mode or forward-mode autodiff code for this, that’d be great. The hint is that

\partial\, \textrm{fft}(x) = \textrm{fft}(\partial x)

Presumably something similar holds for the inverse FFT.

ahartikainen · March 18, 2022, 2:47pm

No, I just wanted to add a note (if there was some numerical differences)

WardBrian · March 18, 2022, 4:43pm

Not sure how relevant it would be, but I know a few of the Julia autodiff libraries have custom adjoints for FFT code, e.g.:

github.com

FluxML/Zygote.jl/blob/a133200422e4f12e1d7266f5825154054faf0d9a/src/lib/array.jl#L786-L937


      
          # Abstract FFT
          # ===================
          
          # AbstractFFTs functions do not work with FillArrays, which are needed
          # for some functionality of Zygote. To make it work with FillArrays
          # as well, overload the relevant functions
          AbstractFFTs.fft(x::Fill, dims...) = AbstractFFTs.fft(collect(x), dims...)
          AbstractFFTs.bfft(x::Fill, dims...) = AbstractFFTs.bfft(collect(x), dims...)
          AbstractFFTs.ifft(x::Fill, dims...) = AbstractFFTs.ifft(collect(x), dims...)
          AbstractFFTs.rfft(x::Fill, dims...) = AbstractFFTs.rfft(collect(x), dims...)
          AbstractFFTs.irfft(x::Fill, d, dims...) = AbstractFFTs.irfft(collect(x), d, dims...)
          AbstractFFTs.brfft(x::Fill, d, dims...) = AbstractFFTs.brfft(collect(x), d, dims...)
          
          # the adjoint jacobian of an FFT with respect to its input is the reverse FFT of the
          # gradient of its inputs, but with different normalization factor
          @adjoint function fft(xs)
            return AbstractFFTs.fft(xs), function(Δ)
              return (AbstractFFTs.bfft(Δ),)
            end
          end

This file has been truncated. show original

This is all defined in terms of the AbstractFFT library API · AbstractFFTs.jl

Topic		Replies	Views
2D convolution in stan Developers features , math	5	1057	July 22, 2020
Stan examples of using fft Modeling	2	591	June 21, 2022
CmdStan & Stan 2.30 release candidate Developers	16	1371	August 5, 2022
Difference in `fft` implementation between stan and numpy/scipy Modeling	7	846	August 24, 2022
1D and 2D convolution functions in stan Modeling features , math , ecology , population-dynamics	1	635	September 19, 2022

Related topics