From a86614a83da9f7353b175b7b93b1efe74b4660eb Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 11 Dec 2024 18:22:10 +0000 Subject: [PATCH 01/16] Restricted Cofunction RHS --- firedrake/assemble.py | 2 +- firedrake/functionspaceimpl.py | 4 ++++ firedrake/variational_solver.py | 8 ++++++-- .../regression/test_restricted_function_space.py | 9 ++++++--- 4 files changed, 17 insertions(+), 6 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index f451b3f596..6b1bb7d3e4 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -1201,7 +1201,7 @@ def _apply_dirichlet_bc(self, tensor, bc): bc.zero(tensor) def _check_tensor(self, tensor): - if tensor.function_space() != self._form.arguments()[0].function_space(): + if tensor.function_space() != self._form.arguments()[0].function_space().dual(): raise ValueError("Form's argument does not match provided result tensor") @staticmethod diff --git a/firedrake/functionspaceimpl.py b/firedrake/functionspaceimpl.py index 8fc81244f7..7e7ecfbcbf 100644 --- a/firedrake/functionspaceimpl.py +++ b/firedrake/functionspaceimpl.py @@ -13,6 +13,7 @@ import ufl import finat.ufl +from ufl.duals import is_dual, is_primal from pyop2 import op2, mpi from pyop2.utils import as_tuple @@ -296,6 +297,9 @@ def restore_work_function(self, function): cache[function] = False def __eq__(self, other): + if is_primal(self) != is_primal(other) or \ + is_dual(self) != is_dual(other): + return False try: return self.topological == other.topological and \ self.mesh() is other.mesh() diff --git a/firedrake/variational_solver.py b/firedrake/variational_solver.py index 4a1ac396c5..488b42d55c 100644 --- a/firedrake/variational_solver.py +++ b/firedrake/variational_solver.py @@ -9,7 +9,7 @@ PETSc, OptionsManager, flatten_parameters, DEFAULT_KSP_PARAMETERS, DEFAULT_SNES_PARAMETERS ) -from firedrake.function import Function +from firedrake.function import Function, Cofunction from firedrake.ufl_expr import TrialFunction, TestFunction from firedrake.bcs import DirichletBC, EquationBC, extract_subdomain_ids, restricted_function_space from firedrake.adjoint_utils import NonlinearVariationalProblemMixin, NonlinearVariationalSolverMixin @@ -92,7 +92,11 @@ def __init__(self, F, u, bcs=None, J=None, self.u_restrict = Function(V_res).interpolate(u) v_res, u_res = TestFunction(V_res), TrialFunction(V_res) F_arg, = F.arguments() - self.F = replace(F, {F_arg: v_res, self.u: self.u_restrict}) + replace_F = {F_arg: v_res, self.u: self.u_restrict} + for c in F.coefficients(): + if c.function_space() == V.dual(): + replace_F[c] = Cofunction(V_res.dual()).interpolate(c) + self.F = replace(F, replace_F) v_arg, u_arg = self.J.arguments() self.J = replace(self.J, {v_arg: v_res, u_arg: u_res, self.u: self.u_restrict}) if self.Jp: diff --git a/tests/firedrake/regression/test_restricted_function_space.py b/tests/firedrake/regression/test_restricted_function_space.py index dc9a2ecc64..6dd093f58a 100644 --- a/tests/firedrake/regression/test_restricted_function_space.py +++ b/tests/firedrake/regression/test_restricted_function_space.py @@ -146,7 +146,8 @@ def test_poisson_inhomogeneous_bcs_2(j): @pytest.mark.parallel(nprocs=3) -def test_poisson_inhomogeneous_bcs_high_level_interface(): +@pytest.mark.parametrize("assembled_rhs", [False, True], ids=("Form", "Cofunction")) +def test_poisson_inhomogeneous_bcs_high_level_interface(assembled_rhs): mesh = UnitSquareMesh(8, 8) V = FunctionSpace(mesh, "CG", 2) bc1 = DirichletBC(V, 0., 1) @@ -155,9 +156,11 @@ def test_poisson_inhomogeneous_bcs_high_level_interface(): v = TestFunction(V) a = inner(grad(u), grad(v)) * dx u = Function(V) - L = inner(Constant(0), v) * dx + L = inner(Constant(-2), v) * dx + if assembled_rhs: + L = assemble(L) solve(a == L, u, bcs=[bc1, bc2], restrict=True) - assert errornorm(SpatialCoordinate(mesh)[0], u) < 1.e-12 + assert errornorm(SpatialCoordinate(mesh)[0]**2, u) < 1.e-12 @pytest.mark.parametrize("j", [1, 2, 5]) From aef788674fd7d405127eff7245b61ec67bd29bfe Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 11 Dec 2024 21:16:30 +0000 Subject: [PATCH 02/16] Fix BCs on Cofunction --- demos/netgen/netgen_mesh.py.rst | 3 +-- .../adjoint_utils/blocks/dirichlet_bc.py | 6 +++--- firedrake/adjoint_utils/blocks/solving.py | 19 +++++++------------ tests/firedrake/regression/test_netgen.py | 6 ++---- 4 files changed, 13 insertions(+), 21 deletions(-) diff --git a/demos/netgen/netgen_mesh.py.rst b/demos/netgen/netgen_mesh.py.rst index 47fe769e70..3c1bd7301b 100755 --- a/demos/netgen/netgen_mesh.py.rst +++ b/demos/netgen/netgen_mesh.py.rst @@ -380,8 +380,7 @@ We will now show how to solve the Poisson problem on a high-order mesh, of order bc = DirichletBC(V, 0.0, [1]) A = assemble(a, bcs=bc) - b = assemble(l) - bc.apply(b) + b = assemble(l, bcs=bc, zero_bc_nodes=True) solve(A, sol, b, solver_parameters={"ksp_type": "cg", "pc_type": "lu"}) VTKFile("output/Sphere.pvd").write(sol) diff --git a/firedrake/adjoint_utils/blocks/dirichlet_bc.py b/firedrake/adjoint_utils/blocks/dirichlet_bc.py index b06d367da1..a918fb8a9e 100644 --- a/firedrake/adjoint_utils/blocks/dirichlet_bc.py +++ b/firedrake/adjoint_utils/blocks/dirichlet_bc.py @@ -51,7 +51,7 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, adj_output = None for adj_input in adj_inputs: if isconstant(c): - adj_value = firedrake.Function(self.parent_space.dual()) + adj_value = firedrake.Function(self.parent_space) adj_input.apply(adj_value) if self.function_space != self.parent_space: vec = extract_bc_subvector( @@ -88,11 +88,11 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, # you can even use the Function outside its domain. # For now we will just assume the FunctionSpace is the same for # the BC and the Function. - adj_value = firedrake.Function(self.parent_space.dual()) + adj_value = firedrake.Function(self.parent_space) adj_input.apply(adj_value) r = extract_bc_subvector( adj_value, c.function_space(), bc - ) + ).riesz_representation("l2") if adj_output is None: adj_output = r else: diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py index e4664665b0..0e6adce8df 100644 --- a/firedrake/adjoint_utils/blocks/solving.py +++ b/firedrake/adjoint_utils/blocks/solving.py @@ -197,14 +197,12 @@ def _assemble_dFdu_adj(self, dFdu_adj_form, **kwargs): def _assemble_and_solve_adj_eq(self, dFdu_adj_form, dJdu, compute_bdy): dJdu_copy = dJdu.copy() - kwargs = self.assemble_kwargs.copy() # Homogenize and apply boundary conditions on adj_dFdu and dJdu. bcs = self._homogenize_bcs() - kwargs["bcs"] = bcs - dFdu = self._assemble_dFdu_adj(dFdu_adj_form, **kwargs) + dFdu = firedrake.assemble(dFdu_adj_form, bcs=bcs, **self.assemble_kwargs) for bc in bcs: - bc.apply(dJdu) + bc.zero(dJdu) adj_sol = firedrake.Function(self.function_space) firedrake.solve( @@ -219,10 +217,8 @@ def _assemble_and_solve_adj_eq(self, dFdu_adj_form, dJdu, compute_bdy): return adj_sol, adj_sol_bdy def _compute_adj_bdy(self, adj_sol, adj_sol_bdy, dFdu_adj_form, dJdu): - adj_sol_bdy = firedrake.Function( - self.function_space.dual(), dJdu.dat - firedrake.assemble( - firedrake.action(dFdu_adj_form, adj_sol)).dat) - return adj_sol_bdy + adj_sol_bdy = firedrake.assemble(dJdu - firedrake.action(dFdu_adj_form, adj_sol)) + return adj_sol_bdy.riesz_representation("l2") def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepared=None): @@ -654,9 +650,8 @@ def _forward_solve(self, lhs, rhs, func, bcs, **kwargs): def _adjoint_solve(self, dJdu, compute_bdy): dJdu_copy = dJdu.copy() # Homogenize and apply boundary conditions on adj_dFdu and dJdu. - bcs = self._homogenize_bcs() - for bc in bcs: - bc.apply(dJdu) + for bc in self.bcs: + bc.zero(dJdu) if ( self._ad_solvers["forward_nlvs"]._problem._constant_jacobian @@ -876,7 +871,7 @@ def __init__(self, source, target_space, target, bcs=[], **kwargs): self.add_dependency(bc, no_duplicates=True) def apply_mixedmass(self, a): - b = firedrake.Function(self.target_space) + b = firedrake.Function(self.target_space.dual()) with a.dat.vec_ro as vsrc, b.dat.vec_wo as vrhs: self.mixed_mass.mult(vsrc, vrhs) return b diff --git a/tests/firedrake/regression/test_netgen.py b/tests/firedrake/regression/test_netgen.py index 904e9a6819..9f401f1675 100644 --- a/tests/firedrake/regression/test_netgen.py +++ b/tests/firedrake/regression/test_netgen.py @@ -51,8 +51,7 @@ def poisson(h, degree=2): # Assembling matrix A = assemble(a, bcs=bc) - b = assemble(l) - bc.apply(b) + b = assemble(l, bcs=bc, zero_bc_nodes=True) # Solving the problem solve(A, u, b, solver_parameters={"ksp_type": "preonly", "pc_type": "lu"}) @@ -95,8 +94,7 @@ def poisson3D(h, degree=2): # Assembling matrix A = assemble(a, bcs=bc) - b = assemble(l) - bc.apply(b) + b = assemble(l, bcs=bc, zero_bc_nodes=True) # Solving the problem solve(A, u, b, solver_parameters={"ksp_type": "preonly", "pc_type": "lu"}) From 8e5603eb0f068c79b404241fdb38a4cf0dde428f Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 11 Dec 2024 23:52:40 +0000 Subject: [PATCH 03/16] LinearSolver: check function spaces --- firedrake/linear_solver.py | 5 +++++ tests/firedrake/regression/test_assemble_baseform.py | 1 + 2 files changed, 6 insertions(+) diff --git a/firedrake/linear_solver.py b/firedrake/linear_solver.py index c1dfbcc07e..7a9f7d807d 100644 --- a/firedrake/linear_solver.py +++ b/firedrake/linear_solver.py @@ -147,6 +147,11 @@ def solve(self, x, b): if not isinstance(b, (function.Function, cofunction.Cofunction)): raise TypeError("Provided RHS is a '%s', not a Function or Cofunction" % type(b).__name__) + if x.function_space() != self.trial_space or b.function_space() != self.test_space.dual(): + # When solving `Ax = b`, with A: V x U -> R, or equivalently A: V -> U*, + # we need to make sure that x and b belong to V and U*, respectively. + raise ValueError("Mismatching function spaces.") + if len(self.trial_space) > 1 and self.nullspace is not None: self.nullspace._apply(self.trial_space.dof_dset.field_ises) if len(self.test_space) > 1 and self.transpose_nullspace is not None: diff --git a/tests/firedrake/regression/test_assemble_baseform.py b/tests/firedrake/regression/test_assemble_baseform.py index 063a33bdd9..219549faa3 100644 --- a/tests/firedrake/regression/test_assemble_baseform.py +++ b/tests/firedrake/regression/test_assemble_baseform.py @@ -155,6 +155,7 @@ def test_zero_form(M, f, one): assert abs(zero_form - 0.5 * np.prod(f.ufl_shape)) < 1.0e-12 +@pytest.mark.xfail(reason="action(M, M) causes primal-dual error") def test_preprocess_form(M, a, f): from ufl.algorithms import expand_indices, expand_derivatives From 2dd4f762c5f06cb3a3d17ac44ce5d8731dfea154 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 12 Dec 2024 12:25:52 +0000 Subject: [PATCH 04/16] assemble(form, zero_bc_nodes=True) as default --- demos/netgen/netgen_mesh.py.rst | 2 +- firedrake/assemble.py | 18 ++++++++---------- firedrake/matrix_free/operators.py | 2 +- firedrake/slate/static_condensation/scpc.py | 2 +- firedrake/solving_utils.py | 2 +- tests/firedrake/multigrid/test_poisson_gmg.py | 7 +++---- tests/firedrake/regression/test_assemble.py | 2 +- .../regression/test_assemble_baseform.py | 2 +- tests/firedrake/regression/test_netgen.py | 4 ++-- 9 files changed, 19 insertions(+), 22 deletions(-) diff --git a/demos/netgen/netgen_mesh.py.rst b/demos/netgen/netgen_mesh.py.rst index 3c1bd7301b..3947cc828b 100755 --- a/demos/netgen/netgen_mesh.py.rst +++ b/demos/netgen/netgen_mesh.py.rst @@ -380,7 +380,7 @@ We will now show how to solve the Poisson problem on a high-order mesh, of order bc = DirichletBC(V, 0.0, [1]) A = assemble(a, bcs=bc) - b = assemble(l, bcs=bc, zero_bc_nodes=True) + b = assemble(l, bcs=bc) solve(A, sol, b, solver_parameters={"ksp_type": "cg", "pc_type": "lu"}) VTKFile("output/Sphere.pvd").write(sol) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 6b1bb7d3e4..db38430eb3 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -83,7 +83,7 @@ def assemble(expr, *args, **kwargs): zero_bc_nodes : bool If `True`, set the boundary condition nodes in the output tensor to zero rather than to the values prescribed by the - boundary condition. Default is `False`. + boundary condition. Default is `True`. diagonal : bool If assembling a matrix is it diagonal? weight : float @@ -156,7 +156,7 @@ def get_assembler(form, *args, **kwargs): return ZeroFormAssembler(form, form_compiler_parameters=fc_params) elif len(form.arguments()) == 1 or diagonal: return OneFormAssembler(form, *args, bcs=bcs, form_compiler_parameters=fc_params, needs_zeroing=kwargs.get('needs_zeroing', True), - zero_bc_nodes=kwargs.get('zero_bc_nodes', False), diagonal=diagonal) + zero_bc_nodes=kwargs.get('zero_bc_nodes', True), diagonal=diagonal) elif len(form.arguments()) == 2: return TwoFormAssembler(form, *args, **kwargs) else: @@ -308,7 +308,7 @@ def __init__(self, sub_mat_type=None, options_prefix=None, appctx=None, - zero_bc_nodes=False, + zero_bc_nodes=True, diagonal=False, weight=1.0, allocation_integral_types=None): @@ -1190,13 +1190,11 @@ def _apply_bc(self, tensor, bc): raise AssertionError def _apply_dirichlet_bc(self, tensor, bc): - if not self._zero_bc_nodes: - tensor_func = tensor.riesz_representation(riesz_map="l2") - if self._diagonal: - bc.set(tensor_func, 1) - else: - bc.apply(tensor_func) - tensor.assign(tensor_func.riesz_representation(riesz_map="l2")) + if self._diagonal: + bc.set(tensor, self._weight) + elif not self._zero_bc_nodes: + # We cannot set primal data on a dual Cofunction, this will throw an error + bc.apply(tensor) else: bc.zero(tensor) diff --git a/firedrake/matrix_free/operators.py b/firedrake/matrix_free/operators.py index 3ee448730e..e31625354a 100644 --- a/firedrake/matrix_free/operators.py +++ b/firedrake/matrix_free/operators.py @@ -147,7 +147,7 @@ def __init__(self, a, row_bcs=[], col_bcs=[], self._assemble_action = get_assembler(self.action, bcs=self.bcs_action, form_compiler_parameters=self.fc_params, - zero_bc_nodes=True).assemble + ).assemble # For assembling action(adjoint(f), self._y) # Sorted list of equation bcs diff --git a/firedrake/slate/static_condensation/scpc.py b/firedrake/slate/static_condensation/scpc.py index 35fa4742eb..c4bda6dc27 100644 --- a/firedrake/slate/static_condensation/scpc.py +++ b/firedrake/slate/static_condensation/scpc.py @@ -102,7 +102,7 @@ def initialize(self, pc): r_expr = reduced_sys.rhs # Construct the condensed right-hand side - self._assemble_Srhs = get_assembler(r_expr, bcs=bcs, zero_bc_nodes=True, form_compiler_parameters=self.cxt.fc_params).assemble + self._assemble_Srhs = get_assembler(r_expr, bcs=bcs, form_compiler_parameters=self.cxt.fc_params).assemble # Allocate and set the condensed operator form_assembler = get_assembler(S_expr, bcs=bcs, form_compiler_parameters=self.cxt.fc_params, mat_type=mat_type, options_prefix=prefix, appctx=self.get_appctx(pc)) diff --git a/firedrake/solving_utils.py b/firedrake/solving_utils.py index 9e843016b5..1da0c0f836 100644 --- a/firedrake/solving_utils.py +++ b/firedrake/solving_utils.py @@ -221,7 +221,7 @@ def __init__(self, problem, mat_type, pmat_type, appctx=None, self._assemble_residual = get_assembler(self.F, bcs=self.bcs_F, form_compiler_parameters=self.fcp, - zero_bc_nodes=True).assemble + ).assemble self._jacobian_assembled = False self._splits = {} diff --git a/tests/firedrake/multigrid/test_poisson_gmg.py b/tests/firedrake/multigrid/test_poisson_gmg.py index 81f56acbfc..577587a393 100644 --- a/tests/firedrake/multigrid/test_poisson_gmg.py +++ b/tests/firedrake/multigrid/test_poisson_gmg.py @@ -195,12 +195,11 @@ def test_baseform_coarsening(solver_type, mixed): a_terms.append(inner(grad(u), grad(v)) * dx) a = sum(a_terms) - assemble_bcs = lambda L: assemble(L, bcs=bcs, zero_bc_nodes=True) # These are equivalent right-hand sides sources = [sum(forms), # purely symbolic linear form - assemble_bcs(sum(forms)), # purely numerical cofunction - sum(assemble_bcs(form) for form in forms), # symbolic combination of numerical cofunctions - forms[0] + assemble_bcs(sum(forms[1:])), # symbolic plus numerical + assemble(sum(forms), bcs=bcs), # purely numerical cofunction + sum(assemble(form, bcs=bcs) for form in forms), # symbolic combination of numerical cofunctions + forms[0] + assemble(sum(forms[1:]), bcs=bcs), # symbolic plus numerical ] solutions = [] for L in sources: diff --git a/tests/firedrake/regression/test_assemble.py b/tests/firedrake/regression/test_assemble.py index bd8f020e60..72394380cb 100644 --- a/tests/firedrake/regression/test_assemble.py +++ b/tests/firedrake/regression/test_assemble.py @@ -225,7 +225,7 @@ def test_one_form_assembler_cache(mesh): assert len(L._cache[_FORM_CACHE_KEY]) == 3 # changing zero_bc_nodes should increase the cache size - assemble(L, zero_bc_nodes=True) + assemble(L, zero_bc_nodes=False) assert len(L._cache[_FORM_CACHE_KEY]) == 4 diff --git a/tests/firedrake/regression/test_assemble_baseform.py b/tests/firedrake/regression/test_assemble_baseform.py index 219549faa3..c6c528aaaa 100644 --- a/tests/firedrake/regression/test_assemble_baseform.py +++ b/tests/firedrake/regression/test_assemble_baseform.py @@ -155,7 +155,7 @@ def test_zero_form(M, f, one): assert abs(zero_form - 0.5 * np.prod(f.ufl_shape)) < 1.0e-12 -@pytest.mark.xfail(reason="action(M, M) causes primal-dual error") +@pytest.mark.xfail(reason="action(M, M) raises primal-dual TypeError") def test_preprocess_form(M, a, f): from ufl.algorithms import expand_indices, expand_derivatives diff --git a/tests/firedrake/regression/test_netgen.py b/tests/firedrake/regression/test_netgen.py index 9f401f1675..8cc164d98c 100644 --- a/tests/firedrake/regression/test_netgen.py +++ b/tests/firedrake/regression/test_netgen.py @@ -51,7 +51,7 @@ def poisson(h, degree=2): # Assembling matrix A = assemble(a, bcs=bc) - b = assemble(l, bcs=bc, zero_bc_nodes=True) + b = assemble(l, bcs=bc) # Solving the problem solve(A, u, b, solver_parameters={"ksp_type": "preonly", "pc_type": "lu"}) @@ -94,7 +94,7 @@ def poisson3D(h, degree=2): # Assembling matrix A = assemble(a, bcs=bc) - b = assemble(l, bcs=bc, zero_bc_nodes=True) + b = assemble(l, bcs=bc) # Solving the problem solve(A, u, b, solver_parameters={"ksp_type": "preonly", "pc_type": "lu"}) From 3c5e64fb2f400870e22694248fe2b997a7398123 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 12 Dec 2024 18:35:48 +0000 Subject: [PATCH 05/16] Fix FunctionAssignBlock --- firedrake/adjoint_utils/blocks/function.py | 1 + 1 file changed, 1 insertion(+) diff --git a/firedrake/adjoint_utils/blocks/function.py b/firedrake/adjoint_utils/blocks/function.py index e31a0c4567..aa035bcf73 100644 --- a/firedrake/adjoint_utils/blocks/function.py +++ b/firedrake/adjoint_utils/blocks/function.py @@ -79,6 +79,7 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, ) diff_expr_assembled = firedrake.Function(adj_input_func.function_space()) diff_expr_assembled.interpolate(ufl.conj(diff_expr)) + diff_expr_assembled = diff_expr_assembled.riesz_representation(riesz_map="l2") adj_output = firedrake.Function( R, val=firedrake.assemble(ufl.Action(diff_expr_assembled, adj_input_func)) ) From e3449f5ccdf90f380711405bd8f5a0b44f323a2a Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 12 Dec 2024 21:00:05 +0000 Subject: [PATCH 06/16] Allow Cofunction.assign take in constants --- firedrake/assemble.py | 9 +++++---- firedrake/cofunction.py | 8 ++++++-- tests/firedrake/regression/test_bcs.py | 2 +- 3 files changed, 12 insertions(+), 7 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index db38430eb3..59ff6b750e 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -406,7 +406,7 @@ def base_form_assembly_visitor(self, expr, tensor, *args): assembler = ZeroFormAssembler(form, form_compiler_parameters=self._form_compiler_params) elif rank == 1 or (rank == 2 and self._diagonal): assembler = OneFormAssembler(form, bcs=self._bcs, form_compiler_parameters=self._form_compiler_params, - zero_bc_nodes=self._zero_bc_nodes, diagonal=self._diagonal) + zero_bc_nodes=self._zero_bc_nodes, diagonal=self._diagonal, weight=self._weight) elif rank == 2: assembler = TwoFormAssembler(form, bcs=self._bcs, form_compiler_parameters=self._form_compiler_params, mat_type=self._mat_type, sub_mat_type=self._sub_mat_type, @@ -1149,14 +1149,15 @@ class OneFormAssembler(ParloopFormAssembler): @classmethod def _cache_key(cls, form, bcs=None, form_compiler_parameters=None, needs_zeroing=True, - zero_bc_nodes=False, diagonal=False): + zero_bc_nodes=False, diagonal=False, weight=1.0): bcs = solving._extract_bcs(bcs) return tuple(bcs), tuplify(form_compiler_parameters), needs_zeroing, zero_bc_nodes, diagonal @FormAssembler._skip_if_initialised def __init__(self, form, bcs=None, form_compiler_parameters=None, needs_zeroing=True, - zero_bc_nodes=False, diagonal=False): + zero_bc_nodes=False, diagonal=False, weight=1.0): super().__init__(form, bcs=bcs, form_compiler_parameters=form_compiler_parameters, needs_zeroing=needs_zeroing) + self._weight = weight self._diagonal = diagonal self._zero_bc_nodes = zero_bc_nodes if self._diagonal and any(isinstance(bc, EquationBCSplit) for bc in self._bcs): @@ -1185,7 +1186,7 @@ def _apply_bc(self, tensor, bc): elif isinstance(bc, EquationBCSplit): bc.zero(tensor) type(self)(bc.f, bcs=bc.bcs, form_compiler_parameters=self._form_compiler_params, needs_zeroing=False, - zero_bc_nodes=self._zero_bc_nodes, diagonal=self._diagonal).assemble(tensor=tensor) + zero_bc_nodes=self._zero_bc_nodes, diagonal=self._diagonal, weight=self._weight).assemble(tensor=tensor) else: raise AssertionError diff --git a/firedrake/cofunction.py b/firedrake/cofunction.py index 4878e6da59..d4b16c0728 100644 --- a/firedrake/cofunction.py +++ b/firedrake/cofunction.py @@ -225,8 +225,12 @@ def assign(self, expr, subset=None, expr_from_assemble=False): return self.assign( assembled_expr, subset=subset, expr_from_assemble=True) - - raise ValueError('Cannot assign %s' % expr) + elif expr == 0: + self.dat.zero(subset=subset) + else: + from firedrake.assign import Assigner + Assigner(self, expr, subset).assign() + return self def riesz_representation(self, riesz_map='L2', **solver_options): """Return the Riesz representation of this :class:`Cofunction` with respect to the given Riesz map. diff --git a/tests/firedrake/regression/test_bcs.py b/tests/firedrake/regression/test_bcs.py index 0e27515890..16cbc669c6 100644 --- a/tests/firedrake/regression/test_bcs.py +++ b/tests/firedrake/regression/test_bcs.py @@ -327,7 +327,7 @@ def test_bcs_rhs_assemble(a, V): b1 = assemble(a) b1_func = b1.riesz_representation(riesz_map="l2") for bc in bcs: - bc.apply(b1_func) + bc.zero(b1_func) b1.assign(b1_func.riesz_representation(riesz_map="l2")) b2 = assemble(a, bcs=bcs) assert np.allclose(b1.dat.data, b2.dat.data) From 40cf6d706653f015824eb5c610d6039dbe889bbc Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 13 Dec 2024 16:59:49 +0000 Subject: [PATCH 07/16] suggestion from code review --- firedrake/assemble.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 59ff6b750e..347f2aca90 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -143,20 +143,18 @@ def get_assembler(form, *args, **kwargs): """ is_base_form_preprocessed = kwargs.pop('is_base_form_preprocessed', False) - bcs = kwargs.get('bcs', None) - fc_params = kwargs.get('form_compiler_parameters', None) if isinstance(form, ufl.form.BaseForm) and not is_base_form_preprocessed: mat_type = kwargs.get('mat_type', None) + fc_params = kwargs.get('form_compiler_parameters', None) # Preprocess the DAG and restructure the DAG # Only pre-process `form` once beforehand to avoid pre-processing for each assembly call form = BaseFormAssembler.preprocess_base_form(form, mat_type=mat_type, form_compiler_parameters=fc_params) if isinstance(form, (ufl.form.Form, slate.TensorBase)) and not BaseFormAssembler.base_form_operands(form): diagonal = kwargs.pop('diagonal', False) if len(form.arguments()) == 0: - return ZeroFormAssembler(form, form_compiler_parameters=fc_params) + return ZeroFormAssembler(form, **kwargs) elif len(form.arguments()) == 1 or diagonal: - return OneFormAssembler(form, *args, bcs=bcs, form_compiler_parameters=fc_params, needs_zeroing=kwargs.get('needs_zeroing', True), - zero_bc_nodes=kwargs.get('zero_bc_nodes', True), diagonal=diagonal) + return OneFormAssembler(form, *args, diagonal=diagonal, **kwargs) elif len(form.arguments()) == 2: return TwoFormAssembler(form, *args, **kwargs) else: @@ -1149,13 +1147,13 @@ class OneFormAssembler(ParloopFormAssembler): @classmethod def _cache_key(cls, form, bcs=None, form_compiler_parameters=None, needs_zeroing=True, - zero_bc_nodes=False, diagonal=False, weight=1.0): + zero_bc_nodes=True, diagonal=False, weight=1.0): bcs = solving._extract_bcs(bcs) - return tuple(bcs), tuplify(form_compiler_parameters), needs_zeroing, zero_bc_nodes, diagonal + return tuple(bcs), tuplify(form_compiler_parameters), needs_zeroing, zero_bc_nodes, diagonal, weight @FormAssembler._skip_if_initialised def __init__(self, form, bcs=None, form_compiler_parameters=None, needs_zeroing=True, - zero_bc_nodes=False, diagonal=False, weight=1.0): + zero_bc_nodes=True, diagonal=False, weight=1.0): super().__init__(form, bcs=bcs, form_compiler_parameters=form_compiler_parameters, needs_zeroing=needs_zeroing) self._weight = weight self._diagonal = diagonal From fe30b4820090b7f22527d1c9f88864ab1dad2547 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 19 Dec 2024 11:17:09 -0600 Subject: [PATCH 08/16] more suggestions from review --- firedrake/assemble.py | 2 +- firedrake/linear_solver.py | 10 ++++++---- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 347f2aca90..875d27862e 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -1192,7 +1192,7 @@ def _apply_dirichlet_bc(self, tensor, bc): if self._diagonal: bc.set(tensor, self._weight) elif not self._zero_bc_nodes: - # We cannot set primal data on a dual Cofunction, this will throw an error + # NOTE this will only work if tensor is a Function and not a Cofunction bc.apply(tensor) else: bc.zero(tensor) diff --git a/firedrake/linear_solver.py b/firedrake/linear_solver.py index 7a9f7d807d..7721675a57 100644 --- a/firedrake/linear_solver.py +++ b/firedrake/linear_solver.py @@ -147,10 +147,12 @@ def solve(self, x, b): if not isinstance(b, (function.Function, cofunction.Cofunction)): raise TypeError("Provided RHS is a '%s', not a Function or Cofunction" % type(b).__name__) - if x.function_space() != self.trial_space or b.function_space() != self.test_space.dual(): - # When solving `Ax = b`, with A: V x U -> R, or equivalently A: V -> U*, - # we need to make sure that x and b belong to V and U*, respectively. - raise ValueError("Mismatching function spaces.") + # When solving `Ax = b`, with A: V x U -> R, or equivalently A: V -> U*, + # we need to make sure that x and b belong to V and U*, respectively. + if x.function_space() != self.trial_space: + raise ValueError(f"x must be a Function in {self.trial_space}.") + if b.function_space() != self.test_space.dual(): + raise ValueError(f"b must be a Cofunction in {self.test_space.dual()}.") if len(self.trial_space) > 1 and self.nullspace is not None: self.nullspace._apply(self.trial_space.dof_dset.field_ises) From 3d49f31ec5967838dd1215c88d0ebedaf9f29c6a Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 19 Dec 2024 14:51:26 -0600 Subject: [PATCH 09/16] remove BaseFormAssembler test --- firedrake/cofunction.py | 2 - .../regression/test_assemble_baseform.py | 48 ++----------------- 2 files changed, 3 insertions(+), 47 deletions(-) diff --git a/firedrake/cofunction.py b/firedrake/cofunction.py index d4b16c0728..436d0b7556 100644 --- a/firedrake/cofunction.py +++ b/firedrake/cofunction.py @@ -225,8 +225,6 @@ def assign(self, expr, subset=None, expr_from_assemble=False): return self.assign( assembled_expr, subset=subset, expr_from_assemble=True) - elif expr == 0: - self.dat.zero(subset=subset) else: from firedrake.assign import Assigner Assigner(self, expr, subset).assign() diff --git a/tests/firedrake/regression/test_assemble_baseform.py b/tests/firedrake/regression/test_assemble_baseform.py index c6c528aaaa..54cc0b183c 100644 --- a/tests/firedrake/regression/test_assemble_baseform.py +++ b/tests/firedrake/regression/test_assemble_baseform.py @@ -1,7 +1,7 @@ import pytest import numpy as np from firedrake import * -from firedrake.assemble import BaseFormAssembler, get_assembler +from firedrake.assemble import get_assembler from firedrake.utils import ScalarType import ufl @@ -43,39 +43,15 @@ def fs(request, mesh): @pytest.fixture def f(fs): f = Function(fs, name="f") - f_split = f.subfunctions x = SpatialCoordinate(fs.mesh())[0] - - # NOTE: interpolation of UFL expressions into mixed - # function spaces is not yet implemented - for fi in f_split: - fs_i = fi.function_space() - if fs_i.rank == 1: - fi.interpolate(as_vector((x,) * fs_i.value_size)) - elif fs_i.rank == 2: - fi.interpolate(as_tensor([[x for i in range(fs_i.mesh().geometric_dimension())] - for j in range(fs_i.rank)])) - else: - fi.interpolate(x) + f.interpolate(as_tensor(np.full(f.ufl_shape, x))) return f @pytest.fixture def one(fs): one = Function(fs, name="one") - ones = one.subfunctions - - # NOTE: interpolation of UFL expressions into mixed - # function spaces is not yet implemented - for fi in ones: - fs_i = fi.function_space() - if fs_i.rank == 1: - fi.interpolate(Constant((1.0,) * fs_i.value_size)) - elif fs_i.rank == 2: - fi.interpolate(Constant([[1.0 for i in range(fs_i.mesh().geometric_dimension())] - for j in range(fs_i.rank)])) - else: - fi.interpolate(Constant(1.0)) + one.interpolate(Constant(np.ones(one.ufl_shape))) return one @@ -155,24 +131,6 @@ def test_zero_form(M, f, one): assert abs(zero_form - 0.5 * np.prod(f.ufl_shape)) < 1.0e-12 -@pytest.mark.xfail(reason="action(M, M) raises primal-dual TypeError") -def test_preprocess_form(M, a, f): - from ufl.algorithms import expand_indices, expand_derivatives - - expr = action(action(M, M), f) - A = BaseFormAssembler.preprocess_base_form(expr) - B = action(expand_derivatives(M), action(M, f)) - - assert isinstance(A, ufl.Action) - try: - # Need to expand indices to be able to match equal (different MultiIndex used for both). - assert expand_indices(A.left()) == expand_indices(B.left()) - assert expand_indices(A.right()) == expand_indices(B.right()) - except KeyError: - # Index expansion doesn't seem to play well with tensor elements. - pass - - def test_tensor_copy(a, M): # 1-form tensor From df04f4b90d962f0047f7b4d88f5e512833cec682 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 19 Dec 2024 19:16:03 -0600 Subject: [PATCH 10/16] only supply relevant kwargs to OneFormAssembler --- firedrake/assemble.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 875d27862e..afd7076114 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -143,18 +143,24 @@ def get_assembler(form, *args, **kwargs): """ is_base_form_preprocessed = kwargs.pop('is_base_form_preprocessed', False) + fc_params = kwargs.get('form_compiler_parameters', None) if isinstance(form, ufl.form.BaseForm) and not is_base_form_preprocessed: mat_type = kwargs.get('mat_type', None) - fc_params = kwargs.get('form_compiler_parameters', None) # Preprocess the DAG and restructure the DAG # Only pre-process `form` once beforehand to avoid pre-processing for each assembly call form = BaseFormAssembler.preprocess_base_form(form, mat_type=mat_type, form_compiler_parameters=fc_params) if isinstance(form, (ufl.form.Form, slate.TensorBase)) and not BaseFormAssembler.base_form_operands(form): diagonal = kwargs.pop('diagonal', False) if len(form.arguments()) == 0: - return ZeroFormAssembler(form, **kwargs) + return ZeroFormAssembler(form, form_compiler_parameters=fc_params) elif len(form.arguments()) == 1 or diagonal: - return OneFormAssembler(form, *args, diagonal=diagonal, **kwargs) + return OneFormAssembler(form, *args, + bcs=kwargs.get("bcs", None), + form_compiler_parameters=fc_params, + needs_zeroing=kwargs.get("needs_zeroing", True), + zero_bc_nodes=kwargs.get("zero_bc_nodes", True), + diagonal=diagonal, + weight=kwargs.get("weight", 1.0)) elif len(form.arguments()) == 2: return TwoFormAssembler(form, *args, **kwargs) else: @@ -1192,7 +1198,7 @@ def _apply_dirichlet_bc(self, tensor, bc): if self._diagonal: bc.set(tensor, self._weight) elif not self._zero_bc_nodes: - # NOTE this will only work if tensor is a Function and not a Cofunction + # NOTE this only works if tensor is a Function and not a Cofunction bc.apply(tensor) else: bc.zero(tensor) From 950e42d1130f8454eecebf07744821c9649507fa Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 20 Dec 2024 21:15:56 -0600 Subject: [PATCH 11/16] Only interpolate the residual, not every cofunction in the RHS --- firedrake/assemble.py | 12 +++++++++--- firedrake/variational_solver.py | 18 +++++++++--------- 2 files changed, 18 insertions(+), 12 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index afd7076114..a1d0009eb1 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -385,6 +385,12 @@ def visitor(e, *operands): visited = {} result = BaseFormAssembler.base_form_postorder_traversal(self._form, visitor, visited) + # Apply BCs after assembly + rank = len(self._form.arguments()) + if rank == 1: + for bc in self._bcs: + bc.zero(result) + if tensor: BaseFormAssembler.update_tensor(result, tensor) return tensor @@ -409,7 +415,7 @@ def base_form_assembly_visitor(self, expr, tensor, *args): if rank == 0: assembler = ZeroFormAssembler(form, form_compiler_parameters=self._form_compiler_params) elif rank == 1 or (rank == 2 and self._diagonal): - assembler = OneFormAssembler(form, bcs=self._bcs, form_compiler_parameters=self._form_compiler_params, + assembler = OneFormAssembler(form, form_compiler_parameters=self._form_compiler_params, zero_bc_nodes=self._zero_bc_nodes, diagonal=self._diagonal, weight=self._weight) elif rank == 2: assembler = TwoFormAssembler(form, bcs=self._bcs, form_compiler_parameters=self._form_compiler_params, @@ -811,9 +817,9 @@ def restructure_base_form(expr, visited=None): return ufl.action(expr, ustar) # -- Case (6) -- # - if isinstance(expr, ufl.FormSum) and all(isinstance(c, ufl.core.base_form_operator.BaseFormOperator) for c in expr.components()): + if isinstance(expr, ufl.FormSum) and all(not isinstance(c, ufl.form.BaseForm) for c in expr.components()): # Return ufl.Sum - return sum([c for c in expr.components()]) + return sum(w*c for w, c in zip(expr.weights(), expr.components())) return expr @staticmethod diff --git a/firedrake/variational_solver.py b/firedrake/variational_solver.py index 488b42d55c..3c8fc8b930 100644 --- a/firedrake/variational_solver.py +++ b/firedrake/variational_solver.py @@ -9,11 +9,12 @@ PETSc, OptionsManager, flatten_parameters, DEFAULT_KSP_PARAMETERS, DEFAULT_SNES_PARAMETERS ) -from firedrake.function import Function, Cofunction -from firedrake.ufl_expr import TrialFunction, TestFunction +from firedrake.function import Function +from firedrake.ufl_expr import TrialFunction, TestFunction, action from firedrake.bcs import DirichletBC, EquationBC, extract_subdomain_ids, restricted_function_space from firedrake.adjoint_utils import NonlinearVariationalProblemMixin, NonlinearVariationalSolverMixin -from ufl import replace +from firedrake.__future__ import interpolate +from ufl import replace, Form __all__ = ["LinearVariationalProblem", "LinearVariationalSolver", @@ -91,12 +92,11 @@ def __init__(self, F, u, bcs=None, J=None, bcs = [bc.reconstruct(V=V_res, indices=bc._indices) for bc in bcs] self.u_restrict = Function(V_res).interpolate(u) v_res, u_res = TestFunction(V_res), TrialFunction(V_res) - F_arg, = F.arguments() - replace_F = {F_arg: v_res, self.u: self.u_restrict} - for c in F.coefficients(): - if c.function_space() == V.dual(): - replace_F[c] = Cofunction(V_res.dual()).interpolate(c) - self.F = replace(F, replace_F) + if isinstance(F, Form): + F_arg, = F.arguments() + self.F = replace(F, {F_arg: v_res, self.u: self.u_restrict}) + else: + self.F = action(replace(F, {self.u: self.u_restrict}), interpolate(v_res, V)) v_arg, u_arg = self.J.arguments() self.J = replace(self.J, {v_arg: v_res, u_arg: u_res, self.u: self.u_restrict}) if self.Jp: From a86f3f56f42847e3bd481af9beb381a19b435b85 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 20 Dec 2024 21:21:27 -0600 Subject: [PATCH 12/16] DROP BEFORE MERGE --- .github/workflows/build.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 0eb616c24d..51bb17c2f1 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -84,6 +84,7 @@ jobs: --install defcon \ --install gadopt \ --install asQ \ + --package-branch ufl pbrubeck/fix/formsum-weights \ || (cat firedrake-install.log && /bin/false) - name: Install test dependencies run: | From 337d0872148a93c31d5f9402b1857f019d21ec7c Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 21 Dec 2024 08:49:32 -0600 Subject: [PATCH 13/16] Fix tests --- firedrake/assemble.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index a1d0009eb1..de7baf4f6d 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -817,8 +817,8 @@ def restructure_base_form(expr, visited=None): return ufl.action(expr, ustar) # -- Case (6) -- # - if isinstance(expr, ufl.FormSum) and all(not isinstance(c, ufl.form.BaseForm) for c in expr.components()): - # Return ufl.Sum + if isinstance(expr, ufl.FormSum) and all(ufl.duals.is_dual(a.function_space()) for a in expr.arguments()): + # Return ufl.Sum if we are assembling a FormSum with Coarguments (a primal expression) return sum(w*c for w, c in zip(expr.weights(), expr.components())) return expr From ed34164e29498a991c158159625a43be2d63c895 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 21 Dec 2024 19:05:28 -0600 Subject: [PATCH 14/16] Fix adjoint utils --- firedrake/adjoint_utils/blocks/solving.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py index 0e6adce8df..2cf6d9fd36 100644 --- a/firedrake/adjoint_utils/blocks/solving.py +++ b/firedrake/adjoint_utils/blocks/solving.py @@ -260,8 +260,11 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, return dFdm dFdm = -firedrake.derivative(F_form, c_rep, trial_function) - dFdm = firedrake.adjoint(dFdm) - dFdm = dFdm * adj_sol + if isinstance(dFdm, ufl.Form): + dFdm = firedrake.adjoint(dFdm) + dFdm = firedrake.action(dFdm, adj_sol) + else: + dFdm = dFdm(adj_sol) dFdm = firedrake.assemble(dFdm, **self.assemble_kwargs) return dFdm From 027ad371a2e53e232202c41b7cf88084dad25a91 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 21 Dec 2024 21:28:25 -0600 Subject: [PATCH 15/16] More robust test for (unrestricted) Cofunction RHS --- .../regression/test_solving_interface.py | 20 ++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/tests/firedrake/regression/test_solving_interface.py b/tests/firedrake/regression/test_solving_interface.py index f32e6f3214..ed47d9e82e 100644 --- a/tests/firedrake/regression/test_solving_interface.py +++ b/tests/firedrake/regression/test_solving_interface.py @@ -222,22 +222,24 @@ def test_constant_jacobian_lvs(): def test_solve_cofunction_rhs(): - mesh = UnitSquareMesh(10, 10) + mesh = UnitIntervalMesh(10) V = FunctionSpace(mesh, "CG", 1) + x, = SpatialCoordinate(mesh) u = TrialFunction(V) v = TestFunction(V) - a = inner(u, v) * dx - + a = inner(grad(u), grad(v)) * dx L = Cofunction(V.dual()) - L.vector()[:] = 1. + bcs = [DirichletBC(V, x, "on_boundary")] + # Set the wrong BCs on the RHS + for bc in bcs: + bc.set(L, 888) + Lold = L.copy() w = Function(V) - solve(a == L, w) - - Aw = assemble(action(a, w)) - assert isinstance(Aw, Cofunction) - assert np.allclose(Aw.dat.data_ro, L.dat.data_ro) + solve(a == L, w, bcs=bcs) + assert errornorm(x, w) < 1E-10 + assert np.allclose(L.dat.data, Lold.dat.data) def test_solve_empty_form_rhs(): From d68113f9721a60c7a9d460b9d83338a4a6ebb5fa Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 3 Jan 2025 11:31:17 -0600 Subject: [PATCH 16/16] set bcs directly on diagonal Cofunction --- firedrake/matrix_free/operators.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/firedrake/matrix_free/operators.py b/firedrake/matrix_free/operators.py index e31625354a..e9ecf63a30 100644 --- a/firedrake/matrix_free/operators.py +++ b/firedrake/matrix_free/operators.py @@ -183,11 +183,9 @@ def _assemble_diagonal(self): def getDiagonal(self, mat, vec): self._assemble_diagonal(tensor=self._diagonal) - diagonal_func = self._diagonal.riesz_representation(riesz_map="l2") for bc in self.bcs: # Operator is identity on boundary nodes - bc.set(diagonal_func, 1) - self._diagonal.assign(diagonal_func.riesz_representation(riesz_map="l2")) + bc.set(self._diagonal, 1) with self._diagonal.dat.vec_ro as v: v.copy(vec)