AuxiliaryRealBlockPC and AuxiliaryComplexBlockPC

asQ/preconditioners/__init__.py

@@ -1,2 +1,3 @@
 from asQ.preconditioners.circulantpc import * # noqa F401
 from asQ.preconditioners.jacobipc import * # noqa F401
+from asQ.preconditioners.auxiliary_blockpc import * # noqa F401

asQ/preconditioners/auxiliary_blockpc.py

@@ -0,0 +1,152 @@
+import firedrake as fd
+from firedrake.petsc import PETSc
+
+from functools import partial
+
+__all__ = ['AuxiliaryRealBlockPC', 'AuxiliaryComplexBlockPC']
+
+
+class AuxiliaryBlockPCBase(fd.AuxiliaryOperatorPC):
+ def _setup(self, pc, v, u):
+ appctx = self.get_appctx(pc)
+ self.appctx = appctx
+
+ self.prefix = pc.getOptionsPrefix() + self._prefix
+ self.options = PETSc.Options(self.prefix)
+
+ self.uref = appctx.get('uref')
+ assert self.uref.function_space() == u.function_space()
+
+ self.bcs = appctx['bcs']
+ self.tref = appctx['tref']
+
+ form_mass = appctx['form_mass']
+ form_function = appctx['form_function']
+
+ self.form_mass = appctx.get('aux_form_mass', form_mass)
+ self.form_function = appctx.get('aux_form_function', form_function)
+
+
+class AuxiliaryRealBlockPC(AuxiliaryBlockPCBase):
+ """
+ A preconditioner for the serial-in-time problem that builds a PC using a specified form.
+
+ This preconditioner is analogous to firedrake.AuxiliaryOperatorPC. Given
+ `form_mass` and `form_function` functions (with the usual call-signatures),
+ it constructs an AuxiliaryOperatorPC.
+
+ Required appctx entries (usually filled by the all-at-once preconditioner):
+
+ 'uref': Firedrake Function around which to linearise the form_function.
+ 'tref': The time at which to linearise the form_function.
+ 'bcs': A list of the boundary conditions.
+ 'form_mass': The function to generate the mass matrix.
+ 'form_function': The function to generate the stiffness matrix.
+ 'dt': The timestep size.
+ 'theta': The implicit parameter.
+
+ Optional appctx entries. Used instead of the required entries if present.
+
+ 'aux_form_mass': Alternative function used to generate the mass matrix.
+ 'aux_form_function': Alternative function used to generate the stiffness matrix.
+
+ PETSc options. Used instead of the appctx entries if present.
+
+ 'aux_dt': <float>
+ Alternative timestep size.
+ 'aux_theta': <float>
+ Alternative implicit theta parameter.
+ """
+ def form(self, pc, v, u):
+ self._setup(pc, v, u)
+
+ dt = self.appctx['dt']
+ theta = self.appctx['theta']
+
+ dt = self.options.getReal('dt', default=dt)
+ theta = self.options.getReal('theta', default=theta)
+
+ us = fd.split(self.uref)
+ vs = fd.split(v)
+
+ dt1 = fd.Constant(1/dt)
+ thet = fd.Constant(theta)
+
+ M = self.form_mass(*fd.split(u), *vs)
+
+ F = self.form_function(*us, *vs, self.tref)
+ K = fd.derivative(F, self.uref)
+
+ a = dt1*M + thet*K
+
+ return (a, self.bcs)
+
+
+class AuxiliaryComplexBlockPC(AuxiliaryBlockPCBase):
+ """
+ A preconditioner for the complex blocks that builds a PC using a specified form.
+
+ This preconditioner is analogous to firedrake.AuxiliaryOperatorPC. Given
+ `form_mass` and `form_function` functions on the real function space (with the
+ usual call-signatures), it constructs an AuxiliaryOperatorPC on the complex
+ block function space.
+
+ Required appctx entries (usually filled by the all-at-once preconditioner).
+
+ 'uref': Firedrake Function around which to linearise the form_function.
+ 'tref': The time at which to linearise the form_function.
+ 'bcs': A list of the boundary conditions on the real space.
+ 'form_mass': The function to generate the mass matrix.
+ 'form_function': The function to generate the stiffness matrix.
+ 'd1': The complex coefficient on the mass matrix.
+ 'd2': The complex coefficient on the stiffness matrix.
+ 'cpx': The complex_proxy submodule to generate the complex-valued forms.
+
+ Optional appctx entries. Used instead of the required entries if present.
+
+ 'aux_form_mass': Alternative function used to generate the mass matrix.
+ 'aux_form_function': Alternative function used to generate the stiffness matrix.
+
+ PETSc options. Used instead of the appctx entries if present.
+
+ 'aux_d1r': <float>
+ Real part of an alternative complex coefficient on the mass matrix.
+ 'aux_d1i': <float>
+ Imaginary part of an alternative complex coefficient on the mass matrix.
+ 'aux_d2r': <float>
+ Real part of an alternative complex coefficient on the stiffness matrix.
+ 'aux_d2i': <float>
+ Imaginary part of an alternative complex coefficient on the stiffness matrix.
+ """
+ def form(self, pc, v, u):
+ self._setup(pc, v, u)
+
+ cpx = self.appctx['cpx']
+
+ d1 = self.appctx['d1']
+ d2 = self.appctx['d2']
+
+ # PETScScalar is real so we can't get a complex number directly from options
+ d1r = self.options.getReal('d1r', default=d1.real)
+ d1i = self.options.getReal('d1i', default=d1.imag)
+
+ d2r = self.options.getReal('d2r', default=d2.real)
+ d2i = self.options.getReal('d2i', default=d2.imag)
+
+ d1 = complex(d1r, d1i)
+ d2 = complex(d2r, d2i)
+
+ # complex and real valued function spaces
+ W = v.function_space()
+ V = W.sub(0)
+
+ bcs = tuple((cb
+ for bc in self.bcs
+ for cb in cpx.DirichletBC(W, V, bc, 0*bc.function_arg)))
+
+ M = cpx.BilinearForm(W, d1, self.form_mass)
+ K = cpx.derivative(d2, partial(self.form_function, t=self.tref), self.uref)
+
+ a = M + K
+
+ return (a, bcs)

