Merge remote-tracking branch 'origin/fix_array2d' into coefficient-ar…

…rray-dev
mfem · Oct 3, 2024 · dd805be · dd805be
2 parents 4056b07 + da23550
commit dd805be
Show file tree

Hide file tree

Showing 5 changed files with 151 additions and 10 deletions.
diff --git a/mfem/_par/array.i b/mfem/_par/array.i
@@ -169,13 +169,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
@@ -180,14 +180,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/mfem/common/array_setitem_typemap.i b/mfem/common/array_setitem_typemap.i
@@ -0,0 +1,36 @@
+//
+//   this typemap is used together with %extend, which
+//   adds a class member taking int *pymfem_size 
+//
+%define ARRAY_SETITEM(T)
+  %typemap(in) (int i, const T v) {
+  // generated by ARRAY_LISTTUPLE_INPUT(XXX, YYY)
+  if (!PyList_Check($input)) {
+      if (!PyTuple_Check($input)) {
+         PyErr_SetString(PyExc_ValueError, "Expecting a list/tuple");
+         return NULL;
+      } else {
+	is_tuple = true;
+      }
+  }
+  size = (is_tuple) ? PyTuple_Size($input) : PyList_Size($input);
+  $1 = (void *) & size;
+}
+
+%typemap(argout) (void *List_or_Tuple, XXX *_unused ) {
+  for (int i = 0; i < size$argnum; i++) {
+      PyObject *s = (is_tuple$argnum) ? PyTuple_GetItem($input, i) : PyList_GetItem($input,i);
+      (* result)[i] =  (XXX)YYY(s);
+  }
+}
+
+%typemap(typecheck) (void *List_or_Tuple, XXX *_unused ) {
+  $1 = 0;
+  if (PyList_Check($input)){
+    $1 = 1;
+  }
+  if (PyTuple_Check($input)){
+    $1 = 1;
+  }
+}
+%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())