Start to implement PETSc comms

Addresses usnistgov#644

fipy/solvers/__init__.py

@@ -46,11 +46,18 @@ def __init__(self, solver):
 
 if solver is None and _desired_solver in ["petsc", None]:
     try:
-        if _Nproc > 1:
-            raise  SerialSolverError('petsc')
         from fipy.solvers.petsc import *
         __all__.extend(petsc.__all__)
 
+        from fipy.solvers.petsc.comms.serialPETScCommWrapper import SerialPETScCommWrapper
+        serialComm = SerialPETScCommWrapper()
+
+        if _Nproc > 1:
+            from fipy.solvers.petsc.comms.parallelPETScCommWrapper import ParallelPETScCommWrapper
+            parallelComm = ParallelPETScCommWrapper()
+        else:
+            parallelComm = SerialPETScCommWrapper()
+
         from fipy.matrices.petscMatrix import _PETScMeshMatrix
         _MeshMatrix =  _PETScMeshMatrix
         solver = "petsc"

fipy/solvers/petsc/comms/parallelPETScCommWrapper.py

@@ -4,7 +4,7 @@
  # ###################################################################
  #  FiPy - Python-based finite volume PDE solver
  # 
- #  FILE: "parallelEpetraCommWrapper.py"
+ #  FILE: "parallelPETScCommWrapper.py"
  #
  #  Author: Jonathan Guyer <guyer@nist.gov>
  #  Author: Daniel Wheeler <daniel.wheeler@nist.gov>
@@ -34,27 +34,36 @@
  # ###################################################################
  ##
 
-from PyTrilinos import Epetra
+from petsc4py import PETSc
 from mpi4py import MPI
 
 from fipy.tools import numerix
-from fipy.solvers.trilinos.comms.epetraCommWrapper import EpetraCommWrapper
+from fipy.solvers.petsc.comms.petscCommWrapper import PETScCommWrapper
 
-__all__ = ["ParallelEpetraCommWrapper"]
+__all__ = ["ParallelPETScCommWrapper"]
 
-class ParallelEpetraCommWrapper(EpetraCommWrapper):
+class ParallelPETScCommWrapper(PETScCommWrapper):
     """MPI Communicator wrapper
     
-    Encapsulates capabilities needed for both Epetra and mpi4py.
+    Encapsulates capabilities needed for PETSc.
     """
 
     def __init__(self):
-        self.mpi4py_comm = MPI.COMM_WORLD
-        super(ParallelEpetraCommWrapper, self).__init__()
+        self.petsc4py_comm = PETSc.COMM_WORLD
+        self.mpi4py_comm = PETSc.COMM_WORLD.tompi4py()
+        super(ParallelPETScCommWrapper, self).__init__()
 
-    def __setstate__(self, dict):
-        self.__init__()
+    @property
+    def procID(self):
+        return self.petsc4py_comm.rank
 
+    @property
+    def Nproc(self):
+        return self.petsc4py_comm.size
+
+    def Barrier(self):
+        self.petsc4py_comm.Barrier()
+
     def all(self, a, axis=None):
         return self.mpi4py_comm.allreduce(a.all(axis=axis), op=MPI.LAND)
 
@@ -72,3 +81,12 @@ def bcast(self, obj=None, root=0):
 
     def allgather(self, sendobj=None, recvobj=None):
         return self.mpi4py_comm.allgather(sendobj=sendobj, recvobj=recvobj)
+
+    def sum(self, a, axis=None):
+        return self.mpi4py_comm.allreduce(numerix.array(a).sum(axis=axis), op=MPI.SUM)
+
+    def MaxAll(self, vec):
+        return self.mpi4py_comm.allreduce(max(vec), op=MPI.MAX)
+
+    def MinAll(self, vec):
+        return self.mpi4py_comm.allreduce(min(vec), op=MPI.MIN)

fipy/solvers/petsc/comms/petscCommWrapper.py

@@ -4,7 +4,7 @@
  # ###################################################################
  #  FiPy - Python-based finite volume PDE solver
  # 
- #  FILE: "epetraCommWrapper.py"
+ #  FILE: "petscCommWrapper.py"
  #
  #  Author: Jonathan Guyer <guyer@nist.gov>
  #  Author: Daniel Wheeler <daniel.wheeler@nist.gov>
