matrix_op.hpp

Go to the documentation of this file.

// -------------------------------------------------------//
//
// SHAMROCK code for hydrodynamics
// Copyright (c) 2021-2026 Timothée David--Cléris <tim.shamrock@proton.me>
// SPDX-License-Identifier: CeCILL Free Software License Agreement v2.1
// Shamrock is licensed under the CeCILL 2.1 License, see LICENSE for more information
//
// -------------------------------------------------------//
 
#pragma once

 
#include "shambase/assert.hpp"
#include "shambase/type_traits.hpp"
#include "shambackends/math.hpp"
#include "shambackends/sycl.hpp"
#include <experimental/mdspan>
#include <array>
 
namespace shammath {

    template<class T, class Extents, class Layout, class Accessor, class Func>
    inline void mat_set_vals(const std::mdspan<T, Extents, Layout, Accessor> &input, Func &&func) {
 
        shambase::check_functor_signature<T, int, int>(func);
 
        for (int i = 0; i < input.extent(0); i++) {
            for (int j = 0; j < input.extent(1); j++) {
                input(i, j) = func(i, j);
            }
        }
    }

    template<class T, class Extents, class Layout, class Accessor, class Func>
    inline void vec_set_vals(const std::mdspan<T, Extents, Layout, Accessor> &input, Func &&func) {
 
        shambase::check_functor_signature<T, int>(func);
 
        for (int i = 0; i < input.extent(0); i++) {
            input(i) = func(i);
        }
    }

    template<class T, class Extents, class Layout, class Accessor, class Func>
    inline void mat_update_vals(
        const std::mdspan<T, Extents, Layout, Accessor> &input, Func &&func) {
 
        shambase::check_functor_signature<void, T &, int, int>(func);
 
        for (int i = 0; i < input.extent(0); i++) {
            for (int j = 0; j < input.extent(1); j++) {
                func(input(i, j), i, j);
            }
        }
    }

    template<class T, class Extents, class Layout, class Accessor>
    inline void mat_set_identity(const std::mdspan<T, Extents, Layout, Accessor> &input1) {
 
        SHAM_ASSERT(input1.extent(0) == input1.extent(1));
 
        mat_set_vals(input1, [](auto i, auto j) -> T {
            return (i == j) ? 1 : 0;
        });
    }

    template<class T, class Extents, class Layout, class Accessor>
    inline void mat_mul_scalar(
        const std::mdspan<T, Extents, Layout, Accessor> &input, const T &scalar) {
        mat_update_vals(input, [&](T &v, auto i, auto j) {
            v *= scalar;
        });
    }

    template<class T, class Extents, class Layout, class Accessor>
    inline void mat_copy(
        const std::mdspan<T, Extents, Layout, Accessor> &input,
        const std::mdspan<T, Extents, Layout, Accessor> &output) {
 
        SHAM_ASSERT(input.extent(0) == output.extent(0));
        SHAM_ASSERT(input.extent(1) == output.extent(1));
 
        for (int i = 0; i < input.extent(0); i++) {
            for (int j = 0; j < input.extent(1); j++) {
                output(i, j) = input(i, j);
            }
        }
    }

    template<
        class T,
        class Extents1,
        class Extents2,
        class Extents3,
        class Layout1,
        class Layout2,
        class Layout3,
        class Accessor1,
        class Accessor2,
        class Accessor3>
    inline void mat_plus(
        const std::mdspan<T, Extents1, Layout1, Accessor1> &input1,
        const std::mdspan<T, Extents2, Layout2, Accessor2> &input2,
        const std::mdspan<T, Extents3, Layout3, Accessor3> &output) {
 
        SHAM_ASSERT(input1.extent(0) == output.extent(0));
        SHAM_ASSERT(input1.extent(1) == output.extent(1));
        SHAM_ASSERT(input1.extent(0) == input2.extent(0));
        SHAM_ASSERT(input1.extent(1) == input2.extent(1));
 
        for (int i = 0; i < input1.extent(0); i++) {
            for (int j = 0; j < input1.extent(1); j++) {
                output(i, j) = input1(i, j) + input2(i, j);
            }
        }
    }

    template<
        class T,
        class Extents1,
        class Extents2,
        class Layout1,
        class Layout2,
        class Accessor1,
        class Accessor2>
    inline void mat_plus_equal(
        const std::mdspan<T, Extents1, Layout1, Accessor1> &inout,
        const std::mdspan<T, Extents2, Layout2, Accessor2> &matb) {
 
        SHAM_ASSERT(inout.extent(0) == inout.extent(0));
        SHAM_ASSERT(inout.extent(1) == inout.extent(1));
 
        for (int i = 0; i < inout.extent(0); i++) {
            for (int j = 0; j < inout.extent(1); j++) {
                inout(i, j) += matb(i, j);
            }
        }
    }

