From e265f0a77ef12124d2f9558324343eb96d99e1d1 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Fri, 21 Mar 2025 12:09:59 +0000 Subject: [PATCH 1/9] initial python type FieldsplitSNES implementation --- firedrake/preconditioners/__init__.py | 1 + firedrake/preconditioners/fieldsplit_snes.py | 82 ++++++++++++++++++++ 2 files changed, 83 insertions(+) create mode 100644 firedrake/preconditioners/fieldsplit_snes.py diff --git a/firedrake/preconditioners/__init__.py b/firedrake/preconditioners/__init__.py index 491a73657b..1a74065aa1 100644 --- a/firedrake/preconditioners/__init__.py +++ b/firedrake/preconditioners/__init__.py @@ -13,3 +13,4 @@ from firedrake.preconditioners.hiptmair import * # noqa: F401 from firedrake.preconditioners.facet_split import * # noqa: F401 from firedrake.preconditioners.bddc import * # noqa: F401 +from firedrake.preconditioners.fieldsplit_snes import * # noqa: F401 diff --git a/firedrake/preconditioners/fieldsplit_snes.py b/firedrake/preconditioners/fieldsplit_snes.py new file mode 100644 index 0000000000..579318594f --- /dev/null +++ b/firedrake/preconditioners/fieldsplit_snes.py @@ -0,0 +1,82 @@ +from firedrake.preconditioners.base import SNESBase +from firedrake.petsc import PETSc +from firedrake.dmhooks import get_appctx +from firedrake.function import Function + +__all__ = ("FieldsplitSNES",) + + +class FieldsplitSNES(SNESBase): + prefix = "fieldsplit_" + + # TODO: + # - Allow setting field grouping/ordering like fieldsplit + + def initialize(self, snes): + from firedrake.variational_solver import NonlinearVariationalSolver # ImportError if we do this at file level + ctx = get_appctx(snes.dm) + self.sol = ctx._problem.u_restrict + + # buffer to save solution to outer problem during solve + self.sol_outer = Function(self.sol.function_space()) + + # buffers for shuffling solutions during solve + self.sol_current = Function(self.sol.function_space()) + self.sol_new = Function(self.sol.function_space()) + + # options for setting up the fieldsplit + snes_prefix = snes.getOptionsPrefix() + 'snes_' + self.prefix + # options for each field + sub_prefix = snes.getOptionsPrefix() + self.prefix + + snes_options = PETSc.Options(snes_prefix) + self.fieldsplit_type = snes_options.getString('type', 'additive') + if self.fieldsplit_type not in ('additive', 'multiplicative'): + raise ValueError( + 'FieldsplitSNES option snes_fieldsplit_type must be' + ' "additive" or "multiplicative"') + + split_ctxs = ctx.split([(i,) for i in range(len(self.sol))]) + + self.solvers = tuple( + NonlinearVariationalSolver( + ctx._problem, appctx=ctx.appctx, + options_prefix=sub_prefix+str(i)) + for i, ctx in enumerate(split_ctxs) + ) + + def update(self, snes): + pass + + def step(self, snes, x, f, y): + # store current value of outer solution + self.sol_outer.assign(self.sol) + + # the full form in ctx now has the most up to date solution + with self.sol_current.dat.vec_wo as vec: + x.copy(vec) + self.sol.assign(self.sol_current) + + # The current snes solution x is held in sol_current, and we + # will place the new solution in sol_new. + # The solvers evaluate forms containing sol, so for each + # splitting type sol needs to hold: + # - additive: all fields need to hold sol_current values + # - multiplicative: fields need to hold sol_current before + # they are are solved for, and keep the updated sol_new + # values afterwards. + for solver, u, ucurr, unew in zip(self.solvers, + self.sol.subfunctions, + self.sol_current.subfunctions, + self.sol_new.subfunctions): + solver.solve() + unew.assign(u) + if self.fieldsplit_type == 'additive': + u.assign(ucurr) + + with self.sol_new.dat.vec_ro as vec: + vec.copy(y) + y.aypx(-1, x) + + # restore outer solution + self.sol.assign(self.sol_outer) From 136eb7bbd1baca628a5ecd1ebffd513ba8611c1a Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Fri, 21 Mar 2025 12:23:56 +0000 Subject: [PATCH 2/9] initial test to check python type FieldsplitSNES doesn't crash --- .../regression/test_fieldsplit_snes.py | 95 +++++++++++++++++++ 1 file changed, 95 insertions(+) create mode 100644 tests/firedrake/regression/test_fieldsplit_snes.py diff --git a/tests/firedrake/regression/test_fieldsplit_snes.py b/tests/firedrake/regression/test_fieldsplit_snes.py new file mode 100644 index 0000000000..d489fea805 --- /dev/null +++ b/tests/firedrake/regression/test_fieldsplit_snes.py @@ -0,0 +1,95 @@ +from firedrake import * + + +def test_fieldsplit_snes(): + re = Constant(100) + nu = Constant(1/re) + + nx = 50 + dt = Constant(0.1) # CFL = dt*nx + + mesh = PeriodicUnitIntervalMesh(nx) + x, = SpatialCoordinate(mesh) + + Vu = VectorFunctionSpace(mesh, "CG", 2) + Vq = FunctionSpace(mesh, "DG", 1) + W = Vu*Vq + + w0 = Function(W) + u0, q0 = w0.subfunctions + u0.project(as_vector([0.5 + 1.0*sin(2*pi*x)])) + q0.interpolate(cos(2*pi*x)) + + def M(u, v): + return inner(u, v)*dx + + def Aburgers(u, v, nu): + return ( + inner(dot(u, nabla_grad(u)), v)*dx + + nu*inner(grad(u), grad(v))*dx + ) + + def Ascalar(q, p, u): + n = FacetNormal(mesh) + un = 0.5*(dot(u, n) + abs(dot(u, n))) + return (- q*div(p*u)*dx + + jump(p)*jump(un*q)*dS) + + # current and next timestep + w = Function(W) + wn = Function(W) + + u, q = split(w) + un, qn = split(wn) + + v, p = TestFunctions(W) + + # Trapezium rule + F = ( + M(un - u, v) + 0.5*dt*(Aburgers(un, v, nu) + Aburgers(u, v, nu)) + + M(qn - q, p) + 0.5*dt*(Ascalar(qn, p, un) + Ascalar(q, p, u)) + ) + + common_params = { + 'snes_converged_reason': None, + 'snes_monitor': None, + 'snes_rtol': 1e-8, + 'snes_atol': 1e-12, + 'ksp_converged_reason': None, + 'ksp_monitor': None, + } + + newton_params = { + 'snes_type': 'newtonls', + 'mat_type': 'aij', + 'ksp_type': 'preonly', + 'pc_type': 'lu', + } + + uparams = common_params | newton_params + qparams = common_params | newton_params | {'snes_type': 'ksponly'} + + python_params = { + 'snes_type': 'nrichardson', + 'npc_snes_type': 'python', + 'npc_snes_python_type': 'firedrake.FieldsplitSNES', + 'npc_snes_fieldsplit_type': 'additive', + 'npc_fieldsplit_0': uparams, + 'npc_fieldsplit_1': qparams, + } + + params = common_params | python_params + + w.assign(w0) + wn.assign(w0) + u, q = w.subfunctions + un, qn = wn.subfunctions + solver = NonlinearVariationalSolver( + NonlinearVariationalProblem(F, wn), + solver_parameters=params, + options_prefix="") + + nsteps = 2 + for i in range(nsteps): + w.assign(wn) + solver.solve() From 3645598fc4d3a9c1c73280593c82a477aaede411 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Wed, 16 Apr 2025 11:52:14 -0600 Subject: [PATCH 3/9] snesfieldsplit fix for vector function spaces --- firedrake/preconditioners/fieldsplit_snes.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/firedrake/preconditioners/fieldsplit_snes.py b/firedrake/preconditioners/fieldsplit_snes.py index 579318594f..0a514cb29b 100644 --- a/firedrake/preconditioners/fieldsplit_snes.py +++ b/firedrake/preconditioners/fieldsplit_snes.py @@ -1,6 +1,6 @@ from firedrake.preconditioners.base import SNESBase from firedrake.petsc import PETSc -from firedrake.dmhooks import get_appctx +from firedrake.dmhooks import get_appctx, get_function_space from firedrake.function import Function __all__ = ("FieldsplitSNES",) @@ -15,6 +15,7 @@ class FieldsplitSNES(SNESBase): def initialize(self, snes): from firedrake.variational_solver import NonlinearVariationalSolver # ImportError if we do this at file level ctx = get_appctx(snes.dm) + W = get_function_space(snes.dm) self.sol = ctx._problem.u_restrict # buffer to save solution to outer problem during solve @@ -36,7 +37,7 @@ def initialize(self, snes): 'FieldsplitSNES option snes_fieldsplit_type must be' ' "additive" or "multiplicative"') - split_ctxs = ctx.split([(i,) for i in range(len(self.sol))]) + split_ctxs = ctx.split([(i,) for i in range(len(W))]) self.solvers = tuple( NonlinearVariationalSolver( From e0e042540dfb4db51ffdc5c81e291d0475cc3c47 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Wed, 16 Apr 2025 15:45:18 -0600 Subject: [PATCH 4/9] Update tests/firedrake/regression/test_fieldsplit_snes.py --- tests/firedrake/regression/test_fieldsplit_snes.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/firedrake/regression/test_fieldsplit_snes.py b/tests/firedrake/regression/test_fieldsplit_snes.py index d489fea805..2078699294 100644 --- a/tests/firedrake/regression/test_fieldsplit_snes.py +++ b/tests/firedrake/regression/test_fieldsplit_snes.py @@ -32,8 +32,8 @@ def Aburgers(u, v, nu): def Ascalar(q, p, u): n = FacetNormal(mesh) un = 0.5*(dot(u, n) + abs(dot(u, n))) - return (- q*div(p*u)*dx - + jump(p)*jump(un*q)*dS) + return (- q*div(u*p)*dx + + jump(un*q)*jump(p)*dS) # current and next timestep w = Function(W) From 72aa2c186cfbe999d795d64300b224ca87e4cf85 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Mon, 12 May 2025 22:48:01 +0100 Subject: [PATCH 5/9] profile fieldsplit snes methods --- firedrake/preconditioners/fieldsplit_snes.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/firedrake/preconditioners/fieldsplit_snes.py b/firedrake/preconditioners/fieldsplit_snes.py index 0a514cb29b..179ce8aa47 100644 --- a/firedrake/preconditioners/fieldsplit_snes.py +++ b/firedrake/preconditioners/fieldsplit_snes.py @@ -12,6 +12,7 @@ class FieldsplitSNES(SNESBase): # TODO: # - Allow setting field grouping/ordering like fieldsplit + @PETSc.Log.EventDecorator() def initialize(self, snes): from firedrake.variational_solver import NonlinearVariationalSolver # ImportError if we do this at file level ctx = get_appctx(snes.dm) @@ -49,6 +50,7 @@ def initialize(self, snes): def update(self, snes): pass + @PETSc.Log.EventDecorator() def step(self, snes, x, f, y): # store current value of outer solution self.sol_outer.assign(self.sol) From b118805c446a6dd5059e33ae0eade66e8785afd3 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Mon, 12 May 2025 22:48:18 +0100 Subject: [PATCH 6/9] initial auxiliary snes impl --- firedrake/preconditioners/auxiliary_snes.py | 48 +++++++++++++++++++++ 1 file changed, 48 insertions(+) create mode 100644 firedrake/preconditioners/auxiliary_snes.py diff --git a/firedrake/preconditioners/auxiliary_snes.py b/firedrake/preconditioners/auxiliary_snes.py new file mode 100644 index 0000000000..a80aefe591 --- /dev/null +++ b/firedrake/preconditioners/auxiliary_snes.py @@ -0,0 +1,48 @@ +from firedrake.preconditioners.base import SNESBase +from firedrake.petsc import PETSc +from firedrake.dmhooks import get_appctx, get_function_space +from firedrake.function import Function, TestFunction + +__all__ = ("AuxiliaryOperatorSNES",) + + +class AuxiliaryOperatorSNES(SNESBase): + prefix = "aux_" + + @PETSc.Log.EventDecorator() + def initialize(self, snes): + from firedrake.variational_solver import ( # ImportError if this is at file level + NonlinearVariationalSolver, NonlinearVariationalProblem) + + appctx = get_appctx(snes.dm).appctx + V = get_function_space(snes.dm).collapse() + + fcp = appctx.get("form_compiler_parameters") + + self.u = Function(V) + v = TestFunction(V) + + F, bcs, J, Jp = self.form(snes, self.u, v) + + prefix = snes.getOptionsPrefix() + self.prefix + + self.solver = NonlinearVariationalSolver( + NonlinearVariationalProblem( + F, self.u, bcs=bcs, J=J, Jp=Jp, + form_compiler_parameters=fcp), + appctx=appctx, options_prefix=prefix) + + def update(self, snes): + pass + + @PETSc.Log.EventDecorator() + def step(self, snes, x, f, y): + with self.u.dat.vec_wo as vec: + x.copy(vec) + self.solver.solve() + with self.u.dat.vec_ro as vec: + vec.copy(y) + y.aypx(-1, x) + + def form(self, snes, u, v): + raise NotImplementedError From cba2545eef521e9a2cad8d9ef7dec83d8a1a71e4 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Tue, 13 May 2025 09:48:51 +0100 Subject: [PATCH 7/9] import AuxiliaryOperatorSNES in preconditioners.__init__ --- firedrake/preconditioners/__init__.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/firedrake/preconditioners/__init__.py b/firedrake/preconditioners/__init__.py index 1a74065aa1..aa9ba1f3d7 100644 --- a/firedrake/preconditioners/__init__.py +++ b/firedrake/preconditioners/__init__.py @@ -13,4 +13,5 @@ from firedrake.preconditioners.hiptmair import * # noqa: F401 from firedrake.preconditioners.facet_split import * # noqa: F401 from firedrake.preconditioners.bddc import * # noqa: F401 -from firedrake.preconditioners.fieldsplit_snes import * # noqa: F401 +from firedrake.preconditioners.fieldsplit_snes import * # noqa: F401 +from firedrake.preconditioners.auxiliary_snes import * # noqa: F401 From dbba932c96acc937c32718de3c486ba7e54b63ea Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Wed, 14 May 2025 10:24:57 +0100 Subject: [PATCH 8/9] auxsnes form update --- firedrake/preconditioners/auxiliary_snes.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/firedrake/preconditioners/auxiliary_snes.py b/firedrake/preconditioners/auxiliary_snes.py index a80aefe591..89815c8345 100644 --- a/firedrake/preconditioners/auxiliary_snes.py +++ b/firedrake/preconditioners/auxiliary_snes.py @@ -1,7 +1,6 @@ from firedrake.preconditioners.base import SNESBase from firedrake.petsc import PETSc from firedrake.dmhooks import get_appctx, get_function_space -from firedrake.function import Function, TestFunction __all__ = ("AuxiliaryOperatorSNES",) @@ -11,8 +10,10 @@ class AuxiliaryOperatorSNES(SNESBase): @PETSc.Log.EventDecorator() def initialize(self, snes): - from firedrake.variational_solver import ( # ImportError if this is at file level - NonlinearVariationalSolver, NonlinearVariationalProblem) + from firedrake import ( # ImportError if this is at file level + NonlinearVariationalSolver, + NonlinearVariationalProblem, + Function, TestFunction) appctx = get_appctx(snes.dm).appctx V = get_function_space(snes.dm).collapse() @@ -22,13 +23,13 @@ def initialize(self, snes): self.u = Function(V) v = TestFunction(V) - F, bcs, J, Jp = self.form(snes, self.u, v) + F, bcs, self.u = self.form(snes, self.u, v) prefix = snes.getOptionsPrefix() + self.prefix self.solver = NonlinearVariationalSolver( NonlinearVariationalProblem( - F, self.u, bcs=bcs, J=J, Jp=Jp, + F, self.u, bcs=bcs, form_compiler_parameters=fcp), appctx=appctx, options_prefix=prefix) From c34c20fed802a179f5ddcb5d8d53e1f23369ad5f Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Thu, 15 May 2025 12:11:20 +0100 Subject: [PATCH 9/9] auxiliary snes wip --- firedrake/preconditioners/auxiliary_snes.py | 38 ++++++- .../regression/test_fieldsplit_snes.py | 105 ++++++++++++++++++ 2 files changed, 139 insertions(+), 4 deletions(-) diff --git a/firedrake/preconditioners/auxiliary_snes.py b/firedrake/preconditioners/auxiliary_snes.py index 89815c8345..21b50beb02 100644 --- a/firedrake/preconditioners/auxiliary_snes.py +++ b/firedrake/preconditioners/auxiliary_snes.py @@ -13,17 +13,21 @@ def initialize(self, snes): from firedrake import ( # ImportError if this is at file level NonlinearVariationalSolver, NonlinearVariationalProblem, - Function, TestFunction) + Function, TestFunction, Cofunction) - appctx = get_appctx(snes.dm).appctx + ctx = get_appctx(snes.dm) V = get_function_space(snes.dm).collapse() + appctx = ctx.appctx fcp = appctx.get("form_compiler_parameters") - self.u = Function(V) + u = Function(V) v = TestFunction(V) - F, bcs, self.u = self.form(snes, self.u, v) + F, bcs, self.u = self.form(snes, u, v) + + self.b = Cofunction(V.dual()) + F += self.b prefix = snes.getOptionsPrefix() + self.prefix @@ -32,18 +36,44 @@ def initialize(self, snes): F, self.u, bcs=bcs, form_compiler_parameters=fcp), appctx=appctx, options_prefix=prefix) + outer_snes = snes + inner_snes = self.solver.snes + inner_snes.incrementTabLevel(1, parent=outer_snes) + inner_snes.ksp.incrementTabLevel(1, parent=outer_snes) + inner_snes.ksp.pc.incrementTabLevel(1, parent=outer_snes) def update(self, snes): pass @PETSc.Log.EventDecorator() def step(self, snes, x, f, y): + from firedrake import errornorm with self.u.dat.vec_wo as vec: x.copy(vec) + # PETSc.Sys.Print(f"{x.norm() = }") + if f is not None: + with self.b.dat.vec_wo as vec: + f.copy(vec) + else: + self.b.zero() + # self.b.zero() + + # PETSc.Sys.Print(f"Before: {errornorm(self.un, self.u) = :.5e}") + PETSc.Sys.Print(f"Before: {errornorm(self.un1, self.u) = :.5e}") self.solver.solve() + # PETSc.Sys.Print(f"After: {errornorm(self.un, self.u) = :.5e}") + PETSc.Sys.Print(f"After: {errornorm(self.un1, self.u) = :.5e}") with self.u.dat.vec_ro as vec: + # PETSc.Sys.Print(f"{vec.norm() = }") vec.copy(y) y.aypx(-1, x) + # PETSc.Sys.Print(f"{y.norm() = }") def form(self, snes, u, v): raise NotImplementedError + + def view(self, snes, viewer=None): + super().view(snes, viewer) + if hasattr(self, "solver"): + viewer.printfASCII("SNES to apply auxiliary inverse\n") + self.solver.snes.view(viewer) diff --git a/tests/firedrake/regression/test_fieldsplit_snes.py b/tests/firedrake/regression/test_fieldsplit_snes.py index 2078699294..14f9530dd1 100644 --- a/tests/firedrake/regression/test_fieldsplit_snes.py +++ b/tests/firedrake/regression/test_fieldsplit_snes.py @@ -93,3 +93,108 @@ def Ascalar(q, p, u): for i in range(nsteps): w.assign(wn) solver.solve() + + +def M(u, v): + return inner(u, v)*dx + + +def A(u, v, nu): + return ( + inner(dot(u, nabla_grad(u)), v)*dx + + nu*inner(grad(u), grad(v))*dx + ) + +class AuxiliaryBurgersSNES(AuxiliaryOperatorSNES): + def form(self, snes, u, v): + appctx = self.get_appctx(snes) + nu = appctx["nu"] + dt = appctx["dt"] + un = appctx["un"] + un1 = appctx["un1"] + uh = (u + un)/2 + F = M(u - un, v) + dt*A(uh, v, nu) + self.un = un + self.un1 = un1 + return F, None, u + + +def test_auxiliary_snes(): + re = Constant(100) + re_aux = Constant(50) + + nu = Constant(1/re) + nu_aux = Constant(1/re_aux) + + nx = 50 + dt = Constant(0.1) # CFL = dt*nx + + mesh = PeriodicUnitIntervalMesh(nx) + x, = SpatialCoordinate(mesh) + + V = VectorFunctionSpace(mesh, "CG", 2) + + # current and next timestep + ic = as_vector([1.0 + 0.5*sin(2*pi*x)]) + un = Function(V).project(ic) + un1 = Function(V).project(ic) + + v = TestFunction(V) + + # Implicit midpoint rule + uh = (un + un1)/2 + F = M(un1 - un, v) + dt*A(uh, v, nu) + + solver_parameters = { + 'snes': { + 'view': ':snes_view.log', + 'converged_reason': None, + 'monitor': None, + 'rtol': 1e-8, + 'atol': 0, + 'max_it': 3, + 'convergence_test': 'skip', + 'linesearch_type': 'l2', + 'linesearch_damping': 1.0, + 'linesearch_monitor': None, + }, + 'snes_type': 'nrichardson', + 'npc_snes_type': 'python', + 'npc_snes_python_type': f'{__name__}.AuxiliaryBurgersSNES', + 'npc_aux': { + 'snes': { + 'converged_reason': None, + 'monitor': None, + 'rtol': 1e-4, + 'atol': 0, + 'max_it': 2, + 'convergence_test': 'skip', + }, + 'snes_type': 'newtonls', + 'mat_type': 'aij', + 'ksp_type': 'preonly', + 'pc_type': 'lu', + 'pc_factor_mat_solver_type': 'petsc', + }, + } + + appctx = { + "nu": nu_aux, + "dt": dt, + "un": un, + "un1": un1, + } + + solver = NonlinearVariationalSolver( + NonlinearVariationalProblem(F, un1), + solver_parameters=solver_parameters, + options_prefix="fd", appctx=appctx) + + nsteps = 1 + for i in range(nsteps): + solver.solve() + un.assign(un1) + + +if __name__ == "__main__": + test_auxiliary_snes()