matrix.hpp

Go to the documentation of this file.

// SPDX-FileCopyrightText: 2023 - 2026 NeoN authors
//
// SPDX-License-Identifier: MIT
 
#pragma once
 
#include "NeoN/core/vector/vector.hpp"
#include "NeoN/linearAlgebra/cooSparsityPattern.hpp"
#include "NeoN/linearAlgebra/csrSparsityPattern.hpp"
#include "NeoN/linearAlgebra/faceToMatrixAddress.hpp"
 
namespace NeoN::la
{
 
template<typename ValueType, typename SparsityViewType>
struct MatrixView
{
    MatrixView(const View<ValueType> valueView, SparsityViewType sparsityView)
        : values(valueView), sparsity(sparsityView) {};
 
    ~MatrixView() = default;
 
    KOKKOS_INLINE_FUNCTION
    ValueType& entry(const localIdx i, const localIdx j) const
    {
        return values[sparsity.entry(i, j)];
    }
 
    KOKKOS_INLINE_FUNCTION
    ValueType& entry(const localIdx offset) const { return values[offset]; }
 
    View<ValueType> values; 
    SparsityViewType sparsity;
};
 
template<typename ValueType, typename SparsityType>
class Matrix : public NeoN::SupportsCopyTo<Matrix<ValueType, SparsityType>>
{
 
    void validate()
    {
        NF_ASSERT(values_.exec() == sparsityPattern_->exec(), "Executors are not the same");
        // TODO this is not necessarily true for matrix types with padding like ELL
        NF_ASSERT(values_.size() == sparsityPattern_->nnz(), "Matrix values and columns mismatch");
    }
 
public:
 
    using MatrixValueType = ValueType;
    using MatrixSparsityType = SparsityType;
 
    Matrix(const Vector<ValueType>& values, std::shared_ptr<const SparsityType> sp)
        : values_(values), sparsityPattern_(sp)
    {
        validate();
    }
 
    Matrix(
        const Vector<ValueType>& values,
        const Vector<typename SparsityType::SparsityIndexType>& colIdxs,
        const Vector<typename SparsityType::SparsityIndexType>& rowOffs,
        Dimensions dimensions
    )
        : values_(values),
          sparsityPattern_(
              std::make_shared<const SparsityType>(Vector(colIdxs), Vector(rowOffs), dimensions)
          )
    {
        validate();
    }
 
    Matrix(
        const Vector<ValueType>& values,
        std::shared_ptr<const SparsityType> sparsity,
        std::shared_ptr<const FaceToMatrixAddress> faceToMatrixAddress
    )
        requires std::is_same_v<typename SparsityType::SparsityIndexType, localIdx>
        : values_(values), sparsityPattern_(sparsity), faceToMatrixAddress_(faceToMatrixAddress)
    {
        validate();
    }
 
    ~Matrix() = default;
 
    [[nodiscard]] const Executor& exec() const { return values_.exec(); }
 
    [[nodiscard]] localIdx nRows() const { return sparsityPattern_->rows(); }
 
    [[nodiscard]] localIdx nNonZeros() const { return sparsityPattern_->nnz(); }
 
    void reset() { values_ = ValueType {}; }
 
 
    [[nodiscard]] Vector<ValueType>& values() { return values_; }
 
    [[nodiscard]] const Vector<typename SparsityType::SparsityIndexType>& colIdxs() const
    {
        return sparsityPattern_->colIdxs();
    }
 
    [[nodiscard]] const Vector<typename SparsityType::SparsityIndexType>& rowOffs() const
    {
        return sparsityPattern_->rowOffs();
    }
 
    [[nodiscard]] const Vector<ValueType>& values() const { return values_; }
 
    [[nodiscard]] Matrix<ValueType, SparsityType> copyToExecutor(Executor dstExec) const override;
 
    [[nodiscard]] std::shared_ptr<const SparsityType> sparsity() const { return sparsityPattern_; }
 
