Skip to content

Commit

Permalink
docs: Explain the meanings of axes
Browse files Browse the repository at this point in the history
  • Loading branch information
Yuuichi Asahi committed Oct 1, 2024
1 parent 7b8ad76 commit 74beb82
Showing 1 changed file with 41 additions and 1 deletion.
42 changes: 41 additions & 1 deletion docs/intro/using.rst
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ If the rank of Views is higher than the dimension of FFT, a batched FFT plan is
APIs start from ``i`` represent inverse transforms.
For Real FFTs, users have to pay attention to the input and output data types as well as their extents.
Inconsistent data types are suppressed by compilation errors. If extents are inconsistent,
it will raise runtime errors (C++ exceptions or assertions).
it will raise runtime errors (C++ ``std::runtime_error``).
The following listing shows good and bad examples of Real FFTs.

.. code-block:: C++
Expand Down Expand Up @@ -130,3 +130,43 @@ In some backend, FFT plan creation leads to some overhead, wherein we need this
.. note::

Input and Output Views used to call FFT APIs must have the same types and extents as the ones used for plan creation.

Axes parameters
---------------

As well as ``numpy.fft``, you can specify negative axes to perform FFT over chosen axes, which is not common in C++.
Actually for FFT APIs, default axes are set as ``{-DIM, -(DIM-1), ...}`` where ``DIM`` is the rank of the FFT dimensions,
corresponding to the FFTs over last ``DIM`` axes. For example, ``KokkosFFT::fft2(execution_space(), in, out)`` is equivalent to ``KokkosFFT::fft2(execution_space(), in, out, axis_type<2>({-2, -1}))``.
Negative axes are counted from the last axis, which is the same as ``numpy.fft``.
For example, ``-1`` means the last axis, ``-2`` means the second last axis, and so on.
Negative axes ``-1`` and ``-2`` respectively correspond to ``rank-1`` and ``rank-2``, where the ``rank`` is the rank of the Views.

The following listing shows examples of axes parameters with negative or positive values.

.. code-block:: C++

template <typename T> using View2D = Kokkos::View<T**, Kokkos::LayoutLeft, execution_space>;
template <typename T> using View3D = Kokkos::View<T***, Kokkos::LayoutLeft, execution_space>;
constexpr int n0 = 4, n1 = 8, n2 = 5;

View2D<double> x2("x2", n0, n1);
View3D<double> x3("x3", n0, n1, n2);
View2D<Kokkos::complex<double> > x2_hat("x2_hat", n0/2+1, n1);
View3D<Kokkos::complex<double> > x3_hat("x3_hat", n0, n1/2+1, n2);

// Follwoing codes are all equivalent to np.fft(np.rfft(axis=0), axis=1)
KokkosFFT::rfft2(execution_space(), x2, x2_hat, /*axes=*/{-1, -2});
KokkosFFT::rfft2(execution_space(), x2, x2_hat, /*axes=*/{-1, 0});
KokkosFFT::rfft2(execution_space(), x2, x2_hat, /*axes=*/{1, -2});
KokkosFFT::rfft2(execution_space(), x2, x2_hat, /*axes=*/{1, 0});

// Follwoing codes are all equivalent to np.fft(np.rfft(axis=1), axis=2)
KokkosFFT::rfft2(execution_space(), x3, x3_hat, /*axes=*/{-1, -2});
KokkosFFT::rfft2(execution_space(), x3, x3_hat, /*axes=*/{-1, 1});
KokkosFFT::rfft2(execution_space(), x3, x3_hat, /*axes=*/{2, -2});
KokkosFFT::rfft2(execution_space(), x3, x3_hat, /*axes=*/{2, 1});

.. note::

If you rely on negative axes, the corresponding positve axes are different depending on the rank of Views.
Thus, it is recommended to use negative axes for simplicity.

0 comments on commit 74beb82

Please sign in to comment.