Array wrapping enhancement

mfem/_par/array.i

@@ -86,6 +86,8 @@ XXXPTR_SIZE_IN(bool *data_, int asize, bool)
 namespace mfem{
 %pythonprepend Array::__setitem__ %{
     i = int(i)
+    if hasattr(v, "thisown"):
+        v.thisown = False
 %}
 %pythonprepend Array::Append %{
     if isinstance(args[0], list):
@@ -167,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(DenseMatrixArray2D) mfem::Array2D<mfem::DenseMatrix *>;
+%template(SparseMatrixArray2D) mfem::Array2D<mfem::SparseMatrix *>;
+%template(HypreParMatrixArray2D) mfem::Array2D<mfem::HypreParMatrix *>;
+

mfem/_par/bilininteg.i

@@ -40,6 +40,7 @@ import_array();
 
 %ignore  mfem::MassIntegrator::SetupPA;
 
+%include "../common/kernel_dispatch.i"
 %include "fem/bilininteg.hpp"
 
 %feature("director") mfem::PyBilinearFormIntegrator;

mfem/_par/coefficient.i

@@ -110,6 +110,23 @@ namespace mfem {
 
 %include "../common/typemap_macros.i"
 LIST_TO_MFEMOBJ_POINTERARRAY_IN(mfem::IntegrationRule const *irs[],  mfem::IntegrationRule *, 0)
+LIST_TO_MFEMOBJ_ARRAY_IN(const mfem::Array<mfem::Coefficient*> & coefs,  mfem::Coefficient *)
+LIST_TO_MFEMOBJ_ARRAY_IN(const mfem::Array<mfem::VectorCoefficient*> & coefs,  mfem::VectorCoefficient *)
+LIST_TO_MFEMOBJ_ARRAY_IN(const mfem::Array<mfem::MatrixCoefficient*> & coefs,  mfem::MatrixCoefficient *)
+
+/* define CoefficientPtrArray, VectorCoefficientPtrArray, MatrixCoefficientPtrArray */
+%import "../common/array_listtuple_typemap.i"
+ARRAY_LISTTUPLE_INPUT_SWIGOBJ(mfem::Coefficient *, 1)
+ARRAY_LISTTUPLE_INPUT_SWIGOBJ(mfem::VectorCoefficient *, 1)
+ARRAY_LISTTUPLE_INPUT_SWIGOBJ(mfem::MatrixCoefficient *, 1)  
+
+%import "../common/array_instantiation_macro.i"
+IGNORE_ARRAY_METHODS(mfem::Coefficient *)
+INSTANTIATE_ARRAY0(Coefficient *, Coefficient, 1)
+IGNORE_ARRAY_METHODS(mfem::VectorCoefficient *)
+INSTANTIATE_ARRAY0(VectorCoefficient *, VectorCoefficient, 1)
+IGNORE_ARRAY_METHODS(mfem::MatrixCoefficient *)
+INSTANTIATE_ARRAY0(MatrixCoefficient *, MatrixCoefficient, 1)
 
 %include "fem/coefficient.hpp"
 %include "../common/numba_coefficient.i"

mfem/_par/quadinterpolator.i

@@ -25,6 +25,7 @@ import_array();
 
 %import "../common/numpy_int_typemap.i"
 
+%include "../common/kernel_dispatch.i"
 %include "fem/quadinterpolator.hpp"
 
 #endif

mfem/_ser/array.i

@@ -40,17 +40,6 @@ XXXPTR_SIZE_IN(int *data_, int asize, int)
 XXXPTR_SIZE_IN(double *data_, int asize, double)
 XXXPTR_SIZE_IN(bool *data_, int asize, bool)
 
-%pythonappend mfem::Array::Array %{
-  if len(args) == 1 and isinstance(args[0], list):
-      if (len(args[0]) == 2 and hasattr(args[0][0], 'disown') and
-          not hasattr(args[0][1], 'disown')):
-          ## first element is SwigObject, like <Swig Object of type 'int *'>
-          ## We do not own data in this case.
-          pass
-      else:
-          self.MakeDataOwner()
-%}
-
 %ignore mfem::Array::operator[];
 %ignore mfem::Array2D::operator[];
 %ignore mfem::BlockArray::operator[];
@@ -92,8 +81,20 @@ XXXPTR_SIZE_IN(bool *data_, int asize, bool)
 
 };
 namespace mfem{
+%pythonappend Array::Array %{
+  if len(args) == 1 and isinstance(args[0], list):
+      if (len(args[0]) == 2 and hasattr(args[0][0], 'disown') and
+          not hasattr(args[0][1], 'disown')):
+          ## first element is SwigObject, like <Swig Object of type 'int *'>
+          ## We do not own data in this case.
+          pass
+      else:
+          self.MakeDataOwner()
+%}
 %pythonprepend Array::__setitem__ %{
     i = int(i)
+    if hasattr(v, "thisown"):
+        v.thisown = False
 %}
 %pythonprepend Array::Append %{
     if isinstance(args[0], list):
@@ -179,14 +180,60 @@ 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(DenseMatrixArray2D) mfem::Array2D<mfem::DenseMatrix *>;
+%template(SparseMatrixArray2D) mfem::Array2D<mfem::SparseMatrix *>;

mfem/_ser/bilininteg.i

@@ -38,6 +38,7 @@ import_array();
 
 %ignore  mfem::MassIntegrator::SetupPA;
 
+%include "../common/kernel_dispatch.i"
 %include "fem/bilininteg.hpp"
 
 %feature("director") mfem::PyBilinearFormIntegrator;

mfem/_ser/coefficient.i

@@ -110,6 +110,23 @@ namespace mfem {
 */
 %include "../common/typemap_macros.i"
 LIST_TO_MFEMOBJ_POINTERARRAY_IN(mfem::IntegrationRule const *irs[],  mfem::IntegrationRule *, 0)
+LIST_TO_MFEMOBJ_ARRAY_IN(const mfem::Array<mfem::Coefficient*> & coefs,  mfem::Coefficient *)
+LIST_TO_MFEMOBJ_ARRAY_IN(const mfem::Array<mfem::VectorCoefficient*> & coefs,  mfem::VectorCoefficient *)
+LIST_TO_MFEMOBJ_ARRAY_IN(const mfem::Array<mfem::MatrixCoefficient*> & coefs,  mfem::MatrixCoefficient *)
+
+/* define CoefficientPtrArray, VectorCoefficientPtrArray, MatrixCoefficientPtrArray */
+%import "../common/array_listtuple_typemap.i"
+ARRAY_LISTTUPLE_INPUT_SWIGOBJ(mfem::Coefficient *, 1)
+ARRAY_LISTTUPLE_INPUT_SWIGOBJ(mfem::VectorCoefficient *, 1)
+ARRAY_LISTTUPLE_INPUT_SWIGOBJ(mfem::MatrixCoefficient *, 1)  
+
+%import "../common/array_instantiation_macro.i"
+IGNORE_ARRAY_METHODS(mfem::Coefficient *)
+INSTANTIATE_ARRAY0(Coefficient *, Coefficient, 1)
+IGNORE_ARRAY_METHODS(mfem::VectorCoefficient *)
+INSTANTIATE_ARRAY0(VectorCoefficient *, VectorCoefficient, 1)
+IGNORE_ARRAY_METHODS(mfem::MatrixCoefficient *)
+INSTANTIATE_ARRAY0(MatrixCoefficient *, MatrixCoefficient, 1)
 
 %include "fem/coefficient.hpp"
 %include "../common/numba_coefficient.i"

mfem/_ser/lininteg.i

@@ -19,10 +19,12 @@
 import_array();
 %}
 
+
+
 %include "exception.i"
 %import "globals.i"
 
-//%include "fem/coefficient.hpp"
+
 %import "fe.i"
 %import "vector.i"
 %import "eltrans.i"

mfem/_ser/quadinterpolator.i

@@ -25,6 +25,7 @@ import_array();
 
 %import "../common/numpy_int_typemap.i"
 
+%include "../common/kernel_dispatch.i"
 %include "fem/quadinterpolator.hpp"
 
 #endif

mfem/common/array_instantiation_macro.i

@@ -1,6 +1,7 @@
 %define INSTANTIATE_ARRAY2(XXX, YYY, ZZZ, USEPTR)
 #if USEPTR == 1
-%template(##ZZZ##Ptr##Array) mfem::Array<mfem::XXX>;
+ //%template(##ZZZ##Ptr##Array) mfem::Array<mfem::XXX>;
+%template(##ZZZ##Array) mfem::Array<mfem::XXX>;
 #else
 %template(##ZZZ##Array) mfem::Array<mfem::XXX>;
 #endif
@@ -325,6 +326,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
 
 

mfem/common/array_listtuple_typemap.i

@@ -71,6 +71,7 @@
       if (!SWIG_IsOK(SWIG_ConvertPtr(s, (void **) &temp_ptr$argnum, ty, 0|0))) {
    	  PyErr_SetString(PyExc_ValueError, "List items must be XXX");
       } else {
+          PyObject_SetAttrString(s, "thisown", Py_False);
 	  #if USEPTR==1
           (* result)[i] =  temp_ptr$argnum;
 	  #else

mfem/common/kernel_dispatch.i

@@ -0,0 +1,6 @@
+/*
+   kernel_dispatch.i
+
+   We set a macro simpily ingnore the MFEM_REGISTER_KERNELS macro
+ */
+#define MFEM_REGISTER_KERNELS(...) //

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())

test/test_coefficient.py

@@ -0,0 +1,67 @@
+import sys
+import numpy as np
+import gc
+from scipy.sparse import csr_matrix
+
+if len(sys.argv) > 1 and sys.argv[1] == '-p':
+    import mfem.par as mfem
+    use_parallel = True
+
+    from mfem.common.mpi_debug import nicePrint as print
+
+    from mpi4py import MPI
+    myid = MPI.COMM_WORLD.rank
+
+else:
+    import mfem.ser as mfem
+    use_parallel = False
+    myid = 0
+
+def test():
+    dim = 3
+    max_attr = 5
+    sigma_attr_coefs = mfem.MatrixCoefficientArray(max_attr)
+    sigma_attr = mfem.intArray(max_attr)
+
+
+    tensors =  [mfem.DenseMatrix(np.ones((3,3))*i) for i in range(max_attr)]
+    tensor_coefs =  [mfem.MatrixConstantCoefficient(mat) for mat in tensors]
+    sigma_array = mfem.MatrixCoefficientArray(tensor_coefs)
+    sigma_attr = mfem.intArray([1,2,3,4,5])
+    sigmaCoef = mfem.PWMatrixCoefficient(dim, sigma_attr, sigma_array, False)
+
+    for ti, tensor in enumerate(tensors): 
+        # add matrix coefficient to list
+        xx = mfem.MatrixConstantCoefficient(tensor)
+        sigma_attr_coefs[ti] = xx
+        sigma_attr[ti] = ti+1   
+
+    # Create PW Matrix Coefficient
+    sigmaCoef = mfem.PWMatrixCoefficient(dim, sigma_attr, sigma_attr_coefs, False)
+
+    sigmaCoef = mfem.PWMatrixCoefficient(dim, sigma_attr, tensor_coefs, False)
+    tensor_coefs = mfem.MatrixCoefficientArray([mfem.MatrixConstantCoefficient(mat) for mat in tensors])
+    sigmaCoef = mfem.PWMatrixCoefficient(dim, sigma_attr, tensor_coefs, False)
+
+    data = tensor_coefs.GetData()
+    tensor_coefs2 = mfem.MatrixCoefficientArray(data, 5, False)
+    sigmaCoef = mfem.PWMatrixCoefficient(dim, sigma_attr, tensor_coefs2, False)
+
+    print("exiting")
+
+
+if __name__ == '__main__':
+    import tracemalloc
+
+    tracemalloc.start()
+
+    print(tracemalloc.get_traced_memory())
+
+    for i in range(3000):
+        test()
+        print(tracemalloc.get_traced_memory())
+    snapshot = tracemalloc.take_snapshot()
+    top_stats = snapshot.statistics('lineno')
+    for stat in top_stats[:10]: 
+        print(stat)
+