polar_spline.hpp

#pragma once
 
#include <ddc/ddc.hpp>
 
template <class PolarBSplinesType, class MemSpace = Kokkos::DefaultExecutionSpace::memory_space>
struct PolarSplineSpan;
 
template <class PolarBSplinesType, class MemSpace = Kokkos::DefaultExecutionSpace::memory_space>
struct PolarSplineView;
 
 
template <class PolarBSplinesType, class MemSpace = Kokkos::DefaultExecutionSpace::memory_space>
struct PolarSpline
{
public:
    using BSplinesR = typename PolarBSplinesType::BSplinesR_tag;
    using BSplinesTheta = typename PolarBSplinesType::BSplinesTheta_tag;
 
public:
    ddc::Chunk<
            double,
            ddc::DiscreteDomain<BSplinesR, BSplinesTheta>,
            ddc::KokkosAllocator<double, MemSpace>>
            spline_coef;
 
    ddc::Chunk<
            double,
            ddc::DiscreteDomain<PolarBSplinesType>,
            ddc::KokkosAllocator<double, MemSpace>>
            singular_spline_coef;
 
public:
    PolarSpline<PolarBSplinesType, MemSpace>(ddc::DiscreteDomain<BSplinesR, BSplinesTheta> domain)
        : spline_coef(ddc::DiscreteDomain<BSplinesR, BSplinesTheta>(
                ddc::select<BSplinesR>(domain).remove_first(
                        ddc::DiscreteVector<BSplinesR>(PolarBSplinesType::continuity + 1)),
                ddc::select<BSplinesTheta>(domain)))
        , singular_spline_coef(ddc::DiscreteDomain<PolarBSplinesType>(
                  ddc::DiscreteElement<PolarBSplinesType>(0),
                  ddc::DiscreteVector<PolarBSplinesType>(PolarBSplinesType::n_singular_basis())))
    {
    }
 
    PolarSpline<PolarBSplinesType, MemSpace>(
            ddc::DiscreteDomain<PolarBSplinesType> singular_domain,
            ddc::DiscreteDomain<BSplinesR, BSplinesTheta> domain)
        : spline_coef(domain)
        , singular_spline_coef(singular_domain.take_first(
                  ddc::DiscreteVector<PolarBSplinesType>(PolarBSplinesType::n_singular_basis())))
    {
    }
 
    PolarSplineSpan<PolarBSplinesType, MemSpace> span_view()
    {
        return PolarSplineSpan<PolarBSplinesType, MemSpace>(*this);
    }
 
    PolarSplineView<PolarBSplinesType, MemSpace> span_cview() const
    {
        return PolarSplineView<PolarBSplinesType, MemSpace>(*this);
    }
};
 
template <class PolarBSplinesType, class MemSpace>
struct PolarSplineSpan
{
public:
    using BSplinesR = typename PolarBSplinesType::BSplinesR_tag;
    using BSplinesTheta = typename PolarBSplinesType::BSplinesTheta_tag;
 
public:
    ddc::ChunkSpan<
            double,
            ddc::DiscreteDomain<BSplinesR, BSplinesTheta>,
            Kokkos::layout_right,
            MemSpace>
            spline_coef;
 
    ddc::ChunkSpan<double, ddc::DiscreteDomain<PolarBSplinesType>, Kokkos::layout_right, MemSpace>
            singular_spline_coef;
 
public:
    PolarSplineSpan<PolarBSplinesType, MemSpace>(PolarSpline<PolarBSplinesType, MemSpace>& spl)
        : spline_coef(spl.spline_coef.span_view())
        , singular_spline_coef(spl.singular_spline_coef.span_view())
    {
    }
 
    PolarSplineSpan<PolarBSplinesType, MemSpace> span_view()
    {
        return *this;
    }
 
    PolarSplineView<PolarBSplinesType, MemSpace> span_cview() const
    {
        return PolarSplineView<PolarBSplinesType, MemSpace>(*this);
    }
};
 
template <class PolarBSplinesType, class MemSpace>
struct PolarSplineView
{
public:
    using BSplinesR = typename PolarBSplinesType::BSplinesR_tag;
    using BSplinesTheta = typename PolarBSplinesType::BSplinesTheta_tag;
 
public:
    ddc::ChunkSpan<
            double const,
            ddc::DiscreteDomain<BSplinesR, BSplinesTheta>,
            Kokkos::layout_right,
            MemSpace> const spline_coef;
 
    ddc::ChunkSpan<
            double const,
            ddc::DiscreteDomain<PolarBSplinesType>,
            Kokkos::layout_right,
            MemSpace> const singular_spline_coef;
 
public:
    PolarSplineView<PolarBSplinesType, MemSpace>(
            PolarSpline<PolarBSplinesType, MemSpace> const& spl)
        : spline_coef(spl.spline_coef.span_cview())
        , singular_spline_coef(spl.singular_spline_coef.span_cview())
    {
    }
 