    template<
        class T,
        class Extents1,
        class Extents2,
        class Extents3,
        class Layout1,
        class Layout2,
        class Layout3,
        class Accessor1,
        class Accessor2,
        class Accessor3>
    inline void mat_sub(
        const std::mdspan<T, Extents1, Layout1, Accessor1> &input1,
        const std::mdspan<T, Extents2, Layout2, Accessor2> &input2,
        const std::mdspan<T, Extents3, Layout3, Accessor3> &output) {
 
        SHAM_ASSERT(input1.extent(0) == output.extent(0));
        SHAM_ASSERT(input1.extent(1) == output.extent(1));
        SHAM_ASSERT(input1.extent(0) == input2.extent(0));
        SHAM_ASSERT(input1.extent(1) == input2.extent(1));
 
        for (int i = 0; i < input1.extent(0); i++) {
            for (int j = 0; j < input1.extent(1); j++) {
                output(i, j) = input1(i, j) - input2(i, j);
            }
        }
    }

    template<
        class T,
        class Extents1,
        class Extents2,
        class Layout1,
        class Layout2,
        class Accessor1,
        class Accessor2>
    inline void mat_sub_equal(
        const std::mdspan<T, Extents1, Layout1, Accessor1> &inout,
        const std::mdspan<T, Extents2, Layout2, Accessor2> &matb) {
 
        SHAM_ASSERT(inout.extent(0) == inout.extent(0));
        SHAM_ASSERT(inout.extent(1) == inout.extent(1));
 
        for (int i = 0; i < inout.extent(0); i++) {
            for (int j = 0; j < inout.extent(1); j++) {
                inout(i, j) -= matb(i, j);
            }
        }
    }

    template<
        class Ta,
        class Tb,
        class Extents1,
        class Extents2,
        class Extents3,
        class Layout1,
        class Layout2,
        class Layout3,
        class Accessor1,
        class Accessor2,
        class Accessor3>
    inline void mat_prod(
        const std::mdspan<Ta, Extents1, Layout1, Accessor1> &input1,
        const std::mdspan<Ta, Extents2, Layout2, Accessor2> &input2,
        const std::mdspan<Tb, Extents3, Layout3, Accessor3> &output) {
 
        SHAM_ASSERT(input1.extent(0) == output.extent(0));
        SHAM_ASSERT(input1.extent(1) == input2.extent(0));
        SHAM_ASSERT(input2.extent(1) == output.extent(1));
 
        // output_ij = mat1_ik mat2_jk
        for (int i = 0; i < input1.extent(0); i++) {
            for (int j = 0; j < input2.extent(1); j++) {
                Tb sum = 0;
                for (int k = 0; k < input1.extent(1); k++) {
                    sum += input1(i, k) * input2(k, j);
                }
                output(i, j) = sum;
            }
        }
    }

    template<class T, class SizeType, class Layout, class Accessor>
    inline void mat_inv_33(
        const std::mdspan<T, std::extents<SizeType, 3, 3>, Layout, Accessor> &input,
        const std::mdspan<T, std::extents<SizeType, 3, 3>, Layout, Accessor> &output) {
 
        T &a00 = input(0, 0);
        T &a10 = input(1, 0);
        T &a20 = input(2, 0);
 
        T &a01 = input(0, 1);
        T &a11 = input(1, 1);
        T &a21 = input(2, 1);
 
        T &a02 = input(0, 2);
        T &a12 = input(1, 2);
        T &a22 = input(2, 2);
 
        T det
            = (-a02 * a11 * a20 + a01 * a12 * a20 + a02 * a10 * a21 - a00 * a12 * a21
               - a01 * a10 * a22 + a00 * a11 * a22);
 
        output(0, 0) = (-a12 * a21 + a11 * a22) / det;
        output(1, 0) = (a12 * a20 - a10 * a22) / det;
        output(2, 0) = (-a11 * a20 + a10 * a21) / det;
 
        output(0, 1) = (a02 * a21 - a01 * a22) / det;
        output(1, 1) = (-a02 * a20 + a00 * a22) / det;
        output(2, 1) = (a01 * a20 - a00 * a21) / det;
 
        output(0, 2) = (-a02 * a11 + a01 * a12) / det;
        output(1, 2) = (a02 * a10 - a00 * a12) / det;
        output(2, 2) = (-a01 * a10 + a00 * a11) / det;
    }