asQ/preconditioners/circulantpc.py

@@ -14,61 +14,7 @@
 
 from functools import partial
 
-__all__ = ['CirculantPC', 'AuxiliaryBlockPC']
-
-
-class AuxiliaryBlockPC(fd.AuxiliaryOperatorPC):
- """
- A preconditioner for the complex blocks that builds a PC using a specified form.
-
- This preconditioner is analogous to firedrake.AuxiliaryOperatorPC. Given
- `form_mass` and `form_function` functions on the real function space (with the
- usual call-signatures), it constructs an AuxiliaryOperatorPC on the complex
- block function space.
-
- By default, the circulant eigenvalues and the form_mass and form_function of the
- circulant preconditioner are used (i.e. exactly the same operator as the block).
-
- User-defined `form_mass` and `form_function` functions and complex coeffiecients
- can be passed through the block_appctx using the following keys:
- 'aux_form_mass': function used to build the mass matrix.
- 'aux_form_function': function used to build the stiffness matrix.
- 'aux_%d_d1': complex coefficient on the mass matrix of the %d'th block.
- 'aux_%d_d2': complex coefficient on the stiffness matrix of the %d'th block.
- """
- def form(self, pc, v, u):
- appctx = self.get_appctx(pc)
-
- cpx = appctx['cpx']
-
- u0 = appctx['u0']
- assert u0.function_space() == v.function_space()
-
- bcs = appctx['bcs']
- t0 = appctx['t0']
-
- d1 = appctx['d1']
- d2 = appctx['d2']
-
- blockid = appctx.get('blockid', None)
- blockid_str = f'{blockid}_' if blockid is not None else ''
-
- aux_d1 = appctx.get(f'aux_{blockid_str}d1', d1)
- aux_d2 = appctx.get(f'aux_{blockid_str}d2', d2)
-
- form_mass = appctx['form_mass']
- form_function = appctx['form_function']
-
- aux_form_mass = appctx.get('aux_form_mass', form_mass)
- aux_form_function = appctx.get('aux_form_function', form_function)
-
- Vc = v.function_space()
- M = cpx.BilinearForm(Vc, aux_d1, aux_form_mass)
- K = cpx.derivative(aux_d2, partial(aux_form_function, t=t0), u0)
-
- A = M + K
-
- return (A, bcs)
+__all__ = ['CirculantPC']
 
 
 class CirculantPC(AllAtOnceBlockPCBase):
@@ -130,12 +76,11 @@ class CirculantPC(AllAtOnceBlockPCBase):
  If the AllAtOnceSolver's appctx contains a 'block_appctx' dictionary, this is
  added to the appctx of each block solver. The appctx of each block solver also
  contains the following:
- 'blockid': index of the block.
  'd1': circulant eigenvalue of the mass matrix.
  'd2': circulant eigenvalue of the stiffness matrix.
  'cpx': complex-proxy module implementation set with 'diagfft_complex_proxy'.
- 'u0': state around which the blocks are linearised.
- 't0': time at which the blocks are linearised.
+ 'uref': state around which the blocks are linearised.
+ 'tref': time at which the blocks are linearised.
  'bcs': block boundary conditions.
  'form_mass': function used to build the block mass matrix.
  'form_function': function used to build the block stiffness matrix.
