diff --git a/include/bout/arraynd.hxx b/include/bout/arraynd.hxx new file mode 100644 index 0000000000..87a965d11c --- /dev/null +++ b/include/bout/arraynd.hxx @@ -0,0 +1,283 @@ +#ifndef __ARRAYND_HXX__ +#define __ARRAYND_HXX__ + +#include "output.hxx" +#include "bout/array.hxx" + +template +class ArrayND { +public: + using data_type = T; + using slice_type = ArrayND; + using size_type = int; + using shape_type = decltype( + std::tuple_cat(std::tuple{}, typename slice_type::shape_type{})); + + constexpr static int ndims = ndim; + + ArrayND() : len(0){}; + + template + ArrayND(size_type len, allSizes... sizes) : len(len) { + static_assert(sizeof...(sizes) == ndim - 1, + "Incorrect number of dimension sizes passed as arguments to ArrayND."); + + data = Array(len); + for (auto& i : data) { + i = slice_type(sizes...); + } + } + ArrayND(const ArrayND& other) : len(other.len), data(other.data) { + } + + ArrayND& operator=(const ArrayND& other) { + len = other.len; + data = other.data; + return *this; + } + + template + inline T& operator()(size_type i1, allSizes... sizes) { + ASSERT2(0 <= i1 && i1 < len); + return data[i1](sizes...); + } + inline slice_type& operator[](size_type i1) { + ASSERT2(0 <= i1 && i1 < len); + return data[i1]; + } + inline slice_type& operator()(size_type i1) { + ASSERT2(0 <= i1 && i1 < len); + return data[i1]; + } + + template + inline const T& operator()(size_type i1, allSizes... sizes) const { + ASSERT2(0 <= i1 && i1 < len); + return data[i1](sizes...); + } + inline const slice_type& operator[](size_type i1) const { + ASSERT2(0 <= i1 && i1 < len); + return data[i1]; + } + inline const slice_type& operator()(size_type i1) const { + ASSERT2(0 <= i1 && i1 < len); + return data[i1]; + } + + ArrayND& operator=(const T& val) { + for (auto& i : data) { + i = val; + }; + return *this; + }; + + slice_type* begin() { return std::begin(data); }; + const slice_type* begin() const { return std::begin(data); }; + slice_type* end() { return std::end(data); }; + const slice_type* end() const { return std::end(data); }; + + // Note we assume a non-jagged array, i.e. we assume + // data[0].shape() == data[j].shape() for 0 <= j < len + constexpr auto shape() const -> shape_type { + return std::tuple_cat(std::make_tuple(len), data[0].shape()); + } + + // Note we assume a non-jagged array, i.e. we assume + // data[0].size() == data[j].size() for 0 <= j < len + size_type size() const { return len * data[0].size(); } + + bool empty() const { return size() == 0; } + + bool unique() const { + bool res = data.unique(); + for (const auto& i : data) { + res = res && i.unique(); + // Try to return early if we can. + if (!res) + return res; + }; + return res; + } + + void ensureUnique() { + data.ensureUnique(); + for (auto& i : data) { + i.ensureUnique(); + }; + } + + void pack(ArrayND& flat) { + ASSERT1(size() == flat.size()); + size_type currentPosition = 0; + currentPosition = setFromFlatArray(flat, currentPosition); + ASSERT1(currentPosition == size()); + }; + + void pack(Array& flat) { + ASSERT1(size() == flat.size()); + size_type currentPosition = 0; + currentPosition = setFromFlatArray(flat, currentPosition); + ASSERT1(currentPosition == size()); + }; + + size_type setFromFlatArray(Array& flat, size_type currentPos = 0) { + ArrayND tmp(flat.size()); + tmp.data = flat; + return this->setFromFlatArray(tmp, currentPos); + } + + size_type setFromFlatArray(ArrayND& flat, size_type currentPos = 0) { + ASSERT1(flat.size() >= size()); + + for (auto& i : data) { + currentPos = i.setFromFlatArray(flat, currentPos); + } + + return currentPos; + }; + + ArrayND flatten() { + ArrayND result(size()); + size_type start = 0; + for (auto& i : data) { + result.setSpan(start, i.size(), i.flatten()); + start += i.size(); + } + + return result; + }; + + Array data; + +private: + size_type len; +}; + +template +class ArrayND { +public: + using data_type = T; + using size_type = int; + using shape_type = std::tuple; + constexpr static int ndim = 1; + + constexpr static int ndims = ndim; + + ArrayND() : len(0){}; + ArrayND(size_type len) : len(len) { data = Array(len); } + // Should only end up calling this if we pass too many dimension sizes + template + ArrayND(size_type len, allSizes... sizes) : len(len) { + static_assert( + sizeof...(sizes) == ndim - 1, + "Incorrect number of dimension sizes passed as arguments to ArrayND."); + } + + ArrayND(const ArrayND& other) : len(other.len), data(other.data) { + } + + ArrayND& operator=(const ArrayND& other) { + len = other.len; + data = other.data; + return *this; + } + + inline T& operator()(size_type i1) { + ASSERT2(0 <= i1 && i1 < len); + return data[i1]; + } + inline T& operator[](size_type i1) { + ASSERT2(0 <= i1 && i1 < len); + return data[i1]; + } + inline const T& operator()(size_type i1) const { + ASSERT2(0 <= i1 && i1 < len); + return data[i1]; + } + inline const T& operator[](size_type i1) const { + ASSERT2(0 <= i1 && i1 < len); + return data[i1]; + } + + ArrayND& operator=(const T& val) { + for (auto& i : data) { + i = val; + }; + return *this; + }; + + T* begin() { return std::begin(data); }; + const T* begin() const { return std::begin(data); }; + T* end() { return std::end(data); }; + const T* end() const { return std::end(data); }; + + constexpr std::tuple shape() const { return std::make_tuple(len); } + + size_type size() const { return len; } + + bool empty() const { return size() == 0; } + + bool unique() const { return data.unique(); } + + /*! + * Ensures that this ArrayND does not share data with another + * This should be called before performing any write operations + * on the data. + */ + void ensureUnique() { data.ensureUnique(); } + + size_type setFromFlatArray(Array& flat, size_type currentPos = 0) { + ArrayND tmp(flat.size()); + tmp.data = flat; + return this->setFromFlatArray(tmp, currentPos); + } + + size_type setFromFlatArray(ArrayND& flat, size_type currentPos = 0) { + ASSERT1(flat.size() >= size()); + setSpan(0, size(), flat.getSpan(currentPos, size())); + return currentPos + size(); + }; + + ArrayND flatten() { return *this; } + + void setSpan(size_type start, size_type count, ArrayND source) { + this->setSpan(start, count, source.data); + } + + void setSpan(size_type start, size_type count, Array source) { + ASSERT1(start >= 0); + ASSERT1(start < len); + ASSERT1(start + count - 1 <= len); + ASSERT1(source.size() >= count); + for (size_type i = 0; i < count; i++) { + data[start + i] = source[i]; + } + }; + + Array getSpan(size_type start, size_type count) { + ASSERT1(start >= 0); + ASSERT1(start < len); + ASSERT1(start + count - 1 <= len); + Array result{count}; + for (size_type i = 0; i < count; i++) { + result[i] = data[start + i]; + } + return result; + }; + + Array data; + +private: + size_type len; +}; + +template +using Array1D = ArrayND; + +template +using Array2D = ArrayND; + +template +using Array3D = ArrayND; + +#endif diff --git a/tests/unit/include/bout/test_arraynd.cxx b/tests/unit/include/bout/test_arraynd.cxx new file mode 100644 index 0000000000..37b8248015 --- /dev/null +++ b/tests/unit/include/bout/test_arraynd.cxx @@ -0,0 +1,47 @@ +#include "gtest/gtest.h" + +#include "bout/arraynd.hxx" + +class ArrayNDTest : public ::testing::Test { +public: + ArrayNDTest() {} + ~ArrayNDTest() {} +}; + +TEST_F(ArrayNDTest, ArrayNDSize1D) { + ArrayND a(5); + + ASSERT_FALSE(a.empty()); + EXPECT_EQ(a.size(), 5); + EXPECT_EQ(std::tuple_size::value, 1); + EXPECT_EQ(std::get<0>(a.shape()), 5); + EXPECT_TRUE(a.unique()); +} + +TEST_F(ArrayNDTest, ArrayNDSize2D) { + ArrayND a(5, 10); + + ASSERT_FALSE(a.empty()); + EXPECT_EQ(a.size(), 5 * 10); + EXPECT_EQ(std::tuple_size::value, 2); + EXPECT_EQ(std::get<0>(a.shape()), 5); + EXPECT_EQ(std::get<1>(a.shape()), 10); + EXPECT_TRUE(a.unique()); +} + +TEST_F(ArrayNDTest, ArrayNDUnique2D) { + ArrayND a(5, 10); + EXPECT_TRUE(a.unique()); + auto b = a[0]; + EXPECT_FALSE(a.unique()); + EXPECT_FALSE(b.unique()); + b.ensureUnique(); + EXPECT_TRUE(b.unique()); + EXPECT_TRUE(a.unique()); + auto c = a.data[0]; + EXPECT_FALSE(c.unique()); + EXPECT_FALSE(a.unique()); + a.ensureUnique(); + EXPECT_TRUE(c.unique()); + EXPECT_TRUE(a.unique()); +}