fft.hpp

// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_MATRIX_FFT_HPP_
#define GKO_PUBLIC_CORE_MATRIX_FFT_HPP_
 
 
#include <ginkgo/core/base/array.hpp>
#include <ginkgo/core/base/lin_op.hpp>
 
 
namespace gko {
namespace matrix {
 
 
class Fft : public EnableLinOp<Fft>,
            public WritableToMatrixData<std::complex<float>, int32>,
            public WritableToMatrixData<std::complex<float>, int64>,
            public WritableToMatrixData<std::complex<double>, int32>,
            public WritableToMatrixData<std::complex<double>, int64>,
            public Transposable {
    friend class EnablePolymorphicObject<Fft, LinOp>;
 
public:
    using EnableLinOp<Fft>::convert_to;
    using EnableLinOp<Fft>::move_to;
 
    using value_type = std::complex<double>;
    using index_type = int64;
    using transposed_type = Fft;
 
    std::unique_ptr<LinOp> transpose() const override;
 
    std::unique_ptr<LinOp> conj_transpose() const override;
 
    void write(matrix_data<std::complex<float>, int32>& data) const override;
 
    void write(matrix_data<std::complex<float>, int64>& data) const override;
 
    void write(matrix_data<std::complex<double>, int32>& data) const override;
 
    void write(matrix_data<std::complex<double>, int64>& data) const override;
 
    dim<1> get_fft_size() const;
 
    bool is_inverse() const;
 
    static std::unique_ptr<Fft> create(std::shared_ptr<const Executor> exec);
 
    static std::unique_ptr<Fft> create(std::shared_ptr<const Executor> exec,
                                       size_type size = 0,
                                       bool inverse = false);
 
protected:
    Fft(std::shared_ptr<const Executor> exec, size_type size = 0,
        bool inverse = false);
 
    void apply_impl(const LinOp* b, LinOp* x) const override;
 
    void apply_impl(const LinOp* alpha, const LinOp* b, const LinOp* beta,
                    LinOp* x) const override;
 
private:
    mutable array<char> buffer_;
    bool inverse_;
};
 
 
class Fft2 : public EnableLinOp<Fft2>,
             public WritableToMatrixData<std::complex<float>, int32>,
             public WritableToMatrixData<std::complex<float>, int64>,
             public WritableToMatrixData<std::complex<double>, int32>,
             public WritableToMatrixData<std::complex<double>, int64>,
             public Transposable {
    friend class EnablePolymorphicObject<Fft2, LinOp>;
 
public:
    using EnableLinOp<Fft2>::convert_to;
    using EnableLinOp<Fft2>::move_to;
 
    using value_type = std::complex<double>;
    using index_type = int64;
    using transposed_type = Fft2;
 
    std::unique_ptr<LinOp> transpose() const override;
 
    std::unique_ptr<LinOp> conj_transpose() const override;
 
    void write(matrix_data<std::complex<float>, int32>& data) const override;
 
    void write(matrix_data<std::complex<float>, int64>& data) const override;
 
    void write(matrix_data<std::complex<double>, int32>& data) const override;
 
    void write(matrix_data<std::complex<double>, int64>& data) const override;
 
    dim<2> get_fft_size() const;
 
    bool is_inverse() const;
 
    static std::unique_ptr<Fft2> create(std::shared_ptr<const Executor> exec);
 
    static std::unique_ptr<Fft2> create(std::shared_ptr<const Executor> exec,
                                        size_type size);
 
    static std::unique_ptr<Fft2> create(std::shared_ptr<const Executor> exec,
                                        size_type size1, size_type size2,
                                        bool inverse = false);
 
protected:
    Fft2(std::shared_ptr<const Executor> exec, size_type size1 = 0,
         size_type size2 = 0, bool inverse = false);
 
    void apply_impl(const LinOp* b, LinOp* x) const override;
 
    void apply_impl(const LinOp* alpha, const LinOp* b, const LinOp* beta,
                    LinOp* x) const override;
 
private:
    mutable array<char> buffer_;
    dim<2> fft_size_;
    bool inverse_;
};
 
 
class Fft3 : public EnableLinOp<Fft3>,
             public WritableToMatrixData<std::complex<float>, int32>,
             public WritableToMatrixData<std::complex<float>, int64>,
             public WritableToMatrixData<std::complex<double>, int32>,
             public WritableToMatrixData<std::complex<double>, int64>,
             public Transposable {
    friend class EnablePolymorphicObject<Fft3, LinOp>;
 
public:
    using EnableLinOp<Fft3>::convert_to;
    using EnableLinOp<Fft3>::move_to;
 
    using value_type = std::complex<double>;
    using index_type = int64;
    using transposed_type = Fft3;
 
    std::unique_ptr<LinOp> transpose() const override;
 
    std::unique_ptr<LinOp> conj_transpose() const override;
 
    void write(matrix_data<std::complex<float>, int32>& data) const override;
 
    void write(matrix_data<std::complex<float>, int64>& data) const override;
 
    void write(matrix_data<std::complex<double>, int32>& data) const override;
 
    void write(matrix_data<std::complex<double>, int64>& data) const override;
 
    dim<3> get_fft_size() const;
 
    bool is_inverse() const;
 
    static std::unique_ptr<Fft3> create(std::shared_ptr<const Executor> exec);
 
    static std::unique_ptr<Fft3> create(std::shared_ptr<const Executor> exec,
                                        size_type size);
 
    static std::unique_ptr<Fft3> create(std::shared_ptr<const Executor> exec,
                                        size_type size1, size_type size2,
                                        size_type size3, bool inverse = false);
 
protected:
    Fft3(std::shared_ptr<const Executor> exec, size_type size1 = 0,
         size_type size2 = 0, size_type size3 = 0, bool inverse = false);
 
    void apply_impl(const LinOp* b, LinOp* x) const override;
 
    void apply_impl(const LinOp* alpha, const LinOp* b, const LinOp* beta,
                    LinOp* x) const override;
 
private:
    mutable array<char> buffer_;
    dim<3> fft_size_;
    bool inverse_;
};
 
 
}  // namespace matrix
}  // namespace gko
 
 
#endif  // GKO_PUBLIC_CORE_MATRIX_FFT_HPP_