    template<class T, class U, class Extents, class Layout, class Accessor>
    inline void mat_L1_norm(const std::mdspan<T, Extents, Layout, Accessor> &input, U &res) {
        res = 0;
        for (auto i = 0; i < input.extent(0); i++) {
            T sum = 0;
            for (auto j = 0; j < input.extent(1); j++) {
                sum += sham::abs(input(i, j));
            }
            res = sham::max(res, sum);
        }
    }

    template<class T, class Extents, class Layout, class Accessor>
    inline void mat_set_nul(const std::mdspan<T, Extents, Layout, Accessor> &input) {
        mat_set_vals(input, [](auto i, auto j) -> T {
            return 0;
        });
    }

    template<class T, class Extents, class Layout, class Accessor>
    inline void vec_set_nul(const std::mdspan<T, Extents, Layout, Accessor> &input) {
        for (auto i = 0; i < input.extent(0); i++) {
            input(i) = 0;
        }
    }

    template<class T, class Extents, class Layout, class Accessor>
    inline void vec_copy(
        const std::mdspan<T, Extents, Layout, Accessor> &input,
        const std::mdspan<T, Extents, Layout, Accessor> &output) {
        SHAM_ASSERT(input.extent(0) == output.extent(0));
 
        for (int i = 0; i < input.extent(0); i++) {
            output(i) = input(i);
        }
    }

    template<
        class T,
        class U,
        class Extents1,
        class Extents2,
        class Layout1,
        class Layout2,
        class Accessor1,
        class Accessor2>
    inline void vec_axpy_beta(
        const U alpha,
        const std::mdspan<T, Extents1, Layout1, Accessor1> &input,
        const U beta,
        const std::mdspan<T, Extents2, Layout2, Accessor2> &output) {
 
        SHAM_ASSERT(input.extent(0) == output.extent(0));
 
        for (int i = 0; i < input.extent(0); i++) {
            output(i) = alpha * input(i) + beta * output(i);
        }
    }

    template<
        class T,
        class U,
        class Extents1,
        class Extents2,
        class Layout1,
        class Layout2,
        class Accessor1,
        class Accessor2>
    inline void vec_axpy(
        const U alpha,
        const std::mdspan<T, Extents1, Layout1, Accessor1> &input,
        const std::mdspan<T, Extents2, Layout2, Accessor2> &output) {
 
        vec_axpy_beta(alpha, input, U{1}, output);
    }

    template<
        class T,
        class U,
        class Extents1,
        class Extents2,
        class Layout1,
        class Layout2,
        class Accessor1,
        class Accessor2>
    inline void mat_axpy_beta(
        const U alpha,
        const std::mdspan<T, Extents1, Layout1, Accessor1> &input,
        const U beta,
        const std::mdspan<T, Extents2, Layout2, Accessor2> &output) {
 
        SHAM_ASSERT(input.extent(0) == output.extent(0));
        SHAM_ASSERT(input.extent(1) == output.extent(1));
 
        for (int i = 0; i < input.extent(0); i++) {
            for (int j = 0; j < input.extent(1); j++) {
                output(i, j) = alpha * input(i, j) + beta * output(i, j);
            }
        }
    }

    template<
        class T,
        class U,
        class Extents1,
        class Extents2,
        class Layout1,
        class Layout2,
        class Accessor1,
        class Accessor2>
    inline void mat_axpy(
        const U alpha,
        const std::mdspan<T, Extents1, Layout1, Accessor1> &input,
        const std::mdspan<T, Extents2, Layout2, Accessor2> &output) {
 
        mat_axpy_beta(alpha, input, U{1}, output);
    }

    template<
        class T,
        class U,
        class Extents1,
        class Extents2,
        class Extents3,
        class Layout1,
        class Layout2,
        class Layout3,
        class Accessor1,
        class Accessor2,
        class Accessor3>
    inline void mat_gemm(
        const U alpha,
        const std::mdspan<T, Extents1, Layout1, Accessor1> &input1,
        const std::mdspan<T, Extents2, Layout2, Accessor2> &input2,
        const U beta,
        const std::mdspan<T, Extents3, Layout3, Accessor3> &output) {
 
        SHAM_ASSERT(input1.extent(0) == output.extent(0));
        SHAM_ASSERT(input1.extent(1) == input2.extent(0));
        SHAM_ASSERT(input2.extent(1) == output.extent(1));
 
        for (int i = 0; i < input1.extent(0); i++) {
            for (int j = 0; j < input2.extent(1); j++) {
                T sum = 0;
                for (int k = 0; k < input1.extent(1); k++) {
                    sum += input1(i, k) * input2(k, j);
                }
                output(i, j) = alpha * sum + beta * output(i, j);
            }
        }
    }

 
    template<class T, class U, class Extents1, class Layout1, class Accessor1>
    inline void mat_plus_equal_scalar_id(
        const std::mdspan<T, Extents1, Layout1, Accessor1> &inout, const U beta) {
        for (int i = 0; i < inout.extent(0); i++) {
            inout(i, i) = inout(i, i) + beta;
        }
    }