@@ -281,12 +226,11 @@ def initialize(self, pc):
 
  # pass parameters into PC:
  appctx_h = {
- "blockid": i,
  "d1": d1,
  "d2": d2,
  "cpx": cpx,
- "u0": self.u0,
- "t0": self.t_average,
+ "uref": self.u0,
+ "tref": self.t_average,
  "bcs": self.block_bcs,
  "form_mass": self.form_mass,
  "form_function": self.form_function,

asQ/preconditioners/jacobipc.py

@@ -112,11 +112,10 @@ def initialize(self, pc):
 
  # pass parameters into PC:
  appctx_h = {
- "blockid": i,
  "dt": self.dt,
  "theta": self.theta,
- "t0": t0,
- "u0": u0,
+ "tref": t0,
+ "uref": u0,
  "bcs": self.block_bcs,
  "form_mass": self.form_mass,
  "form_function": self.form_function,

examples/complex_proxy/advection.py

@@ -123,12 +123,10 @@ def form_function(q, phi, t=None):
  'converged_reason': None,
  'rtol': rtol,
  },
- 'ksp_type': 'gmres',
+ 'ksp_type': 'preonly',
  'pc_type': 'python',
- 'pc_python_type': 'asQ.AuxiliaryBlockPC',
- 'aux': {
- 'pc_type': 'ilu'
- }
+ 'pc_python_type': 'asQ.AuxiliaryComplexBlockPC',
+ 'aux': lu_pc
 }
 
 # All of these are given to the block solver by the diagpc
@@ -139,8 +137,8 @@ def form_function(q, phi, t=None):
 
 appctx = {
  'cpx': cpx,
- 'u0': w,
- 't0': None,
+ 'uref': w,
+ 'tref': None,
  'd1': D1pc,
  'd2': D2pc,
  'bcs': [],

examples/complex_proxy/heat.py

@@ -0,0 +1,64 @@
+import firedrake as fd
+from asQ.complex_proxy import vector as cpx
+
+import numpy as np
+
+mesh = fd.UnitSquareMesh(4, 4)
+
+V = fd.FunctionSpace(mesh, "CG", 1)
+W = cpx.FunctionSpace(V)
+
+
+def form_mass(u, v):
+ return u*v*fd.dx
+
+
+def form_function(u, v, t=None):
+ return fd.inner(fd.grad(u), fd.grad(v))*fd.dx
+
+
+d1 = 1.5 + 0.4j
+d2 = 0.3 - 0.2j
+
+u = fd.Function(W)
+
+M = cpx.BilinearForm(W, d1, form_mass)
+K = cpx.derivative(d2, form_function, u)
+
+A = M + K
+
+L = fd.Cofunction(W.dual())
+np.random.seed(12345)
+for dat in L.dat:
+ dat.data[:] = np.random.rand(*(dat.data.shape))
+
+sparams = {
+ 'ksp': {
+ 'monitor': None,
+ 'converged_rate': None,
+ },
+ 'ksp_type': 'preonly',
+ 'pc_type': 'python',
+ 'pc_python_type': 'asQ.AuxiliaryComplexBlockPC',
+ 'aux': {
+ 'pc_type': 'lu',
+ 'd1r': d1.real,
+ 'd1i': d1.imag,
+ }
+}
+
+appctx = {
+ 'cpx': cpx,
+ 'uref': u,
+ 'tref': None,
+ 'd1': d1,
+ 'd2': d2,
+ 'bcs': [],
+ 'form_mass': form_mass,
+ 'form_function': form_function
+}
+
+problem = fd.LinearVariationalProblem(A, L, u)
+solver = fd.LinearVariationalSolver(problem, appctx=appctx,
+ solver_parameters=sparams)
+solver.solve()

examples/heat/heat_variable_diffusion.py

@@ -49,9 +49,9 @@ def aux_form_function(u, v, t):
 }
 
 aux_parameters = {
- 'ksp_type': 'richardson',
+ 'ksp_type': 'gmres',
  'pc_type': 'python',
- 'pc_python_type': 'asQ.AuxiliaryBlockPC',
+ 'pc_python_type': 'asQ.AuxiliaryComplexBlockPC',
  'aux': {
  'pc_type': 'ilu'
  }
@@ -68,15 +68,16 @@ def aux_form_function(u, v, t):
  'snes_type': 'ksponly',
  'ksp': {
  'monitor': None,
- 'converged_reason': None,
+ 'converged_rate': None,
  'rtol': 1e-8,
  },
  'mat_type': 'matfree',
- 'ksp_type': 'gmres',
+ 'ksp_type': 'richardson',
+ 'ksp_norm_type': 'unpreconditioned',
  'pc_type': 'python',
  'pc_python_type': 'asQ.CirculantPC',
- 'diagfft_alpha': 1e-3,
- 'diagfft_block': block_parameters
+ 'circulant_alpha': 1e-3,
+ 'circulant_block': block_parameters
 }
 
 appctx = {'block_appctx': block_appctx}