Radon3D.h

Go to the documentation of this file.

#pragma once
 
#include "Common.h"
 
namespace reg23 {
 
at::Tensor Radon3D_CPU(const at::Tensor &volume, const at::Tensor &volumeSpacing, const at::Tensor &phiValues,
                       const at::Tensor &thetaValues, const at::Tensor &rValues, int64_t samplesPerDirection);
 
at::Tensor DRadon3DDR_CPU(const at::Tensor &volume, const at::Tensor &volumeSpacing, const at::Tensor &phiValues,
                          const at::Tensor &thetaValues, const at::Tensor &rValues, int64_t samplesPerDirection);
 
__host__ at::Tensor Radon3D_CUDA(const at::Tensor &volume, const at::Tensor &volumeSpacing, const at::Tensor &phiValues,
                                 const at::Tensor &thetaValues, const at::Tensor &rValues, int64_t samplesPerDirection);
 
__host__ at::Tensor Radon3D_CUDA_V2(const at::Tensor &volume, const at::Tensor &volumeSpacing,
                                    const at::Tensor &phiValues, const at::Tensor &thetaValues,
                                    const at::Tensor &rValues, int64_t samplesPerDirection);
 
__host__ at::Tensor DRadon3DDR_CUDA(const at::Tensor &volume, const at::Tensor &volumeSpacing,
                                    const at::Tensor &phiValues, const at::Tensor &thetaValues,
                                    const at::Tensor &rValues, int64_t samplesPerDirection);
 
__host__ at::Tensor DRadon3DDR_CUDA_V2(const at::Tensor &volume, const at::Tensor &volumeSpacing,
                                       const at::Tensor &phiValues, const at::Tensor &thetaValues,
                                       const at::Tensor &rValues, int64_t samplesPerDirection);
 
template <typename texture_t> struct Radon3D {
 
    static_assert(texture_t::DIMENSIONALITY == 3);
 
    using IntType = typename texture_t::IntType;
    using FloatType = typename texture_t::FloatType;
    using SizeType = typename texture_t::SizeType;
    using VectorType = typename texture_t::VectorType;
    using AddressModeType = typename texture_t::AddressModeType;
 
    struct CommonData {
        texture_t inputTexture{};
        Linear<VectorType> mappingIndexToOffset{};
        FloatType scaleFactor{};
        at::Tensor flatOutput{};
    };
 
    __host__ static CommonData Common(const at::Tensor &volume, const at::Tensor &volumeSpacing,
                                      const at::Tensor &phiValues, const at::Tensor &thetaValues,
                                      const at::Tensor &rValues, int64_t samplesPerDirection, at::DeviceType device) {
        // volume should be a 3D array of floats on the chosen device
        TORCH_CHECK(volume.sizes().size() == 3);
        TORCH_CHECK(volume.dtype() == at::kFloat);
        TORCH_INTERNAL_ASSERT(volume.device().type() == device);
        // volumeSpacing should be a 1D tensor of 3 doubles
        TORCH_CHECK(volumeSpacing.sizes() == at::IntArrayRef{3});
        TORCH_CHECK(volumeSpacing.dtype() == at::kDouble);
        // phiValues, thetaValues and rValues should have matching sizes, contain floats, and be on the chosen device
        TORCH_CHECK(phiValues.sizes() == thetaValues.sizes());
        TORCH_CHECK(thetaValues.sizes() == rValues.sizes());
        TORCH_CHECK(phiValues.dtype() == at::kFloat);
        TORCH_CHECK(thetaValues.dtype() == at::kFloat);
        TORCH_CHECK(rValues.dtype() == at::kFloat);
        TORCH_INTERNAL_ASSERT(phiValues.device().type() == device);
        TORCH_INTERNAL_ASSERT(thetaValues.device().type() == device);
        TORCH_INTERNAL_ASSERT(rValues.device().type() == device);
 
        CommonData ret{};
        ret.inputTexture = texture_t::FromTensor(volume, VectorType::FromTensor(volumeSpacing));
        const FloatType planeSize = sqrt(
            ret.inputTexture.SizeWorld().template Apply<FloatType>(&Square<FloatType>).Sum());
        ret.mappingIndexToOffset = GetMappingIToOffset(planeSize, samplesPerDirection);
        const FloatType rootScaleFactor = planeSize / static_cast<FloatType>(samplesPerDirection);
        ret.scaleFactor = rootScaleFactor * rootScaleFactor;
        ret.flatOutput = torch::zeros(at::IntArrayRef({phiValues.numel()}), volume.contiguous().options());
        return ret;
    }
 
    __host__ __device__ [[nodiscard]] static Linear<VectorType> GetMappingIToOffset(FloatType planeSize,
        int64_t samplesPerDirection) {
 
        return {VectorType::Full(-.5 * planeSize),
                VectorType::Full(planeSize / static_cast<FloatType>(samplesPerDirection - 1))};
    }
 
    __host__ __device__ [[nodiscard]] static Linear2<VectorType>
    GetMappingIndexToTexCoord(const texture_t &textureIn, FloatType phi, FloatType theta, FloatType r,
                              const Linear<VectorType> &mappingIToOffset) {
        const FloatType sp = sin(phi);
        const FloatType cp = cos(phi);
        const FloatType st = sin(theta);
        const FloatType ct = cos(theta);
        const Linear2<VectorType> mappingOffsetToWorld{{r * ct * cp, r * ct * sp, r * st}, {-sp, cp, 0.f},
                                                       {-st * cp, -st * sp, ct}};
        return textureIn.MappingWorldToTexCoord()(mappingOffsetToWorld(mappingIToOffset));
    }
 
    __host__ __device__ [[nodiscard]] static VectorType GetDTexCoordDR(const texture_t &textureIn, FloatType phi,
                                                                       FloatType theta, FloatType r) {
        const FloatType sign = static_cast<FloatType>(r > 0.) - static_cast<FloatType>(r < 0.);
        const FloatType sp = sin(phi);
        const FloatType cp = cos(phi);
        const FloatType st = sin(theta);
        const FloatType ct = cos(theta);
        return textureIn.MappingWorldToTexCoord().gradient * sign * VectorType{ct * cp, ct * sp, st};
    }
 
    __host__ __device__ [[nodiscard]] static float
    IntegrateLooped(const texture_t &texture, const Linear2<VectorType> &mappingIndexToTexCoord,
                    int64_t samplesPerDirection) {
        float ret = 0.f;
        for (int64_t j = 0; j < samplesPerDirection; ++j) {
            for (int64_t i = 0; i < samplesPerDirection; ++i) {
                const FloatType iF = static_cast<FloatType>(i);
                const FloatType jF = static_cast<FloatType>(j);
                ret += texture.Sample(mappingIndexToTexCoord(VectorType::Full(iF), VectorType::Full(jF)));
            }
        }
        return ret;
    }
 
    __host__ __device__ [[nodiscard]] static float
    DIntegrateLoopedDMappingParameter(const texture_t &texture, const Linear2<VectorType> &mappingIndexToTexCoord,
                                      const VectorType &dTexCoordDMappingParameter, int64_t samplesPerDirection) {
        float ret = 0.f;
        for (int64_t j = 0; j < samplesPerDirection; ++j) {
            for (int64_t i = 0; i < samplesPerDirection; ++i) {
                const FloatType iF = static_cast<FloatType>(i);
                const FloatType jF = static_cast<FloatType>(j);
                const VectorType texCoord = mappingIndexToTexCoord(VectorType::Full(iF), VectorType::Full(jF));
                ret += VecDot(texture.DSampleDTexCoord(texCoord), dTexCoordDMappingParameter);
            }
        }
        return ret;
    }
};
 
} // namespace reg23