block matrix can be used in petsc by passing the "-block_matrix" flag

through solver_options

python/_tools.py

@@ -2,6 +2,8 @@ import time
 import atexit
 import gmsh
 import signal
+import numpy as np
+import ctypes
 
 gmsh.initialize()
 gmsh.option.setNumber("General.Terminal",1)
@@ -10,6 +12,13 @@ signal.signal(signal.SIGINT, signal.SIG_DFL)
 
 timers = {}
 
+def _np2c(a,dtype=float) :
+    tmp = np.require(a,dtype,"C")
+    r = ctypes.c_void_p(tmp.ctypes.data)
+    r.tmp = tmp
+    return r
+
+
 def timeit(func):
     def wrapper(*args, **kwargs):
         self = args[0]

python/fluid.py

@@ -33,7 +33,7 @@ import numpy as np
 import os
 import sys
 from . import VTK
-from ._tools import gmsh, get_linear_system_package
+from ._tools import gmsh, get_linear_system_package, _np2c
 
 dir_path = os.path.dirname(os.path.realpath(__file__))
 lib2 = np.ctypeslib.load_library("libmbfluid2",dir_path)
@@ -60,12 +60,6 @@ def _is_string(s) :
     else :
         return isinstance(s, basestring)
 
-def _np2c(a,dtype=float) :
-    tmp = np.require(a,dtype,"C")
-    r = c_void_p(tmp.ctypes.data)
-    r.tmp = tmp
-    return r
-
 def _load_msh(mesh_file_name, lib, dim) :
     """
     mesh_file_name -- Name of the mesh.msh file containing information about the domain

python/petsc4pylsys.py

@@ -6,88 +6,59 @@ import numpy as np
 from ._tools import timeit
 from .csr import CSR
 
-class LinearSystemAIJ:
+class LinearSystem:
 
     def __init__(self, elements, n_fields, options, constraints) :
-        idx = ((elements*n_fields)[:,:,None]+np.arange(n_fields)[None,None,:]).reshape([elements.shape[0],-1])
-        self.localsize = elements.shape[1]*n_fields
-        self.constraints = list(c[:,0]*n_fields + c[:,1] for c in constraints)
-        self.idx = (elements[:,None,:]*n_fields+np.arange(n_fields)[None,:,None]).reshape([-1])
         PETSc.Options().prefixPush("fluid_")
-        PETSc.Options().insertString("-pc_type lu -pc_factor_shift_type NONZERO -pc_factor_shift_amount 1e-16")
+        PETSc.Options().insertString("-pc_type lu -pc_factor_shift_type INBLOCKS -pc_factor_shift_amount 1e-16")
         PETSc.Options().insertString(options)
         PETSc.Options().prefixPop()
+        self.block = PETSc.Options().getBool("fluid_block_matrix",False)
         self.ksp = PETSc.KSP().create()
         self.ksp.setOptionsPrefix(b"fluid_")
         self.ksp.setFromOptions()
-        self.csr = CSR(idx, self.idx,self.constraints)
-        self.val = np.zeros(self.csr.col.size)
-        self.mat = PETSc.Mat().createAIJWithArrays(self.csr.size,(self.csr.row,self.csr.col,self.val),bsize=1)
-
-    @timeit
-    def local_to_global(self, localv, localm, u, constraints_value):
-        self.csr.assemble_mat(localm, constraints_value, self.val)
-        self.mat.assemble()
-        return self.csr.assemble_rhs(localv, u, constraints_value)
-
-    @timeit
-    def solve(self,rhs) :
-        x = np.ndarray(rhs.shape)
-        prhs = PETSc.Vec().createWithArray(rhs.reshape([-1]))
-        px = PETSc.Vec().createWithArray(x.reshape([-1]))
-        self.ksp.setOperators(self.mat)
-        self.ksp.solve(prhs,px)
-        return x[:self.csr.ndof]
-
-
-class LinearSystemBAIJ :
-    def __init__(self, elements, n_fields, options, constraints = []) :
-        nnodes = np.max(elements)+1
-        nn = elements.shape[1]
-        pairs = np.ndarray((elements.shape[0],nn*(nn-1)),dtype=([('i0',np.int32),('i1',np.int32)]))
-        pairs['i0'] = elements[:,np.repeat(range(nn),nn-1)]
-        idx = np.hstack([[i for i in range(nn) if i!=j] for j in range(nn)]) 
-        pairs['i1'] = elements[:,idx]
-        pairs = np.unique(pairs)
-        nnz = (np.bincount(pairs["i0"])+1).astype(np.int32)
-        PETSc.Options().prefixPush("fluid_")
-        PETSc.Options().insertString(options)
-        PETSc.Options().prefixPop()
-        self.mat = PETSc.Mat().createBAIJ(nnodes*n_fields,n_fields,nnz)
-        self.ksp = PETSc.KSP().create()
-        self.ksp.setOptionsPrefix(b"fluid_")
-        self.ksp.setFromOptions()
-        self.mat.zeroEntries()
-        self.constraints = list(c[:,0]*n_fields + c[:,1] for c in constraints)
-        self.ndof = (np.max(elements) + 1)*n_fields
-        ###
         self.localsize = elements.shape[1]*n_fields
-        self.elements = elements
-        self.idx = (self.elements[:,None,:]*n_fields+np.arange(n_fields)[None,:,None]).reshape([-1])
+        if self.block:
+            if len(constraints) != 0 :
+                raise ValueError("petsc block matrices not implemented with constraints")
+            nn = elements.shape[1]
+            self.csr = CSR(elements,elements,[])
+            self.mat = PETSc.Mat().createBAIJ(self.csr.size*n_fields, n_fields, csr=(self.csr.row, self.csr.col))
+            self.idx = (elements[:,None,:]*n_fields+np.arange(n_fields)[None,:,None]).reshape([-1])
+            self.n_fields = n_fields
+            csrmapl = np.arange(n_fields*n_fields)[None,:].reshape(n_fields,n_fields)
+            csrmap = (self.csr.map[0]*n_fields*n_fields)[:,None,None]+csrmapl[None,:,:]
+            self.csrmap = np.swapaxes(csrmap.reshape([elements.shape[0],nn,n_fields,nn,n_fields]),2,3).reshape([-1])
+        else:
+            idx = ((elements*n_fields)[:,:,None]+np.arange(n_fields)[None,None,:]).reshape([elements.shape[0],-1])
+            self.constraints = list(c[:,0]*n_fields + c[:,1] for c in constraints)
+            idx2 = (elements[:,None,:]*n_fields+np.arange(n_fields)[None,:,None]).reshape([-1])
+            self.csr = CSR(idx, idx2, self.constraints)
+            self.val = np.zeros(self.csr.col.size)
+            self.mat = PETSc.Mat().createAIJWithArrays(self.csr.size,(self.csr.row,self.csr.col,self.val))
 
     @timeit
-    def local_to_global2(self,localv,localm, u, constraints_value = []):
-        self.mat.zeroEntries()
-        for e,m in zip(self.elements,localm.reshape([-1,self.localsize**2])) :
-            self.mat.setValuesBlocked(e,e,m,PETSc.InsertMode.ADD)
-        self.mat.assemble()
-
-    def local_to_global(self,localv,localm, u, constraints_value = []):
-        rhs = np.zeros((self.ndof + len(constraints_value,)))
-        np.add.at(rhs.reshape([-1]),self.idx,localv)
-        self.local_to_global2(localv,localm,u,constraints_value)
-        return rhs
+    def local_to_global(self, localv, localm, u, constraints_value):
+        if self.block:
+            rhs = np.bincount(self.idx,localv,self.csr.size*self.n_fields)
+            val = np.bincount(self.csrmap,localm,self.csr.col.size*self.n_fields**2)
+            self.mat.zeroEntries()
+            self.mat.setValuesBlockedCSR(self.csr.row, self.csr.col, val)
+            self.mat.assemble()
+            return rhs
+        else:
+            self.csr.assemble_mat(localm, constraints_value, self.val)
+            self.mat.assemble()
+            return self.csr.assemble_rhs(localv, u, constraints_value)
 
     @timeit
     def solve(self,rhs) :
-        viewer = PETSc.Viewer.ASCII("matinfo")
-        viewer.pushFormat(5)
-        self.mat.view(viewer)
         x = np.ndarray(rhs.shape)
         prhs = PETSc.Vec().createWithArray(rhs.reshape([-1]))
         px = PETSc.Vec().createWithArray(x.reshape([-1]))
         self.ksp.setOperators(self.mat)
         self.ksp.solve(prhs,px)
-        return x
-
-LinearSystem = LinearSystemAIJ
+        if self.block:
+            return x
+        else:
+            return x[:self.csr.ndof]

python/petsclsys.py

@@ -22,7 +22,7 @@ import numpy as np
 import atexit
 import weakref
 
-from ._tools import timeit
+from ._tools import timeit, _np2c
 from .csr import CSR
 
 import ctypes
@@ -36,12 +36,6 @@ lib.PetscInitialize(None, None, None)
 COMM_WORLD = ctypes.c_void_p.in_dll(lib,"PETSC_COMM_WORLD")
 #atexit.register(lib.PetscFinalize)
 
-def _np2c(a,dtype=float) :
-    tmp = np.require(a,dtype,"C")
-    r = ctypes.c_void_p(tmp.ctypes.data)
-    r.tmp = tmp
-    return r
-
 
 class PetscObj:
     
@@ -58,113 +52,88 @@ class PetscObj:
         return self._p
 
 
-class Viewer(PetscObj) :
-    def __init__(self):
-        super().__init__("PetscViewerASCIIOpen","PetscViewerDestroy",COMM_WORLD,"matinfo".encode())
-        lib.PetscViewerPushFormat(self,4)
-
-
-class KSP(PetscObj) :
-
-    def __init__(self, options) :
-        super().__init__("KSPCreate","KSPDestroy",COMM_WORLD)
-
-        lib.PetscOptionsPrefixPush(None, b"fluid_")
-        lib.PetscOptionsInsertString(None, b"-pc_type lu -pc_factor_shift_type NONZERO -pc_factor_shift_amount 1e-16")
-        lib.PetscOptionsInsertString(None, options.encode())
-        lib.PetscOptionsPrefixPop(None)
-        lib.KSPSetOptionsPrefix(self,b"fluid_")
-        lib.KSPSetFromOptions(self)
-
-    def solve(self, mat, rhs):
-        x = np.zeros_like(rhs)
-        #lib.KSPSetReusePreconditioner(self, 0)
-        lib.KSPSetOperators(self, mat, mat)
-        lib.KSPSolve(self, Vec(rhs), Vec(x))
-        return x
-
-
-class MatSeqBAIJ(PetscObj) :
-
-    def __init__(self, ms, bs, csrrow, csrcol) :
-        super().__init__("MatCreate","MatDestroy",COMM_WORLD)
-        lib.MatSetType(self, b"seqbaij")
-        lib.MatSetSizes(self,ctypes.c_int(ms),ctypes.c_int(ms),ctypes.c_int(ms),ctypes.c_int(ms))
-        lib.MatSeqBAIJSetPreallocationCSR(self,
-                ctypes.c_int(bs), _np2c(csrrow, np.int32), _np2c(csrcol, np.int32), None)
-        self.ms = ms
-
-    def set(self, v) :
-        aptr = ctypes.POINTER(ctypes.c_double)()
-        lib.MatSeqBAIJGetArray(self, ctypes.byref(aptr))
-        np.ctypeslib.as_array(aptr,shape=(v.size,))[:] = v
-        lib.MatSeqBAIJRestoreArray(self, ctypes.byref(aptr))
-        lib.MatAssemblyBegin(self,0)
-        lib.MatAssemblyEnd(self,0)
-
-
 class Vec(PetscObj) :
 
     def __init__(self, v) :
         super().__init__("VecCreateSeqWithArray","VecDestroy",COMM_WORLD,
                 ctypes.c_int(v.shape[-1]), ctypes.c_int(v.size),ctypes.c_void_p(v.ctypes.data))
 
+def insert_options(prefix,options):
+    lib.PetscOptionsPrefixPush(None, prefix.encode())
+    lib.PetscOptionsInsertString(None, options.encode())
+    lib.PetscOptionsPrefixPop(None)
 
-class LinearSystemAIJ:
+def get_option_bool(prefix, name, default):
+    v = ctypes.c_int(0)
+    lib.PetscOptionsGetBool(None, prefix.encode(), name.encode(), ctypes.c_int(default), ctypes.byref(v))
+    return v.value
+
+class LinearSystem:
 
     def __init__(self, elements, n_fields, options, constraints = []) :
-        idx = ((elements*n_fields)[:,:,None]+np.arange(n_fields)[None,None,:]).reshape([elements.shape[0],-1])
-        self.localsize = elements.shape[1]*n_fields
-        self.constraints = list(c[:,0]*n_fields + c[:,1] for c in constraints)
-        self.idx = (elements[:,None,:]*n_fields+np.arange(n_fields)[None,:,None]).reshape([-1])
-        self.ksp = KSP(options)
-        self.csr = CSR(idx, self.idx,self.constraints)
-        self.val = np.zeros(self.csr.col.size)
-        self.mat = PetscObj("MatCreateSeqAIJWithArrays", "MatDestroy", COMM_WORLD,
-                ctypes.c_int(self.csr.size), ctypes.c_int(self.csr.size), 
-                _np2c(self.csr.row,np.int32), _np2c(self.csr.col,np.int32), _np2c(self.val))
+        insert_options("fluid", "-pc_type lu -pc_factor_shift_type NONZERO -pc_factor_shift_amount 1e-16")
+        insert_options("fluid", options)
+        self.block = get_option_bool("fluid","-block_matrix",0)
+        self.ksp = PetscObj("KSPCreate", "KSPDestroy", COMM_WORLD)
+        lib.KSPSetOptionsPrefix(self.ksp,b"fluid_")
+        lib.KSPSetFromOptions(self.ksp)
+        if self.block:
+            if len(constraints) != 0 :
+                raise ValueError("petsc BAIJ not implemented with constraints")
+            nnodes = np.max(elements)+1
+            nn = elements.shape[1]
+            self.csr = CSR(elements,elements,[])
+            self.localsize = elements.shape[1]*n_fields
+            self.elements = elements
+            self.idx = (self.elements[:,None,:]*n_fields+np.arange(n_fields)[None,:,None]).reshape([-1])
+            self.size = nnodes*n_fields
+            self.n_fields = n_fields
+            csrmapl = np.arange(nn*nn)[None,:].reshape(nn,nn).T
+            csrmap = (self.csr.map[0]*nn*nn)[:,None,None]+csrmapl[None,:,:]
+            self.csrmap = np.copy(np.swapaxes(csrmap.reshape([elements.shape[0],nn,nn, n_fields,n_fields]),2,3).reshape([-1]))
+            self.vsize = np.max(self.csrmap)+1
+            self.mat = PetscObj("MatCreate", "MatDestroy", COMM_WORLD)
+            lib.MatSetType(self.mat, b"seqbaij")
+            ms = self.csr.size*n_fields
+            lib.MatSetSizes(self.mat,ctypes.c_int(ms),ctypes.c_int(ms),ctypes.c_int(ms),ctypes.c_int(ms))
+            lib.MatSeqBAIJSetPreallocationCSR(self.mat,
+                    ctypes.c_int(n_fields), _np2c(self.csr.row, np.int32), _np2c(self.csr.col, np.int32), None)
+        else:
+            idx = ((elements*n_fields)[:,:,None]+np.arange(n_fields)[None,None,:]).reshape([elements.shape[0],-1])
+            self.localsize = elements.shape[1]*n_fields
+            self.constraints = list(c[:,0]*n_fields + c[:,1] for c in constraints)
+            self.idx = (elements[:,None,:]*n_fields+np.arange(n_fields)[None,:,None]).reshape([-1])
+            self.csr = CSR(idx, self.idx,self.constraints)
+            self.val = np.zeros(self.csr.col.size)
+            self.mat = PetscObj("MatCreateSeqAIJWithArrays", "MatDestroy", COMM_WORLD,
+                    ctypes.c_int(self.csr.size), ctypes.c_int(self.csr.size), 
+                    _np2c(self.csr.row,np.int32), _np2c(self.csr.col,np.int32), _np2c(self.val))
 
     @timeit
     def local_to_global(self,localv,localm, u, constraints_value = []):
-        self.csr.assemble_mat(localm, constraints_value, self.val)
+        if self.block:
+            rhs = np.bincount(self.idx,localv,self.csr.size*self.n_fields)
+            aptr = ctypes.POINTER(ctypes.c_double)()
+            lib.MatSeqBAIJGetArray(self.mat, ctypes.byref(aptr))
+            a = np.ctypeslib.as_array(aptr,shape=(self.vsize,))
+            a[:] = np.bincount(self.csrmap,localm,self.csr.col.size*self.n_fields**2)
+            lib.MatSeqBAIJRestoreArray(self.mat, ctypes.byref(aptr))
+        else:
+            self.csr.assemble_mat(localm, constraints_value, self.val)
+            rhs = self.csr.assemble_rhs(localv, u , constraints_value)
         lib.MatAssemblyBegin(self.mat,0)
         lib.MatAssemblyEnd(self.mat,0)
-        return self.csr.assemble_rhs(localv, u , constraints_value)
-
-    @timeit
-    def solve(self,rhs) :
-        return self.ksp.solve(self.mat, rhs.reshape([-1])).reshape(rhs.shape)[:self.csr.ndof]
-
-
-class LinearSystemBAIJ :
-
-    def __init__(self, elements, n_fields, options, constraints = []) :
-        if len(constraints) != 0 :
-            raise ValueError("petsc BAIJ not implemented with constraints")
-        nnodes = np.max(elements)+1
-        nn = elements.shape[1]
-        self.ksp = KSP(options)
-        self.csr = CSR(elements,elements,[])
-        self.mat = MatSeqBAIJ(self.csr.size*n_fields, n_fields, self.csr.row, self.csr.col)
-        self.localsize = elements.shape[1]*n_fields
-        self.elements = elements
-        self.idx = (self.elements[:,None,:]*n_fields+np.arange(n_fields)[None,:,None]).reshape([-1])
-        self.size = nnodes*n_fields
-        self.n_fields = n_fields
-        csrmapl = np.arange(nn*nn)[None,:].reshape(nn,nn).T
-        csrmap = (self.csr.map[0]*nn*nn)[:,None,None]+csrmapl[None,:,:]
-        self.csrmap = np.copy(np.swapaxes(csrmap.reshape([elements.shape[0],nn,nn, n_fields,n_fields]),2,3).reshape([-1]))
-        self.val = np.zeros(self.csr.col.size)
-
-    @timeit
-    def local_to_global(self,localv,localm, u, constraints_value = []):
-        rhs = np.bincount(self.idx,localv,self.csr.size*self.n_fields)
-        self.mat.set(np.bincount(self.csrmap,localm,self.csr.col.size*self.n_fields**2))
         return rhs
 
     @timeit
     def solve(self,rhs) :
-        return self.ksp.solve(self.mat, rhs.reshape([-1,self.n_fields])).reshape(rhs.shape)
-
-LinearSystem = LinearSystemAIJ
+        x = np.zeros_like(rhs)
+        #lib.KSPSetReusePreconditioner(self, 0)
+        lib.KSPSetOperators(self.ksp, self.mat, self.mat)
+        if self.block:
+            lib.KSPSolve(self.ksp, Vec(rhs.reshape(-1,self.n_fields)), Vec(x))
+            return x.reshape(rhs.shape)
+        else:
+            lib.KSPSolve(self.ksp, Vec(rhs.reshape(-1)), Vec(x))
+            return x.reshape(rhs.shape)[:self.csr.ndof]