diff --git a/asQ/preconditioners/__init__.py b/asQ/preconditioners/__init__.py index f51f3d5b..a68b2480 100644 --- a/asQ/preconditioners/__init__.py +++ b/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 diff --git a/asQ/preconditioners/auxiliary_blockpc.py b/asQ/preconditioners/auxiliary_blockpc.py new file mode 100644 index 00000000..4444b970 --- /dev/null +++ b/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': + Alternative timestep size. + 'aux_theta': + 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': + Real part of an alternative complex coefficient on the mass matrix. + 'aux_d1i': + Imaginary part of an alternative complex coefficient on the mass matrix. + 'aux_d2r': + Real part of an alternative complex coefficient on the stiffness matrix. + 'aux_d2i': + 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) diff --git a/asQ/preconditioners/circulantpc.py b/asQ/preconditioners/circulantpc.py index e6cc6272..10475834 100644 --- a/asQ/preconditioners/circulantpc.py +++ b/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, diff --git a/asQ/preconditioners/jacobipc.py b/asQ/preconditioners/jacobipc.py index 4c2e46c6..c4b475c7 100644 --- a/asQ/preconditioners/jacobipc.py +++ b/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, diff --git a/examples/complex_proxy/advection.py b/examples/complex_proxy/advection.py index c0a291a8..bed48394 100644 --- a/examples/complex_proxy/advection.py +++ b/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': [], diff --git a/examples/complex_proxy/heat.py b/examples/complex_proxy/heat.py new file mode 100644 index 00000000..24257d44 --- /dev/null +++ b/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() diff --git a/examples/heat/heat_variable_diffusion.py b/examples/heat/heat_variable_diffusion.py index d27c0c37..c1ad996d 100644 --- a/examples/heat/heat_variable_diffusion.py +++ b/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} diff --git a/tests/preconditioners/test_auxiliarypc.py b/tests/preconditioners/test_auxiliarypc.py new file mode 100644 index 00000000..04d17999 --- /dev/null +++ b/tests/preconditioners/test_auxiliarypc.py @@ -0,0 +1,172 @@ +import firedrake as fd +import numpy as np +import pytest + + +def create_complex_solver(cpx, mesh, V, W, bcs, sparams): + 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) + + cpx_bcs = tuple((cb for bc in bcs + for cb in cpx.DirichletBC(W, V, bc, 0*bc.function_arg))) + + 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)) + + appctx = { + 'cpx': cpx, + 'uref': u, + 'tref': None, + 'd1': d1, + 'd2': d2, + 'bcs': cpx_bcs, + 'form_mass': form_mass, + 'form_function': form_function + } + + problem = fd.LinearVariationalProblem(A, L, u, bcs=cpx_bcs) + solver = fd.LinearVariationalSolver(problem, appctx=appctx, + solver_parameters=sparams) + return solver, u + + +bcopts = ['nobc', 'dirichlet'] +cpxopts = ['cpx_vector', 'cpx_mixed'] + + +@pytest.mark.parametrize("bcopt", bcopts) +@pytest.mark.parametrize("cpxopt", cpxopts) +def test_complex_blockpc(bcopt, cpxopt): + if cpxopt == 'cpx_vector': + from asQ.complex_proxy import vector as cpx + elif cpxopt == 'cpx_mixed': + from asQ.complex_proxy import mixed as cpx + else: + assert False and "complex_proxy type not recognised" + + mesh = fd.UnitSquareMesh(4, 4) + + V = fd.FunctionSpace(mesh, "CG", 1) + W = cpx.FunctionSpace(V) + + if bcopt == 'nobc': + bcs = [] + elif bcopt == 'dirichlet': + bcs = [fd.DirichletBC(V, 0, sub_domain=1)] + else: + assert False and "boundary condition option not recognised" + + pc_type = 'ilu' + + aux_sparams = { + 'ksp_type': 'preonly', + 'pc_type': 'python', + 'pc_python_type': 'asQ.AuxiliaryComplexBlockPC', + 'aux_pc_type': pc_type + } + + direct_sparams = { + 'ksp_type': 'preonly', + 'pc_type': pc_type, + } + + direct_solver, udirect = create_complex_solver(cpx, mesh, V, W, bcs, direct_sparams) + aux_solver, uaux = create_complex_solver(cpx, mesh, V, W, bcs, aux_sparams) + + direct_solver.solve() + aux_solver.solve() + + assert fd.errornorm(udirect, uaux) < 1e-12 + + +def create_real_solver(mesh, V, bcs, sparams): + 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 + + dt = 0.05 + theta = 0.75 + + dt1 = fd.Constant(1/0.05) + thet = fd.Constant(0.75) + + u = fd.Function(V) + v = fd.TestFunction(V) + + M = form_mass(u, v) + K = form_function(u, v) + F = dt1*M + thet*K + A = fd.derivative(F, u) + + L = fd.Cofunction(V.dual()) + np.random.seed(12345) + for dat in L.dat: + dat.data[:] = np.random.rand(*(dat.data.shape)) + + appctx = { + 'uref': u, + 'tref': None, + 'dt': dt, + 'theta': theta, + 'bcs': bcs, + 'form_mass': form_mass, + 'form_function': form_function + } + + problem = fd.LinearVariationalProblem(A, L, u, bcs=bcs) + solver = fd.LinearVariationalSolver(problem, appctx=appctx, + solver_parameters=sparams) + return solver, u + + +@pytest.mark.parametrize("bcopt", bcopts) +def test_real_blockpc(bcopt): + mesh = fd.UnitSquareMesh(4, 4) + + V = fd.FunctionSpace(mesh, "CG", 1) + + if bcopt == 'nobc': + bcs = [] + elif bcopt == 'dirichlet': + bcs = [fd.DirichletBC(V, 0, sub_domain=1)] + else: + assert False and "boundary condition option not recognised" + + pc_type = 'ilu' + + aux_sparams = { + 'ksp_type': 'preonly', + 'pc_type': 'python', + 'pc_python_type': 'asQ.AuxiliaryRealBlockPC', + 'aux_pc_type': pc_type + } + + direct_sparams = { + 'ksp_type': 'preonly', + 'pc_type': pc_type, + } + + direct_solver, udirect = create_real_solver(mesh, V, bcs, direct_sparams) + aux_solver, uaux = create_real_solver(mesh, V, bcs, aux_sparams) + + direct_solver.solve() + aux_solver.solve() + + assert fd.errornorm(udirect, uaux) < 1e-12