firedrakeproject / firedrake

Firedrake is an automated system for the portable solution of partial differential equations using the finite element method (FEM)
https://firedrakeproject.org
Other
495 stars 157 forks source link

Commutative property for adjoint and boundary conditions? #1957

Open salazardetroya opened 3 years ago

salazardetroya commented 3 years ago

I am interacting with TSAdjoint from firedrake. TSAdjoint takes the jacobian of your PDE and calls KSPSolveTranspose to get the adjoint function. The jacobian is already assembled with the boundary conditions. I believe this is messing up with the preconditioners. For instance, in the following mwe, I calculate the adjoint function using solve and using KSPSolveTranspose. With a direct solver, they both return the same solution. With the iterative solver, KSPSolveTranspose does not converge in 100 iterations.

from firedrake import *
import petsc4py

n_elem = 50
m = UnitSquareMesh(n_elem, n_elem)
char_size = 1.0 / 50
mesh = ExtrudedMesh(m, n_elem, layer_height=char_size)
V = FunctionSpace(mesh, "CG", 1)
u, v = TrialFunction(V), TestFunction(V)

a = inner(grad(u), grad(v)) * dx + u * v * dx
L = Constant(1.0) * v * dx
bcs = DirichletBC(V, Constant(1e3), "bottom")

u_sol = Function(V)
direct_parameters = {
    "ksp_type": "preonly",
    "pc_type": "lu",
    "pc_factor_mat_solver_type": "mumps",
    "mat_mumps_icntl_14": 200,
    "mat_mumps_icntl_24": 1,
}
iterative_parameters = {
    "ksp_type": "gmres",
    "ksp_converged_reason": None,
    "ksp_rtol": 1e-7,
    "ksp_max_it": 100,
    "snes_rtol": 1e-7,
    "snes_monitor": None,
    "pc_type": "hypre",
    "pc_hypre_type": "boomeramg",
    "pc_hypre_boomeramg_max_iter": 5,
    "pc_hypre_boomeramg_coarsen_type": "PMIS",
    "pc_hypre_boomeramg_agg_nl": 2,
    "pc_hypre_boomeramg_strong_threshold": 0.95,
    "pc_hypre_boomeramg_interp_type": "ext+i",
    "pc_hypre_boomeramg_P_max": 2,
    "pc_hypre_boomeramg_relax_type_all": "sequential-Gauss-Seidel",
    "pc_hypre_boomeramg_grid_sweeps_all": 1,
    "pc_hypre_boomeramg_truncfactor": 0.3,
    "pc_hypre_boomeramg_max_levels": 6,
}

problem = LinearVariationalProblem(a, L, u_sol, bcs=bcs)
solver = LinearVariationalSolver(problem, solver_parameters=iterative_parameters)
solver.solve()

J = u_sol * u_sol * dx

lmbda = Function(V)
dJdu = derivative(J, u_sol)
solve(
    adjoint(a) == dJdu,
    lmbda,
    bcs=homogenize(bcs),
    solver_parameters=iterative_parameters,
)
File("lmbd.pvd").write(lmbda)

bcs_hom = homogenize(bcs)
dJdu_vec = assemble(dJdu)
bcs_hom.apply(dJdu_vec)

with dJdu_vec.dat.vec_ro as b:
    with lmbda.dat.vec as x:
        solver.snes.ksp.solveTranspose(b, x)

File("lmbd_ksp_transpose.pvd").write(lmbda)

It seems that the main difference is that in the first case, the adjoint/transpose operation precedes the application of the boundary conditions (I believe), whereas the opposite happens in the second case.

Is there any way to hand KSPSolveTranspose() a jacobian or solver parameters to make it converge reasonably? Bear in mind that KSPSolveTranspose() is deep in the TSAdjoint implementation. I do not know how easy it will be to replace with something that looks like the first adjoint solve.

salazardetroya commented 3 years ago

Waiting to hear from PETSc as the issue might be on their end (or hypre's)

wence- commented 3 years ago

This sounds like another reason (possibly) to condense constraints out. Let me try and scope what would be involved in an issue.