diff --git a/include/openmc/tensor.h b/include/openmc/tensor.h index 9d428aaeb..9e443e72f 100644 --- a/include/openmc/tensor.h +++ b/include/openmc/tensor.h @@ -762,6 +762,24 @@ public: data_[i] += o.data_[i]; return *this; } + Tensor& operator-=(const Tensor& o) + { + for (size_t i = 0; i < data_.size(); ++i) + data_[i] -= o.data_[i]; + return *this; + } + Tensor& operator*=(const Tensor& o) + { + for (size_t i = 0; i < data_.size(); ++i) + data_[i] *= o.data_[i]; + return *this; + } + Tensor& operator/=(const Tensor& o) + { + for (size_t i = 0; i < data_.size(); ++i) + data_[i] /= o.data_[i]; + return *this; + } Tensor operator+(const Tensor& o) const { @@ -784,6 +802,13 @@ public: r.data_[i] = data_[i] / o.data_[i]; return r; } + Tensor operator*(const Tensor& o) const + { + Tensor r(shape_); + for (size_t i = 0; i < data_.size(); ++i) + r.data_[i] = data_[i] * o.data_[i]; + return r; + } Tensor operator+(T val) const { diff --git a/tests/cpp_unit_tests/test_tensor.cpp b/tests/cpp_unit_tests/test_tensor.cpp index 6eae1cea6..0936d866a 100644 --- a/tests/cpp_unit_tests/test_tensor.cpp +++ b/tests/cpp_unit_tests/test_tensor.cpp @@ -350,6 +350,9 @@ TEST_CASE("Tensor element-wise arithmetic") c = a / b; REQUIRE(c[0] == 0.25); + + c = a * b; + REQUIRE(c[0] == 4.0); } TEST_CASE("Tensor scalar arithmetic") @@ -378,7 +381,7 @@ TEST_CASE("Tensor scalar arithmetic") REQUIRE(b[0] == 11.0); } -TEST_CASE("Tensor compound addition with tensor") +TEST_CASE("Tensor compound arithmetic with tensor") { Tensor a({3}, 0.0); Tensor b({3}, 0.0); @@ -388,6 +391,24 @@ TEST_CASE("Tensor compound addition with tensor") REQUIRE(a[0] == 11.0); REQUIRE(a[1] == 22.0); REQUIRE(a[2] == 33.0); + + a = {1.0, 2.0, 3.0}; + b -= a; + REQUIRE(b[0] == 9.0); + REQUIRE(b[1] == 18.0); + REQUIRE(b[2] == 27.0); + + b = {10.0, 20.0, 30.0}; + a *= b; + REQUIRE(a[0] == 10.0); + REQUIRE(a[1] == 40.0); + REQUIRE(a[2] == 90.0); + + a = {1.0, 2.0, 3.0}; + b /= a; + REQUIRE(b[0] == 10.0); + REQUIRE(b[1] == 10.0); + REQUIRE(b[2] == 10.0); } TEST_CASE("Tensor comparison operators")