prepare some infrastructure for next steps

eisvogel-project · Jan 11, 2024 · ed22cb0 · ed22cb0
1 parent 0f58a7f
commit ed22cb0
Show file tree

Hide file tree

Showing 3 changed files with 78 additions and 24 deletions.
diff --git a/include/Eisvogel/DistributedNDArray.hh b/include/Eisvogel/DistributedNDArray.hh
@@ -55,19 +55,26 @@ public:
   ~DistributedNDArray();
 
   using chunk_t = DenseNDArray<T, dims>;
+  using shape_t = std::array<std::size_t, dims>;
 
   // For assembling a distributed array
   void RegisterChunk(const chunk_t& chunk, const IndexVector start_ind);
   void Flush();
 
   // For accessing a distributed array
   T& operator()(IndexVector& inds);
+  const shape_t& shape();
 
 private:
 
   bool chunkContainsInds(const ChunkMetadata& chunk_meta, const IndexVector& inds);
   std::size_t getChunkIndex(const IndexVector& inds);
   chunk_t& retrieveChunk(std::size_t chunk_ind);
+  void calculateShape();
+  bool isContiguous();
+
+protected:
+  const shape_t m_shape{};
 
 private:
 
@@ -124,7 +131,7 @@ void DistributedNDArray<T, dims>::RegisterChunk(const DenseNDArray<T, dims>& chu
   std::string chunk_filename = "chunk_" + std::to_string(m_chunk_index.size()) + ".bin";
   ChunkMetadata meta(chunk_filename, start_ind, stop_ind);
   m_chunk_index.push_back(meta);
-
+  
   // write chunk data to disk --> this is where fancy sparsification and compression would happen
   std::string chunk_path = m_dirpath + "/" + chunk_filename;
   std::fstream ofs;
@@ -153,6 +160,8 @@ T& DistributedNDArray<T, dims>::operator()(IndexVector& inds) {
   // retrieve chunk from cache or load from file
   chunk_t& found_chunk = retrieveChunk(chunk_ind);
 
+  std::cout << "found element in chunk " << std::to_string(chunk_ind) << std::endl;
+
   // index and return element
   return found_chunk(inds);
 }
@@ -203,4 +212,24 @@ DistributedNDArray<T, dims>::chunk_t& DistributedNDArray<T, dims>::retrieveChunk
   return m_chunk_cache.find(chunk_ind) -> second;
 }
 
+template <class T, std::size_t dims>
+const DistributedNDArray<T, dims>::shape_t& DistributedNDArray<T, dims>::shape() {
+  return m_shape;
+}
+
+template <class T, std::size_t dims>
+void DistributedNDArray<T, dims>::calculateShape() {
+
+  // go through chunks and determine global start and end inds
+
+}
+
+template <class T, std::size_t dims>
+bool DistributedNDArray<T, dims>::isContiguous() {
+
+  // determine global (distributed-array wide) start and stop inds
+  // determine total number of elements
+  // check if sum of elements of chunks give teh same value  
+}
+
 #endif
diff --git a/include/Eisvogel/NDArray.hh b/include/Eisvogel/NDArray.hh
@@ -48,6 +48,13 @@ private:
 
 private:
   friend struct stor::Traits<DenseNDArray<T, dims>>;
+
+  template <typename T0, std::size_t dims0, typename T1, typename T2>
+  friend inline DenseNDArray<T0, dims0> operator_binary(const DenseNDArray<T1, dims0>& lhs, const DenseNDArray<T2, dims0>& rhs, auto binary_op);
+
+  template <typename T0, std::size_t dims0, typename T1>
+  friend inline DenseNDArray<T0, dims0> operator_unary(const DenseNDArray<T1, dims0>& arg, auto unary_op);
+
   template <typename T1> friend class DenseNDArray<T, dims>;
 
 public:
