Add routines to create Eigen::tensor objects from numpy.ndarray objec…

…ts; needed to set coefficient arrays from pyEXP
EXP-code · Jul 13, 2024 · b71934f · b71934f
1 parent f9af374
commit b71934f
Show file tree

Hide file tree

Showing 2 changed files with 94 additions and 8 deletions.
diff --git a/pyEXP/CoefWrappers.cc b/pyEXP/CoefWrappers.cc
@@ -1321,7 +1321,12 @@ void CoefficientClasses(py::module &m) {
          SphFldCoefs instance
          )")
     .def("__call__",
-	 &CoefClasses::SphFldCoefs::getMatrix,
+	 [](CoefClasses::SphFldCoefs& A, double t)
+	 {
+	   Eigen::Tensor<std::complex<double>, 3> M = A.getMatrix(t);
+	   py::array_t<std::complex<double>> ret = make_ndarray<std::complex<double>>(M);
+	   return ret;
+	 },
          R"(
          Return the coefficient tensor for the desired time.
 
@@ -1342,16 +1347,21 @@ void CoefficientClasses(py::module &m) {
          )",
          py::arg("time"))
     .def("setMatrix",
-	 &CoefClasses::SphFldCoefs::setMatrix,
+	 [](CoefClasses::SphFldCoefs& A, double t,
+	    py::array_t<std::complex<double>> array)
+	 {
+	   auto M = make_tensor3<std::complex<double>>(array);
+	   A.setMatrix(t, M);
+	 },
          R"(
          Enter and/or rewrite the coefficient tensor at the provided time
 
          Parameters
          ----------
          time : float
              snapshot time corresponding to the the coefficient matrix
-             mat : numpy.ndarray
-                 the new coefficient array.
+         mat : numpy.ndarray
+             the new coefficient array.
 
          Returns
          -------
@@ -1396,7 +1406,12 @@ void CoefficientClasses(py::module &m) {
          CylFldCoefs instance
          )")
     .def("__call__",
-	 &CoefClasses::CylFldCoefs::getMatrix,
+	 [](CoefClasses::CylFldCoefs& A, double t)
+	 {
+	   Eigen::Tensor<std::complex<double>, 3> M = A.getMatrix(t);
+	   py::array_t<std::complex<double>> ret = make_ndarray<std::complex<double>>(M);
+	   return ret;
+	 },
          R"(
          Return the coefficient tensor for the desired time.
 
@@ -1417,16 +1432,21 @@ void CoefficientClasses(py::module &m) {
          )",
          py::arg("time"))
     .def("setMatrix",
-	 &CoefClasses::CylFldCoefs::setMatrix,
+	 [](CoefClasses::CylFldCoefs& A, double t,
+	    py::array_t<std::complex<double>> array)
+	 {
+	   auto M = make_tensor3<std::complex<double>>(array);
+	   A.setMatrix(t, M);
+	 },
          R"(
          Enter and/or rewrite the coefficient tensor at the provided time
 
          Parameters
          ----------
          time : float
              snapshot time corresponding to the the coefficient matrix
-             mat : numpy.ndarray
-                 the new coefficient array.
+         mat : numpy.ndarray
+             the new coefficient array.
 
          Returns
          -------

diff --git a/pyEXP/TensorToArray.H b/pyEXP/TensorToArray.H
@@ -27,6 +27,38 @@ py::array_t<T> make_ndarray(Eigen::Tensor<T, 3>& mat)
      );
 }
 
+template <typename T>
+Eigen::Tensor<T, 3> make_tensor3(py::array_t<T> array)
+{
+  // Request a buffer descriptor from Python
+  py::buffer_info buffer_info = array.request();
+
+  // Get the array dimenions
+  T *data = static_cast<T *>(buffer_info.ptr);
+  std::vector<ssize_t> shape = buffer_info.shape;
+
+  // Check rank
+  if (shape.size() != 3) {
+    std::ostringstream sout;
+    sout << "make_tensor3: tensor rank must be 3, found "
+	 << shape.size();
+    throw std::runtime_error(sout.str());
+  }
+
+  // Build result tensor with col-major ordering
+  Eigen::Tensor<T, 3> tensor(shape[0], shape[1], shape[2]);
+  for (int i=0, l=0; i < shape[0]; i++) {
+    for (int j=0; j < shape[1]; j++) {
+      for (int k=0; k < shape[2]; k++) {
+	tensor(i, j, k) = data[l++];
+      }
+    }
+  }
+
+  return tensor;
+}
+
+
 //! Helper function that maps the Eigen::Tensor<T, 4> into an numpy.ndarray
 template <typename T>
 py::array_t<T> make_ndarray4(Eigen::Tensor<T, 4>& mat)
@@ -53,4 +85,38 @@ py::array_t<T> make_ndarray4(Eigen::Tensor<T, 4>& mat)
      );
 }
 
+template <typename T>
+Eigen::Tensor<T, 4> make_tensor4(py::array_t<T> array)
+{
+  // Request a buffer descriptor from Python
+  py::buffer_info buffer_info = array.request();
+
+  // Get the array dimenions
+  T *data = static_cast<T *>(buffer_info.ptr);
+  std::vector<ssize_t> shape = buffer_info.shape;
+
+  // Check rank
+  if (shape.size() != 4) {
+    std::ostringstream sout;
+    sout << "make_tensor4: tensor rank must be 4, found "
+	 << shape.size();
+    throw std::runtime_error(sout.str());
+  }
+
+  // Build result tensor with col-major ordering
+  Eigen::Tensor<T, 4> tensor(shape[0], shape[1], shape[2], shape[3]);
+  for (int i=0, l=0; i < shape[0]; i++) {
+    for (int j=0; j < shape[1]; j++) {
+      for (int k=0; k < shape[2]; k++) {
+	for (int l=0; l < shape[3]; k++) {
+	  tensor(i, j, k, l) = data[l++];
+	}
+      }
+    }
+  }
+
+  return tensor;
+}
+
+
 #endif