matrix_data.hpp

// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_BASE_MATRIX_DATA_HPP_
#define GKO_PUBLIC_CORE_BASE_MATRIX_DATA_HPP_
 
 
#include <algorithm>
#include <numeric>
#include <tuple>
#include <vector>
 
#include <ginkgo/core/base/dim.hpp>
#include <ginkgo/core/base/math.hpp>
#include <ginkgo/core/base/range.hpp>
#include <ginkgo/core/base/range_accessors.hpp>
#include <ginkgo/core/base/types.hpp>
#include <ginkgo/core/base/utils.hpp>
 
 
namespace gko {
 
 
namespace detail {
 
 
// internal structure used to get around explicit constructors in std::tuple
template <typename ValueType, typename IndexType>
struct input_triple {
    IndexType row;
    IndexType col;
    ValueType val;
};
 
 
template <typename ValueType, typename Distribution, typename Generator>
typename std::enable_if<!is_complex_s<ValueType>::value, ValueType>::type
get_rand_value(Distribution&& dist, Generator&& gen)
{
    return dist(gen);
}
 
 
template <typename ValueType, typename Distribution, typename Generator>
typename std::enable_if<is_complex_s<ValueType>::value, ValueType>::type
get_rand_value(Distribution&& dist, Generator&& gen)
{
    return ValueType(dist(gen), dist(gen));
}
 
 
}  // namespace detail
 
 
template <typename ValueType, typename IndexType>
struct matrix_data_entry {
    using value_type = ValueType;
    using index_type = IndexType;
    matrix_data_entry() = default;
 
    GKO_ATTRIBUTES matrix_data_entry(index_type r, index_type c, value_type v)
        : row(r), column(c), value(v)
    {}
 
    bool operator==(const matrix_data_entry& other) const
    {
        return std::tie(this->row, this->column, this->value) ==
               std::tie(other.row, other.column, other.value);
    };
    bool operator!=(const matrix_data_entry& other) const
    {
        return std::tie(this->row, this->column, this->value) !=
               std::tie(other.row, other.column, other.value);
    };
 
#define GKO_DEFINE_DEFAULT_COMPARE_OPERATOR(_op)            \
    bool operator _op(const matrix_data_entry& other) const \
    {                                                       \
        return std::tie(this->row, this->column)            \
            _op std::tie(other.row, other.column);          \
    }
 
    GKO_DEFINE_DEFAULT_COMPARE_OPERATOR(<);
    GKO_DEFINE_DEFAULT_COMPARE_OPERATOR(>);
    GKO_DEFINE_DEFAULT_COMPARE_OPERATOR(<=);
    GKO_DEFINE_DEFAULT_COMPARE_OPERATOR(>=);
 
#undef GKO_DEFINE_DEFAULT_COMPARE_OPERATOR
 
    friend std::ostream& operator<<(std::ostream& os,
                                    const matrix_data_entry& x)
    {
        os << '(' << x.row << ',' << x.column << ',' << x.value << ')';
        return os;
    }
 
    index_type row;
    index_type column;
    value_type value;
};
 
 
template <typename ValueType = default_precision, typename IndexType = int32>
struct matrix_data {
    using value_type = ValueType;
    using index_type = IndexType;
    using nonzero_type = matrix_data_entry<value_type, index_type>;
 
    matrix_data(dim<2> size_ = dim<2>{}, ValueType value = zero<ValueType>())
        : size{size_}
    {
        if (is_zero(value)) {
            return;
        }
        nonzeros.reserve(size[0] * size[1]);
        for (size_type row = 0; row < size[0]; ++row) {
            for (size_type col = 0; col < size[1]; ++col) {
                nonzeros.emplace_back(row, col, value);
            }
        }
    }
 
    template <typename RandomDistribution, typename RandomEngine>
    matrix_data(dim<2> size_, RandomDistribution&& dist, RandomEngine&& engine)
        : size{size_}
    {
        nonzeros.reserve(size[0] * size[1]);
        for (size_type row = 0; row < size[0]; ++row) {
            for (size_type col = 0; col < size[1]; ++col) {
                const auto value =
                    detail::get_rand_value<ValueType>(dist, engine);
                if (is_nonzero(value)) {
                    nonzeros.emplace_back(row, col, value);
                }
            }
        }
    }
 
