JacobiPC and SliceJacobiPC for all-at-once-system

asQ/__init__.py

@@ -1,6 +1,7 @@
+from asQ.ensemble import *  # noqa: F401
 from asQ.parallel_arrays import *  # noqa: F401
 from asQ.paradiag import *  # noqa: F401
-from asQ.diag_preconditioner import *  # noqa: F401
+from asQ.preconditioners import *  # noqa: F401
 from asQ.allatonce import *  # noqa: F401
 from asQ.post import *  # noqa: F401
 from asQ import complex_proxy  # noqa: F401

asQ/allatonce/jacobian.py

@@ -1,4 +1,6 @@
 import firedrake as fd
+from firedrake.petsc import PETSc
+
 from functools import partial
 
 from asQ.profiling import profiler
@@ -29,29 +31,26 @@ class AllAtOnceJacobian(TimePartitionMixin):
     prefix = "aaos_jacobian_"
 
     @profiler()
-    def __init__(self, aaoform, current_state,
+    def __init__(self, aaoform,
                  reference_state=None,
                  options_prefix="",
                  appctx={}):
         """
         Python context for a PETSc Mat for the Jacobian of an AllAtOnceForm.
 
         :arg aaoform: The AllAtOnceForm object to linearise.
-        :arg current_state: The AllAtOnceFunction being solved for.
         :arg reference_state: A firedrake.Function for a single timestep.
             Only needed if 'aaos_jacobian_state' is 'reference'.
         :arg options_prefix: string prefix for the Jacobian PETSc options.
         :arg appctx: the appcontext for the Jacobian and the preconditioner.
         """
         self._time_partition_setup(aaoform.ensemble, aaoform.time_partition)
-        prefix = self.prefix + options_prefix
+        prefix = options_prefix + self.prefix
 
         aaofunc = aaoform.aaofunc
         self.aaoform = aaoform
         self.aaofunc = aaofunc
 
-        self.current_state = current_state
-
         self.appctx = appctx
 
         # function the Jacobian acts on, and contribution from timestep at end of previous slice
@@ -110,34 +109,32 @@ def update(self, X=None):
         Update the state to linearise around according to aaos_jacobian_state.
 
         :arg X: an optional AllAtOnceFunction or global PETSc Vec.
-            If X is not None and aaos_jacobian_state = 'current' then the state
-            is updated from X instead of self.current_state.
+            If X is not None then the state is updated from X.
         """
 
         aaofunc = self.aaofunc
         jacobian_state = self.jacobian_state()
 
-        if jacobian_state == 'linear':
-            pass
+        if jacobian_state in ('linear', 'user'):
+            return
 
-        elif jacobian_state == 'current':
-            if X is None:
-                X = self.current_state
+        if X is not None:
             self.aaofunc.assign(X)
 
+        if jacobian_state == 'current':
+            return
+
         elif jacobian_state in ('window', 'slice'):
-            time_average(self.current_state, self.ureduce, self.uwrk, average=jacobian_state)
+            time_average(self.aaofunc, self.ureduce, self.uwrk,
+                         average=jacobian_state)
             aaofunc.assign(self.ureduce)
 
         elif jacobian_state == 'initial':
-            aaofunc.assign(self.current_state.initial_condition)
+            aaofunc.assign(self.aaofunc.initial_condition)
 
         elif jacobian_state == 'reference':
             aaofunc.assign(self.reference_state)
 
-        elif jacobian_state == 'user':
-            pass
-
         return
 
     @profiler()
@@ -168,3 +165,16 @@ def mult(self, mat, X, Y):
 
         with self.F.global_vec_ro() as v:
             v.copy(Y)
+
+    @profiler()
+    def petsc_mat(self):
+        """
+        Return a petsc4py.PETSc.Mat with this AllAtOnceJacobian as the python context.
+        """
+        mat = PETSc.Mat().create(comm=self.ensemble.global_comm)
+        mat.setType("python")
+        sizes = (self.aaofunc.nlocal_dofs, self.aaofunc.nglobal_dofs)
+        mat.setSizes((sizes, sizes))
+        mat.setPythonContext(self)
+        mat.setUp()
+        return mat

asQ/allatonce/solver.py

@@ -1,10 +1,11 @@
 from firedrake.petsc import PETSc, OptionsManager, flatten_parameters
 
 from asQ.profiling import profiler
-from asQ.allatonce import AllAtOnceJacobian, AllAtOnceCofunction
+from asQ.allatonce import (AllAtOnceCofunction, AllAtOnceFunction,
+                           AllAtOnceJacobian)
 from asQ.allatonce.mixin import TimePartitionMixin
 
-__all__ = ['AllAtOnceSolver']
+__all__ = ['AllAtOnceSolver', 'LinearSolver']
 
 
 class AllAtOnceSolver(TimePartitionMixin):
@@ -20,7 +21,10 @@ def __init__(self, aaoform, aaofunc,
                  pre_jacobian_callback=lambda solver, X, J: None,
                  post_jacobian_callback=lambda solver, X, J: None):
         """
-        Solves an all-at-once form over an all-at-once function.
+        Solve an all-at-once form over an all-at-once function.
+
+        This is used to solve for a timeseries defined by the all-at-once form
+        and the initial condition in the all-at-once function.
 
         :arg aaoform: the AllAtOnceForm to solve.
         :arg aaofunc: the AllAtOnceFunction solution.
@@ -99,18 +103,11 @@ def assemble_function(snes, X, F):
         # Jacobian
         with self.options.inserted_options():
             self.jacobian = AllAtOnceJacobian(self.jacobian_form,
-                                              current_state=aaofunc,
                                               reference_state=jacobian_reference_state,
                                               options_prefix=options_prefix,
                                               appctx=appctx)
 
-        jacobian_mat = PETSc.Mat().create(comm=self.ensemble.global_comm)
-        jacobian_mat.setType("python")
-        sizes = (aaofunc.nlocal_dofs, aaofunc.nglobal_dofs)
-        jacobian_mat.setSizes((sizes, sizes))
-        jacobian_mat.setPythonContext(self.jacobian)
-        jacobian_mat.setUp()
-        self.jacobian_mat = jacobian_mat
+        self.jacobian_mat = self.jacobian.petsc_mat()
 
         def form_jacobian(snes, X, J, P):
             self.pre_jacobian_callback(self, X, J)
@@ -119,7 +116,9 @@ def form_jacobian(snes, X, J, P):
             J.assemble()
             P.assemble()
 
-        self.snes.setJacobian(form_jacobian, J=jacobian_mat, P=jacobian_mat)
+        self.snes.setJacobian(form_jacobian,
+                              J=self.jacobian_mat,
+                              P=self.jacobian_mat)
 
         # complete the snes setup
         self.options.set_from_options(self.snes)
@@ -140,3 +139,73 @@ def solve(self, rhs=None):
                     raise TypeError(msg)
                 with rhs.global_vec_ro() as rvec:
                     self.snes.solve(rvec, gvec)
+
+
+class LinearSolver(TimePartitionMixin):
+    @profiler()
+    def __init__(self, aaoform,
+                 solver_parameters={},
+                 appctx={},
+                 options_prefix=""):
+        """
+        Solve a linear system where the matrix is an all-at-once Jacobian.
+
+        This does not solve for a timeseries (use an AllAtOnceSolver if this
+        is what you need), but simply uses the AllAtOnceJacobian as the Mat
+        for a KSP.
+
+        :arg aaoform: the AllAtOnceForm to form the Jacobian from.
+        :arg solver_parameters: solver parameters to pass to PETSc.
+            This should be a dict mapping PETSc options to values.
+        :arg appctx: A dictionary containing application context that is
+            passed to the preconditioner if matrix-free.
+        :arg options_prefix: an optional prefix used to distinguish PETSc options.
+            Use this option if you want to pass options to the solver from the
+            command line in addition to through the solver_parameters dict.
+        """
+        self._time_partition_setup(aaoform.ensemble, aaoform.time_partition)
+
+        self.aaoform = aaoform
+        self.appctx = appctx
+
+        # manage options from both dict and command line
+        self.solver_parameters = solver_parameters
+        self.flat_solver_parameters = flatten_parameters(solver_parameters)
+        self.options = OptionsManager(self.flat_solver_parameters, options_prefix)
+        options_prefix = self.options.options_prefix
+
+        # the solver
+        self.ksp = PETSc.KSP().create(comm=self.ensemble.global_comm)
+        self.ksp.setOptionsPrefix(options_prefix)
+
+        # create the all-at-once jacobian
+        with self.options.inserted_options():
+            self.jacobian = AllAtOnceJacobian(aaoform, appctx=appctx,
+                                              options_prefix=options_prefix)
+
+        # create petsc matrix
+        self.jacobian_mat = self.jacobian.petsc_mat()
+
+        # finish setting up the ksp
+        self.ksp.setOperators(self.jacobian_mat)
+        self.options.set_from_options(self.ksp)
+
+    @profiler()
+    def solve(self, b, x):
+        """
+        Solve the all-at-once matrix Ax=b.
+
+        :arg b: AllAtOnceCofunction right hand side vector.
+        :arg x: AllAtOnceFunction solution vector.
+        """
+        if not isinstance(x, AllAtOnceFunction):
+            msg = f"Solution of all-at-once problem must be AllAtOnceFunction not {type(x)}."
+            raise TypeError(msg)
+
+        if not isinstance(b, AllAtOnceCofunction):
+            msg = f"Right hand side of all-at-once problem must be AllAtOnceCofunction not {type(b)}."
+            raise TypeError(msg)
+
+        with x.global_vec() as xvec, b.global_vec_ro() as bvec:
+            with self.options.inserted_options():
+                self.ksp.solve(bvec, xvec)

asQ/ensemble.py

@@ -0,0 +1,82 @@
+from firedrake import COMM_WORLD, Ensemble
+from pyop2.mpi import internal_comm, decref
+
+__all__ = ['create_ensemble', 'split_ensemble', 'EnsembleConnector']
+
+
+def create_ensemble(time_partition, comm=COMM_WORLD):
+    '''
+    Create an Ensemble for the given slice partition.
+    Checks that the number of slices and the size of the communicator are compatible.
+
+    :arg time_partition: a list of integers, the number of timesteps on each time-rank
+    :arg comm: the global communicator for the ensemble
+    '''
+    nslices = 1 if type(time_partition) is int else len(time_partition)
+    nranks = comm.size
+
+    if nranks % nslices != 0:
+        raise ValueError("Number of time slices must be exact factor of number of MPI ranks")
+
+    nspatial_domains = nranks//nslices
+
+    return Ensemble(comm, nspatial_domains)
+
+
+def split_ensemble(ensemble, split_size):
+    """
+    Split an Ensemble into multiple smaller Ensembles which share the same
+    spatial communicators `ensemble.comm`.
+
+    Each smaller Ensemble returned is defined over a contiguous subset of the
+    members of the large Ensemble.
+
+    :arg ensemble: the large Ensemble to split.
+    :arg split_size: the number of members in each smaller Ensemble.
+    """
+    if (ensemble.ensemble_comm.size % split_size) != 0:
+        msg = "Ensemble size must be integer multiple of split_size"
+        raise ValueError(msg)
+
+    # which split are we part of?
+    split_rank = ensemble.ensemble_comm.rank // split_size
+
+    # create split_ensemble.global_comm
+    split_comm = ensemble.global_comm.Split(color=split_rank,
+                                            key=ensemble.global_comm.rank)
+
+    return EnsembleConnector(split_comm, ensemble.comm, split_size)
+
+
+class EnsembleConnector(Ensemble):
+    def __init__(self, global_comm, local_comm, nmembers):
+        """
+        An Ensemble created from provided spatial communicators (ensemble.comm).
+
+        :arg global_comm: global communicator the Ensemble is defined over.
+        :arg local_comm: communicator to use for the Ensemble.comm member.
+        :arg nmembers: number of Ensemble members (ensemble.ensemble_comm.size).
+        """
+        if nmembers*local_comm.size != global_comm.size:
+            msg = "The global ensemble must have the same number of ranks as the sum of the local comms"
+            raise ValueError(msg)
+
+        self.global_comm = global_comm
+        self._global_comm = internal_comm(self.global_comm)
+
+        self.comm = local_comm
+        self._comm = internal_comm(self.comm)
+
+        self.ensemble_comm = self.global_comm.Split(color=self.comm.rank,
+                                                    key=global_comm.rank)
+
+        self._ensemble_comm = internal_comm(self.ensemble_comm)
+
+    def __del__(self):
+        if hasattr(self, "ensemble_comm"):
+            self.ensemble_comm.Free()
+            del self.ensemble_comm
+        for comm_name in ["_global_comm", "_comm", "_ensemble_comm"]:
+            if hasattr(self, comm_name):
+                comm = getattr(self, comm_name)
+                decref(comm)

asQ/paradiag.py

@@ -1,32 +1,11 @@
-import firedrake as fd
 from firedrake.petsc import PETSc
 from asQ.profiling import profiler
 
 from asQ.allatonce import AllAtOnceFunction, AllAtOnceForm, AllAtOnceSolver
 from asQ.allatonce.mixin import TimePartitionMixin
 from asQ.parallel_arrays import SharedArray
 
-__all__ = ['create_ensemble', 'Paradiag']
-
-
-@profiler()
-def create_ensemble(time_partition, comm=fd.COMM_WORLD):
-    '''
-    Create an Ensemble for the given slice partition.
-    Checks that the number of slices and the size of the communicator are compatible.
-
-    :arg time_partition: a list of integers, the number of timesteps on each time-rank
-    :arg comm: the global communicator for the ensemble
-    '''
-    nslices = 1 if type(time_partition) is int else len(time_partition)
-    nranks = comm.size
-
-    if nranks % nslices != 0:
-        raise ValueError("Number of time slices must be exact factor of number of MPI ranks")
-
-    nspatial_domains = nranks//nslices
-
-    return fd.Ensemble(comm, nspatial_domains)
+__all__ = ['Paradiag']
 
 
 class Paradiag(TimePartitionMixin):
@@ -167,9 +146,11 @@ def sync_diagnostics(self):
 
         Until this method is called, diagnostic information is not guaranteed to be valid.
         """
-        pc_block_iterations = self.solver.jacobian.pc.block_iterations
-        pc_block_iterations.synchronise()
-        self.block_iterations.data(deepcopy=False)[:] = pc_block_iterations.data(deepcopy=False)
+        jacobian = self.solver.jacobian
+        if hasattr(jacobian, "pc") and hasattr(jacobian.pc, "block_iterations"):
+            pc_block_iterations = self.solver.jacobian.pc.block_iterations
+            pc_block_iterations.synchronise()
+            self.block_iterations.data(deepcopy=False)[:] = pc_block_iterations.data(deepcopy=False)
 
     @profiler()
     def solve(self,