Dham/adj kwargs

firedrake/adjoint_utils/variational_solver.py

@@ -1,5 +1,6 @@
 import copy
 from functools import wraps
+import warnings
 from pyadjoint.tape import get_working_tape, stop_annotating, annotate_tape, no_annotations
 from firedrake.adjoint_utils.blocks import NonlinearVariationalSolveBlock
 from firedrake.ufl_expr import derivative, adjoint
@@ -73,6 +74,17 @@ def wrapper(self, **kwargs):
                 tape = get_working_tape()
                 problem = self._ad_problem
                 sb_kwargs = NonlinearVariationalSolveBlock.pop_kwargs(kwargs)
+                if (
+                    "adj_kwargs" in self._ad_kwargs
+                    and "adj_kwargs" in sb_kwargs
+                ):
+                    raise TypeError("Cannot pass adj_kwargs to both "
+                                    "solver constructor and solve().")
+                if "adj_kwargs" in sb_kwargs:
+                    warnings.warn(
+                        "Passing adj_kwargs to the solve method has been"
+                        " deprecated, pass it to the"
+                        " NonlinearVariationalProblem instead.", FutureWarning)
                 sb_kwargs.update(kwargs)
 
                 block = NonlinearVariationalSolveBlock(problem._ad_F == 0,

firedrake/variational_solver.py

@@ -194,49 +194,53 @@ def __init__(self, problem, *, solver_parameters=None,
                  post_jacobian_callback=None,
                  pre_function_callback=None,
                  post_function_callback=None,
-                 pre_apply_bcs=True):
+                 pre_apply_bcs=True,
+                 adj_kwargs=None):
         r"""
         :arg problem: A :class:`NonlinearVariationalProblem` to solve.
         :kwarg nullspace: an optional :class:`.VectorSpaceBasis` (or
-               :class:`.MixedVectorSpaceBasis`) spanning the null
-               space of the operator.
-        :kwarg transpose_nullspace: as for the nullspace, but used to
-               make the right hand side consistent.
-        :kwarg near_nullspace: as for the nullspace, but used to
-               specify the near nullspace (for multigrid solvers).
-        :kwarg solver_parameters: Solver parameters to pass to PETSc.
-               This should be a dict mapping PETSc options to values.
-        :kwarg appctx: A dictionary containing application context that
-               is passed to the preconditioner if matrix-free.
-        :kwarg options_prefix: an optional prefix used to distinguish
-               PETSc options.  If not provided a unique prefix will be
-               created.  Use this option if you want to pass options
-               to the solver from the command line in addition to
-               through the ``solver_parameters`` dict.
-        :kwarg pre_jacobian_callback: A user-defined function that will
-               be called immediately before Jacobian assembly. This can
-               be used, for example, to update a coefficient function
-               that has a complicated dependence on the unknown solution.
-        :kwarg post_jacobian_callback: As above, but called after the
-               Jacobian has been assembled.
-        :kwarg pre_function_callback: As above, but called immediately
-               before residual assembly.
-        :kwarg post_function_callback: As above, but called immediately
-               after residual assembly.
+               :class:`.MixedVectorSpaceBasis`) spanning the null space of the
+               operator.
+        :kwarg transpose_nullspace: as for the nullspace, but used to make the
+               right hand side consistent.
+        :kwarg near_nullspace: as for the nullspace, but used to specify the
+               near nullspace (for multigrid solvers).
+        :kwarg solver_parameters: Solver parameters to pass to PETSc. This
+               should be a dict mapping PETSc options to values.
+        :kwarg appctx: A dictionary containing application context that is
+               passed to the preconditioner if matrix-free.
+        :kwarg options_prefix: an optional prefix used to distinguish PETSc
+               options.  If not provided a unique prefix will be created.  Use
+               this option if you want to pass options to the solver from the
+               command line in addition to through the ``solver_parameters``
+               dict.
+        :kwarg pre_jacobian_callback: A user-defined function that will be
+               called immediately before Jacobian assembly. This can be used,
+               for example, to update a coefficient function that has a
+               complicated dependence on the unknown solution.
+        :kwarg post_jacobian_callback: As above, but called after the Jacobian
+               has been assembled.
+        :kwarg pre_function_callback: As above, but called immediately before
+               residual assembly.
+        :kwarg post_function_callback: As above, but called immediately after
+               residual assembly.
         :kwarg pre_apply_bcs: If `True`, the bcs are applied before the solve.
                Otherwise, the problem is linearised around the initial guess
-               before imposing bcs, and the bcs are appended to the nonlinear system.
+               before imposing bcs, and the bcs are appended to the nonlinear
+               system.
+        :kwarg adj_kwargs: A dictionary of keyword arguments for the adjoint
+               LinearVariationalSolver. Arguments not set here default to the
+               values given for the forward solve.
 
-        Example usage of the ``solver_parameters`` option: to set the
-        nonlinear solver type to just use a linear solver, use
+        Example usage of the ``solver_parameters`` option: to set the nonlinear
+        solver type to just use a linear solver, use
 
         .. code-block:: python3
 
             {'snes_type': 'ksponly'}
 
-        PETSc flag options (where the presence of the option means something) should
-        be specified with ``None``.
-        For example:
+        PETSc flag options (where the presence of the option means something)
+        should be specified with ``None``. For example:
 
         .. code-block:: python3
 
@@ -251,12 +255,15 @@ def __init__(self, problem, *, solver_parameters=None,
             def update_diffusivity(current_solution):
                 with cursol.dat.vec_wo as v:
                     current_solution.copy(v)
-                solve(trial*test*dx == dot(grad(cursol), grad(test))*dx, diffusivity)
+                solve(trial*test*dx == dot(grad(cursol), grad(test))*dx,
+                diffusivity)
 
             solver = NonlinearVariationalSolver(problem,
                                                 pre_jacobian_callback=update_diffusivity)
 
         """