    matrix_data(std::initializer_list<std::initializer_list<ValueType>> values)
        : size{values.size(), 0}
    {
        for (size_type row = 0; row < values.size(); ++row) {
            const auto row_data = begin(values)[row];
            size[1] = std::max(size[1], row_data.size());
            for (size_type col = 0; col < row_data.size(); ++col) {
                const auto& val = begin(row_data)[col];
                if (is_nonzero(val)) {
                    nonzeros.emplace_back(row, col, val);
                }
            }
        }
    }
 
    matrix_data(
        dim<2> size_,
        std::initializer_list<detail::input_triple<ValueType, IndexType>>
            nonzeros_)
        : size{size_}, nonzeros()
    {
        nonzeros.reserve(nonzeros_.size());
        for (const auto& elem : nonzeros_) {
            nonzeros.emplace_back(elem.row, elem.col, elem.val);
        }
    }
 
    matrix_data(dim<2> size_, const matrix_data& block)
        : size{size_ * block.size}
    {
        nonzeros.reserve(size_[0] * size_[1] * block.nonzeros.size());
        for (size_type row = 0; row < size_[0]; ++row) {
            for (size_type col = 0; col < size_[1]; ++col) {
                for (const auto& elem : block.nonzeros) {
                    nonzeros.emplace_back(row * block.size[0] + elem.row,
                                          col * block.size[1] + elem.column,
                                          elem.value);
                }
            }
        }
        this->sort_row_major();
    }
 
    template <typename Accessor>
    matrix_data(const range<Accessor>& data)
        : size{data.length(0), data.length(1)}
    {
        for (gko::size_type row = 0; row < size[0]; ++row) {
            for (gko::size_type col = 0; col < size[1]; ++col) {
                if (is_nonzero(data(row, col))) {
                    nonzeros.emplace_back(row, col, data(row, col));
                }
            }
        }
    }
 
    static matrix_data diag(dim<2> size_, ValueType value)
    {
        matrix_data res(size_);
        if (is_nonzero(value)) {
            const auto num_nnz = std::min(size_[0], size_[1]);
            res.nonzeros.reserve(num_nnz);
            for (size_type i = 0; i < num_nnz; ++i) {
                res.nonzeros.emplace_back(i, i, value);
            }
        }
        return res;
    }
 
    static matrix_data diag(dim<2> size_,
                            std::initializer_list<ValueType> nonzeros_)
    {
        matrix_data res(size_);
        res.nonzeros.reserve(nonzeros_.size());
        int pos = 0;
        for (auto value : nonzeros_) {
            res.nonzeros.emplace_back(pos, pos, value);
            ++pos;
        }
        return res;
    }
 
    static matrix_data diag(dim<2> size_, const matrix_data& block)
    {
        matrix_data res(size_ * block.size);
        const auto num_blocks = std::min(size_[0], size_[1]);
        res.nonzeros.reserve(num_blocks * block.nonzeros.size());
        for (size_type b = 0; b < num_blocks; ++b) {
            for (const auto& elem : block.nonzeros) {
                res.nonzeros.emplace_back(b * block.size[0] + elem.row,
                                          b * block.size[1] + elem.column,
                                          elem.value);
            }
        }
        return res;
    }
 
    template <typename ForwardIterator>
    static matrix_data diag(ForwardIterator begin, ForwardIterator end)
    {
        matrix_data res(std::accumulate(
            begin, end, dim<2>{}, [](dim<2> s, const matrix_data& d) {
                return dim<2>{s[0] + d.size[0], s[1] + d.size[1]};
            }));
 
        size_type row_offset{};
        size_type col_offset{};
        for (auto it = begin; it != end; ++it) {
            for (const auto& elem : it->nonzeros) {
                res.nonzeros.emplace_back(row_offset + elem.row,
                                          col_offset + elem.column, elem.value);
            }
            row_offset += it->size[0];
            col_offset += it->size[1];
        }
 
        return res;
    }
 
    static matrix_data diag(std::initializer_list<matrix_data> blocks)
    {
        return diag(begin(blocks), end(blocks));
    }
 
    template <typename RandomDistribution, typename RandomEngine>
    static matrix_data cond(size_type size,
                            remove_complex<ValueType> condition_number,
                            RandomDistribution&& dist, RandomEngine&& engine,
                            size_type num_reflectors)
    {
        using range = range<accessor::row_major<ValueType, 2>>;
        std::vector<ValueType> mtx_data(size * size, zero<ValueType>());
        std::vector<ValueType> ref_data(size);
        std::vector<ValueType> work(size);
        range matrix(mtx_data.data(), size, size, size);
        range reflector(ref_data.data(), size, 1u, 1u);
 
        initialize_diag_with_cond(condition_number, matrix);
        for (size_type i = 0; i < num_reflectors; ++i) {
            generate_random_reflector(dist, engine, reflector);
            reflect_domain(reflector, matrix, work.data());
            generate_random_reflector(dist, engine, reflector);
            reflect_range(reflector, matrix, work.data());
        }
        return matrix;
    }
 