    template<
        class T,
        class U,
        class Extents1,
        class Extents2,
        class Extents3,
        class Layout1,
        class Layout2,
        class Layout3,
        class Accessor1,
        class Accessor2,
        class Accessor3>
    inline void mat_gemv(
        const U alpha,
        const std::mdspan<T, Extents1, Layout1, Accessor1> &M,
        const std::mdspan<T, Extents2, Layout2, Accessor2> &x,
        const U beta,
        const std::mdspan<T, Extents3, Layout3, Accessor3> &y) {
 
        SHAM_ASSERT(M.extent(1) == x.extent(0));
        SHAM_ASSERT(M.extent(0) == y.extent(0));
 
        for (int i = 0; i < M.extent(0); i++) {
            T sum = 0;
            for (int j = 0; j < M.extent(1); j++) {
                sum += M(i, j) * x(j);
            }
            y(i) = alpha * sum + beta * y(i);
        }
    }

    template<
        class T,
        class Extents1,
        class Extents2,
        class Layout1,
        class Layout2,
        class Accessor1,
        class Accessor2>
    inline void mat_transpose(
        const std::mdspan<T, Extents1, Layout1, Accessor1> &input,
        const std::mdspan<T, Extents2, Layout2, Accessor2> &output) {
 
        SHAM_ASSERT(input.extent(0) == output.extent(1));
        SHAM_ASSERT(input.extent(1) == output.extent(0));
 
        for (int i = 0; i < output.extent(0); i++) {
            for (int j = 0; j < output.extent(1); j++) {
                output(i, j) = input(j, i);
            }
        }
    }

    template<
        class T,
        class Extents1,
        class Extents2,
        class Layout1,
        class Layout2,
        class Accessor1,
        class Accessor2>
    inline void Cholesky_decomp(
        const std::mdspan<T, Extents1, Layout1, Accessor1> &M,
        const std::mdspan<T, Extents2, Layout2, Accessor2> &L) {
 
        SHAM_ASSERT(M.extent(1) == M.extent(0));
        SHAM_ASSERT(M.extent(0) == L.extent(0));
        SHAM_ASSERT(L.extent(1) == L.extent(0));
 
        for (int i = 0; i < M.extent(0); i++) {
            T sum_ik = 0.0;
            for (int k = 0; k < i; k++) {
                sum_ik += L(i, k) * L(i, k);
            }
            L(i, i) = sycl::sqrt(M(i, i) - sum_ik);
            for (int j = i + 1; j < M.extent(1); j++) {
                T sum_ikjk = 0.0;
                for (int k = 0; k < i; k++) {
                    sum_ikjk += L(i, k) * L(j, k);
                }
                L(j, i) = (M(i, j) - sum_ikjk) / L(i, i);
                L(i, j) = 0.0;
            }
        }
    }

    template<
        class T,
        class Extents1,
        class Extents2,
        class Extents3,
        class Layout1,
        class Layout2,
        class Layout3,
        class Accessor1,
        class Accessor2,
        class Accessor3>
    inline void Cholesky_solve(
        const std::mdspan<T, Extents1, Layout1, Accessor1> &M,
        const std::mdspan<T, Extents2, Layout2, Accessor2> &y,
        const std::mdspan<T, Extents3, Layout3, Accessor3> &x) {
 
        SHAM_ASSERT(M.extent(1) == M.extent(0));
        SHAM_ASSERT(M.extent(1) == x.extent(0));
        SHAM_ASSERT(M.extent(0) == y.extent(0));
 
        std::vector<T> a(M.extent(0));
        std::vector<T> L_storage(M.extent(0) * M.extent(1));
 
        std::mdspan<T, Extents1> L{L_storage.data(), M.extent(0), M.extent(1)};
        Cholesky_decomp(M, L);
 
        for (int i = 0; i < M.extent(0); i++) {
            T sum = 0.0;
            for (int k = 0; k < i; k++) {
                sum += L(i, k) * a[k];
            }
            a[i] = (y(i) - sum) / L(i, i);
        }
        for (int i = M.extent(0) - 1; i >= 0; i--) {
            T sum = 0.0;
            for (int k = i + 1; k < M.extent(0); k++) {
                sum += L(k, i) * x(k);
            }
            x(i) = (a[i] - sum) / L(i, i);
        }
    }
 
} // namespace shammath