@@ -36,53 +36,17 @@
 
 __docformat__ = 'restructuredtext'
 
-from PyTrilinos import Epetra
-
 from fipy.tools import numerix
 from fipy.tools.comms.abstractCommWrapper import AbstractCommWrapper
 
-__all__ = ["EpetraCommWrapper"]
+__all__ = ["PETScCommWrapper"]
 
-class EpetraCommWrapper(AbstractCommWrapper):
+class PETScCommWrapper(AbstractCommWrapper):
     """MPI Communicator wrapper
     
-    Encapsulates capabilities needed for Epetra. 
+    Encapsulates capabilities needed for PETSc. 
     Some capabilities are not parallel.
     """
 
-    def __init__(self):
-        self.epetra_comm = Epetra.PyComm()
-        super(EpetraCommWrapper, self).__init__() 
-
-    @property
-    def procID(self):
-        return self.epetra_comm.MyPID()
-
-    @property
-    def Nproc(self):
-        return self.epetra_comm.NumProc()
-
-    def Barrier(self):
-        self.epetra_comm.Barrier()
-
-    def sum(self, a, axis=None):
-        summed = numerix.array(a).sum(axis=axis)
-        shape = summed.shape
-        if shape == ():
-            summed = summed.reshape((1,))
-        parallelSummed = self.epetra_comm.SumAll(summed)
-        if shape == ():
-            parallelSummed = parallelSummed.reshape(())
-        return parallelSummed
-
-    def __setstate__(self, dict):
-        self.__init__()
-
     def Norm2(self, vec):
-        return vec.Norm2()
-
-    def MaxAll(self, vec):
-        return self.epetra_comm.MaxAll(numerix.array(vec))
-
-    def MinAll(self, vec):
-        return self.epetra_comm.MinAll(numerix.array(vec))
+        return vec.norm(norm_type=1)

fipy/solvers/petsc/comms/serialPETScCommWrapper.py

@@ -4,7 +4,7 @@
  # ###################################################################
  #  FiPy - Python-based finite volume PDE solver
  # 
- #  FILE: "serialEpetraCommWrapper.py"
+ #  FILE: "serialPETScCommWrapper.py"
  #
  #  Author: Jonathan Guyer <guyer@nist.gov>
  #  Author: Daniel Wheeler <daniel.wheeler@nist.gov>
@@ -34,20 +34,16 @@
  # ###################################################################
  ##
 
-from fipy.tools import numerix
-from fipy.solvers.trilinos.comms.epetraCommWrapper import EpetraCommWrapper
+from fipy.solvers.petsc.comms.petscCommWrapper import PETScCommWrapper
 
-__all__ = ["SerialEpetraCommWrapper"]
+__all__ = ["SerialPETScCommWrapper"]
 
-class SerialEpetraCommWrapper(EpetraCommWrapper):
+class SerialPETScCommWrapper(PETScCommWrapper):
     @property
     def procID(self):
         return 0
 
     @property
     def Nproc(self):
         return 1
-
-    def Norm2(self, vec):
-        return numerix.L2norm(vec)
 

fipy/tools/comms/abstractCommWrapper.py

@@ -48,7 +48,7 @@ def allgather(self, obj):
         return obj
 
     def sum(self, a, axis=None):
-        raise NotImplementedError
+        return a.sum(axis=axis)
 
     def __getstate__(self):
         return {'dummy': 0}
@@ -57,10 +57,10 @@ def __setstate__(self, dict):
         self.__init__()
 
     def Norm2(self, vec):
-        raise NotImplementedError
+        return numerix.L2norm(vec)
 
     def MaxAll(self, vec):
-        raise NotImplementedError
+        return max(numerix.array(vec))
 
     def MinAll(self, vec):
-        raise NotImplementedError
+        return min(numerix.array(vec))

fipy/tools/comms/dummyComm.py

@@ -14,12 +14,3 @@ def procID(self):
     @property
     def Nproc(self):
         return 1
-
-    def sum(self, a, axis=None):
-        return a.sum(axis=axis)
-
-    def MaxAll(self, vec):
-        return max(numerix.array(vec))
-
-    def MinAll(self, vec):
-        return min(numerix.array(vec))