fixing Array2D wrapping

mfem · Aug 16, 2024 · 1259886 · 1259886
1 parent 0054c54
commit 1259886
Show file tree

Hide file tree

Showing 4 changed files with 115 additions and 10 deletions.
diff --git a/mfem/_par/array.i b/mfem/_par/array.i
@@ -167,13 +167,62 @@ INSTANTIATE_ARRAY_BOOL
 IGNORE_ARRAY_METHODS_PREMITIVE(unsigned int)
 INSTANTIATE_ARRAY_NUMPYARRAY(uint, unsigned int, NPY_UINT)       // 32bit
 
+/* Array< Array<int> *> */
+IGNORE_ARRAY_METHODS(mfem::Array<int> *)
+INSTANTIATE_ARRAY2(Array<int> *, Array<int>, intArray, 1)
+
 /*
-   for these Array2D, we instantiate it. But we dont extend it since, Array2D<T> does not
-   expose the interanl pointer to array1d.
+   Array2D:: Assign and data access
 */
+
+%extend mfem::Array2D{
+  void Assign(const T &a){
+     *self = a;
+  }
+  void Assign(const mfem::Array2D<T> &a){
+     *self = a;
+  }
+
+  void __setitem__(PyObject* param, const T v) {
+    if (PyTuple_Check(param)) {
+        PyErr_Clear();
+        int i = PyInt_AsLong(PyTuple_GetItem(param, 0));
+        int j = PyInt_AsLong(PyTuple_GetItem(param, 1));
+
+        if (PyErr_Occurred()) {
+           PyErr_SetString(PyExc_ValueError, "Argument must be i, j");
+           return;
+        }
+	T *arr =  self -> GetRow(i);
+        arr[j] = v;
+    }
+  }
+  T __getitem__(PyObject* param) {
+    if (PyTuple_Check(param)) {
+        PyErr_Clear();
+        int i = PyInt_AsLong(PyTuple_GetItem(param, 0));
+        int j = PyInt_AsLong(PyTuple_GetItem(param, 1));
+
+        if (PyErr_Occurred()) {
+           PyErr_SetString(PyExc_ValueError, "Argument must be i, j");
+	   i = 0;
+	   j = 0;
+        }
+	T *arr = self -> GetRow(i);
+        return arr[j];
+    }
+  }
+}
 %template(intArray2D) mfem::Array2D<int>;
 %template(doubleArray2D) mfem::Array2D<double>;
 
-/* Array< Array<int> *> */
-IGNORE_ARRAY_METHODS(mfem::Array<int> *)
-INSTANTIATE_ARRAY2(Array<int> *, Array<int>, intArray, 1)
+/* Array2D<* DenseMatrix>, Array2D<* SparseMatrix>,  Array2D<* HypreParMatrix> */
+
+IGNORE_ARRAY_METHODS(mfem::DenseMatrix *)
+IGNORE_ARRAY_METHODS(mfem::SparseMatrix *)
+IGNORE_ARRAY_METHODS(mfem::HypreParMatrix *)
+
+%template(densematArray2D) mfem::Array2D<mfem::DenseMatrix *>;
+%template(sparsematArray2D) mfem::Array2D<mfem::SparseMatrix *>;
+%template(hypreparmatArray2D) mfem::Array2D<mfem::HypreParMatrix *>;
+
diff --git a/mfem/_ser/array.i b/mfem/_ser/array.i
@@ -179,14 +179,64 @@ INSTANTIATE_ARRAY_BOOL
 IGNORE_ARRAY_METHODS_PREMITIVE(unsigned int)
 INSTANTIATE_ARRAY_NUMPYARRAY(uint, unsigned int, NPY_UINT)       // 32bit
 
+/* Array< Array<int> *> */
+IGNORE_ARRAY_METHODS(mfem::Array<int> *)
+INSTANTIATE_ARRAY2(Array<int> *, Array<int>, intArray, 1)
+
 /*
-   for these Array2D, we instantiate it. But we dont extend it since, Array2D<T> does not
-   expose the interanl pointer to array1d.
+   Array2D:: Assign and data access
 */
+
+%extend mfem::Array2D{
+  void Assign(const T &a){
+     *self = a;
+  }
+  void Assign(const mfem::Array2D<T> &a){
+     *self = a;
+  }
+
+  void __setitem__(PyObject* param, const T v) {
+    if (PyTuple_Check(param)) {
+        PyErr_Clear();
+        int i = PyInt_AsLong(PyTuple_GetItem(param, 0));
+        int j = PyInt_AsLong(PyTuple_GetItem(param, 1));
+
+        if (PyErr_Occurred()) {
+           PyErr_SetString(PyExc_ValueError, "Argument must be i, j");
+           return;
+        }
+	T *arr =  self -> GetRow(i);
+        arr[j] = v;
+    }
+  }
+  T __getitem__(PyObject* param) {
+    if (PyTuple_Check(param)) {
+        PyErr_Clear();
+        int i = PyInt_AsLong(PyTuple_GetItem(param, 0));
+        int j = PyInt_AsLong(PyTuple_GetItem(param, 1));
+
+        if (PyErr_Occurred()) {
+           PyErr_SetString(PyExc_ValueError, "Argument must be i, j");
+	   i = 0;
+	   j = 0;
+        }
+	T *arr = self -> GetRow(i);
+        return arr[j];
+    }
+  }
+}
+
 %template(intArray2D) mfem::Array2D<int>;
 %template(doubleArray2D) mfem::Array2D<double>;
 
-/* Array< Array<int> *> */
-IGNORE_ARRAY_METHODS(mfem::Array<int> *)
-INSTANTIATE_ARRAY2(Array<int> *, Array<int>, intArray, 1)
+/* Array2D<* DenseMatrix>, Array2D<* SparseMatrix>,  Array2D<* HypreParMatrix> */
+
+IGNORE_ARRAY_METHODS(mfem::DenseMatrix *)
+IGNORE_ARRAY_METHODS(mfem::SparseMatrix *)
+
+%template(densematArray2D) mfem::Array2D<mfem::DenseMatrix *>;
+%template(sparsematArray2D) mfem::Array2D<mfem::SparseMatrix *>;
+//%template(hypreparmatArray2D) mfem::Array2D<mfem::HypreParMatrix *>;
+
+
 
diff --git a/mfem/common/array_instantiation_macro.i b/mfem/common/array_instantiation_macro.i
@@ -325,6 +325,8 @@ INSTANTIATE_ARRAY2(XXX, YYY, YYY, USEPTR)
 %ignore mfem::Array2D<XXX>::Print;
 %ignore mfem::Array2D<XXX>::PrintGZ;
 %ignore mfem::Array2D<XXX>::SaveGZ;
+%ignore mfem::Array2D<XXX>::Load;
+%ignore mfem::Array2D<XXX>::Save;
 %enddef
 
 

diff --git a/test/test_blockmatrix.py b/test/test_blockmatrix.py
@@ -35,3 +35,7 @@
 m.SetBlock(1, 0, mmat)
 print(m._offsets)
 print(m._linked_mat)
+
+from mfem.common.sparse_utils import sparsemat_to_scipycsr
+print(m.CreateMonolithic())
+print(sparsemat_to_scipycsr(m.CreateMonolithic(), np.float64).todense())