#include "xtensor/containers/xarray.hpp"

#include "xtensor/generators/xbuilder.hpp"

#include "xtensor/views/xstrided_view.hpp"

#include "xtensor/views/xview.hpp"

#include "doctest/doctest.h"

#include "xtensor-blas/xlinalg.hpp"

namespace xt

{

TEST_SUITE("xdot")

{

TEST_CASE("matrix_times_vector")

{

xarray<float> a = xt::ones<float>({1, 4});

xarray<float> b = {{1, 2, 3}, {4, 5, 6}, {7, 8, 9}, {1, 1, 1}};

xarray<float> e1 = {{13, 16, 19}};

auto r1 = linalg::dot(a, b);

CHECK_EQ(e1, r1);

xarray<float> c = xt::ones<float>({3, 1});

auto r2 = linalg::dot(b, c);

xarray<float> e2 = {{6, 15, 24, 3}};

e2.reshape({4, 1});

CHECK_EQ(e2, r2);

CHECK_THROWS_AS(linalg::dot(b, a), std::runtime_error);

CHECK_THROWS_AS(linalg::dot(c, b), std::runtime_error);

}

TEST_CASE("matrix_transpose_times_column")

{

xarray<double, layout_type::row_major> a = xt::ones<double>({2, 4});

xarray<double, layout_type::row_major> b = xt::ones<double>({2, 1});

auto r1 = linalg::dot(xt::transpose(a), b);

CHECK(all(equal(r1, 2.0)));

}

TEST_CASE("matrix_transpose_times_column_cm")

{

xarray<double, layout_type::column_major> a = xt::ones<double>({2, 4});

xarray<double, layout_type::column_major> b = xt::ones<double>({2, 1});

auto r1 = linalg::dot(xt::transpose(a), b);

CHECK(all(equal(r1, 2.0)));

}

TEST_CASE("square_matrix_times_vector")

{

xarray<float> a = {{1, 1, 1}};

xarray<float> b = {{1, 2, 3}, {4, 5, 6}, {7, 8, 9}};

auto r1 = linalg::dot(a, b);

xarray<float> e1 = {{12, 15, 18}};

CHECK_EQ(r1, e1);

auto r2 = linalg::dot(b, xt::transpose(a));

xarray<float> e2 = xarray<float>::from_shape({3, 1});

e2(0, 0) = 6.f;

e2(1, 0) = 15.f;

e2(2, 0) = 24.f;

CHECK_EQ(r2, e2);

CHECK_THROWS_AS(linalg::dot(b, a), std::runtime_error);

}

TEST_CASE("vector_times_vector")

{

xarray<float> a = xt::ones<float>({1, 3});

xarray<float> b = xt::ones<float>({3, 1});

auto r1 = linalg::dot(a, b);

xarray<float> e1 = xarray<float>::from_shape({1, 1});

e1(0, 0) = 3;

CHECK_EQ(e1, r1);

auto r2 = linalg::dot(b, a);

xarray<float> e2 = xt::ones<float>({3, 3});

CHECK_EQ(e2, r2);

auto r3 = linalg::dot(b, e1);

CHECK_EQ(b * 3.f, r3);

}

TEST_CASE("matrix_times_1d")

{

xarray<float> a = xt::ones<float>({5, 3});

xarray<float> b = xt::ones<float>({5});

xarray<float> c = xt::ones<float>({3});

auto r1 = linalg::dot(xt::transpose(a), b);

CHECK(all(equal(r1, 5.0)));

auto r2 = linalg::dot(c, xt::transpose(a));

CHECK(all(equal(r2, 3.0)));

auto r3 = linalg::dot(a, c);

CHECK(all(equal(r3, 3.0)));

auto r4 = linalg::dot(c, xt::ones<float>({3, 5}));

CHECK(all(equal(r3, 3.0)));

}

TEST_CASE("A_times_A_T")

{

xarray<float> a = xt::ones<float>({5, 3});

auto r1 = linalg::dot(a, xt::transpose(a));

CHECK(all(equal(r1, 3.0)));

auto r2 = linalg::dot(xt::transpose(a), a);

CHECK(all(equal(r2, 5.0)));

}

TEST_CASE("matrix_times_vector_cm")

{

xarray<float, layout_type::column_major> a = xt::ones<float>({1, 4});

xarray<float, layout_type::column_major> b = {{1, 2, 3}, {4, 5, 6}, {7, 8, 9}, {1, 1, 1}};

xarray<float, layout_type::column_major> e1 = {{13, 16, 19}};

auto r1 = linalg::dot(a, b);

CHECK_EQ(e1, r1);

xarray<float, layout_type::column_major> c = xt::ones<float>({3, 1});

auto r2 = linalg::dot(b, c);

xarray<float, layout_type::column_major> e2 = {{6, 15, 24, 3}};

e2.reshape({4, 1});

CHECK_EQ(e2, r2);

CHECK_THROWS_AS(linalg::dot(b, a), std::runtime_error);

CHECK_THROWS_AS(linalg::dot(c, b), std::runtime_error);

}

TEST_CASE("square_matrix_times_vector_cm")

{

xarray<float, layout_type::column_major> a = {{1, 1, 1}};

xarray<float, layout_type::column_major> b = {{1, 2, 3}, {4, 5, 6}, {7, 8, 9}};

auto r1 = linalg::dot(a, b);

xarray<float, layout_type::column_major> e1 = {{12, 15, 18}};

CHECK_EQ(r1, e1);

auto r2 = linalg::dot(b, xt::transpose(a));

xarray<float, layout_type::column_major> e2 = xarray<float, layout_type::column_major>::from_shape(

{3, 1}

);

e2(0, 0) = 6.f;

e2(1, 0) = 15.f;

e2(2, 0) = 24.f;

CHECK_EQ(r2, e2);

CHECK_THROWS_AS(linalg::dot(b, a), std::runtime_error);

}

TEST_CASE("vector_times_vector_cm")

{

xarray<float, layout_type::column_major> a = xt::ones<float>({1, 3});

xarray<float, layout_type::column_major> b = xt::ones<float>({3, 1});

auto r1 = linalg::dot(a, b);

xarray<float, layout_type::column_major> e1 = xarray<float, layout_type::column_major>::from_shape(

{1, 1}

);

e1(0, 0) = 3;

CHECK_EQ(e1, r1);

auto r2 = linalg::dot(b, a);

xarray<float, layout_type::column_major> e2 = xt::ones<float>({3, 3});

CHECK_EQ(e2, r2);

auto r3 = linalg::dot(b, e1);

CHECK_EQ(b * 3.f, r3);

}

TEST_CASE("on_view")

{

xt::xarray<int> a = xt::reshape_view(xt::arange<int>(10 * 10 * 10), {10, 10, 10});

xt::xarray<int> b = xt::reshape_view(xt::arange<int>(10 * 10 * 10), {10, 10, 10});

auto res = xt::linalg::dot(view(a, 0, 0), view(b, 0, 0));

auto res1 = xt::linalg::dot(view(a, 0, range(0, 3)), transpose(view(b, 0, range(0, 3))));

CHECK_EQ(res(0), 285.);

CHECK_EQ(res1(0, 0), 285.);

CHECK_EQ(res1(1, 2), 3635.);

CHECK_EQ(res1.dimension(), 2u);

CHECK_EQ(res1.shape()[0], 3u);

CHECK_EQ(res1.shape()[1], 3u);

auto res2 = xt::linalg::dot(strided_view(a, {0, 0}), strided_view(b, {0, 0}));

auto res3 = xt::linalg::dot(

strided_view(a, {0, range(0, 3)}),

transpose(strided_view(b, {0, range(0, 3)}))

);

CHECK_EQ(res2(0), 285.);

CHECK_EQ(res3(0, 0), 285.);

CHECK_EQ(res3(1, 2), 3635.);

CHECK_EQ(res3.dimension(), 2u);

CHECK_EQ(res3.shape()[0], 3u);

CHECK_EQ(res3.shape()[1], 3u);

}

}

} // namespace xt