    [[nodiscard]] std::shared_ptr<const FaceToMatrixAddress> faceToMatrixAddress() const
    {
        return faceToMatrixAddress_;
    }
 
    [[nodiscard]] MatrixView<ValueType, SparsityView<typename SparsityType::SparsityIndexType>>
    view()
    {
        return MatrixView<ValueType, SparsityView<typename SparsityType::SparsityIndexType>>(
            values_.view(), sparsityPattern_->view()
        );
    }
 
    [[nodiscard]] MatrixView<
        const ValueType,
        SparsityView<typename SparsityType::SparsityIndexType>>
    view() const
    {
        return MatrixView<const ValueType, SparsityView<typename SparsityType::SparsityIndexType>>(
            View<const ValueType>(values_.view()), sparsityPattern_->view()
        );
    }
 
    [[nodiscard]] Vector<ValueType> diag() const;
 
 
private:
 
    Vector<ValueType> values_; 
 
    std::shared_ptr<const SparsityType> sparsityPattern_;
 
    std::shared_ptr<const FaceToMatrixAddress> faceToMatrixAddress_;
};
 
 
template<typename ValueType, typename IndexType>
using CSRMatrix = Matrix<ValueType, la::CsrSparsityPattern<IndexType>>;
 
template<typename ValueType, typename IndexType>
using COOMatrix = Matrix<ValueType, la::CooSparsityPattern<IndexType>>;
 
template<typename ValueType, typename IndexType>
[[nodiscard]] Vector<ValueType> upper(const CSRMatrix<ValueType, IndexType>&);
 
[[nodiscard]] Vector<scalar>
scaledInverseDiag(const CSRMatrix<Vec3, localIdx>&, const Vector<scalar>&);
 
void scaledInverseDiag(
    const CSRMatrix<Vec3, localIdx>& mtx, const Vector<scalar>& a, Vector<scalar>& out
);
 
[[nodiscard]] Vector<scalar>
scaledInverseDiag(const CSRMatrix<Vec3, localIdx>&, const FaceToMatrixAddress& mi, const Vector<scalar>&);
 
void scaledInverseDiag(
    const CSRMatrix<Vec3, localIdx>& mtx,
    const FaceToMatrixAddress& mi,
    const Vector<scalar>& a,
    Vector<scalar>& out
);
 
[[nodiscard]] Vector<scalar>
scaledInverseDiag(const CSRMatrix<scalar, localIdx>&, const FaceToMatrixAddress& mi, const Vector<scalar>&);
 
void scaledInverseDiag(
    const CSRMatrix<scalar, localIdx>& mtx,
    const FaceToMatrixAddress& mi,
    const Vector<scalar>& a,
    Vector<scalar>& out
);
 
/* @brief given Matrix<Vec3> this function returns a component Matrix<scalar>*/
template<unsigned int I>
[[nodiscard]] auto getComponent(const CSRMatrix<Vec3, localIdx>& in)
{
    auto sparsity = in.sparsity();
    return CSRMatrix<scalar, localIdx>(getComponent<I>(in.values()), sparsity);
}
 
void negLUx(
    const CSRMatrix<Vec3, localIdx>& mtx,
    const Vector<Vec3>& a,
    const Vector<Vec3>& b,
    const Vector<scalar>& rAU,
    const Vector<scalar>& vol,
    Vector<Vec3>& out
);
 
void scaledInvDiagNegLUx(
    const CSRMatrix<Vec3, localIdx>& mtx,
    const Vector<Vec3>& a,
    const Vector<Vec3>& b,
    const Vector<scalar>& vol,
    Vector<scalar>& rAU,
    Vector<Vec3>& out
);
 
void scaledInvDiagNegLUx(
    const CSRMatrix<scalar, localIdx>& mtx,
    const Vector<Vec3>& a,
    const Vector<Vec3>& b,
    const Vector<scalar>& vol,
    Vector<scalar>& rAU,
    Vector<Vec3>& out
);
 
} // namespace NeoN