    template <typename RandomDistribution, typename RandomEngine>
    static matrix_data cond(size_type size,
                            remove_complex<ValueType> condition_number,
                            RandomDistribution&& dist, RandomEngine&& engine)
    {
        return cond(size, condition_number,
                    std::forward<RandomDistribution>(dist),
                    std::forward<RandomEngine>(engine), size - 1);
    }
 
    dim<2> size;
 
    std::vector<nonzero_type> nonzeros;
 
    void sort_row_major()
    {
        std::sort(
            begin(nonzeros), end(nonzeros), [](nonzero_type x, nonzero_type y) {
                return std::tie(x.row, x.column) < std::tie(y.row, y.column);
            });
    }
 
    GKO_DEPRECATED("Use sort_row_major() instead") void ensure_row_major_order()
    {
        this->sort_row_major();
    }
 
    void remove_zeros()
    {
        nonzeros.erase(
            std::remove_if(begin(nonzeros), end(nonzeros),
                           [](nonzero_type nz) { return is_zero(nz.value); }),
            end(nonzeros));
    }
 
    void sum_duplicates()
    {
        this->sort_row_major();
        std::vector<nonzero_type> new_nonzeros;
        if (!nonzeros.empty()) {
            new_nonzeros.emplace_back(nonzeros.front().row,
                                      nonzeros.front().column,
                                      zero<ValueType>());
            for (auto entry : nonzeros) {
                if (entry.row != new_nonzeros.back().row ||
                    entry.column != new_nonzeros.back().column) {
                    new_nonzeros.emplace_back(entry.row, entry.column,
                                              zero<ValueType>());
                }
                new_nonzeros.back().value += entry.value;
            }
            nonzeros = std::move(new_nonzeros);
        }
    }
 
private:
    template <typename Accessor>
    static void initialize_diag_with_cond(
        remove_complex<ValueType> condition_number,
        const range<Accessor>& matrix)
    {
        using sigma_type = remove_complex<ValueType>;
        const auto size = matrix.length(0);
        const auto min_sigma = one(condition_number) / sqrt(condition_number);
        const auto max_sigma = sqrt(condition_number);
 
        matrix = zero(matrix);
        for (gko::size_type i = 0; i < size; ++i) {
            matrix(i, i) = max_sigma * static_cast<sigma_type>(size - i - 1) /
                               static_cast<sigma_type>(size - 1) +
                           min_sigma * static_cast<sigma_type>(i) /
                               static_cast<sigma_type>(size - 1);
        }
    }
 
    template <typename RandomDistribution, typename RandomEngine,
              typename Accessor>
    static void generate_random_reflector(RandomDistribution&& dist,
                                          RandomEngine&& engine,
                                          const range<Accessor>& reflector)
    {
        for (gko::size_type i = 0; i < reflector.length(0); ++i) {
            reflector(i, 0) = detail::get_rand_value<ValueType>(dist, engine);
        }
    }
 
    template <typename Accessor>
    static void reflect_domain(const range<Accessor>& reflector,
                               const range<Accessor>& matrix,
                               ValueType* work_data)
    {
        const auto two = one<ValueType>() + one<ValueType>();
        range<accessor::row_major<ValueType, 2>> work(work_data,
                                                      matrix.length(0), 1u, 1u);
        work = mmul(matrix, reflector);
        const auto ct_reflector = conj(transpose(reflector));
        const auto scale = two / mmul(ct_reflector, reflector)(0, 0);
        matrix = matrix - scale * mmul(work, ct_reflector);
    }
 
    template <typename Accessor>
    static void reflect_range(const range<Accessor>& reflector,
                              const range<Accessor>& matrix,
                              ValueType* work_data)
    {
        const auto two = one<ValueType>() + one<ValueType>();
        range<accessor::row_major<ValueType, 2>> work(
            work_data, 1u, matrix.length(0), matrix.length(0));
        const auto ct_reflector = conj(transpose(reflector));
        work = mmul(ct_reflector, matrix);
        const auto scale = two / mmul(ct_reflector, reflector)(0, 0);
        matrix = matrix - scale * mmul(reflector, work);
    }
};
 
 
}  // namespace gko
 
 
#endif  // GKO_PUBLIC_CORE_BASE_MATRIX_DATA_HPP_