custom_index_array.h

Go to the documentation of this file.

/* 
* Copyright (C) 2020-2026 MEmilio
*
* Authors: Daniel Abele, Khoa Nguyen
*
* Contact: Martin J. Kuehn <Martin.Kuehn@DLR.de>
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
*     http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#ifndef MIO_UTILS_CUSTOM_INDEX_ARRAY_H
#define MIO_UTILS_CUSTOM_INDEX_ARRAY_H
 
#include "memilio/math/eigen_util.h"
#include "memilio/utils/index.h"
#include "memilio/utils/metaprogramming.h"
 
#include <concepts>
 
namespace
{
 
//calculate the product of tuple elements
// TODO std::apply or fold expression in C++17; current version required by CustomIndexArray::Slice
template <int I, template <class...> class Index, class... Ts>
size_t product(Index<Ts...> const& t)
{
    if constexpr (I < sizeof...(Ts)) {
        return static_cast<size_t>(mio::get<I>(t)) * product<I + 1, Index, Ts...>(t);
    }
    else {
        return 1;
    }
}
 
template <template <class...> class Index, class... Ts>
size_t product(Index<Ts...> const& t)
{
    return product<0, Index, Ts...>(t);
}
 
} // namespace
 
namespace mio
{
 
namespace details
{
 
// Internal implementation for flatten_index
template <size_t I, typename Index>
std::pair<size_t, size_t> flatten_index(Index const& indices, Index const& dimensions)
{
    assert(get<I>(indices) < get<I>(dimensions));
    if constexpr (I < (Index::size - 1)) {
        auto [val, prod] = flatten_index<I + 1>(indices, dimensions);
        return {val + (size_t)mio::get<I>(indices) * prod, prod * (size_t)mio::get<I>(dimensions)};
    }
    else {
        return {(size_t)mio::get<I>(indices), (size_t)mio::get<I>(dimensions)};
    }
}
 
template <typename T>
struct is_random_access_iterator
    : std::is_base_of<typename std::iterator_traits<T>::iterator_category, std::random_access_iterator_tag> {
};
 
} // namespace details
 
template <typename MultiIndex>
size_t flatten_index(MultiIndex const& indices, MultiIndex const& dimensions)
{
    return details::flatten_index<0>(indices, dimensions).first;
}
 
template <class Typ, class... Tags>
class CustomIndexArray
{
public:
    using Type              = Typ;
    using Index             = ::mio::Index<Tags...>;
    using InternalArrayType = Eigen::Array<Type, Eigen::Dynamic, 1>;
 
    explicit CustomIndexArray()
        : m_dimensions(Index::Zero())
        , m_numel(0)
        , m_y()
    {
    }
 
    template <class... Ts>
        requires(std::is_constructible_v<Type, Ts...>)
    CustomIndexArray(Index const& dims, Ts&&... args)
        : m_dimensions{dims}
        , m_numel(product(dims))
        , m_y(InternalArrayType::Constant(m_numel, 1, Type{std::forward<Ts>(args)...}))
    {
    }
 
    template <std::input_or_output_iterator Iter>
    CustomIndexArray(Index const& dims, Iter b, Iter e)
        : m_dimensions(dims)
        , m_numel(product(dims))
        , m_y(m_numel, 1)
    {
        std::copy(b, e, begin());
    }
 
    size_t constexpr numel() const
    {
        return m_numel;
    }
 
    template <typename Tag>
    mio::Index<Tag> size() const
    {
        return get<Tag>(m_dimensions);
    }
 
    Index size() const
    {
        return m_dimensions;
    }
 
    void resize(Index new_dims)
    {
        m_dimensions = new_dims;
        m_numel      = product(m_dimensions);
        m_y.resize(m_numel);
    }
 
    template <class Tag>
    void resize(mio::Index<Tag> new_dim)
    {
        std::get<mio::Index<Tag>>(m_dimensions.indices) = new_dim;
        m_numel                                         = product(m_dimensions);
        m_y.resize(m_numel);
    }
 
    auto const& array() const
    {
        return m_y;
    }
    auto& array()
    {
        return m_y;
    }
 
    Type& operator[](Index const& index)
    {
        return m_y[get_flat_index(index)];
    }
 
    Type const& operator[](Index const& index) const
    {
        return m_y[get_flat_index(index)];
    }
 
    CustomIndexArray& operator=(const Type& scalar)
    {
        m_y = scalar;
        return *this;
    }
 
    bool operator==(const CustomIndexArray& other) const
    {
        return this->m_dimensions == other.m_dimensions && (this->m_y.array() == other.m_y.array()).all();
    }
    bool operator!=(const CustomIndexArray& other) const
    {
        return !(*this == other);
    }
    size_t get_flat_index(Index const& index) const
    {
        return (Eigen::Index)flatten_index(index, m_dimensions);
    }
 
    void set_multiple(const std::vector<typename CustomIndexArray<Typ, Tags...>::Index>& indices, const Typ& value)
    {
        for (const auto& index : indices) {
            m_y[get_flat_index(index)] = value;
        }
    }
 
    template <class OtherType>
        requires std::convertible_to<Type, OtherType>
    CustomIndexArray<OtherType, Tags...> convert() const
    {
        CustomIndexArray<OtherType, Tags...> other;
        other.resize(m_dimensions);
        other.array() = m_y.template cast<OtherType>();
        return other;
    }
 
private:
    template <typename Tag, typename iter_type>
        requires details::is_random_access_iterator<iter_type>::value
    class Slice
    {
        using difference_type = typename iter_type::difference_type;
 
        template <typename T>
        class Iterator
        {
        public:
            using iterator_category = std::random_access_iterator_tag;
            using difference_type   = std::ptrdiff_t;
            using value_type        = T;
            using pointer           = value_type*;
            using reference         = value_type&;
 
            Iterator(iter_type begin_, size_t di_, size_t dr_, Seq<size_t> const& seq_, difference_type offset = 0)
                : data_begin(begin_)
                , di(di_)
                , dr(dr_)
                , seq(seq_)
                , inner_offset(offset)
            {
            }
 
            Iterator& operator=(size_t rhs)
            {
                inner_offset = rhs;
                return *this;
            }
            Iterator& operator=(const Iterator& rhs)
            {
                inner_offset = rhs.inner_offset;
                return *this;
            }
            Iterator& operator+=(const int& rhs)
            {
                inner_offset += rhs;
                return *this;
            }
            Iterator& operator-=(const int& rhs)
            {
                inner_offset -= rhs;
                return *this;
            }
 
            reference operator[](const difference_type& rhs)
            {
                return data_begin[outer_offset(inner_offset + rhs)];
            }
            value_type const& operator[](const difference_type& rhs) const
            {
                return data_begin[outer_offset(inner_offset + rhs)];
            }
            reference operator*()
            {
                return data_begin[outer_offset(inner_offset)];
            }
            value_type const& operator*() const
            {
                return *(data_begin[outer_offset(inner_offset)]);
            }
            pointer operator->()
            {
                return data_begin + outer_offset(inner_offset);
            }
 
            Iterator& operator++()
            {
                inner_offset++;
                return *this;
            }
            Iterator operator++(int)
            {
                Iterator tmp = *this;
                ++(*this);
                return tmp;
            }
            Iterator& operator--()
            {
                --inner_offset;
                return *this;
            }
            Iterator operator--(int)
            {
                Iterator tmp = *this;
                --(*this);
                return tmp;
            }
 
            Iterator operator+(const int& rhs) const
            {
                return Iterator(data_begin, di, dr, seq, inner_offset + rhs);
            }
            Iterator operator-(const int& rhs) const
            {
                return Iterator(data_begin, di, dr, seq, inner_offset - rhs);
            }
 
            friend bool operator==(const Iterator& a, const Iterator& b)
            {
                return a.inner_offset == b.inner_offset && a.data_begin == b.data_begin && a.di == b.di &&
                       a.dr == b.dr && a.seq.start == b.seq.start && a.seq.n == b.seq.n && a.seq.stride == b.seq.stride;
            }
            friend bool operator!=(const Iterator& a, const Iterator& b)
            {
                return !(a == b);
            }
            friend bool operator<(const Iterator& a, const Iterator& b)
            {
                return a.inner_offset < b.inner_offset;
            }
            friend bool operator<=(const Iterator& a, const Iterator& b)
            {
                return a.inner_offset <= b.inner_offset;
            }
            friend bool operator>(const Iterator& a, const Iterator& b)
            {
                return a.inner_offset > b.inner_offset;
            }
            friend bool operator>=(const Iterator& a, const Iterator& b)
            {
                return a.inner_offset >= b.inner_offset;
            }
 
        private:
            inline Slice::difference_type outer_offset(difference_type const& inner) const
            {
 
                // calculate the outer offset from the inner offset
 
                // first unravel the inner index into an index (i,j,k) for a 3-dim array with dims (dl, idx_sequence.n, dr)
                auto dv           = std::div(inner, seq.n * dr);
                difference_type i = dv.quot;
                dv                = std::div(dv.rem, dr);
                difference_type j = dv.quot * seq.stride + seq.start;
                difference_type k = dv.rem;
 
                // then flatten the index for a 3-dim array with dims (dl, di, dr)
                return i * di * dr + j * dr + k;
            }
 
            iter_type data_begin;
            size_t di, dr;
            Seq<size_t> seq;
            difference_type inner_offset;
        };
 
    public:
        using value_type     = Type;
        using reference      = Type&;
        using iterator       = Iterator<Type>;
        using const_iterator = Iterator<Type const>;
 
        Slice(Index const& dimensions, iter_type const& start_iter, Seq<size_t> idx_sequence_)
            : data_begin(start_iter)
            , idx_sequence(idx_sequence_)
            , m_dimensions(dimensions)
            , di(mio::get<Tag>(dimensions))
            , dr(product<index_of_type_v<Tag, Index>>(dimensions) / di)
            , dl(product(dimensions) / (di * dr))
        {
            assert((size_t)idx_sequence.start + idx_sequence.n <= di);
 
            mio::get<Tag>(m_dimensions) = mio::Index<Tag>(idx_sequence.n);
        }
 
        // returns the number of elements in a slice
        size_t numel() const
        {
            return dl * dr * (idx_sequence.n);
        }
 
        // returns an stl-compatible random access iterator into the slice
        iterator begin()
        {
            return iterator(data_begin, di, dr, idx_sequence, 0);
        }
 
        // returns an stl-compatible random access iterator into the slice
        const_iterator begin() const
        {
            return const_iterator(data_begin, di, dr, idx_sequence, 0);
        }
 
        // returns an stl-compatible end random access iterator into the slice
        iterator end()
        {
            return iterator(data_begin, di, dr, idx_sequence, numel());
        }
 
        // returns an stl-compatible end random access iterator into the slice
        const_iterator end() const
        {
            return const_iterator(data_begin, di, dr, idx_sequence, numel());
        }
 
        // copies the slice elements into a CustomIndexArray of appropriate dimension
        CustomIndexArray<Type, Tags...> as_array()
        {
            CustomIndexArray<Type, Tags...> array(m_dimensions);
            Eigen::Index idx = 0;
            for (auto& v : *this) {
                array.array()[idx++] = v;
            };
            return array;
        }
 
        // comparison operators
        friend bool operator==(const Slice& a, const Slice& b)
        {
            return a.data_begin == b.data_begin && a.idx_sequence.start == b.idx_sequence.start &&
                   a.idx_sequence.n == b.idx_sequence.n && a.idx_sequence.stride == b.idx_sequence.stride &&
                   a.m_dimensions == b.m_dimensions;
        }
 
        friend bool operator!=(const Slice& a, const Slice& b)
        {
            return !(a == b);
        }
 
        template <typename OtherTag>
        Slice<OtherTag, iterator> slice(Seq<size_t> idx_sequence_)
        {
            return Slice<OtherTag, iterator>(m_dimensions, begin(), idx_sequence_);
        }
 
        template <typename OtherTag>
        Slice<OtherTag, const_iterator> slice(Seq<size_t> idx_sequence_) const
        {
            return Slice<OtherTag, const_iterator>(m_dimensions, begin(), idx_sequence_);
        }
 
        Slice& operator=(const Type& scalar)
        {
            for (auto&& e : *this) {
                e = scalar;
            }
            return *this;
        }
 
    private:
        iter_type data_begin;
        Seq<size_t> idx_sequence;
        Index m_dimensions;
        size_t di, dr, dl;
    };
 
    // An array storying the size of each category
    Index m_dimensions;
 
    // number of elements stored
    size_t m_numel;
 
    // An array containing the elements
    InternalArrayType m_y{};
 
public:
    using value_type     = Type;
    using reference      = Type&;
    using iterator       = typename InternalArrayType::iterator;
    using const_iterator = typename InternalArrayType::const_iterator;
 
    iterator begin()
    {
        return array().begin();
    }
 
    const_iterator begin() const
    {
        return array().cbegin();
    }
 
    iterator end()
    {
        return array().end();
    }
 
    const_iterator end() const
    {
        return array().cend();
    }
 
    template <typename Tag>
    Slice<Tag, iterator> slice(Seq<size_t> idx_seq)
    {
        return Slice<Tag, iterator>(m_dimensions, begin(), idx_seq);
    }
    template <typename Tag>
    Slice<Tag, const_iterator> slice(Seq<size_t> idx_seq) const
    {
        return Slice<Tag, const_iterator>(m_dimensions, begin(), idx_seq);
    }
    template <typename Tag>
    Slice<Tag, iterator> slice(mio::Index<Tag> idx)
    {
        return slice<Tag>({(size_t)idx, 1});
    }
    template <typename Tag>
    Slice<Tag, const_iterator> slice(mio::Index<Tag> idx) const
    {
        return slice<Tag>({(size_t)idx, 1});
    }
    template <typename Tag>
        requires std::is_enum_v<Tag>
    Slice<Tag, iterator> slice(Tag idx)
    {
        return slice<Tag>(mio::Index<Tag>(idx));
    }
    template <typename Tag>
        requires std::is_enum_v<Tag>
    Slice<Tag, const_iterator> slice(Tag idx) const
    {
        return slice<Tag>(mio::Index<Tag>(idx));
    }
    template <class IOContext>
    void serialize(IOContext& io) const
    {
        auto obj = io.create_object("Array");
        obj.add_element("Dimensions", m_dimensions);
        obj.add_list("Elements", begin(), end());
    }
 
protected:
    template <class IOContext, class Derived>
    static IOResult<Derived> deserialize(IOContext& io, Tag<Derived>)
    {
        auto obj  = io.expect_object("Array");
        auto dims = obj.expect_element("Dimensions", Tag<Index>{});
        auto els  = obj.expect_list("Elements", Tag<Type>{});
        return apply(
            io,
            [](auto&& d, auto&& e) -> IOResult<Derived> {
                auto a = Derived(d);
                if (a.numel() != e.size())
                    return failure(StatusCode::OutOfRange, "Dimensions of Array don't match the number of elements.");
                std::copy(e.begin(), e.end(), mio::begin(a.m_y));
                return success(a);
            },
            dims, els);
    }
 
public:
    template <class IOContext>
    static IOResult<CustomIndexArray> deserialize(IOContext& io)
    {
        return deserialize(io, Tag<CustomIndexArray>{});
    }
};
 
} // namespace mio
 
#endif // MIO_UTILS_CUSTOM_INDEX_ARRAY_H