    PolarSplineView<PolarBSplinesType, MemSpace>(
            PolarSplineSpan<PolarBSplinesType, MemSpace> const& spl)
        : spline_coef(spl.spline_coef.span_cview())
        , singular_spline_coef(spl.singular_spline_coef.span_cview())
    {
    }
 
    PolarSplineSpan<PolarBSplinesType, MemSpace> span_view() const
    {
        return *this;
    }
 
    PolarSplineView<PolarBSplinesType, MemSpace> span_cview() const
    {
        return *this;
    }
};
 
template <class T>
inline constexpr bool is_polar_spline_v = false;
 
template <class PolarBSplinesType, class MemSpace>
inline constexpr bool is_polar_spline_v<PolarSpline<PolarBSplinesType, MemSpace>> = true;
 
template <class PolarBSplinesType, class MemSpace>
inline constexpr bool is_polar_spline_v<PolarSplineSpan<PolarBSplinesType, MemSpace>> = true;
 
template <class PolarBSplinesType, class MemSpace>
inline constexpr bool is_polar_spline_v<PolarSplineView<PolarBSplinesType, MemSpace>> = true;
 
template <class ExecSpace, class PolarBSplinesType, class MemSpace>
PolarSpline<PolarBSplinesType, typename ExecSpace::memory_space> create_mirror(
        ExecSpace const& exec_space,
        PolarSplineSpan<PolarBSplinesType, MemSpace> const& src)
{
    PolarSpline<PolarBSplinesType, typename ExecSpace::memory_space>
            dst(src.singular_spline_coef.domain(), src.spline_coef.domain());
    return dst;
}
 
template <class PolarBSplinesType, class MemSpace>
PolarSpline<PolarBSplinesType, Kokkos::HostSpace> create_mirror(
        PolarSplineSpan<PolarBSplinesType, MemSpace> const& src)
{
    return create_mirror(Kokkos::DefaultHostExecutionSpace(), src);
}
 
template <class ExecSpace, class PolarBSplinesType, class MemSpace>
PolarSpline<PolarBSplinesType, typename ExecSpace::memory_space> create_mirror_and_copy(
        ExecSpace const& exec_space,
        PolarSplineSpan<PolarBSplinesType, MemSpace> const& src)
{
    PolarSpline<PolarBSplinesType, typename ExecSpace::memory_space> dst
            = create_mirror(exec_space, src);
    ddc::parallel_deepcopy(dst.spline_coef, src.spline_coef);
    ddc::parallel_deepcopy(dst.singular_spline_coef, src.singular_spline_coef);
    return dst;
}
 
template <class PolarBSplinesType, class MemSpace>
PolarSpline<PolarBSplinesType, Kokkos::HostSpace> create_mirror_and_copy(
        PolarSplineSpan<PolarBSplinesType, MemSpace> const& src)
{
    return create_mirror_and_copy(Kokkos::DefaultHostExecutionSpace(), src);
}
 
template <class ExecSpace, class PolarBSplinesType, class MemSpace>
auto create_mirror_view(
        ExecSpace const& exec_space,
        PolarSplineSpan<PolarBSplinesType, MemSpace> const& src)
{
    if constexpr (std::is_same_v<MemSpace, typename ExecSpace::memory_space>) {
        return src;
    } else {
        PolarSpline<PolarBSplinesType, typename ExecSpace::memory_space>
                dst(src.singular_spline_coef.domain(), src.spline_coef.domain());
        return dst;
    }
}
 
template <class PolarBSplinesType, class MemSpace>
auto create_mirror_view(PolarSplineSpan<PolarBSplinesType, MemSpace> const& src)
{
    return create_mirror_view(Kokkos::DefaultHostExecutionSpace(), src);
}
 
template <class ExecSpace, class PolarBSplinesType, class MemSpace>
auto create_mirror_view_and_copy(
        ExecSpace const& exec_space,
        PolarSplineSpan<PolarBSplinesType, MemSpace> const& src)
{
    auto dst = create_mirror_view(exec_space, src);
    ddc::parallel_deepcopy(dst.spline_coef, src.spline_coef);
    ddc::parallel_deepcopy(dst.singular_spline_coef, src.singular_spline_coef);
    return dst;
}
 
template <class PolarBSplinesType, class MemSpace>
auto create_mirror_view_and_copy(PolarSplineSpan<PolarBSplinesType, MemSpace> const& src)
{
    return create_mirror_view_and_copy(Kokkos::DefaultHostExecutionSpace(), src);
}