diff --git a/firedrake/__init__.py b/firedrake/__init__.py index 70dec1d6ca..07ff341863 100644 --- a/firedrake/__init__.py +++ b/firedrake/__init__.py @@ -3,7 +3,7 @@ # the specific version, here we are more permissive. This is to catch the # case where users don't update their PETSc for a really long time or # accidentally install a too-new release that isn't yet supported. -PETSC_SUPPORTED_VERSIONS = ">=3.23.0" +PETSC_SUPPORTED_VERSIONS = ">=3.24.2" def init_petsc(): diff --git a/firedrake/dmhooks.py b/firedrake/dmhooks.py index 046852b2e6..41b71661d2 100644 --- a/firedrake/dmhooks.py +++ b/firedrake/dmhooks.py @@ -363,43 +363,36 @@ def create_subdm(dm, fields, *args, **kwargs): :arg fields: The fields in the new sub-DM. """ W = get_function_space(dm) - ctx = get_appctx(dm) - coarsen = get_ctx_coarsener(dm) - parent = get_parent(dm) if len(fields) == 1: # Subspace is just a single FunctionSpace. idx, = fields - subdm = W[idx].dm + subspace = W[idx] iset = W._ises[idx] - add_hook(parent, setup=partial(push_parent, subdm, parent), teardown=partial(pop_parent, subdm, parent), - call_setup=True) - - if ctx is not None: - ctx, = ctx.split([(idx, )]) - add_hook(parent, setup=partial(push_appctx, subdm, ctx), teardown=partial(pop_appctx, subdm, ctx), - call_setup=True) - add_hook(parent, setup=partial(push_ctx_coarsener, subdm, coarsen), teardown=partial(pop_ctx_coarsener, subdm, coarsen), - call_setup=True) - return iset, subdm else: # Need to build an MFS for the subspace subspace = firedrake.MixedFunctionSpace([W[f] for f in fields]) - add_hook(parent, setup=partial(push_parent, subspace.dm, parent), teardown=partial(pop_parent, subspace.dm, parent), - call_setup=True) # Index set mapping from W into subspace. - iset = PETSc.IS().createGeneral(numpy.concatenate([W._ises[f].indices - for f in fields]), + iset = PETSc.IS().createGeneral(numpy.concatenate([W.dof_dset.field_ises[f].indices for f in fields]), comm=W._comm) - if ctx is not None: - ctx, = ctx.split([fields]) - add_hook(parent, setup=partial(push_appctx, subspace.dm, ctx), - teardown=partial(pop_appctx, subspace.dm, ctx), - call_setup=True) - add_hook(parent, setup=partial(push_ctx_coarsener, subspace.dm, coarsen), - teardown=partial(pop_ctx_coarsener, subspace.dm, coarsen), - call_setup=True) - return iset, subspace.dm + + subdm = subspace.dm + parent = get_parent(dm) + add_hook(parent, setup=partial(push_parent, subdm, parent), + teardown=partial(pop_parent, subdm, parent), + call_setup=True) + + ctx = get_appctx(dm) + coarsen = get_ctx_coarsener(dm) + if ctx is not None: + ctx, = ctx.split([fields]) + add_hook(parent, setup=partial(push_appctx, subdm, ctx), + teardown=partial(pop_appctx, subdm, ctx), + call_setup=True) + add_hook(parent, setup=partial(push_ctx_coarsener, subdm, coarsen), + teardown=partial(pop_ctx_coarsener, subdm, coarsen), + call_setup=True) + return iset, subdm @PETSc.Log.EventDecorator() diff --git a/firedrake/formmanipulation.py b/firedrake/formmanipulation.py index eb830493e1..0ee187b8bc 100644 --- a/firedrake/formmanipulation.py +++ b/firedrake/formmanipulation.py @@ -1,9 +1,8 @@ - import numpy import collections from ufl import as_tensor, as_vector, split -from ufl.classes import Zero, FixedIndex, ListTensor, ZeroBaseForm +from ufl.classes import Form, Zero, FixedIndex, ListTensor, ZeroBaseForm from ufl.algorithms.map_integrands import map_integrand_dags from ufl.algorithms import expand_derivatives from ufl.corealg.map_dag import MultiFunction, map_expr_dags @@ -161,27 +160,58 @@ def cofunction(self, o): return Cofunction(W, val=MixedDat(o.dat[i] for i in indices)) def matrix(self, o): + from firedrake.bcs import DirichletBC, EquationBCSplit ises = [] args = [] + argument_indices = [] for a in o.arguments(): V = a.function_space() iset = PETSc.IS() if a.number() in self.blocks: + fields = self.blocks[a.number()] asplit = self._subspace_argument(a) - for f in self.blocks[a.number()]: + for f in fields: fset = V.dof_dset.field_ises[f] iset = iset.expand(fset) else: + fields = tuple(range(len(V))) asplit = a for fset in V.dof_dset.field_ises: iset = iset.expand(fset) ises.append(iset) args.append(asplit) + argument_indices.append(fields) + + if isinstance(o.a, Form): + form = self.split(o.a, argument_indices=argument_indices) + if isinstance(form, ZeroBaseForm): + return form + else: + form = None submat = o.petscmat.createSubMatrix(*ises) - bcs = () - return AssembledMatrix(tuple(args), bcs, submat) + + bcs = [] + spaces = [a.function_space() for a in o.arguments()] + for bc in o.bcs: + W = bc.function_space() + while W.parent is not None: + W = W.parent + + number = spaces.index(W) + V = args[number].function_space() + field = self.blocks[number] + if isinstance(bc, DirichletBC): + bc_temp = bc.reconstruct(field=field, V=V, g=bc.function_arg, use_split=True) + elif isinstance(bc, EquationBCSplit): + row_field, col_field = argument_indices + bc_temp = bc.reconstruct(field=field, V=V, row_field=row_field, col_field=col_field, use_split=True) + bc_temp = None + if bc_temp is not None: + bcs.append(bc_temp) + + return AssembledMatrix(form or tuple(args), tuple(bcs), submat) def zero_base_form(self, o): return ZeroBaseForm(tuple(map(self, o.arguments()))) diff --git a/firedrake/matrix.py b/firedrake/matrix.py index 2f33841289..d2c84451f7 100644 --- a/firedrake/matrix.py +++ b/firedrake/matrix.py @@ -170,9 +170,7 @@ def __init__(self, a, bcs, mat_type, *args, **kwargs): self.mat_type = mat_type def assemble(self): - raise NotImplementedError("API compatibility to apply bcs after 'assemble(a)'\ - has been removed. Use 'assemble(a, bcs=bcs)', which\ - now returns an assembled matrix.") + self.M.assemble() class ImplicitMatrix(MatrixBase): @@ -250,3 +248,9 @@ def __init__(self, a, bcs, petscmat, *args, **kwargs): def mat(self): return self.petscmat + + def assemble(self): + # Bump petsc matrix state by assembling it. + # Ensures that if the matrix changed, the preconditioner is + # updated if necessary. + self.petscmat.assemble() diff --git a/firedrake/solving_utils.py b/firedrake/solving_utils.py index 661f3ae218..523faf38d1 100644 --- a/firedrake/solving_utils.py +++ b/firedrake/solving_utils.py @@ -353,9 +353,17 @@ def split(self, fields): splits = [] problem = self._problem splitter = ExtractSubBlock() + + Fbig = problem.F + # Reuse the submatrices if we are splitting a MatNest + Jbig = self._jac if self.mat_type == "nest" else problem.J + Jpbig = self._pjac if self.pmat_type == "nest" else problem.Jp + if Jpbig is Jbig: + Jpbig = None + for field in fields: - F = splitter.split(problem.F, argument_indices=(field, )) - J = splitter.split(problem.J, argument_indices=(field, field)) + F = splitter.split(Fbig, argument_indices=(field, )) + J = splitter.split(Jbig, argument_indices=(field, field)) us = problem.u_restrict.subfunctions V = F.arguments()[0].function_space() # Exposition: @@ -397,7 +405,6 @@ def split(self, fields): # solving for, and some spaces that have just become # coefficients in the new form. u = as_vector(vec) - J = replace(J, {problem.u_restrict: u}) if problem.is_linear and isinstance(J, MatrixBase): # The BC lifting term is action(MatrixBase, u). # We cannot replace u with the split solution, as action expects a Function. @@ -407,23 +414,35 @@ def split(self, fields): F += problem.compute_bc_lifting(J, subu) else: F = replace(F, {problem.u_restrict: u}) - if problem.Jp is not None: - Jp = splitter.split(problem.Jp, argument_indices=(field, field)) + + J = replace(J, {problem.u_restrict: u}) + if Jpbig is not None: + Jp = splitter.split(Jpbig, argument_indices=(field, field)) Jp = replace(Jp, {problem.u_restrict: u}) else: Jp = None - bcs = [] - for bc in problem.bcs: - if isinstance(bc, DirichletBC): - bc_temp = bc.reconstruct(field=field, V=V, g=bc.function_arg, sub_domain=bc.sub_domain) - elif isinstance(bc, EquationBC): - bc_temp = bc.reconstruct(V, subu, u, field, problem.is_linear) - if bc_temp is not None: - bcs.append(bc_temp) + + if isinstance(J, MatrixBase) and J.has_bcs: + # The BCs of the problem are already encoded in the Jacobian + bcs = None + else: + bcs = [] + for bc in problem.bcs: + if isinstance(bc, DirichletBC): + bc_temp = bc.reconstruct(field=field, V=V, g=bc.function_arg) + elif isinstance(bc, EquationBC): + bc_temp = bc.reconstruct(V, subu, u, field, problem.is_linear) + if bc_temp is not None: + bcs.append(bc_temp) + new_problem = NLVP(F, subu, bcs=bcs, J=J, Jp=Jp, is_linear=problem.is_linear, form_compiler_parameters=problem.form_compiler_parameters) new_problem._constant_jacobian = problem._constant_jacobian - splits.append(type(self)(new_problem, mat_type=self.mat_type, pmat_type=self.pmat_type, + splits.append(type(self)(new_problem, + mat_type=self.mat_type, + pmat_type=self.pmat_type, + sub_mat_type=self.sub_mat_type, + sub_pmat_type=self.sub_pmat_type, appctx=self.appctx, transfer_manager=self.transfer_manager, pre_apply_bcs=self.pre_apply_bcs)) diff --git a/pyproject.toml b/pyproject.toml index 6376531907..06536d0a3c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,7 +30,7 @@ dependencies = [ "loopy>2024.1", "numpy", "packaging", - "petsc4py==3.24.0", + "petsc4py==3.24.2", "petsctools", "pkgconfig", "progress", @@ -150,7 +150,7 @@ requires = [ "pkgconfig", "pybind11", "setuptools>=77.0.3", - "petsc4py==3.24.0", + "petsc4py==3.24.2", "rtree>=1.2", ] build-backend = "setuptools.build_meta" diff --git a/scripts/firedrake-configure b/scripts/firedrake-configure index 1b2ffcb42b..c3d52639af 100755 --- a/scripts/firedrake-configure +++ b/scripts/firedrake-configure @@ -39,7 +39,7 @@ ARCH_DEFAULT = FiredrakeArch.DEFAULT ARCH_COMPLEX = FiredrakeArch.COMPLEX -SUPPORTED_PETSC_VERSION = "v3.24.0" +SUPPORTED_PETSC_VERSION = "v3.24.2" def main(): diff --git a/setup.py b/setup.py index 27015d23cc..76ea9bf8fd 100644 --- a/setup.py +++ b/setup.py @@ -20,7 +20,7 @@ # Ensure that the PETSc getting linked against is compatible -petsctools.init(version_spec=">=3.23.0") +petsctools.init(version_spec=">=3.24.2") import petsc4py diff --git a/tests/firedrake/equation_bcs/test_equation_bcs.py b/tests/firedrake/equation_bcs/test_equation_bcs.py index 53daf97a54..c3151e46e5 100644 --- a/tests/firedrake/equation_bcs/test_equation_bcs.py +++ b/tests/firedrake/equation_bcs/test_equation_bcs.py @@ -296,8 +296,8 @@ def test_EquationBC_mixedpoisson_matrix(eq_type): assert abs(math.log2(err[0][0]) - math.log2(err[1][0]) - (porder+1)) < 0.05 -def test_EquationBC_mixedpoisson_matrix_fieldsplit(): - mat_type = "aij" +@pytest.mark.parametrize("mat_type", ("aij", "nest")) +def test_EquationBC_mixedpoisson_matrix_fieldsplit(mat_type): eq_type = "linear" porder = 0 # Mixed poisson with EquationBCs diff --git a/tests/firedrake/regression/test_nested_fieldsplit_solves.py b/tests/firedrake/regression/test_nested_fieldsplit_solves.py index 35de398216..e061a810ad 100644 --- a/tests/firedrake/regression/test_nested_fieldsplit_solves.py +++ b/tests/firedrake/regression/test_nested_fieldsplit_solves.py @@ -152,6 +152,49 @@ def test_nested_fieldsplit_solve_parallel(W, A, b, expect): assert norm(f) < 1e-11 +@pytest.mark.parametrize("mat_type,pmat_type", [("nest", "nest"), ("matfree", "nest"), ("matfree", "aij")]) +def test_nonlinear_fieldsplit(mat_type, pmat_type): + mesh = UnitIntervalMesh(1) + V = FunctionSpace(mesh, "DG", 0) + Z = V * V * V + + u = Function(Z) + u0, u1, u2 = split(u) + v0, v1, v2 = TestFunctions(Z) + + F = inner(u0, v0) * dx + F += inner(0.5*u1**2 + u1, v1) * dx + F += inner(u2, v2) * dx + u.subfunctions[1].assign(Constant(1)) + + sp = { + "mat_type": mat_type, + "pmat_type": pmat_type, + "snes_max_it": 10, + "ksp_type": "fgmres", + "pc_type": "fieldsplit", + "pc_fieldsplit_type": "additive", + "pc_fieldsplit_0_fields": "0", + "pc_fieldsplit_1_fields": "1,2", + "fieldsplit_1_ksp_view_eigenvalues": None, + "fieldsplit": { + "ksp_type": "gmres", + "pc_type": "jacobi", + }, + } + problem = NonlinearVariationalProblem(F, u) + solver = NonlinearVariationalSolver(problem, solver_parameters=sp) + + def mymonitor(snes, it, fnorm): + if it == 0: + # This call happens before the first linear solve + return + assert np.allclose(snes.ksp.pc.getFieldSplitSubKSP()[1].computeEigenvalues(), 1) + + solver.snes.setMonitor(mymonitor) + solver.solve() + + def test_matrix_types(W): a = inner(TrialFunction(W), TestFunction(W))*dx diff --git a/tests/firedrake/regression/test_split.py b/tests/firedrake/regression/test_split.py index fcf8f221f8..3a49db9c07 100644 --- a/tests/firedrake/regression/test_split.py +++ b/tests/firedrake/regression/test_split.py @@ -91,6 +91,31 @@ def test_assemble_split_mixed_derivative(): assert np.allclose(actual.M.values, expect.M.values) +@pytest.mark.parametrize("mat_type", ("aij", "nest")) +@pytest.mark.parametrize("bcs", (True, False)) +def test_split_assembled_matrix(mat_type, bcs): + mesh = UnitSquareMesh(2, 2) + V = FunctionSpace(mesh, "CG", 1) + Q = FunctionSpace(mesh, "DG", 0) + Z = V * Q + bcs = [DirichletBC(Z.sub(0), 0, "on_boundary")] if bcs else [] + + test = TestFunction(Z) + trial = TrialFunction(Z) + + a = inner(test, trial)*dx + A = assemble(a, bcs=bcs, mat_type=mat_type) + + splitter = ExtractSubBlock() + actual = splitter.split(A, (0, 0)) + + bcs = [bc.reconstruct(V=V) for bc in bcs] + expect = assemble(splitter.split(a, (0, 0)), bcs=bcs) + + expect.petscmat.axpy(-1, actual.petscmat) + assert np.allclose(expect.petscmat[:, :], 0) + + def test_split_coordinate_derivative(): mesh = UnitSquareMesh(1, 1) V = FunctionSpace(mesh, "P", 1)