segmentedField.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: MIT
// SPDX-FileCopyrightText: 2023 NeoFOAM authors
#pragma once
 
#include "NeoFOAM/core/primitives/label.hpp"
#include "NeoFOAM/fields/field.hpp"
 
namespace NeoFOAM
{
 
template<typename IndexType>
IndexType segmentsFromIntervals(const Field<IndexType>& intervals, Field<IndexType>& offsets)
{
    IndexType finalValue = 0;
    auto inSpan = intervals.span();
    auto offsSpan = offsets.span();
    NF_ASSERT_EQUAL(inSpan.size() + 1, offsSpan.size());
    NeoFOAM::parallelScan(
        intervals.exec(),
        {1, offsSpan.size()},
        KOKKOS_LAMBDA(const std::size_t i, NeoFOAM::localIdx& update, const bool final) {
            update += inSpan[i - 1];
            if (final)
            {
                offsSpan[i] = update;
            }
        },
        finalValue
    );
    return finalValue;
}
 
template<typename ValueType, typename IndexType = NeoFOAM::localIdx>
class SegmentedFieldView
{
public:
 
    std::span<ValueType> values;
 
    std::span<IndexType> segments;
 
    KOKKOS_INLINE_FUNCTION
    Kokkos::pair<IndexType, IndexType> bounds(std::size_t segI) const
    {
        return Kokkos::pair<IndexType, IndexType> {segments[segI], segments[segI + 1]};
    }
 
    KOKKOS_INLINE_FUNCTION
    Kokkos::pair<IndexType, IndexType> range(std::size_t segI) const
    {
        return Kokkos::pair<IndexType, IndexType> {
            segments[segI], segments[segI + 1] - segments[segI]
        };
    }
 
    KOKKOS_INLINE_FUNCTION std::span<ValueType> span(std::size_t segI) const
    {
        auto [start, length] = range(segI);
        return values.subspan(start, length);
    }
 
    KOKKOS_INLINE_FUNCTION
    IndexType operator[](std::size_t i) const { return segments[i]; }
};
 
template<typename ValueType, typename IndexType>
class SegmentedField
{
public:
 
 
    SegmentedField(const Executor& exec, size_t size, size_t numSegments)
        : values_(exec, size), segments_(exec, numSegments + 1)
    {}
 
    /*
     * @brief Create a segmented field from intervals.
     * @param intervals The intervals to create the segmented field from.
     * @note The intervals are the lengths of each segment
     */
    SegmentedField(const Field<IndexType>& intervals)
        : values_(intervals.exec(), 0),
          segments_(intervals.exec(), intervals.size() + 1, IndexType(0))
    {
        IndexType valueSize = segmentsFromIntervals(intervals, segments_);
        values_ = Field<ValueType>(intervals.exec(), valueSize);
    }
 
 
    SegmentedField(const Field<ValueType>& values, const Field<IndexType>& segments)
        : values_(values), segments_(segments)
    {
        NF_ASSERT(values.exec() == segments.exec(), "Executors are not the same.");
    }
 
 
    const Executor& exec() const { return values_.exec(); }
 
    size_t size() const { return values_.size(); }
 
    size_t numSegments() const { return segments_.size() - 1; }
 
 
    [[nodiscard]] SegmentedFieldView<ValueType, IndexType> view() &
    {
        return SegmentedFieldView<ValueType, IndexType> {values_.span(), segments_.span()};
    }
 
    // ensures no return a span of a temporary object --> invalid memory access
    [[nodiscard]] SegmentedFieldView<ValueType, IndexType> view() && = delete;
 
    [[nodiscard]] std::pair<std::span<ValueType>, std::span<IndexType>> spans() &
    {
        return {values_.span(), segments_.span()};
    }
 
    // ensures not to return a span of a temporary object --> invalid memory access
    [[nodiscard]] std::pair<std::span<ValueType>, std::span<IndexType>> spans() && = delete;
 
    const Field<ValueType>& values() const { return values_; }
 
    const Field<IndexType>& segments() const { return segments_; }
 
private:
 
    Field<ValueType> values_;
    Field<IndexType> segments_; 
};
 
} // namespace NeoFOAM