@@ -142,7 +149,7 @@ public:
 
   // printing
   void print() {
-    for(T& cur: m_data) {
+    for(const T& cur: m_data) {
       std::cout << cur << " ";
     }
     std::cout << std::endl;
@@ -157,59 +164,66 @@ public:
   const std::size_t size() const requires(dims == 1) {return m_data.size();}
 
   friend DenseNDArray<T, dims> operator+(const DenseNDArray<T, dims>& lhs, const DenseNDArray<T, dims>& rhs) {
-    return operator_binary(lhs, rhs, std::plus<>());
+    return operator_binary<T, dims, T, T>(lhs, rhs, std::plus<>());
   }
 
   friend DenseNDArray<T, dims> operator+(const DenseNDArray<T, dims>& lhs, const T& rhs) {
     auto plus_rhs = [&](const T& el){return std::plus<>()(el, rhs);};
-    return operator_unary(lhs, plus_rhs);
+    return operator_unary<T, dims, T>(lhs, plus_rhs);
   }
 
   friend DenseNDArray<T, dims> operator-(const DenseNDArray<T, dims>& lhs, const DenseNDArray<T, dims>& rhs) {
-    return operator_binary(lhs, rhs, std::minus<>());
+    return operator_binary<T, dims, T, T>(lhs, rhs, std::minus<>());
   }
 
   friend DenseNDArray<T, dims> operator-(const DenseNDArray<T, dims>& lhs, const T& rhs) {
     auto minus_rhs = [&](const T& el){return std::minus<>()(el, rhs);};
-    return operator_unary(lhs, minus_rhs);
+    return operator_unary<T, dims, T>(lhs, minus_rhs);
   }
 
   friend DenseNDArray<T, dims> operator*(const DenseNDArray<T, dims>& lhs, const DenseNDArray<T, dims>& rhs) {
-    return operator_binary(lhs, rhs, std::multiplies<>());
+    return operator_binary<T, dims, T, T>(lhs, rhs, std::multiplies<>());
   }
 
   friend DenseNDArray<T, dims> operator*(const DenseNDArray<T, dims>& lhs, const T& rhs) {
     auto multiplies_rhs = [&](const T& el){return std::multiplies<>()(el, rhs);};
-    return operator_unary(lhs, multiplies_rhs);
+    return operator_unary<T, dims, T>(lhs, multiplies_rhs);
   }
 
   friend DenseNDArray<T, dims> operator/(const DenseNDArray<T, dims>& lhs, const DenseNDArray<T, dims>& rhs) {
-    return operator_binary(lhs, rhs, std::divides<>());
+    return operator_binary<T, dims, T, T>(lhs, rhs, std::divides<>());
   }
 
   friend DenseNDArray<T, dims> operator/(const DenseNDArray<T, dims>& lhs, const T& rhs) {
     auto divides_rhs = [&](const T& el){return std::divides<>()(el, rhs);};
-    return operator_unary(lhs, divides_rhs);
+    return operator_unary<T, dims, T>(lhs, divides_rhs);
   }
 
+  friend DenseNDArray<bool, dims> operator<(const DenseNDArray<T, dims>& lhs, const DenseNDArray<T, dims>& rhs) {
+    auto elementwise_lt = [&](const T& el_lhs, const T& el_rhs){return el_lhs < el_rhs;};
+    return operator_binary<bool, dims, T, T>(lhs, rhs, elementwise_lt);
+  }
+
 private:
 
   stride_t m_strides = {};
   data_t m_data = {};
+};
 
-  friend inline DenseNDArray<T, dims> operator_binary(const DenseNDArray<T, dims>& lhs, const DenseNDArray<T, dims>& rhs,
-						      auto binary_op) {
-    DenseNDArray<T, dims> result(lhs.m_shape, T());
-    std::transform(lhs.m_data.begin(), lhs.m_data.end(), rhs.m_data.begin(), result.m_data.begin(), binary_op);
-    return result;
-  }
+template <typename T0, std::size_t dims0, typename T1, typename T2>
+inline DenseNDArray<T0, dims0> operator_binary(const DenseNDArray<T1, dims0>& lhs, const DenseNDArray<T2, dims0>& rhs,
+					       auto binary_op) {
+  DenseNDArray<T0, dims0> result(lhs.m_shape, T0());
+  std::transform(lhs.m_data.begin(), lhs.m_data.end(), rhs.m_data.begin(), result.m_data.begin(), binary_op);
+  return result;
+}
 
-  friend inline DenseNDArray<T, dims> operator_unary(const DenseNDArray<T, dims>& arg, auto unary_op) {
-    DenseNDArray<T, dims> result(arg.m_shape, T());
-    std::transform(arg.m_data.begin(), arg.m_data.end(), result.m_data.begin(), unary_op);
-    return result;
-  }
-};
+template <typename T0, std::size_t dims0, typename T1>
+inline DenseNDArray<T0, dims0> operator_unary(const DenseNDArray<T1, dims0>& arg, auto unary_op) {
+  DenseNDArray<T0, dims0> result(arg.m_shape, T0());
+  std::transform(arg.m_data.begin(), arg.m_data.end(), result.m_data.begin(), unary_op);
+  return result;
+}
 
 namespace stor {
 

diff --git a/tests/io/testDistributedNDArray.cxx b/tests/io/testDistributedNDArray.cxx
@@ -11,6 +11,12 @@ int main(int argc, char* argv[]) {
   chunk1(1, 1) = 4.0;
   chunk1.print();
 
+  DenseNDArray<float, 2> chunk1_1({2, 2}, 0.0);
+  chunk1_1.print();
+
+  DenseNDArray<bool, 2> comp = chunk1_1 < chunk1;
+  comp.print();
+
   IndexVector start_ind1({0, 0});
   start_ind1.print();
 
@@ -34,7 +40,12 @@ int main(int argc, char* argv[]) {
   std::cout << darr_load(acc_ind1) << std::endl;
   std::cout << darr_load(acc_ind1) << std::endl;
 
-  IndexVector acc_ind2 = {10,10};
+  IndexVector acc_ind2 = {2,1};
   std::cout << darr_load(acc_ind2) << std::endl;
-  std::cout << darr_load(acc_ind2) << std::endl;  
+  std::cout << darr_load(acc_ind2) << std::endl;
+
+  std::cout << "shape:" << std::endl;  
+  for(auto cur : darr_load.shape()) {
+    std::cout << cur << std::endl;
+  }
 }