+        # Note adj_kwargs is picked up by the wrapper it is only included here
+        # for documentation purposes.
         assert isinstance(problem, NonlinearVariationalProblem)
 
         solver_parameters = flatten_parameters(solver_parameters or {})

tests/firedrake/adjoint/test_adj_kwargs.py

@@ -0,0 +1,118 @@
+import pytest
+import warnings
+
+from firedrake import *
+from firedrake.adjoint import *
+
+
+@pytest.fixture(autouse=True)
+def autouse_set_test_tape(set_test_tape):
+    pass
+
+
+@pytest.fixture()
+def forcing():
+    mesh = UnitSquareMesh(20,20)
+    U = FunctionSpace(mesh, "CG", 1)
+    return Function(U)
+
+
+@pytest.fixture()
+def nvp(forcing):
+    U = forcing.function_space()
+    u = Function(U)
+    v = TestFunction(U)
+
+    F = inner(grad(u), grad(v))*dx - inner(forcing, v)*dx
+
+    bcs = [DirichletBC(U, Constant(1), (4,)),
+        DirichletBC(U, Constant(0), (1, 2, 3))]
+
+    return NonlinearVariationalProblem(F, u, bcs=bcs)
+
+
+# Spike the adjoint solve so we can check that the options are actually
+# set.
+sp_lu = {
+        "pc_type": "none",
+        "ksp_type": "cg",
+        "ksp_max_it": 1,
+        "ksp_view": None
+        }
+
+#@pytest.mark.parametrize("mode", ("solver", "solve", "both", "none", "free"))
+
+
+@pytest.mark.skipcomplex
+def test_adj_kwargs_solver(nvp, forcing):
+    nvs = NonlinearVariationalSolver(nvp,
+                                     adj_kwargs={"solver_parameters": sp_lu})
+    nvs.solve()
+    u = nvp.u
+
+    J = assemble(inner(u, u)*dx)
+    Jhat = ReducedFunctional(J, Control(forcing))
+
+    with pytest.raises(ConvergenceError):
+        Jhat.derivative()
+
+
+@pytest.mark.skipcomplex
+def test_adj_kwargs_none(nvp, forcing):
+    # Don't pass adj_kwargs anywhere. Should succeed.
+    nvs = NonlinearVariationalSolver(nvp)
+    nvs.solve()
+    u=nvp.u
+    J = assemble(inner(u, u)*dx)
+    Jhat = ReducedFunctional(J, Control(forcing))
+    Jhat.derivative()
+
+
+@pytest.mark.skipcomplex
+def test_adj_kwargs_both(nvp):
+    nvs = NonlinearVariationalSolver(nvp,
+                                     adj_kwargs={"solver_parameters": sp_lu})
+    # Passing adj_kwargs to both the solver and solve() is an error.
+    with pytest.raises(TypeError):
+        nvs.solve(adj_kwargs={"solver_parameters": sp_lu})
+
+
+@pytest.mark.skipcomplex
+def test_adj_kwargs_solve(nvp, forcing):
+    nvs = NonlinearVariationalSolver(nvp)
+    with warnings.catch_warnings(record=True) as w:
+        # Cause all warnings to always be triggered.
+        warnings.simplefilter("always")
+        # Trigger a warning.
+        nvs.solve(adj_kwargs={"solver_parameters": sp_lu})
+        # Verify some things
+        assert len(w) == 1
+        assert issubclass(w[-1].category, FutureWarning)
+
+    u = nvp.u
+    J = assemble(inner(u, u)*dx)
+    Jhat = ReducedFunctional(J, Control(forcing))
+
+    with pytest.raises(ConvergenceError):
+        Jhat.derivative()
+
+
+# Unclear why this doesn't work.
+@pytest.mark.xfail
+@pytest.mark.skipcomplex
+def test_adj_kwargs_solve_free_function(nvp, forcing):
+    u = nvp.u
+
+    with warnings.catch_warnings(record=True) as w:
+        # Cause all warnings to always be triggered.
+        warnings.simplefilter("always")
+        # Should not trigger warnings.
+        solve(nvp.F == 0, u, adj_kwargs={"solver_parameters": sp_lu})
+        # Verify no warning raised.
+        assert not w
+
+    J = assemble(inner(u, u)*dx)
+    Jhat = ReducedFunctional(J, Control(forcing))
+
+    with pytest.raises(ConvergenceError):
+        Jhat.derivative()