matrix_assembly_data.hpp

// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_BASE_MATRIX_ASSEMBLY_DATA_HPP_
#define GKO_PUBLIC_CORE_BASE_MATRIX_ASSEMBLY_DATA_HPP_
 
 
#include <algorithm>
#include <iterator>
#include <numeric>
#include <tuple>
#include <unordered_map>
 
#include <ginkgo/core/base/dim.hpp>
#include <ginkgo/core/base/math.hpp>
#include <ginkgo/core/base/matrix_data.hpp>
#include <ginkgo/core/base/types.hpp>
#include <ginkgo/core/base/utils.hpp>
 
 
namespace gko {
namespace detail {
 
 
template <typename IndexType>
struct symbolic_nonzero_hash {
    symbolic_nonzero_hash() = default;
 
    explicit symbolic_nonzero_hash(size_type num_cols) noexcept
        : num_cols_{num_cols}
    {}
 
    std::size_t operator()(std::pair<IndexType, IndexType> nnz) const noexcept
    {
        return static_cast<std::size_t>(nnz.first) * num_cols_ + nnz.second;
    }
 
    size_type num_cols_;
};
 
 
}  // namespace detail
 

template <typename ValueType = default_precision, typename IndexType = int32>
class matrix_assembly_data {
    GKO_ASSERT_SUPPORTED_VALUE_TYPE;
 
public:
    using value_type = ValueType;
    using index_type = IndexType;
 
    explicit matrix_assembly_data(dim<2> size)
        : size_{size},
          nonzeros_(0, detail::symbolic_nonzero_hash<index_type>(size_[1]))
    {}

    void add_value(index_type row, index_type col, value_type val)
    {
        auto ind = std::make_pair(row, col);
        nonzeros_[ind] += val;
    }

    void set_value(index_type row, index_type col, value_type val)
    {
        auto ind = std::make_pair(row, col);
        nonzeros_[ind] = val;
    }

    value_type get_value(index_type row, index_type col)
    {
        const auto it = nonzeros_.find(std::make_pair(row, col));
        if (it == nonzeros_.end()) {
            return zero<value_type>();
        } else {
            return it->second;
        }
    }

    bool contains(index_type row, index_type col)
    {
        return nonzeros_.find(std::make_pair(row, col)) != nonzeros_.end();
    }

    dim<2> get_size() const noexcept { return size_; }

    size_type get_num_stored_elements() const noexcept
    {
        return nonzeros_.size();
    }

    matrix_data<ValueType, IndexType> get_ordered_data() const
    {
        using output_type = matrix_data<ValueType, IndexType>;
        using nonzero_type = typename output_type::nonzero_type;
        using entry_type =
            std::pair<std::pair<index_type, index_type>, value_type>;
        output_type data{size_};
        data.nonzeros.reserve(nonzeros_.size());
        std::transform(nonzeros_.begin(), nonzeros_.end(),
                       std::back_inserter(data.nonzeros), [](entry_type entry) {
                           return nonzero_type{entry.first.first,
                                               entry.first.second,
                                               entry.second};
                       });
        data.sort_row_major();
        return data;
    }
 
private:
    dim<2> size_;

    std::unordered_map<std::pair<index_type, index_type>, value_type,
                       detail::symbolic_nonzero_hash<index_type>>
        nonzeros_;
};
 
 
}  // namespace gko
 
 
#endif  // GKO_PUBLIC_CORE_BASE_MATRIX_ASSEMBLY_DATA_HPP_