Commit 112dcb2e authored by Jonas Jelten's avatar Jonas Jelten 🥕
Browse files

add fourier transform operator

parent a35fe939
......@@ -69,6 +69,27 @@ if(${ELSA_BUILD_WITH_QUICKVEC})
list(APPEND MODULE_PUBLIC_DEPS "Quickvec::quickvec")
endif()
# use FFTW if available
option(WITH_FFTW "Build elsa using fftw for faster fourier transforms" ON)
# workaround for
# https://github.com/FFTW/fftw3/issues/130
# better would be find_package(fftw3)
# fftw3f: float, fftw3: double, fftw3l: 128bit, _omp: OpenMP
if(WITH_FFTW)
find_package(PkgConfig)
if(PKG_CONFIG_FOUND)
pkg_check_modules(FFTW3d IMPORTED_TARGET "fftw3")
pkg_check_modules(FFTW3f IMPORTED_TARGET "fftw3f")
if(FFTW3d_FOUND AND FFTW3f_FOUND)
set(FFTW3_FOUND TRUE)
list(APPEND MODULE_PUBLIC_DEPS "PkgConfig::FFTW3d" "PkgConfig::FFTW3f")
# TODO: also add fftw3_omp if supported
endif()
endif()
endif()
ADD_ELSA_MODULE(
core "${MODULE_HEADERS}" "${MODULE_SOURCES}" INSTALL_DIR PUBLIC_DEPS ${MODULE_PUBLIC_DEPS}
PRIVATE_DEPS ${MODULE_PRIVATE_DEPS}
......@@ -93,6 +114,11 @@ if(${ELSA_BUILD_WITH_QUICKVEC})
target_compile_definitions(${ELSA_MODULE_TARGET_NAME} PUBLIC ELSA_ENABLE_CUDA_VECTOR)
endif()
if(WITH_FFTW AND FFTW3_FOUND)
target_compile_definitions("${ELSA_MODULE_TARGET_NAME}" PRIVATE WITH_FFTW)
endif()
if(ELSA_BUILD_PYTHON_BINDINGS)
set(BINDINGS_SOURCES
${MODULE_SOURCES}
......
......@@ -183,6 +183,18 @@ namespace elsa
return _dataHandler->maxElement();
}
template <typename data_t>
void DataContainer<data_t>::fft() const
{
this->_dataHandler->fft(*this->_dataDescriptor);
}
template <typename data_t>
void DataContainer<data_t>::ifft() const
{
this->_dataHandler->ifft(*this->_dataDescriptor);
}
template <typename data_t>
DataContainer<data_t>& DataContainer<data_t>::operator+=(const DataContainer<data_t>& dc)
{
......
......@@ -29,6 +29,7 @@ namespace elsa
* @author David Frank - added DataHandler concept, iterators
* @author Nikola Dinev - add block support
* @author Jens Petit - expression templates
* @author Jonas Jelten - various enhancements, fft, complex handling, pretty formatting
*
* @tparam data_t - data type that is stored in the DataContainer, defaulting to real_t.
*
......@@ -47,7 +48,8 @@ namespace elsa
DataContainer() = delete;
/**
* @brief Constructor for empty DataContainer, no initialisation is performed
* @brief Constructor for empty DataContainer, no initialisation is performed,
* but the underlying space is allocated.
*
* @param[in] dataDescriptor containing the associated metadata
* @param[in] handlerType the data handler (default: CPU)
......@@ -246,6 +248,12 @@ namespace elsa
/// return the max of all elements of this signal
data_t maxElement() const;
/// convert to the fourier transformed signal
void fft() const;
/// convert to the inverse fourier transformed signal
void ifft() const;
/// if the datacontainer is already complex, return itself.
template <typename _data_t = data_t>
typename std::enable_if_t<isComplex<_data_t>, DataContainer<_data_t>> asComplex() const
......
......@@ -14,6 +14,8 @@
namespace elsa
{
class DataDescriptor;
/**
* @brief Base class encapsulating data handling. The data is stored transparently, for example
* on CPU or GPU.
......@@ -81,6 +83,12 @@ namespace elsa
/// return the max of all elements of the data vector
virtual data_t maxElement() const = 0;
/// in-place create the fourier transformed of the data vector
virtual DataHandler<data_t>& fft(const DataDescriptor& source_desc) = 0;
/// in-place create the inverse fourier transformed of the data vector
virtual DataHandler<data_t>& ifft(const DataDescriptor& source_desc) = 0;
/// compute in-place element-wise addition of another vector v
virtual DataHandler<data_t>& operator+=(const DataHandler<data_t>& v) = 0;
......
#include "DataHandlerCPU.h"
#include "DataHandlerMapCPU.h"
#include "Error.h"
#include "TypeCasts.hpp"
#include "DataDescriptor.h"
#include <iostream>
#if WITH_FFTW
#define EIGEN_FFTW_DEFAULT
#endif
#include <unsupported/Eigen/FFT>
namespace elsa
{
......@@ -130,6 +140,20 @@ namespace elsa
}
}
template <typename data_t>
DataHandler<data_t>& DataHandlerCPU<data_t>::fft(const DataDescriptor& source_desc)
{
this->base_fft<true>(source_desc);
return *this;
}
template <typename data_t>
DataHandler<data_t>& DataHandlerCPU<data_t>::ifft(const DataDescriptor& source_desc)
{
this->base_fft<false>(source_desc);
return *this;
}
template <typename data_t>
DataHandler<data_t>& DataHandlerCPU<data_t>::operator+=(const DataHandler<data_t>& v)
{
......@@ -459,6 +483,94 @@ namespace elsa
return DataMap_t(&(_data->operator[](0)), getSize());
}
template <typename data_t>
template <bool is_forward>
void DataHandlerCPU<data_t>::base_fft(const DataDescriptor& source_desc)
{
if constexpr (isComplex<data_t>) {
// now that we change this datahandler,
// copy-on-write it.
this->detach();
const auto& src_shape = source_desc.getNumberOfCoefficientsPerDimension();
const auto& src_dims = source_desc.getNumberOfDimensions();
typename DataVector_t::Scalar* this_data = this->_data->data();
// TODO: fftw variant
// generalization of an 1D-FFT
// walk over each dimenson and 1d-fft one 'line' of data
for (index_t dim_idx = 0; dim_idx < src_dims; ++dim_idx) {
// jumps in the data for the current dimension's data
// dim_size[0] * dim_size[1] * ...
// 1 for dim_idx == 0.
const index_t stride = src_shape.head(dim_idx).prod();
// number of coefficients for the current dimension
const index_t dim_size = src_shape(dim_idx);
// number of coefficients for the other dimensions
// this is the number of 1d-ffts we'll do
// e.g. shape=[2, 3, 4] and we do dim_idx=2 (=shape 4)
// -> src_shape.prod() == 24 / 4 = 6 == 2*3
const index_t other_dims_size = src_shape.prod() / dim_size;
#ifndef EIGEN_FFTW_DEFAULT
// when using eigen+fftw, this corrupts the memory, so don't parallelize.
// error messages may include:
// * double free or corruption (fasttop)
// * malloc_consolidate(): unaligned fastbin chunk detected
#pragma omp parallel for
#endif
// do all the 1d-ffts along the current dimensions axis
for (index_t i = 0; i < other_dims_size; ++i) {
index_t ray_start = i;
// each time i is a multiple of stride,
// jump forward the current+previous dimensions' shape product
// (draw an indexed 3d cube to visualize this)
ray_start += (stride * (dim_size - 1)) * ((i - (i % stride)) / stride);
// this is one "ray" through the volume
Eigen::Map<DataVector_t, Eigen::AlignmentType::Unaligned, Eigen::InnerStride<>>
input_map(this_data + ray_start, dim_size, Eigen::InnerStride<>(stride));
Eigen::FFT<GetFloatingPointType_t<typename DataVector_t::Scalar>> fft_op;
Eigen::Matrix<data_t, Eigen::Dynamic, 1> fft_in{dim_size};
Eigen::Matrix<data_t, Eigen::Dynamic, 1> fft_out{dim_size};
// eigen internally copies the fwd input matrix anyway if
// it doesn't have stride == 1
fft_in = input_map.block(0, 0, dim_size, 1);
if (unlikely(dim_size == 1)) {
// eigen kiss-fft crashes for size=1...
fft_out = fft_in;
} else {
// arguments for in and out _must not_ be the same matrix!
// they will corrupt wildly otherwise.
if constexpr (is_forward) {
fft_op.fwd(fft_out, fft_in);
} else {
// eigen inv-fft already scales down by dim_size
fft_op.inv(fft_out, fft_in);
}
}
// we can't directly use the map as fft output,
// since Eigen internally just uses the pointer to
// the map's first element, and doesn't respect stride at all..
input_map.block(0, 0, dim_size, 1) = fft_out;
}
}
} else {
throw Error{"fft with non-complex input container not supported"};
}
}
// ------------------------------------------
// explicit template instantiation
template class DataHandlerCPU<float>;
......
......@@ -16,6 +16,7 @@ namespace elsa
// forward declaration for friend test function
template <typename data_t = real_t>
class DataHandlerCPU;
// forward declaration, used for testing and defined in test file (declared as friend)
template <typename data_t>
long useCount(const DataHandlerCPU<data_t>& /*dh*/);
......@@ -129,6 +130,12 @@ namespace elsa
/// return the max of all elements of the data vector
data_t maxElement() const override;
/// create the fourier transformed of the data vector
DataHandler<data_t>& fft(const DataDescriptor& source_desc) override;
/// create the inverse fourier transformed of the data vector
DataHandler<data_t>& ifft(const DataDescriptor& source_desc) override;
/// copy assign another DataHandlerCPU to this, other types handled in assign()
DataHandlerCPU<data_t>& operator=(const DataHandlerCPU<data_t>& v);
......@@ -213,5 +220,8 @@ namespace elsa
/// change the vector being handled (rvalue version)
void attach(std::shared_ptr<DataVector_t>&& data);
template <bool is_forward>
void base_fft(const DataDescriptor& source_desc);
};
} // namespace elsa
#include "DataHandlerGPU.h"
#include "DataContainer.h"
#include "DataHandlerMapGPU.h"
#include "TypeCasts.hpp"
#include "elsaDefines.h"
#include <cublas_v2.h>
......@@ -128,6 +131,35 @@ namespace elsa
return _data->maxElement();
}
DataHandler<data_t>& DataHandlerGPU<data_t>::fft(const DataDescriptor& source_desc)
{
// until we have a gpu fft implementation, use the cpu version.
DataContainer<data_t> tmp{source_desc, DataHandlerType::CPU};
for (index_t i = 0; i < this->getSize(); i++) {
tmp[i] = this->operator[](i);
}
tmp.fft();
for (index_t i = 0; i < this->getSize(); i++) {
this->operator[](i) = tmp[i];
}
return *this;
}
template <typename data_t>
DataHandler<data_t>& DataHandlerGPU<data_t>::ifft(const DataDescriptor& source_desc)
{
// until we have a gpu fft implementation, use the cpu version.
DataContainer<data_t> tmp{source_desc, DataHandlerType::CPU};
for (index_t i = 0; i < this->getSize(); i++) {
tmp[i] = this->operator[](i);
}
tmp.ifft();
for (index_t i = 0; i < this->getSize(); i++) {
this->operator[](i) = tmp[i];
}
return *this;
}
template <typename data_t>
DataHandler<data_t>& DataHandlerGPU<data_t>::operator+=(const DataHandler<data_t>& v)
{
......
......@@ -134,6 +134,12 @@ namespace elsa
/// return the max of all elements of the data vector
data_t maxElement() const override;
/// create the fourier transformed of the data vector
DataHandler<data_t>& fft(const DataDescriptor& source_desc) override;
/// create the inverse fourier transformed of the data vector
DataHandler<data_t>& ifft(const DataDescriptor& source_desc) override;
/// copy assign another DataHandlerGPU
DataHandlerGPU<data_t>& operator=(const DataHandlerGPU<data_t>& v);
......
......@@ -132,6 +132,22 @@ namespace elsa
}
}
template <typename data_t>
DataHandler<data_t>& DataHandlerMapCPU<data_t>::fft(const DataDescriptor& source_desc)
{
// detaches internally
this->_dataOwner->fft(source_desc);
return *this;
}
template <typename data_t>
DataHandler<data_t>& DataHandlerMapCPU<data_t>::ifft(const DataDescriptor& source_desc)
{
// detaches internally
this->_dataOwner->ifft(source_desc);
return *this;
}
template <typename data_t>
DataHandler<data_t>& DataHandlerMapCPU<data_t>::operator+=(const DataHandler<data_t>& v)
{
......
......@@ -105,6 +105,12 @@ namespace elsa
/// return the max of all elements of the data vector
data_t maxElement() const override;
/// create the fourier transformed of the data vector
DataHandler<data_t>& fft(const DataDescriptor& source_desc) override;
/// create the inverse fourier transformed of the data vector
DataHandler<data_t>& ifft(const DataDescriptor& source_desc) override;
/// compute in-place element-wise addition of another vector v
DataHandler<data_t>& operator+=(const DataHandler<data_t>& v) override;
......
......@@ -141,6 +141,18 @@ namespace elsa
return _map.maxElement();
}
template <typename data_t>
DataHandler<data_t>& DataHandlerMapGPU<data_t>::fft(const DataDescriptor& source_desc)
{
throw std::runtime_error{"todo implement"};
}
template <typename data_t>
DataHandler<data_t>& DataHandlerMapGPU<data_t>::ifft(const DataDescriptor& source_desc)
{
throw std::runtime_error{"todo implement"};
}
template <typename data_t>
DataHandler<data_t>& DataHandlerMapGPU<data_t>::operator+=(const DataHandler<data_t>& v)
{
......
......@@ -124,6 +124,12 @@ namespace elsa
/// return the max of all elements of the data vector
data_t maxElement() const override;
/// create the fourier transformed of the data vector
DataHandler<data_t>& fft(const DataDescriptor& source_desc) override;
/// create the inverse fourier transformed of the data vector
DataHandler<data_t>& ifft(const DataDescriptor& source_desc) override;
/// compute in-place element-wise addition of another vector v
DataHandler<data_t>& operator+=(const DataHandler<data_t>& v) override;
......
......@@ -83,3 +83,17 @@ namespace elsa
constexpr bool isComplex = std::is_same<RemoveCvRef_t<T>, std::complex<float>>::value
|| std::is_same<RemoveCvRef_t<T>, std::complex<double>>::value;
} // namespace elsa
/*
* Branch prediction tuning.
* the expression is expected to be true (=likely) or false (=unlikely).
*
* btw, this implementation was taken from the Linux kernel.
*/
#if defined(__GNUC__) || defined(__clang__)
#define likely(x) __builtin_expect(!!(x), 1)
#define unlikely(x) __builtin_expect(!!(x), 0)
#else
#define likely(x) (x)
#define unlikely(x) (x)
#endif
# list all the headers of the module
set(MODULE_HEADERS Identity.h Scaling.h FiniteDifferences.h BlockLinearOperator.h Dictionary.h)
set(MODULE_HEADERS
Identity.h
Scaling.h
FiniteDifferences.h
FourierTransform.h
BlockLinearOperator.h
Dictionary.h)
# list all the code files of the module
set(MODULE_SOURCES Identity.cpp Scaling.cpp FiniteDifferences.cpp BlockLinearOperator.cpp Dictionary.cpp)
set(MODULE_SOURCES
Identity.cpp
Scaling.cpp
FiniteDifferences.cpp
FourierTransform.cpp
BlockLinearOperator.cpp
Dictionary.cpp)
list(APPEND MODULE_PUBLIC_DEPS elsa_core elsa_logging)
list(APPEND MODULE_PRIVATE_DEPS)
......
#include "FourierTransform.h"
#include "Error.h"
#include "ExpressionPredicates.h"
#include "Timer.h"
#include <iostream>
namespace elsa
{
template <typename data_t>
FourierTransform<data_t>::FourierTransform(const DataDescriptor& domainDescriptor)
: B(domainDescriptor, domainDescriptor)
{
}
template <typename data_t>
void FourierTransform<data_t>::applyImpl(const DataContainer<data_t>& x,
DataContainer<data_t>& Ax) const
{
Timer timeguard("FourierTransform", "apply()");
// TODO: if domain and range descriptors dont match,
// resize the datacontainer appropriately.
auto x_values = x.getDataDescriptor().getNumberOfCoefficientsPerDimension();
assert(x_values == this->_domainDescriptor->getNumberOfCoefficientsPerDimension());
// copy the input and fouriertransform it
Ax = x;
Ax.fft();
}
template <typename data_t>
void FourierTransform<data_t>::applyAdjointImpl(const DataContainer<data_t>& x,
DataContainer<data_t>& Atx) const
{
Timer timeguard("FourierTransform", "applyAdjoint()");
// TODO: same as applyImpl
auto x_values = x.getDataDescriptor().getNumberOfCoefficientsPerDimension();
assert(x_values == this->_domainDescriptor->getNumberOfCoefficientsPerDimension());
// copy the input and inverse-fouriertransform it
Atx = x;
Atx.ifft();
}
template <typename data_t>
FourierTransform<data_t>* FourierTransform<data_t>::cloneImpl() const
{
auto& domainDescriptor = static_cast<const DataDescriptor&>(*this->_domainDescriptor);
return new FourierTransform(domainDescriptor);
}
template <typename data_t>
bool FourierTransform<data_t>::isEqual(const B& other) const
{
if (!B::isEqual(other))
return false;
auto otherOP = dynamic_cast<const FourierTransform*>(&other);
if (!otherOP)
return false;
// fourier transforms have no configuration yet
return true;
}
template class FourierTransform<std::complex<float>>;
template class FourierTransform<std::complex<double>>;
} // namespace elsa
#pragma once
#include <complex>
#include "LinearOperator.h"
namespace elsa
{
/**
* @brief Operator for applying multi-dimensional fourier transforms.
*
* @author Jonas Jelten - initial code
*
* @tparam data_t data type for the domain and range of the transformation,
* defaulting to real_t
*
* Implements the n-dimensional signal fourier transformation.
* Can support multiple backends, by default uses Eigen::FFT with FFTW.
*/
template <typename data_t = std::complex<real_t>>
class FourierTransform : public LinearOperator<data_t>
{
private:
using B = LinearOperator<data_t>;
/** working data container for the fft.
like in datacontainer, we operate on the vector in n dimensions. */
using fftvector_t = Eigen::Matrix<data_t, Eigen::Dynamic, 1>;
public:
/**
* @brief create a fourier transform operator
*
* @param[in] domainDescriptor metadata defining the domain and range of the transformation
*/
explicit FourierTransform(const DataDescriptor& domainDescriptor);
~FourierTransform() override = default;
protected:
/**
* @brief perform the fourier transformation
* @param x inputData (image matrix)
* @param Ax outputData (fourier transformed image matrix)
*/
void applyImpl(const DataContainer<data_t>& x, DataContainer<data_t>& Ax) const override;
/**
* @brief perform the inverse fourier transformation
* @param x inputData (image matrix in frequency domain)
* @param Atx outputData (inversely fourier transformed image matrix)
*/
void applyAdjointImpl(const DataContainer<data_t>& y,
DataContainer<data_t>& Aty) const override;
/// implement the polymorphic clone operation
FourierTransform* cloneImpl() const override;
/// implement the polymorphic comparison operation
bool isEqual(const B& other) const override;
};
} // namespace elsa
......@@ -16,5 +16,6 @@ add_custom_target(
ELSA_DOCTEST(Identity)
ELSA_DOCTEST(Scaling)
ELSA_DOCTEST(FiniteDifferences)
ELSA_DOCTEST(FourierTransform)
ELSA_DOCTEST(BlockLinearOperator)
ELSA_DOCTEST(Dictionary)
/**
* @file test_FourierTransform.cpp
*
* @brief Tests for the fourier transform operator
*
* @author Jonas Jelten
*/
#include "ExpressionPredicates.h"
#include "doctest/doctest.h"
#include "elsaDefines.h"
#include "testHelpers.h"
#include "FourierTransform.h"
#include "VolumeDescriptor.h"