From 22ea2c454a0389548ca52266c97b3f75ef56b8fc Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 24 Jun 2025 15:13:25 +0100 Subject: [PATCH 01/72] Support mat_type="is" --- firedrake/assemble.py | 2 +- pyop2/types/mat.py | 23 ++++++++++++++++------- 2 files changed, 17 insertions(+), 8 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 9ca4d17237..32743d2c93 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -1328,7 +1328,7 @@ def _get_mat_type(mat_type, sub_mat_type, arguments): for arg in arguments for V in arg.function_space()): mat_type = "nest" - if mat_type not in {"matfree", "aij", "baij", "nest", "dense"}: + if mat_type not in {"matfree", "aij", "baij", "nest", "dense", "is"}: raise ValueError(f"Unrecognised matrix type, '{mat_type}'") if sub_mat_type is None: sub_mat_type = parameters.parameters["default_sub_matrix_type"] diff --git a/pyop2/types/mat.py b/pyop2/types/mat.py index d0fb9e2404..5aef11aea1 100644 --- a/pyop2/types/mat.py +++ b/pyop2/types/mat.py @@ -618,11 +618,16 @@ def _init_monolithic(self): rset, cset = self.sparsity.dsets rlgmap = rset.unblocked_lgmap clgmap = cset.unblocked_lgmap - mat.createAIJ(size=((self.nrows, None), (self.ncols, None)), - nnz=(self.sparsity.nnz, self.sparsity.onnz), - bsize=1, - comm=self.comm) + if self.mat_type == "is": + create = mat.createIS + elif self.mat_type.endswith("aij"): + create = mat.createAIJ + else: + raise ValueError(f"Unsupported mat_type {mat_type}") + size = ((self.nrows, None), (self.ncols, None)) + create(size, bsize=1, comm=self.comm) mat.setLGMap(rmap=rlgmap, cmap=clgmap) + mat.setPreallocationNNZ((self.sparsity.nnz, self.sparsity.onnzi)) self.handle = mat self._blocks = [] rows, cols = self.sparsity.shape @@ -685,7 +690,10 @@ def _init_block(self): col_lg = cset.lgmap rdim, cdim = self.dims[0][0] - if rdim == cdim and rdim > 1 and self.sparsity._block_sparse: + if self.mat_type == "is": + block_sparse = False + create = mat.createIS + elif rdim == cdim and rdim > 1 and self.sparsity._block_sparse: # Size is total number of rows and columns, but the # /sparsity/ is the block sparsity. block_sparse = True @@ -697,10 +705,10 @@ def _init_block(self): create = mat.createAIJ create(size=((self.nrows, None), (self.ncols, None)), - nnz=(self.sparsity.nnz, self.sparsity.onnz), bsize=(rdim, cdim), comm=self.comm) mat.setLGMap(rmap=row_lg, cmap=col_lg) + mat.setPreallocationNNZ((self.sparsity.nnz, self.sparsity.onnz)) # Stash entries destined for other processors mat.setOption(mat.Option.IGNORE_OFF_PROC_ENTRIES, False) # Any add or insertion that would generate a new entry that has not @@ -716,7 +724,8 @@ def _init_block(self): mat.setOption(mat.Option.KEEP_NONZERO_PATTERN, True) # We completely fill the allocated matrix when zeroing the # entries, so raise an error if we "missed" one. - mat.setOption(mat.Option.UNUSED_NONZERO_LOCATION_ERR, True) + if self.mat_type != "is": + mat.setOption(mat.Option.UNUSED_NONZERO_LOCATION_ERR, True) # Put zeros in all the places we might eventually put a value. with profiling.timed_region("MatZeroInitial"): sparsity.fill_with_zeros(mat, self.sparsity.dims[0][0], From 9042f4502083e3c60d1696431e7967a5da728c31 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 24 Jun 2025 17:04:02 +0100 Subject: [PATCH 02/72] MatIS: support DirichletBC, add a test --- pyop2/parloop.py | 6 ++++-- pyop2/types/mat.py | 10 +++++++--- tests/firedrake/regression/test_assemble.py | 19 +++++++++++++++++++ 3 files changed, 30 insertions(+), 5 deletions(-) diff --git a/pyop2/parloop.py b/pyop2/parloop.py index c70f4c9fb7..6d208570b1 100644 --- a/pyop2/parloop.py +++ b/pyop2/parloop.py @@ -296,7 +296,8 @@ def replace_lgmaps(self): olgmaps = [] for m, lgmaps in zip(pl_arg.data, pl_arg.lgmaps): olgmaps.append(m.handle.getLGMap()) - m.handle.setLGMap(*lgmaps) + if m.handle.type != "is": + m.handle.setLGMap(*lgmaps) orig_lgmaps.append(olgmaps) return tuple(orig_lgmaps) @@ -309,7 +310,8 @@ def restore_lgmaps(self, orig_lgmaps): for arg, d in reversed(list(zip(self.global_kernel.arguments, self.arguments))): if isinstance(arg, (MatKernelArg, MixedMatKernelArg)) and d.lgmaps is not None: for m, lgmaps in zip(d.data, orig_lgmaps.pop()): - m.handle.setLGMap(*lgmaps) + if m.handle.type != "is": + m.handle.setLGMap(*lgmaps) @cached_property def _has_mats(self): diff --git a/pyop2/types/mat.py b/pyop2/types/mat.py index 5aef11aea1..0f60b2b469 100644 --- a/pyop2/types/mat.py +++ b/pyop2/types/mat.py @@ -825,9 +825,13 @@ def set_local_diagonal_entries(self, rows, diag_val=1.0, idx=None): rows = rows.reshape(-1, 1) self.change_assembly_state(Mat.INSERT_VALUES) if len(rows) > 0: - values = np.full(rows.shape, diag_val, dtype=dtypes.ScalarType) - self.handle.setValuesLocalRCV(rows, rows, values, - addv=PETSc.InsertMode.INSERT_VALUES) + if self.handle.type == "is": + self.handle.assemble() + self.handle.zeroRowsColumnsLocal(rows, diag_val) + else: + values = np.full(rows.shape, diag_val, dtype=dtypes.ScalarType) + self.handle.setValuesLocalRCV(rows, rows, values, + addv=PETSc.InsertMode.INSERT_VALUES) @mpi.collective def assemble(self): diff --git a/tests/firedrake/regression/test_assemble.py b/tests/firedrake/regression/test_assemble.py index 94e3439c6e..73471f53f0 100644 --- a/tests/firedrake/regression/test_assemble.py +++ b/tests/firedrake/regression/test_assemble.py @@ -114,6 +114,25 @@ def test_mat_nest_real_block_assembler_correctly_reuses_tensor(mesh): assert A2.M is A1.M +@pytest.mark.parametrize("dirichlet_bcs", [False, True]) +def test_assemble_matis(mesh, dirichlet_bcs): + V = FunctionSpace(mesh, "CG", 1) + u = TrialFunction(V) + v = TestFunction(V) + a = inner(grad(u), grad(v))*dx + if dirichlet_bcs: + bcs = DirichletBC(V, 0, (1, 3)) + else: + bcs = None + + ais = assemble(a, bcs=bcs, mat_type="is").petscmat + aijnew = PETSc.Mat() + ais.convert("aij", aijnew) + + aij = assemble(a, bcs=bcs, mat_type="aij").petscmat + aij.axpy(-1, aijnew) + aij.view() + assert np.allclose(aij[:, :], 0) def test_assemble_diagonal(mesh): V = FunctionSpace(mesh, "P", 3) From 642d18b629a1f694d28b528c0e4aa7c94be387ee Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 24 Jun 2025 23:42:46 +0100 Subject: [PATCH 03/72] MatIS: Fix parallel LGMap --- firedrake/assemble.py | 43 +++++++++++++++++++++ pyop2/types/mat.py | 19 +++++---- tests/firedrake/regression/test_assemble.py | 3 +- 3 files changed, 56 insertions(+), 9 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 32743d2c93..ecab839d66 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -1399,6 +1399,10 @@ def _make_sparsity(test, trial, mat_type, sub_mat_type, maps_and_regions): except SparsityFormatError: raise ValueError("Monolithic matrix assembly not supported for systems " "with R-space blocks") + if mat_type == "is": + rlgmap = unghosted_lgmap(sparsity._dsets[0].lgmap, test.function_space()) + clgmap = unghosted_lgmap(sparsity._dsets[1].lgmap, trial.function_space()) + sparsity._lgmaps = (rlgmap, clgmap) return sparsity def _make_maps_and_regions(self): @@ -2210,3 +2214,42 @@ def index_function_spaces(form, indices): return tuple(a.ufl_function_space()[i] for i, a in zip(indices, form.arguments())) else: raise AssertionError + + +def masked_lgmap(lgmap, mask, block=True): + if block: + indices = lgmap.block_indices.copy() + bsize = lgmap.getBlockSize() + else: + indices = lgmap.indices.copy() + bsize = 1 + indices[mask] = -1 + return PETSc.LGMap().create(indices=indices, bsize=bsize, comm=lgmap.comm) + + +def unghosted_lgmap(lgmap, V, block=True): + block_size = lgmap.getBlockSize() if block else 1 + + mesh_dm = V.mesh().topology_dm + depth = mesh_dm.getDepth() + start, end = mesh_dm.getDepthStratum(depth) + + own = [] + for W in V: + section = W.dm.getDefaultSection() + for seed in range(start, end): + # Do not loop over ghost cells + if mesh_dm.getLabelValue("pyop2_ghost", seed) != -1: + continue + closure, _ = mesh_dm.getTransitiveClosure(seed, useCone=True) + for p in closure: + dof = section.getDof(p) + if dof <= 0: + continue + off = section.getOffset(p) + # Local indices within W + W_indices = range(block_size * off, block_size * (off + dof)) + own.extend(W_indices) + + mask = numpy.setdiff1d(range(len(lgmap.indices)), own) + return masked_lgmap(lgmap, mask, block=block) diff --git a/pyop2/types/mat.py b/pyop2/types/mat.py index 0f60b2b469..7eb00aa057 100644 --- a/pyop2/types/mat.py +++ b/pyop2/types/mat.py @@ -58,6 +58,7 @@ def __init__(self, dsets, maps_and_regions, name=None, nest=None, block_sparse=N if self._initialized: return self._dsets = dsets + self._lgmaps = None self._maps_and_regions = maps_and_regions self._block_sparse = block_sparse self._diagonal_block = diagonal_block @@ -619,15 +620,15 @@ def _init_monolithic(self): rlgmap = rset.unblocked_lgmap clgmap = cset.unblocked_lgmap if self.mat_type == "is": + # TODO monolithic lgmaps + rlgmap, clgmap = self.sparsity._lgmaps create = mat.createIS - elif self.mat_type.endswith("aij"): - create = mat.createAIJ else: - raise ValueError(f"Unsupported mat_type {mat_type}") + create = mat.createAIJ size = ((self.nrows, None), (self.ncols, None)) create(size, bsize=1, comm=self.comm) mat.setLGMap(rmap=rlgmap, cmap=clgmap) - mat.setPreallocationNNZ((self.sparsity.nnz, self.sparsity.onnzi)) + mat.setPreallocationNNZ((self.sparsity.nnz, self.sparsity.onnz)) self.handle = mat self._blocks = [] rows, cols = self.sparsity.shape @@ -691,6 +692,7 @@ def _init_block(self): rdim, cdim = self.dims[0][0] if self.mat_type == "is": + row_lg, col_lg = self.sparsity._lgmaps block_sparse = False create = mat.createIS elif rdim == cdim and rdim > 1 and self.sparsity._block_sparse: @@ -703,10 +705,9 @@ def _init_block(self): # the /dof/ sparsity. block_sparse = False create = mat.createAIJ - create(size=((self.nrows, None), - (self.ncols, None)), - bsize=(rdim, cdim), - comm=self.comm) + size = ((self.nrows, None), (self.ncols, None)) + create(size, bsize=(rdim, cdim), comm=self.comm) + mat.setLGMap(rmap=row_lg, cmap=col_lg) mat.setPreallocationNNZ((self.sparsity.nnz, self.sparsity.onnz)) # Stash entries destined for other processors @@ -826,6 +827,8 @@ def set_local_diagonal_entries(self, rows, diag_val=1.0, idx=None): self.change_assembly_state(Mat.INSERT_VALUES) if len(rows) > 0: if self.handle.type == "is": + own = self.handle.getLocalSize()[0] + rows = rows[rows < own] self.handle.assemble() self.handle.zeroRowsColumnsLocal(rows, diag_val) else: diff --git a/tests/firedrake/regression/test_assemble.py b/tests/firedrake/regression/test_assemble.py index 73471f53f0..a2c8c48125 100644 --- a/tests/firedrake/regression/test_assemble.py +++ b/tests/firedrake/regression/test_assemble.py @@ -114,6 +114,7 @@ def test_mat_nest_real_block_assembler_correctly_reuses_tensor(mesh): assert A2.M is A1.M + @pytest.mark.parametrize("dirichlet_bcs", [False, True]) def test_assemble_matis(mesh, dirichlet_bcs): V = FunctionSpace(mesh, "CG", 1) @@ -131,9 +132,9 @@ def test_assemble_matis(mesh, dirichlet_bcs): aij = assemble(a, bcs=bcs, mat_type="aij").petscmat aij.axpy(-1, aijnew) - aij.view() assert np.allclose(aij[:, :], 0) + def test_assemble_diagonal(mesh): V = FunctionSpace(mesh, "P", 3) u = TrialFunction(V) From 2d06f08a8d1fb3ea90a4f5c6d135b43d9e84c62f Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 26 Jun 2025 08:01:58 +0100 Subject: [PATCH 04/72] MatNest + MatIS --- firedrake/assemble.py | 14 ++++---- pyop2/types/mat.py | 39 +++++++++++++-------- tests/firedrake/regression/test_assemble.py | 14 +++++--- 3 files changed, 43 insertions(+), 24 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index ecab839d66..355ae963ee 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -1332,8 +1332,8 @@ def _get_mat_type(mat_type, sub_mat_type, arguments): raise ValueError(f"Unrecognised matrix type, '{mat_type}'") if sub_mat_type is None: sub_mat_type = parameters.parameters["default_sub_matrix_type"] - if sub_mat_type not in {"aij", "baij"}: - raise ValueError(f"Invalid submatrix type, '{sub_mat_type}' (not 'aij' or 'baij')") + if sub_mat_type not in {"aij", "baij", "is"}: + raise ValueError(f"Invalid submatrix type, '{sub_mat_type}' (not 'aij', 'baij', or 'is')") return mat_type, sub_mat_type @@ -2223,7 +2223,8 @@ def masked_lgmap(lgmap, mask, block=True): else: indices = lgmap.indices.copy() bsize = 1 - indices[mask] = -1 + if len(mask) > 0: + indices[mask] = -1 return PETSc.LGMap().create(indices=indices, bsize=bsize, comm=lgmap.comm) @@ -2235,7 +2236,8 @@ def unghosted_lgmap(lgmap, V, block=True): start, end = mesh_dm.getDepthStratum(depth) own = [] - for W in V: + for i, W in enumerate(V): + W_local_ises_indices = V.dof_dset.local_ises[i].indices section = W.dm.getDefaultSection() for seed in range(start, end): # Do not loop over ghost cells @@ -2248,8 +2250,8 @@ def unghosted_lgmap(lgmap, V, block=True): continue off = section.getOffset(p) # Local indices within W - W_indices = range(block_size * off, block_size * (off + dof)) - own.extend(W_indices) + W_indices = slice(block_size * off, block_size * (off + dof)) + own.extend(W_local_ises_indices[W_indices]) mask = numpy.setdiff1d(range(len(lgmap.indices)), own) return masked_lgmap(lgmap, mask, block=block) diff --git a/pyop2/types/mat.py b/pyop2/types/mat.py index 7eb00aa057..dc9935b573 100644 --- a/pyop2/types/mat.py +++ b/pyop2/types/mat.py @@ -641,7 +641,8 @@ def _init_monolithic(self): mat.setOption(mat.Option.KEEP_NONZERO_PATTERN, True) # We completely fill the allocated matrix when zeroing the # entries, so raise an error if we "missed" one. - mat.setOption(mat.Option.UNUSED_NONZERO_LOCATION_ERR, True) + if self.mat_type != "is": + mat.setOption(mat.Option.UNUSED_NONZERO_LOCATION_ERR, True) mat.setOption(mat.Option.IGNORE_OFF_PROC_ENTRIES, False) mat.setOption(mat.Option.NEW_NONZERO_ALLOCATION_ERR, True) # The first assembly (filling with zeros) sets all possible entries. @@ -824,14 +825,16 @@ def set_local_diagonal_entries(self, rows, diag_val=1.0, idx=None): else: rows = np.dstack([rbs*rows + i for i in range(rbs)]).flatten() rows = rows.reshape(-1, 1) - self.change_assembly_state(Mat.INSERT_VALUES) - if len(rows) > 0: - if self.handle.type == "is": - own = self.handle.getLocalSize()[0] - rows = rows[rows < own] - self.handle.assemble() - self.handle.zeroRowsColumnsLocal(rows, diag_val) - else: + if self.handle.type == "is": + self.handle.assemble() + # PETSc does not properly handle local dofs that map + # to a negative global index + rmap, _ = self.handle.getLGMap() + rows = rows[rmap.apply(rows) > -1] + self.handle.zeroRowsColumnsLocal(rows, diag_val) + else: + self.change_assembly_state(Mat.INSERT_VALUES) + if len(rows) > 0: values = np.full(rows.shape, diag_val, dtype=dtypes.ScalarType) self.handle.setValuesLocalRCV(rows, rows, values, addv=PETSc.InsertMode.INSERT_VALUES) @@ -946,11 +949,19 @@ def set_local_diagonal_entries(self, rows, diag_val=1.0, idx=None): else: rows = np.dstack([rbs*rows + i for i in range(rbs)]).flatten() rows = rows.reshape(-1, 1) - self.change_assembly_state(Mat.INSERT_VALUES) - if len(rows) > 0: - values = np.full(rows.shape, diag_val, dtype=dtypes.ScalarType) - self.handle.setValuesLocalRCV(rows, rows, values, - addv=PETSc.InsertMode.INSERT_VALUES) + if self.handle.type == "is": + self.handle.assemble() + # PETSc does not properly handle local dofs that map + # to a negative global index + rmap, _ = self.handle.getLGMap() + rows = rows[rmap.apply(rows) > -1] + self.handle.zeroRowsColumnsLocal(rows, diag_val) + else: + self.change_assembly_state(Mat.INSERT_VALUES) + if len(rows) > 0: + values = np.full(rows.shape, diag_val, dtype=dtypes.ScalarType) + self.handle.setValuesLocalRCV(rows, rows, values, + addv=PETSc.InsertMode.INSERT_VALUES) def addto_values(self, rows, cols, values): """Add a block of values to the :class:`Mat`.""" diff --git a/tests/firedrake/regression/test_assemble.py b/tests/firedrake/regression/test_assemble.py index a2c8c48125..7612a93c0f 100644 --- a/tests/firedrake/regression/test_assemble.py +++ b/tests/firedrake/regression/test_assemble.py @@ -116,17 +116,23 @@ def test_mat_nest_real_block_assembler_correctly_reuses_tensor(mesh): @pytest.mark.parametrize("dirichlet_bcs", [False, True]) -def test_assemble_matis(mesh, dirichlet_bcs): +@pytest.mark.parametrize("mat_type", ["is", "nest"]) +def test_assemble_matis(mesh, mat_type, dirichlet_bcs): V = FunctionSpace(mesh, "CG", 1) + A = 1 + if mat_type == "nest": + V = V * V + A = as_matrix([[2, -1], [-1, 2]]) + u = TrialFunction(V) v = TestFunction(V) - a = inner(grad(u), grad(v))*dx + a = inner(A*grad(u), grad(v))*dx if dirichlet_bcs: - bcs = DirichletBC(V, 0, (1, 3)) + bcs = [DirichletBC(V.sub(i), 0, (i % 4+1, (i+2) % 4+1)) for i in range(len(V))] else: bcs = None - ais = assemble(a, bcs=bcs, mat_type="is").petscmat + ais = assemble(a, bcs=bcs, mat_type=mat_type, sub_mat_type="is").petscmat aijnew = PETSc.Mat() ais.convert("aij", aijnew) From 7bca3e6da00ae96f6b67c0f8f37d73040e4caf4a Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 26 Jun 2025 08:29:26 +0100 Subject: [PATCH 05/72] More comprehensive tests --- tests/firedrake/regression/test_assemble.py | 20 +++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/tests/firedrake/regression/test_assemble.py b/tests/firedrake/regression/test_assemble.py index 7612a93c0f..6a67bb86a8 100644 --- a/tests/firedrake/regression/test_assemble.py +++ b/tests/firedrake/regression/test_assemble.py @@ -115,20 +115,26 @@ def test_mat_nest_real_block_assembler_correctly_reuses_tensor(mesh): assert A2.M is A1.M +@pytest.mark.parametrize("shape,mat_type", [("scalar", "is"), ("vector", "is"), ("mixed", "is"), ("mixed", "nest")]) @pytest.mark.parametrize("dirichlet_bcs", [False, True]) -@pytest.mark.parametrize("mat_type", ["is", "nest"]) -def test_assemble_matis(mesh, mat_type, dirichlet_bcs): - V = FunctionSpace(mesh, "CG", 1) - A = 1 - if mat_type == "nest": - V = V * V +def test_assemble_matis(mesh, shape, mat_type, dirichlet_bcs): + if shape == "vector": + V = VectorFunctionSpace(mesh, "CG", 1) + else: + V = FunctionSpace(mesh, "CG", 1) + if shape == "mixed": + V = V * V + + if V.value_size == 1: + A = 1 + else: A = as_matrix([[2, -1], [-1, 2]]) u = TrialFunction(V) v = TestFunction(V) a = inner(A*grad(u), grad(v))*dx if dirichlet_bcs: - bcs = [DirichletBC(V.sub(i), 0, (i % 4+1, (i+2) % 4+1)) for i in range(len(V))] + bcs = [DirichletBC(V.sub(i), 0, (i % 4+1, (i+2) % 4+1)) for i in range(V.value_size)] else: bcs = None From d0b6469a66b4ac8035c49c7779d5e3545225fd5a Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 26 Jun 2025 15:15:55 +0100 Subject: [PATCH 06/72] Propagate sub_mat_type, fix MatNest + MatIS --- firedrake/assemble.py | 31 +++++++++++++++--- pyop2/types/mat.py | 36 +++++++++++++++++++-- tests/firedrake/regression/test_assemble.py | 33 +++++++++++++++---- 3 files changed, 86 insertions(+), 14 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 355ae963ee..0c31bd5c63 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -356,7 +356,9 @@ def allocate(self): else: test, trial = self._form.arguments() sparsity = ExplicitMatrixAssembler._make_sparsity(test, trial, self._mat_type, self._sub_mat_type, self.maps_and_regions) - return matrix.Matrix(self._form, self._bcs, self._mat_type, sparsity, ScalarType, options_prefix=self._options_prefix) + return matrix.Matrix(self._form, self._bcs, self._mat_type, sparsity, ScalarType, + sub_mat_type=self._sub_mat_type, + options_prefix=self._options_prefix) else: raise NotImplementedError("Only implemented for rank = 2 and diagonal = False") @@ -1377,6 +1379,7 @@ def allocate(self): self._sub_mat_type, self._make_maps_and_regions()) return matrix.Matrix(self._form, self._bcs, self._mat_type, sparsity, ScalarType, + sub_mat_type=self._sub_mat_type, options_prefix=self._options_prefix, fc_params=self._form_compiler_params) @@ -1399,10 +1402,21 @@ def _make_sparsity(test, trial, mat_type, sub_mat_type, maps_and_regions): except SparsityFormatError: raise ValueError("Monolithic matrix assembly not supported for systems " "with R-space blocks") + + # TODO reconstruct dof_dset with the unghosted lgmap if mat_type == "is": - rlgmap = unghosted_lgmap(sparsity._dsets[0].lgmap, test.function_space()) - clgmap = unghosted_lgmap(sparsity._dsets[1].lgmap, trial.function_space()) - sparsity._lgmaps = (rlgmap, clgmap) + rmap = unghosted_lgmap(sparsity._dsets[0].lgmap, test.function_space()) + cmap = unghosted_lgmap(sparsity._dsets[1].lgmap, trial.function_space()) + sparsity._lgmaps = (rmap, cmap) + + elif mat_type == "nest" and sub_mat_type == "is": + Vrow = test.function_space() + Vcol = trial.function_space() + for i, j in numpy.ndindex((len(Vrow), len(Vcol))): + block = sparsity[i, j] + rmap = unghosted_lgmap(block._dsets[0].lgmap, Vrow[i]) + cmap = unghosted_lgmap(block._dsets[1].lgmap, Vcol[j]) + block._lgmaps = (rmap, cmap) return sparsity def _make_maps_and_regions(self): @@ -1498,6 +1512,15 @@ def _apply_bc(self, tensor, bc, u=None): # block is on the matrix diagonal and its index matches the # index of the function space the bc is defined on. op2tensor[index, index].set_local_diagonal_entries(bc.nodes, idx=component, diag_val=self.weight) + if self._mat_type == "nest" and self._sub_mat_type == "is": + for i in range(len(tensor.arguments()[1].function_space())): + if i == index: + continue + op2tensor[index, i].zero_rows(bc.nodes, idx=component, diag_val=0) + for i in range(len(tensor.arguments()[0].function_space())): + if i == index: + continue + op2tensor[i, index].zero_columns(bc.nodes, idx=component, diag_val=0) # Handle off-diagonal block involving real function space. # "lgmaps" is correctly constructed in _matrix_arg, but diff --git a/pyop2/types/mat.py b/pyop2/types/mat.py index dc9935b573..64cb9a8829 100644 --- a/pyop2/types/mat.py +++ b/pyop2/types/mat.py @@ -561,6 +561,7 @@ class Mat(AbstractMat): def __init__(self, *args, **kwargs): self.mat_type = kwargs.pop("mat_type", None) + self.sub_mat_type = kwargs.pop("sub_mat_type", None) super().__init__(*args, **kwargs) self._init() self.assembly_state = Mat.ASSEMBLED @@ -620,7 +621,7 @@ def _init_monolithic(self): rlgmap = rset.unblocked_lgmap clgmap = cset.unblocked_lgmap if self.mat_type == "is": - # TODO monolithic lgmaps + # FIXME monolithic lgmaps rlgmap, clgmap = self.sparsity._lgmaps create = mat.createIS else: @@ -671,7 +672,8 @@ def _init_nest(self): row = [] for j in range(cols): row.append(Mat(self.sparsity[i, j], self.dtype, - '_'.join([self.name, str(i), str(j)]))) + '_'.join([self.name, str(i), str(j)]), + mat_type=self.sub_mat_type)) self._blocks.append(row) # PETSc Mat.createNest wants a flattened list of Mats mat.createNest([[m.handle for m in row_] for row_ in self._blocks], @@ -794,15 +796,43 @@ def zero(self): self.handle.zeroEntries() @mpi.collective - def zero_rows(self, rows, diag_val=1.0): + def zero_rows(self, rows, diag_val=1.0, idx=None): """Zeroes the specified rows of the matrix, with the exception of the diagonal entry, which is set to diag_val. May be used for applying strong boundary conditions. :param rows: a :class:`Subset` or an iterable""" + rows = rows.indices if isinstance(rows, Subset) else rows + rows = np.asarray(rows, dtype=dtypes.IntType) + rbs, _ = self.dims[0][0] + if rbs > 1: + if idx is not None: + rows = rbs * rows + idx + else: + rows = np.dstack([rbs*rows + i for i in range(rbs)]).flatten() self.assemble() + self.handle.zeroRowsLocal(rows, diag_val) + + @mpi.collective + def zero_columns(self, rows, diag_val=1.0, idx=None): + """Zeroes the specified columns of the matrix, with the exception of the + diagonal entry, which is set to diag_val. May be used for applying + strong boundary conditions. + + :param rows: a :class:`Subset` or an iterable""" rows = rows.indices if isinstance(rows, Subset) else rows + rows = np.asarray(rows, dtype=dtypes.IntType) + rbs, _ = self.dims[0][0] + if rbs > 1: + if idx is not None: + rows = rbs * rows + idx + else: + rows = np.dstack([rbs*rows + i for i in range(rbs)]).flatten() + self.assemble() + # FIXME implement zeroColumnsLocal + self.handle.transpose() self.handle.zeroRowsLocal(rows, diag_val) + self.handle.transpose() def _flush_assembly(self): self.handle.assemble(assembly=PETSc.Mat.AssemblyType.FLUSH) diff --git a/tests/firedrake/regression/test_assemble.py b/tests/firedrake/regression/test_assemble.py index 6a67bb86a8..7a4fb5be9b 100644 --- a/tests/firedrake/regression/test_assemble.py +++ b/tests/firedrake/regression/test_assemble.py @@ -124,7 +124,6 @@ def test_assemble_matis(mesh, shape, mat_type, dirichlet_bcs): V = FunctionSpace(mesh, "CG", 1) if shape == "mixed": V = V * V - if V.value_size == 1: A = 1 else: @@ -132,19 +131,39 @@ def test_assemble_matis(mesh, shape, mat_type, dirichlet_bcs): u = TrialFunction(V) v = TestFunction(V) - a = inner(A*grad(u), grad(v))*dx + a = inner(A * grad(u), grad(v))*dx if dirichlet_bcs: bcs = [DirichletBC(V.sub(i), 0, (i % 4+1, (i+2) % 4+1)) for i in range(V.value_size)] else: bcs = None ais = assemble(a, bcs=bcs, mat_type=mat_type, sub_mat_type="is").petscmat - aijnew = PETSc.Mat() - ais.convert("aij", aijnew) - aij = assemble(a, bcs=bcs, mat_type="aij").petscmat - aij.axpy(-1, aijnew) - assert np.allclose(aij[:, :], 0) + aij = PETSc.Mat() + if ais.type == "nest": + blocks = [] + for i in range(len(V)): + row = [] + for j in range(len(V)): + bis = ais.getNestSubMatrix(i, j) + assert bis.type == "is" + bij = PETSc.Mat() + bis.convert("aij", bij) + row.append(bij) + blocks.append(row) + anest = PETSc.Mat() + anest.createNest(blocks, + isrows=V.dof_dset.field_ises, + iscols=V.dof_dset.field_ises, + comm=ais.comm) + anest.convert("aij", aij) + else: + assert ais.type == "is" + ais.convert("aij", aij) + + aij_ref = assemble(a, bcs=bcs, mat_type="aij").petscmat + aij_ref.axpy(-1, aij) + assert np.allclose(aij_ref[:, :], 0) def test_assemble_diagonal(mesh): From 6b7c7f1945d5bbd7362badfa6d733bdd3998698b Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 26 Jun 2025 15:26:58 +0100 Subject: [PATCH 07/72] Only set sub_mat_type on the diagonal blocks --- firedrake/assemble.py | 10 --------- pyop2/types/mat.py | 24 ++------------------- tests/firedrake/regression/test_assemble.py | 3 ++- 3 files changed, 4 insertions(+), 33 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 0c31bd5c63..25b9e4b815 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -1512,16 +1512,6 @@ def _apply_bc(self, tensor, bc, u=None): # block is on the matrix diagonal and its index matches the # index of the function space the bc is defined on. op2tensor[index, index].set_local_diagonal_entries(bc.nodes, idx=component, diag_val=self.weight) - if self._mat_type == "nest" and self._sub_mat_type == "is": - for i in range(len(tensor.arguments()[1].function_space())): - if i == index: - continue - op2tensor[index, i].zero_rows(bc.nodes, idx=component, diag_val=0) - for i in range(len(tensor.arguments()[0].function_space())): - if i == index: - continue - op2tensor[i, index].zero_columns(bc.nodes, idx=component, diag_val=0) - # Handle off-diagonal block involving real function space. # "lgmaps" is correctly constructed in _matrix_arg, but # is ignored by PyOP2 in this case. diff --git a/pyop2/types/mat.py b/pyop2/types/mat.py index 64cb9a8829..ec54e15414 100644 --- a/pyop2/types/mat.py +++ b/pyop2/types/mat.py @@ -671,9 +671,10 @@ def _init_nest(self): for i in range(rows): row = [] for j in range(cols): + # Only set sub_mat_type on the diagonal blocks row.append(Mat(self.sparsity[i, j], self.dtype, '_'.join([self.name, str(i), str(j)]), - mat_type=self.sub_mat_type)) + mat_type=self.sub_mat_type if i == j else None)) self._blocks.append(row) # PETSc Mat.createNest wants a flattened list of Mats mat.createNest([[m.handle for m in row_] for row_ in self._blocks], @@ -813,27 +814,6 @@ def zero_rows(self, rows, diag_val=1.0, idx=None): self.assemble() self.handle.zeroRowsLocal(rows, diag_val) - @mpi.collective - def zero_columns(self, rows, diag_val=1.0, idx=None): - """Zeroes the specified columns of the matrix, with the exception of the - diagonal entry, which is set to diag_val. May be used for applying - strong boundary conditions. - - :param rows: a :class:`Subset` or an iterable""" - rows = rows.indices if isinstance(rows, Subset) else rows - rows = np.asarray(rows, dtype=dtypes.IntType) - rbs, _ = self.dims[0][0] - if rbs > 1: - if idx is not None: - rows = rbs * rows + idx - else: - rows = np.dstack([rbs*rows + i for i in range(rbs)]).flatten() - self.assemble() - # FIXME implement zeroColumnsLocal - self.handle.transpose() - self.handle.zeroRowsLocal(rows, diag_val) - self.handle.transpose() - def _flush_assembly(self): self.handle.assemble(assembly=PETSc.Mat.AssemblyType.FLUSH) diff --git a/tests/firedrake/regression/test_assemble.py b/tests/firedrake/regression/test_assemble.py index 7a4fb5be9b..c89fef9a75 100644 --- a/tests/firedrake/regression/test_assemble.py +++ b/tests/firedrake/regression/test_assemble.py @@ -146,7 +146,8 @@ def test_assemble_matis(mesh, shape, mat_type, dirichlet_bcs): row = [] for j in range(len(V)): bis = ais.getNestSubMatrix(i, j) - assert bis.type == "is" + if i == j: + assert bis.type == "is" bij = PETSc.Mat() bis.convert("aij", bij) row.append(bij) From f9e0bd48b9b2f8a8f4de045962bc4d706dac942e Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 27 Jun 2025 11:30:00 +0100 Subject: [PATCH 08/72] style --- firedrake/assemble.py | 17 ++++++----------- 1 file changed, 6 insertions(+), 11 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 25b9e4b815..5a9fa1d601 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -1408,14 +1408,11 @@ def _make_sparsity(test, trial, mat_type, sub_mat_type, maps_and_regions): rmap = unghosted_lgmap(sparsity._dsets[0].lgmap, test.function_space()) cmap = unghosted_lgmap(sparsity._dsets[1].lgmap, trial.function_space()) sparsity._lgmaps = (rmap, cmap) - elif mat_type == "nest" and sub_mat_type == "is": - Vrow = test.function_space() - Vcol = trial.function_space() - for i, j in numpy.ndindex((len(Vrow), len(Vcol))): + for i, j in numpy.ndindex(sparsity.shape): block = sparsity[i, j] - rmap = unghosted_lgmap(block._dsets[0].lgmap, Vrow[i]) - cmap = unghosted_lgmap(block._dsets[1].lgmap, Vcol[j]) + rmap = unghosted_lgmap(block._dsets[0].lgmap, test.function_space()[i]) + cmap = unghosted_lgmap(block._dsets[1].lgmap, trial.function_space()[j]) block._lgmaps = (rmap, cmap) return sparsity @@ -2242,14 +2239,12 @@ def masked_lgmap(lgmap, mask, block=True): def unghosted_lgmap(lgmap, V, block=True): - block_size = lgmap.getBlockSize() if block else 1 - mesh_dm = V.mesh().topology_dm - depth = mesh_dm.getDepth() - start, end = mesh_dm.getDepthStratum(depth) + start, end = mesh_dm.getHeightStratum(0) own = [] for i, W in enumerate(V): + bsize = W.block_size W_local_ises_indices = V.dof_dset.local_ises[i].indices section = W.dm.getDefaultSection() for seed in range(start, end): @@ -2263,7 +2258,7 @@ def unghosted_lgmap(lgmap, V, block=True): continue off = section.getOffset(p) # Local indices within W - W_indices = slice(block_size * off, block_size * (off + dof)) + W_indices = slice(bsize * off, bsize * (off + dof)) own.extend(W_local_ises_indices[W_indices]) mask = numpy.setdiff1d(range(len(lgmap.indices)), own) From 5c715e17a57244928563e72ef277769e3af516f7 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 28 Jun 2025 15:16:50 +0100 Subject: [PATCH 09/72] Test BDDC, setCoordinates --- firedrake/assemble.py | 4 +- firedrake/preconditioners/bddc.py | 29 +++++++++--- tests/firedrake/regression/test_bddc.py | 60 ++++++++++++++++++------- 3 files changed, 69 insertions(+), 24 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 5a9fa1d601..80876448f7 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -2239,7 +2239,8 @@ def masked_lgmap(lgmap, mask, block=True): def unghosted_lgmap(lgmap, V, block=True): - mesh_dm = V.mesh().topology_dm + mesh = V._mesh + mesh_dm = mesh.topology_dm start, end = mesh_dm.getHeightStratum(0) own = [] @@ -2260,6 +2261,5 @@ def unghosted_lgmap(lgmap, V, block=True): # Local indices within W W_indices = slice(bsize * off, bsize * (off + dof)) own.extend(W_local_ises_indices[W_indices]) - mask = numpy.setdiff1d(range(len(lgmap.indices)), own) return masked_lgmap(lgmap, mask, block=block) diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index c2b3afffde..7e650c1769 100644 --- a/firedrake/preconditioners/bddc.py +++ b/firedrake/preconditioners/bddc.py @@ -4,6 +4,8 @@ from firedrake.petsc import PETSc from firedrake.dmhooks import get_function_space, get_appctx from firedrake.ufl_expr import TestFunction, TrialFunction +from firedrake.function import Function +from firedrake.functionspace import FunctionSpace, VectorFunctionSpace, TensorFunctionSpace from ufl import curl, div, HCurl, HDiv, inner, dx from pyop2.utils import as_tuple import numpy @@ -28,11 +30,6 @@ class BDDCPC(PCBase): keyword arguments supplied to ``PETSc.PC.setBDDCDiscreteGradient``. - ``'divergence_mat'`` for 3D problems in H(div), this sets the Mat with the assembled bilinear form testing the divergence against an L2 space. - - Notes - ----- - Currently the Mat type IS is only supported by FDMPC. - """ _prefix = "bddc_" @@ -44,6 +41,7 @@ def initialize(self, pc): self.prefix = (pc.getOptionsPrefix() or "") + self._prefix V = get_function_space(dm) + variant = V.ufl_element().variant() # Create new PC object as BDDC type bddcpc = PETSc.PC().create(comm=pc.comm) @@ -53,7 +51,7 @@ def initialize(self, pc): bddcpc.setType(PETSc.PC.Type.BDDC) opts = PETSc.Options(bddcpc.getOptionsPrefix()) - if V.ufl_element().variant() == "fdm" and "pc_bddc_use_local_mat_graph" not in opts: + if variant == "fdm" and "pc_bddc_use_local_mat_graph" not in opts: # Disable computation of disconected components of subdomain interfaces opts["pc_bddc_use_local_mat_graph"] = False @@ -80,6 +78,15 @@ def initialize(self, pc): appctx = self.get_appctx(pc) sobolev_space = V.ufl_element().sobolev_space + # set coordinates + if variant != "fdm" and is_lagrange(V): + degree = V.ufl_element().embedded_superdegree + W = VectorFunctionSpace(V.mesh(), "Lagrange", degree, variant=variant) + coords = Function(W).interpolate(V.mesh().coordinates) + gdim, = coords.ufl_shape + view = (slice(None), *(None for _ in V.value_shape), slice(None)) + bddcpc.setCoordinates(numpy.tile(coords.dat.data_ro[view], (1, *V.value_shape, 1)).reshape(-1, gdim)) + tdim = V.mesh().topological_dimension() degree = max(as_tuple(V.ufl_element().degree())) if tdim >= 2 and V.finat_element.formdegree == tdim-1: @@ -133,3 +140,13 @@ def get_vertex_dofs(V): V.dof_dset.lgmap.apply(indices, result=indices) vertex_dofs = PETSc.IS().createGeneral(indices, comm=V.comm) return vertex_dofs + + +def is_lagrange(V): + nodes = V.finat_element.fiat_equivalent.dual_basis() + for node in nodes: + try: + pt, = node.get_point_dict() + except ValueError: + return False + return True diff --git a/tests/firedrake/regression/test_bddc.py b/tests/firedrake/regression/test_bddc.py index ff82094ba9..c2d701f696 100644 --- a/tests/firedrake/regression/test_bddc.py +++ b/tests/firedrake/regression/test_bddc.py @@ -3,13 +3,14 @@ from firedrake.petsc import DEFAULT_DIRECT_SOLVER -def bddc_params(static_condensation): +def bddc_params(): chol = { "pc_type": "cholesky", "pc_factor_mat_solver_type": "petsc", "pc_factor_mat_ordering_type": "natural", } sp = { + "mat_type": "is", "pc_type": "python", "pc_python_type": "firedrake.BDDCPC", "bddc_pc_bddc_neumann": chol, @@ -19,12 +20,14 @@ def bddc_params(static_condensation): return sp -def solver_parameters(static_condensation=True): +def solver_parameters(static_condensation=False, variant=None): rtol = 1E-8 atol = 1E-12 - sp_bddc = bddc_params(static_condensation) - repeated = True - if static_condensation: + sp_bddc = bddc_params() + if variant != "fdm": + sp = sp_bddc + + elif static_condensation: sp = { "pc_type": "python", "pc_python_type": "firedrake.FacetSplitPC", @@ -33,7 +36,7 @@ def solver_parameters(static_condensation=True): "facet_fdm_static_condensation": True, "facet_fdm_pc_use_amat": False, "facet_fdm_mat_type": "is", - "facet_fdm_mat_is_allow_repeated": repeated, + "facet_fdm_mat_is_allow_repeated": True, "facet_fdm_pc_type": "fieldsplit", "facet_fdm_pc_fieldsplit_type": "symmetric_multiplicative", "facet_fdm_pc_fieldsplit_diag_use_amat": False, @@ -47,22 +50,23 @@ def solver_parameters(static_condensation=True): "pc_type": "python", "pc_python_type": "firedrake.FDMPC", "fdm_pc_use_amat": False, - "fdm_mat_type": "is", - "fdm_mat_is_allow_repeated": repeated, + "fdm_mat_is_allow_repeated": True, "fdm": sp_bddc, } + sp.update({ - "mat_type": "matfree", "ksp_type": "cg", "ksp_norm_type": "natural", "ksp_monitor": None, "ksp_rtol": rtol, "ksp_atol": atol, }) + if variant == "fdm": + sp["mat_type"] = "matfree" return sp -def solve_riesz_map(mesh, family, degree, bcs, condense): +def solve_riesz_map(mesh, family, degree, variant, bcs, condense=False, vector=False): dirichlet_ids = [] if bcs: dirichlet_ids = ["on_boundary"] @@ -74,7 +78,10 @@ def solve_riesz_map(mesh, family, degree, bcs, condense): family = "RTCE" if tdim == 2 else "NCE" if family.endswith("F"): family = "RTCF" if tdim == 2 else "NCF" - V = FunctionSpace(mesh, family, degree, variant="fdm") + + fs = VectorFunctionSpace if vector else FunctionSpace + + V = fs(mesh, family, degree, variant=variant) v = TestFunction(V) u = TrialFunction(V) d = { @@ -95,13 +102,21 @@ def solve_riesz_map(mesh, family, degree, bcs, condense): bcs = [DirichletBC(V, u_exact, sub) for sub in dirichlet_ids] nsp = None if formdegree == 0: - nsp = VectorSpaceBasis([Function(V).interpolate(Constant(1))]) + b = np.zeros(V.value_shape) + expr = Constant(b) + basis = [] + for i in np.ndindex(V.value_shape): + b[...] = 0 + b[i] = 1 + expr.assign(b) + basis.append(Function(V).interpolate(expr)) + nsp = VectorSpaceBasis(basis) nsp.orthonormalize() uh = Function(V, name="solution") problem = LinearVariationalProblem(a, L, uh, bcs=bcs) - sp = solver_parameters(condense) + sp = solver_parameters(condense, variant=variant) solver = LinearVariationalSolver(problem, near_nullspace=nsp, solver_parameters=sp, options_prefix="") @@ -120,11 +135,24 @@ def mesh(request): @pytest.mark.parallel -@pytest.mark.parametrize("family", "Q") @pytest.mark.parametrize("degree", (4,)) -@pytest.mark.parametrize("condense", (False, True)) +@pytest.mark.parametrize("family", "Q") +@pytest.mark.parametrize("condense", (False, True), ids=("full", "condense")) def test_bddc_fdm(mesh, family, degree, condense): + variant = "fdm" + bcs = True + tdim = mesh.topological_dimension() + expected = 6 if tdim == 2 else 11 + assert solve_riesz_map(mesh, family, degree, variant, bcs, condense=condense) <= expected + + +@pytest.mark.parallel +@pytest.mark.parametrize("degree", (4,)) +@pytest.mark.parametrize("family", "Q") +@pytest.mark.parametrize("vector", (False, True), ids=("scalar", "vector")) +def test_bddc_aij(mesh, family, degree, vector): + variant = None bcs = True tdim = mesh.topological_dimension() expected = 7 if tdim == 2 else 11 - assert solve_riesz_map(mesh, family, degree, bcs, condense) <= expected + assert solve_riesz_map(mesh, family, degree, variant, bcs, vector=vector) <= expected From 7956efd60d51cd5a5d9f962f593951bdbb374513 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sun, 29 Jun 2025 17:26:50 +0100 Subject: [PATCH 10/72] Fix VectorFunctionSpace --- firedrake/assemble.py | 19 +++++++++++++------ firedrake/preconditioners/bddc.py | 21 ++++++++++++++------- tests/firedrake/regression/test_assemble.py | 10 +++++++--- 3 files changed, 34 insertions(+), 16 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 80876448f7..abbfa21e46 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -2233,20 +2233,26 @@ def masked_lgmap(lgmap, mask, block=True): else: indices = lgmap.indices.copy() bsize = 1 + if len(mask) > 0: indices[mask] = -1 return PETSc.LGMap().create(indices=indices, bsize=bsize, comm=lgmap.comm) def unghosted_lgmap(lgmap, V, block=True): + if block: + ndofs = lgmap.getBlockIndices().size + else: + ndofs = lgmap.getIndices().size + mask = numpy.arange(ndofs, dtype=PETSc.IntType) + mesh = V._mesh mesh_dm = mesh.topology_dm start, end = mesh_dm.getHeightStratum(0) - - own = [] for i, W in enumerate(V): - bsize = W.block_size - W_local_ises_indices = V.dof_dset.local_ises[i].indices + iset = V.dof_dset.local_ises[i] + W_local_indices = iset.indices + bsize = 1 if block else iset.getBlockSize() section = W.dm.getDefaultSection() for seed in range(start, end): # Do not loop over ghost cells @@ -2260,6 +2266,7 @@ def unghosted_lgmap(lgmap, V, block=True): off = section.getOffset(p) # Local indices within W W_indices = slice(bsize * off, bsize * (off + dof)) - own.extend(W_local_ises_indices[W_indices]) - mask = numpy.setdiff1d(range(len(lgmap.indices)), own) + mask[W_local_indices[W_indices]] = -1 + + mask = mask[mask > -1] return masked_lgmap(lgmap, mask, block=block) diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index 7e650c1769..bee7b9bc79 100644 --- a/firedrake/preconditioners/bddc.py +++ b/firedrake/preconditioners/bddc.py @@ -8,6 +8,7 @@ from firedrake.functionspace import FunctionSpace, VectorFunctionSpace, TensorFunctionSpace from ufl import curl, div, HCurl, HDiv, inner, dx from pyop2.utils import as_tuple +import finat import numpy __all__ = ("BDDCPC",) @@ -78,14 +79,12 @@ def initialize(self, pc): appctx = self.get_appctx(pc) sobolev_space = V.ufl_element().sobolev_space - # set coordinates - if variant != "fdm" and is_lagrange(V): + # Set coordinates + if is_lagrange(V.finat_element): degree = V.ufl_element().embedded_superdegree W = VectorFunctionSpace(V.mesh(), "Lagrange", degree, variant=variant) coords = Function(W).interpolate(V.mesh().coordinates) - gdim, = coords.ufl_shape - view = (slice(None), *(None for _ in V.value_shape), slice(None)) - bddcpc.setCoordinates(numpy.tile(coords.dat.data_ro[view], (1, *V.value_shape, 1)).reshape(-1, gdim)) + bddcpc.setCoordinates(coords.dat.data_ro.repeat(V.block_size, axis=0)) tdim = V.mesh().topological_dimension() degree = max(as_tuple(V.ufl_element().degree())) @@ -142,8 +141,16 @@ def get_vertex_dofs(V): return vertex_dofs -def is_lagrange(V): - nodes = V.finat_element.fiat_equivalent.dual_basis() +def is_lagrange(finat_element): + + if isinstance(finat_element, finat.FlattenedDimensions): + return is_lagrange(finat_element.product) + elif isinstance(finat_element, finat.TensorProductElement): + return all(map(is_lagrange, finat_element.factors)) + elif isinstance(finat_element, finat.EnrichedElement): + return all(map(is_lagrange, finat_element.elements)) + + nodes = finat_element.fiat_equivalent.dual_basis() for node in nodes: try: pt, = node.get_point_dict() diff --git a/tests/firedrake/regression/test_assemble.py b/tests/firedrake/regression/test_assemble.py index c89fef9a75..43d823d883 100644 --- a/tests/firedrake/regression/test_assemble.py +++ b/tests/firedrake/regression/test_assemble.py @@ -115,6 +115,7 @@ def test_mat_nest_real_block_assembler_correctly_reuses_tensor(mesh): assert A2.M is A1.M +@pytest.mark.parallel @pytest.mark.parametrize("shape,mat_type", [("scalar", "is"), ("vector", "is"), ("mixed", "is"), ("mixed", "nest")]) @pytest.mark.parametrize("dirichlet_bcs", [False, True]) def test_assemble_matis(mesh, shape, mat_type, dirichlet_bcs): @@ -148,8 +149,10 @@ def test_assemble_matis(mesh, shape, mat_type, dirichlet_bcs): bis = ais.getNestSubMatrix(i, j) if i == j: assert bis.type == "is" - bij = PETSc.Mat() - bis.convert("aij", bij) + bij = PETSc.Mat() + bis.convert("aij", bij) + else: + bij = bis row.append(bij) blocks.append(row) anest = PETSc.Mat() @@ -164,7 +167,8 @@ def test_assemble_matis(mesh, shape, mat_type, dirichlet_bcs): aij_ref = assemble(a, bcs=bcs, mat_type="aij").petscmat aij_ref.axpy(-1, aij) - assert np.allclose(aij_ref[:, :], 0) + ind, iptr, values = aij_ref.getValuesCSR() + assert np.allclose(values, 0) def test_assemble_diagonal(mesh): From 8248cf97cd8e8c15a5dddb04afac72afec0c9ca9 Mon Sep 17 00:00:00 2001 From: Stefano Zampini Date: Wed, 23 Jul 2025 15:33:08 +0200 Subject: [PATCH 11/72] refine BDDC customization --- firedrake/preconditioners/bddc.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index bee7b9bc79..045fa62c18 100644 --- a/firedrake/preconditioners/bddc.py +++ b/firedrake/preconditioners/bddc.py @@ -26,11 +26,11 @@ class BDDCPC(PCBase): - ``'bddc_pc_bddc_coarse'`` to set the coarse solver KSP. This PC also inspects optional arguments supplied in the application context: - - ``'discrete_gradient'`` for problems in H(curl), this sets the arguments + - ``'discrete_gradient'`` for 3D problems in H(curl), this sets the arguments (a Mat tabulating the gradient of the auxiliary H1 space) and keyword arguments supplied to ``PETSc.PC.setBDDCDiscreteGradient``. - - ``'divergence_mat'`` for 3D problems in H(div), this sets the Mat with the - assembled bilinear form testing the divergence against an L2 space. + - ``'divergence_mat'`` for problems in H(div) (resp. 2D H(curl)), this sets the Mat with the + assembled bilinear form testing the divergence (curl) against an L2 space. """ _prefix = "bddc_" @@ -43,6 +43,7 @@ def initialize(self, pc): V = get_function_space(dm) variant = V.ufl_element().variant() + is_lag = is_lagrange(V.finat_element) # Create new PC object as BDDC type bddcpc = PETSc.PC().create(comm=pc.comm) @@ -52,10 +53,12 @@ def initialize(self, pc): bddcpc.setType(PETSc.PC.Type.BDDC) opts = PETSc.Options(bddcpc.getOptionsPrefix()) - if variant == "fdm" and "pc_bddc_use_local_mat_graph" not in opts: - # Disable computation of disconected components of subdomain interfaces + # Do not use CSR of local matrix to define dofs connectivity unless requested + # Using the CSR only makes sense for H1 problems + if "pc_bddc_use_local_mat_graph" not in opts and (not is_lag or variant == "fdm"): opts["pc_bddc_use_local_mat_graph"] = False + # Handle boundary dofs ctx = get_appctx(dm) bcs = tuple(ctx._problem.bcs) if V.extruded: @@ -80,7 +83,7 @@ def initialize(self, pc): sobolev_space = V.ufl_element().sobolev_space # Set coordinates - if is_lagrange(V.finat_element): + if is_lag: degree = V.ufl_element().embedded_superdegree W = VectorFunctionSpace(V.mesh(), "Lagrange", degree, variant=variant) coords = Function(W).interpolate(V.mesh().coordinates) @@ -95,7 +98,7 @@ def initialize(self, pc): d = {HCurl: curl, HDiv: div}[sobolev_space] Q = V.reconstruct(family="DG", degree=degree-1) b = inner(d(TrialFunction(V)), TestFunction(Q)) * dx(degree=2*(degree-1)) - B = assemble(b, mat_type="matfree") + B = assemble(b, mat_type="is") bddcpc.setBDDCDivergenceMat(B.petscmat) elif sobolev_space == HCurl: gradient = appctx.get("discrete_gradient", None) @@ -154,6 +157,6 @@ def is_lagrange(finat_element): for node in nodes: try: pt, = node.get_point_dict() - except ValueError: + except: return False return True From f6c4778c8140b5d244bf2ccf1039e927a906fabd Mon Sep 17 00:00:00 2001 From: Stefano Zampini Date: Thu, 24 Jul 2025 11:23:51 +0200 Subject: [PATCH 12/72] fix lint --- firedrake/preconditioners/bddc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index 045fa62c18..a79cda4237 100644 --- a/firedrake/preconditioners/bddc.py +++ b/firedrake/preconditioners/bddc.py @@ -157,6 +157,6 @@ def is_lagrange(finat_element): for node in nodes: try: pt, = node.get_point_dict() - except: + except (ValueError, TypeError): return False return True From 32b2c4dfb92fd75b93fa8bc1773d1df7441028fd Mon Sep 17 00:00:00 2001 From: Stefano Zampini Date: Thu, 24 Jul 2025 19:55:08 +0200 Subject: [PATCH 13/72] Simplify coordinates handling of bddcpc driver cleanup test --- firedrake/preconditioners/bddc.py | 44 ++++++++----------------- tests/firedrake/regression/test_bddc.py | 5 ++- 2 files changed, 16 insertions(+), 33 deletions(-) diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index a79cda4237..2b00c71cfe 100644 --- a/firedrake/preconditioners/bddc.py +++ b/firedrake/preconditioners/bddc.py @@ -6,9 +6,8 @@ from firedrake.ufl_expr import TestFunction, TrialFunction from firedrake.function import Function from firedrake.functionspace import FunctionSpace, VectorFunctionSpace, TensorFunctionSpace -from ufl import curl, div, HCurl, HDiv, inner, dx +from ufl import curl, div, H1, H2, HCurl, HDiv, inner, dx from pyop2.utils import as_tuple -import finat import numpy __all__ = ("BDDCPC",) @@ -43,7 +42,7 @@ def initialize(self, pc): V = get_function_space(dm) variant = V.ufl_element().variant() - is_lag = is_lagrange(V.finat_element) + sobolev_space = V.ufl_element().sobolev_space # Create new PC object as BDDC type bddcpc = PETSc.PC().create(comm=pc.comm) @@ -54,8 +53,9 @@ def initialize(self, pc): opts = PETSc.Options(bddcpc.getOptionsPrefix()) # Do not use CSR of local matrix to define dofs connectivity unless requested - # Using the CSR only makes sense for H1 problems - if "pc_bddc_use_local_mat_graph" not in opts and (not is_lag or variant == "fdm"): + # Using the CSR only makes sense for H1/H2 problems + is_h1h2 = sobolev_space in [H1, H2] + if "pc_bddc_use_local_mat_graph" not in opts and (not is_h1h2 or variant == "fdm"): opts["pc_bddc_use_local_mat_graph"] = False # Handle boundary dofs @@ -80,18 +80,18 @@ def initialize(self, pc): bddcpc.setBDDCNeumannBoundaries(neu_bndr) appctx = self.get_appctx(pc) - sobolev_space = V.ufl_element().sobolev_space + degree = max(as_tuple(V.ufl_element().degree())) - # Set coordinates - if is_lag: - degree = V.ufl_element().embedded_superdegree + # Set coordinates only if corner selection is requested + # There's no API to query from PC + if "pc_bddc_corner_selection" in opts: W = VectorFunctionSpace(V.mesh(), "Lagrange", degree, variant=variant) coords = Function(W).interpolate(V.mesh().coordinates) bddcpc.setCoordinates(coords.dat.data_ro.repeat(V.block_size, axis=0)) tdim = V.mesh().topological_dimension() - degree = max(as_tuple(V.ufl_element().degree())) if tdim >= 2 and V.finat_element.formdegree == tdim-1: + # Should we use a callable like for hypre_ams? B = appctx.get("divergence_mat", None) if B is None: from firedrake.assemble import assemble @@ -101,14 +101,16 @@ def initialize(self, pc): B = assemble(b, mat_type="is") bddcpc.setBDDCDivergenceMat(B.petscmat) elif sobolev_space == HCurl: + # Should we use a callable like for hypre_ams? gradient = appctx.get("discrete_gradient", None) if gradient is None: from firedrake.preconditioners.fdm import tabulate_exterior_derivative from firedrake.preconditioners.hiptmair import curl_to_grad Q = V.reconstruct(element=curl_to_grad(V.ufl_element())) gradient = tabulate_exterior_derivative(Q, V) - corners = get_vertex_dofs(Q) - gradient.compose('_elements_corners', corners) + if variant == 'fdm': + corners = get_vertex_dofs(Q) + gradient.compose('_elements_corners', corners) grad_args = (gradient,) grad_kwargs = {'order': degree} else: @@ -142,21 +144,3 @@ def get_vertex_dofs(V): V.dof_dset.lgmap.apply(indices, result=indices) vertex_dofs = PETSc.IS().createGeneral(indices, comm=V.comm) return vertex_dofs - - -def is_lagrange(finat_element): - - if isinstance(finat_element, finat.FlattenedDimensions): - return is_lagrange(finat_element.product) - elif isinstance(finat_element, finat.TensorProductElement): - return all(map(is_lagrange, finat_element.factors)) - elif isinstance(finat_element, finat.EnrichedElement): - return all(map(is_lagrange, finat_element.elements)) - - nodes = finat_element.fiat_equivalent.dual_basis() - for node in nodes: - try: - pt, = node.get_point_dict() - except (ValueError, TypeError): - return False - return True diff --git a/tests/firedrake/regression/test_bddc.py b/tests/firedrake/regression/test_bddc.py index c2d701f696..cc875b7097 100644 --- a/tests/firedrake/regression/test_bddc.py +++ b/tests/firedrake/regression/test_bddc.py @@ -137,13 +137,12 @@ def mesh(request): @pytest.mark.parallel @pytest.mark.parametrize("degree", (4,)) @pytest.mark.parametrize("family", "Q") -@pytest.mark.parametrize("condense", (False, True), ids=("full", "condense")) -def test_bddc_fdm(mesh, family, degree, condense): +def test_bddc_fdm(mesh, family, degree): variant = "fdm" bcs = True tdim = mesh.topological_dimension() expected = 6 if tdim == 2 else 11 - assert solve_riesz_map(mesh, family, degree, variant, bcs, condense=condense) <= expected + assert solve_riesz_map(mesh, family, degree, variant, bcs) <= expected @pytest.mark.parallel From 3f33effc51eb40127ad0a946d4eea141261dc400 Mon Sep 17 00:00:00 2001 From: Stefano Zampini Date: Sun, 27 Jul 2025 23:56:30 +0200 Subject: [PATCH 14/72] Don't know why this fixes the mapping for me --- firedrake/preconditioners/bddc.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index 2b00c71cfe..9aa9b08efe 100644 --- a/firedrake/preconditioners/bddc.py +++ b/firedrake/preconditioners/bddc.py @@ -71,11 +71,11 @@ def initialize(self, pc): dir_nodes = numpy.unique(numpy.concatenate([bcdofs(bc, ghost=False) for bc in bcs])) neu_nodes = numpy.setdiff1d(boundary_nodes, dir_nodes) - V.dof_dset.lgmap.apply(dir_nodes, result=dir_nodes) + dir_nodes = V.dof_dset.lgmap.apply(dir_nodes) dir_bndr = PETSc.IS().createGeneral(dir_nodes, comm=pc.comm) bddcpc.setBDDCDirichletBoundaries(dir_bndr) - V.dof_dset.lgmap.apply(neu_nodes, result=neu_nodes) + neu_nodes = V.dof_dset.lgmap.apply(neu_nodes) neu_bndr = PETSc.IS().createGeneral(neu_nodes, comm=pc.comm) bddcpc.setBDDCNeumannBoundaries(neu_bndr) @@ -141,6 +141,6 @@ def applyTranspose(self, pc, x, y): def get_vertex_dofs(V): W = V.reconstruct(element=restrict(V.ufl_element(), "vertex")) indices = get_restriction_indices(V, W) - V.dof_dset.lgmap.apply(indices, result=indices) + indices = V.dof_dset.lgmap.apply(indices) vertex_dofs = PETSc.IS().createGeneral(indices, comm=V.comm) return vertex_dofs From b945f3835d09a7968cad73884d1b5b4a4425905f Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 28 Jul 2025 09:57:22 +0100 Subject: [PATCH 15/72] test tabulate_exterior_derivative --- tests/firedrake/regression/test_fdm.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/tests/firedrake/regression/test_fdm.py b/tests/firedrake/regression/test_fdm.py index 947ae43fbc..bc6151a1c0 100644 --- a/tests/firedrake/regression/test_fdm.py +++ b/tests/firedrake/regression/test_fdm.py @@ -334,3 +334,21 @@ def test_ipdg_direct_solver(fs): assert uvec.norm() < 1E-8 else: assert norm(u_exact-uh, "H1") < 1.0E-8 + + +@pytest.mark.parallel +@pytest.mark.parametrize("degree", range(1, 3)) +@pytest.mark.parametrize("variant", ("spectral", "integral", "fdm")) +def test_tabulate_exterior_derivative(mesh, variant, degree): + from firedrake.preconditioners.fdm import tabulate_exterior_derivative + family = {1: "DG", 2: "RTCE", 3: "NCE"}[mesh.topological_dimension()] + V0 = FunctionSpace(mesh, "Lagrange", degree, variant=variant) + V1 = FunctionSpace(mesh, family, degree, variant=variant) + + D = tabulate_exterior_derivative(V0, V1) + + u0 = Function(V0).interpolate(1) + u1 = Function(V1) + with u0.dat.vec as x, u1.dat.vec as y: + D.mult(x, y) + assert y.norm() < 1E-12 From e4fc07147e5c3edf2be85251411ee7449105e729 Mon Sep 17 00:00:00 2001 From: Stefano Zampini Date: Mon, 28 Jul 2025 12:33:12 +0200 Subject: [PATCH 16/72] add failing test --- tests/firedrake/regression/test_fdm.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/firedrake/regression/test_fdm.py b/tests/firedrake/regression/test_fdm.py index bc6151a1c0..c7105c4abd 100644 --- a/tests/firedrake/regression/test_fdm.py +++ b/tests/firedrake/regression/test_fdm.py @@ -341,11 +341,13 @@ def test_ipdg_direct_solver(fs): @pytest.mark.parametrize("variant", ("spectral", "integral", "fdm")) def test_tabulate_exterior_derivative(mesh, variant, degree): from firedrake.preconditioners.fdm import tabulate_exterior_derivative + from firedrake.preconditioners.bddc import get_vertex_dofs family = {1: "DG", 2: "RTCE", 3: "NCE"}[mesh.topological_dimension()] V0 = FunctionSpace(mesh, "Lagrange", degree, variant=variant) V1 = FunctionSpace(mesh, family, degree, variant=variant) D = tabulate_exterior_derivative(V0, V1) + v = get_vertex_dofs(V0) u0 = Function(V0).interpolate(1) u1 = Function(V1) From f701ec3b4ab34cda06bf356152a1dd07d826c447 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 28 Jul 2025 12:42:34 +0100 Subject: [PATCH 17/72] Fix vertex dofs --- firedrake/preconditioners/facet_split.py | 2 +- tests/firedrake/regression/test_fdm.py | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/firedrake/preconditioners/facet_split.py b/firedrake/preconditioners/facet_split.py index f4ac80eea0..71fbdf2d2c 100644 --- a/firedrake/preconditioners/facet_split.py +++ b/firedrake/preconditioners/facet_split.py @@ -242,7 +242,7 @@ def restricted_dofs(celem, felem): cdofs = celem.entity_dofs() fdofs = felem.entity_dofs() for dim in sorted(cdofs): - for entity in cdofs[dim]: + for entity in sorted(cdofs[dim]): ndofs = len(cdofs[dim][entity]) indices[cdofs[dim][entity]] = fdofs[dim][entity][:ndofs] return indices diff --git a/tests/firedrake/regression/test_fdm.py b/tests/firedrake/regression/test_fdm.py index c7105c4abd..1100f8e49c 100644 --- a/tests/firedrake/regression/test_fdm.py +++ b/tests/firedrake/regression/test_fdm.py @@ -342,9 +342,10 @@ def test_ipdg_direct_solver(fs): def test_tabulate_exterior_derivative(mesh, variant, degree): from firedrake.preconditioners.fdm import tabulate_exterior_derivative from firedrake.preconditioners.bddc import get_vertex_dofs - family = {1: "DG", 2: "RTCE", 3: "NCE"}[mesh.topological_dimension()] + tdim = mesh.topological_dimension() + family = {1: "DG", 2: "RTCE", 3: "NCE"}[tdim] V0 = FunctionSpace(mesh, "Lagrange", degree, variant=variant) - V1 = FunctionSpace(mesh, family, degree, variant=variant) + V1 = FunctionSpace(mesh, family, degree-(tdim == 1), variant=variant) D = tabulate_exterior_derivative(V0, V1) v = get_vertex_dofs(V0) From baa8da5bae88edd0a5bb6c378f1ea3e593716790 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 28 Jul 2025 17:31:40 +0100 Subject: [PATCH 18/72] Add permutation test --- tests/firedrake/regression/test_assemble.py | 16 ++++++++++++++++ tests/firedrake/regression/test_bddc.py | 11 +++++++++++ tests/firedrake/regression/test_fdm.py | 2 -- 3 files changed, 27 insertions(+), 2 deletions(-) diff --git a/tests/firedrake/regression/test_assemble.py b/tests/firedrake/regression/test_assemble.py index 43d823d883..10cffcc71f 100644 --- a/tests/firedrake/regression/test_assemble.py +++ b/tests/firedrake/regression/test_assemble.py @@ -401,3 +401,19 @@ def test_assemble_tensor_empty_shape(mesh): v = Function(V).assign(1) expected = assemble(inner(v, v)*dx) assert np.allclose(result, expected) + + +@pytest.mark.parallel(nprocs=2) +def test_assemble_integral_quads(): + no_overlap = {"overlap_type": (DistributedMeshOverlapType.NONE, 0)} + mesh0 = UnitSquareMesh(3, 1, quadrilateral=True, distribution_parameters=no_overlap) + mesh1 = UnitSquareMesh(3, 1, quadrilateral=True) + + degree = 3 + variant = "integral" + V0 = FunctionSpace(mesh0, "Lagrange", degree, variant=variant) + V1 = FunctionSpace(mesh1, "Lagrange", degree, variant=variant) + c0 = assemble(conj(TestFunction(V0))*dx, bcs=DirichletBC(V0, 0, "on_boundary")) + c1 = assemble(conj(TestFunction(V1))*dx, bcs=DirichletBC(V1, 0, "on_boundary")) + + assert np.linalg.norm(c0.dat.data) == np.linalg.norm(c1.dat.data) diff --git a/tests/firedrake/regression/test_bddc.py b/tests/firedrake/regression/test_bddc.py index cc875b7097..18fa6ab9b0 100644 --- a/tests/firedrake/regression/test_bddc.py +++ b/tests/firedrake/regression/test_bddc.py @@ -134,6 +134,17 @@ def mesh(request): return msh +#@pytest.mark.parallel +@pytest.mark.parametrize("degree", range(1, 3)) +@pytest.mark.parametrize("variant", ("spectral", "fdm")) +def test_vertex_dofs(mesh, variant, degree): + from firedrake.preconditioners.bddc import get_vertex_dofs + P1 = FunctionSpace(mesh, "Lagrange", 1, variant=variant) + V0 = FunctionSpace(mesh, "Lagrange", degree, variant=variant) + v = get_vertex_dofs(V0) + assert v.getSizes() == P1.dof_dset.layout_vec.getSizes() + + @pytest.mark.parallel @pytest.mark.parametrize("degree", (4,)) @pytest.mark.parametrize("family", "Q") diff --git a/tests/firedrake/regression/test_fdm.py b/tests/firedrake/regression/test_fdm.py index 1100f8e49c..0bc370843a 100644 --- a/tests/firedrake/regression/test_fdm.py +++ b/tests/firedrake/regression/test_fdm.py @@ -341,14 +341,12 @@ def test_ipdg_direct_solver(fs): @pytest.mark.parametrize("variant", ("spectral", "integral", "fdm")) def test_tabulate_exterior_derivative(mesh, variant, degree): from firedrake.preconditioners.fdm import tabulate_exterior_derivative - from firedrake.preconditioners.bddc import get_vertex_dofs tdim = mesh.topological_dimension() family = {1: "DG", 2: "RTCE", 3: "NCE"}[tdim] V0 = FunctionSpace(mesh, "Lagrange", degree, variant=variant) V1 = FunctionSpace(mesh, family, degree-(tdim == 1), variant=variant) D = tabulate_exterior_derivative(V0, V1) - v = get_vertex_dofs(V0) u0 = Function(V0).interpolate(1) u1 = Function(V1) From 245362d3f2bc072f89a7a612daee702c1f256143 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 29 Jul 2025 13:55:23 +0100 Subject: [PATCH 19/72] Trivial case for get_restriction_indices --- firedrake/preconditioners/facet_split.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/firedrake/preconditioners/facet_split.py b/firedrake/preconditioners/facet_split.py index 71fbdf2d2c..ac785af650 100644 --- a/firedrake/preconditioners/facet_split.py +++ b/firedrake/preconditioners/facet_split.py @@ -251,6 +251,9 @@ def restricted_dofs(celem, felem): def get_restriction_indices(V, W): """Return the list of dofs in the space V such that W = V[indices]. """ + if V.cell_node_map() is W.cell_node_map(): + return numpy.arange(V.dof_dset.layout_vec.getSizes()[0], dtype=PETSc.IntType) + vdat = V.make_dat(val=numpy.arange(V.dof_count, dtype=PETSc.IntType)) wdats = [Wsub.make_dat(val=numpy.full((Wsub.dof_count,), -1, dtype=PETSc.IntType)) for Wsub in W] wdat = wdats[0] if len(W) == 1 else op2.MixedDat(wdats) From 5397bcf1cd2ed2a3821b3700b09d99c4bbe63f48 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 29 Jul 2025 14:30:10 +0100 Subject: [PATCH 20/72] DROP BEFORE MERGE --- .github/workflows/core.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/core.yml b/.github/workflows/core.yml index 995cbf71e6..cbbe56cadc 100644 --- a/.github/workflows/core.yml +++ b/.github/workflows/core.yml @@ -221,6 +221,7 @@ jobs: --extra-index-url https://download.pytorch.org/whl/cpu \ "$(echo ./firedrake-repo/dist/firedrake-*.tar.gz)[ci,docs]" + pip install -I "firedrake-fiat @ git+https://github.com/firedrakeproject/fiat.git@pbrubeck/fix/permutations" firedrake-clean pip list From 8979c96907471b5b373447aa666faa3961133327 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 29 Jul 2025 17:39:56 +0100 Subject: [PATCH 21/72] mark parellel test --- tests/firedrake/regression/test_bddc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/firedrake/regression/test_bddc.py b/tests/firedrake/regression/test_bddc.py index 18fa6ab9b0..b22f113ea1 100644 --- a/tests/firedrake/regression/test_bddc.py +++ b/tests/firedrake/regression/test_bddc.py @@ -134,7 +134,7 @@ def mesh(request): return msh -#@pytest.mark.parallel +@pytest.mark.parallel @pytest.mark.parametrize("degree", range(1, 3)) @pytest.mark.parametrize("variant", ("spectral", "fdm")) def test_vertex_dofs(mesh, variant, degree): From 79314b521390cb475a5c3e774aa84c1b1d2f61b1 Mon Sep 17 00:00:00 2001 From: Stefano Zampini Date: Wed, 30 Jul 2025 13:31:28 +0200 Subject: [PATCH 22/72] bddc: use callables for gradient and divergence --- firedrake/preconditioners/bddc.py | 52 +++++++++++++++++++++++-------- 1 file changed, 39 insertions(+), 13 deletions(-) diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index 9aa9b08efe..cee8fd359b 100644 --- a/firedrake/preconditioners/bddc.py +++ b/firedrake/preconditioners/bddc.py @@ -6,6 +6,9 @@ from firedrake.ufl_expr import TestFunction, TrialFunction from firedrake.function import Function from firedrake.functionspace import FunctionSpace, VectorFunctionSpace, TensorFunctionSpace +from firedrake.preconditioners.fdm import tabulate_exterior_derivative +from firedrake.preconditioners.hiptmair import curl_to_grad +from firedrake import assemble from ufl import curl, div, H1, H2, HCurl, HDiv, inner, dx from pyop2.utils import as_tuple import numpy @@ -24,12 +27,13 @@ class BDDCPC(PCBase): - ``'bddc_pc_bddc_dirichlet'`` to set sub-KSPs on subdomain interiors, - ``'bddc_pc_bddc_coarse'`` to set the coarse solver KSP. - This PC also inspects optional arguments supplied in the application context: - - ``'discrete_gradient'`` for 3D problems in H(curl), this sets the arguments - (a Mat tabulating the gradient of the auxiliary H1 space) and + This PC also inspects optional callbacks supplied in the application context: + - ``'get_discrete_gradient'`` for 3D problems in H(curl), this is a callable that + provide the arguments (a Mat tabulating the gradient of the auxiliary H1 space) and keyword arguments supplied to ``PETSc.PC.setBDDCDiscreteGradient``. - - ``'divergence_mat'`` for problems in H(div) (resp. 2D H(curl)), this sets the Mat with the - assembled bilinear form testing the divergence (curl) against an L2 space. + - ``'get_divergence_mat'`` for problems in H(div) (resp. 2D H(curl)), this is + provide the arguments (a Mat with the assembled bilinear form testing the divergence + (curl) against an L2 space) and keyword arguments supplied to ``PETSc.PC.setDivergenceMat``. """ _prefix = "bddc_" @@ -112,14 +116,7 @@ def initialize(self, pc): corners = get_vertex_dofs(Q) gradient.compose('_elements_corners', corners) grad_args = (gradient,) - grad_kwargs = {'order': degree} - else: - try: - grad_args, grad_kwargs = gradient - except ValueError: - grad_args = (gradient,) - grad_kwargs = dict() - + grad_kwargs = dict() bddcpc.setBDDCDiscreteGradient(*grad_args, **grad_kwargs) bddcpc.setFromOptions() @@ -144,3 +141,32 @@ def get_vertex_dofs(V): indices = V.dof_dset.lgmap.apply(indices) vertex_dofs = PETSc.IS().createGeneral(indices, comm=V.comm) return vertex_dofs + + +def get_divergence_mat(V): + sobolev_space = V.ufl_element().sobolev_space + degree = max(as_tuple(V.ufl_element().degree())) + d = {HCurl: curl, HDiv: div}[sobolev_space] + if V.shape == (): + make_function_space = FunctionSpace + elif len(V.shape) == 1: + make_function_space = VectorFunctionSpace + else: + make_function_space = TensorFunctionSpace + Q = make_function_space(V.mesh(), "DG", degree-1) + b = inner(d(TrialFunction(V)), TestFunction(Q)) * dx(degree=2*(degree-1)) + B = assemble.assemble(b, mat_type="is") + return (B.petscmat,), dict() + + +def get_discrete_gradient(V): + degree = max(as_tuple(V.ufl_element().degree())) + variant = V.ufl_element().variant() + Q = FunctionSpace(V.mesh(), curl_to_grad(V.ufl_element())) + gradient = tabulate_exterior_derivative(Q, V) + if variant == 'fdm': + corners = get_vertex_dofs(Q) + gradient.compose('_elements_corners', corners) + grad_args = (gradient,) + grad_kwargs = {'order': degree} + return grad_args, grad_kwargs From 72b782c41182d8432a9fa3acb3ccc3b58e4024b7 Mon Sep 17 00:00:00 2001 From: Stefano Zampini Date: Wed, 30 Jul 2025 13:31:34 +0200 Subject: [PATCH 23/72] minor --- firedrake/preconditioners/fdm.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/firedrake/preconditioners/fdm.py b/firedrake/preconditioners/fdm.py index 1efa612e72..b1622146f8 100644 --- a/firedrake/preconditioners/fdm.py +++ b/firedrake/preconditioners/fdm.py @@ -83,7 +83,7 @@ class FDMPC(PCBase): matrices. The PETSc options inspected by this class are: - - 'fdm_mat_type': can be either 'aij' or 'sbaij' + - 'fdm_mat_type': can be either 'aij', 'sbaij', or 'is' - 'fdm_static_condensation': are we assembling the Schur complement on facets? """ @@ -169,7 +169,7 @@ def initialize(self, pc): self.bc_nodes = numpy.empty(0, dtype=PETSc.IntType) # Internally, we just set up a PC object that the user can configure - # however from the PETSc command line. Since PC allows the user to specify + # however from the PETSc command line. Since PC allows the user to specify # a KSP, we can do iterative by -fdm_pc_type ksp. fdmpc = PETSc.PC().create(comm=pc.comm) fdmpc.incrementTabLevel(1, parent=pc) From 0808ccd669f8a377be52a87ee8ee56bf40b2dee0 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 31 Jul 2025 19:16:39 +0100 Subject: [PATCH 24/72] tabulate_exterior_derivative as MatIS --- firedrake/preconditioners/fdm.py | 51 +++++++++++++++++--------- tests/firedrake/regression/test_fdm.py | 26 ++++++++++++- 2 files changed, 59 insertions(+), 18 deletions(-) diff --git a/firedrake/preconditioners/fdm.py b/firedrake/preconditioners/fdm.py index b1622146f8..4ae835f59e 100644 --- a/firedrake/preconditioners/fdm.py +++ b/firedrake/preconditioners/fdm.py @@ -710,19 +710,7 @@ def insert_mode(self): def assembly_lgmaps(self): if self.mat_type != "is": return {Vsub: Vsub.dof_dset.lgmap for Vsub in self.V} - lgmaps = {} - for Vsub in self.V: - lgmap = Vsub.dof_dset.lgmap - if self.allow_repeated: - indices = broken_function(Vsub, lgmap.indices).dat.data_ro - else: - indices = lgmap.indices.copy() - local_indices = numpy.arange(len(indices), dtype=PETSc.IntType) - cell_node_map = broken_function(Vsub, local_indices).dat.data_ro - ghost = numpy.setdiff1d(local_indices, numpy.unique(cell_node_map), assume_unique=True) - indices[ghost] = -1 - lgmaps[Vsub] = PETSc.LGMap().create(indices, bsize=lgmap.getBlockSize(), comm=lgmap.getComm()) - return lgmaps + return {Vsub: get_matis_lgmap(Vsub, self.allow_repeated) for Vsub in self.V} def setup_block(self, Vrow, Vcol): # Preallocate the auxiliary sparse operator @@ -1661,7 +1649,21 @@ def diff_blocks(tdim, formdegree, A00, A11, A10): return A_blocks -def tabulate_exterior_derivative(Vc, Vf, cbcs=[], fbcs=[], comm=None): +def get_matis_lgmap(Vsub, allow_repeated): + lgmap = Vsub.dof_dset.lgmap + if allow_repeated: + indices = broken_function(Vsub, lgmap.indices).dat.data_ro + else: + indices = lgmap.indices.copy() + local_indices = numpy.arange(len(indices), dtype=PETSc.IntType) + cell_node_map = broken_function(Vsub, local_indices).dat.data_ro + ghost = numpy.setdiff1d(local_indices, numpy.unique(cell_node_map), assume_unique=True) + indices[ghost] = -1 + + return PETSc.LGMap().create(indices, bsize=lgmap.getBlockSize(), comm=lgmap.getComm()) + + +def tabulate_exterior_derivative(Vc, Vf, cbcs=[], fbcs=[], comm=None, mat_type="aij", allow_repeated=False): """ Tabulate exterior derivative: Vc -> Vf as an explicit sparse matrix. Works for any tensor-product basis. These are the same matrices one needs for HypreAMS and friends. @@ -1724,9 +1726,18 @@ def tabulate_exterior_derivative(Vc, Vf, cbcs=[], fbcs=[], comm=None): preallocator.setSizes(sizes) preallocator.setUp() + spaces = (Vf, Vc) + bcs = (fbcs, cbcs) + if mat_type == "is": + assert not any(bcs) + lgmaps = tuple(get_matis_lgmap(V, allow_repeated) for V in spaces) + indices = tuple(mask_local_indices(V, V.dof_dset.lgmap, allow_repeated) for V in spaces) + else: + lgmaps = tuple(V.local_to_global_map(bcs) for V, bcs in zip(spaces, bcs)) + indices = tuple(op2.Dat(V.dof_dset, lgmap.indices)(op2.READ, V.cell_node_map()) + for V, lgmap in zip(spaces, lgmaps)) + kernel = ElementKernel(Dhat, name="exterior_derivative").kernel() - indices = tuple(op2.Dat(V.dof_dset, V.local_to_global_map(bcs).indices)(op2.READ, V.cell_node_map()) - for V, bcs in zip((Vf, Vc), (fbcs, cbcs))) assembler = op2.ParLoop(kernel, Vc.mesh().cell_set, *(op2.PassthroughArg(op2.OpaqueType("Mat"), m.handle) for m in (preallocator, Dhat)), @@ -1736,7 +1747,13 @@ def tabulate_exterior_derivative(Vc, Vf, cbcs=[], fbcs=[], comm=None): nnz = get_preallocation(preallocator, sizes[0][0]) preallocator.destroy() - Dmat = PETSc.Mat().createAIJ(sizes, Vf.block_size, nnz=nnz, comm=comm) + Dmat = PETSc.Mat().create(comm=comm) + Dmat.setType(mat_type) + Dmat.setSizes(sizes) + Dmat.setBlockSize(Vf.block_size) + Dmat.setISAllowRepeated(allow_repeated) + Dmat.setLGMap(*lgmaps) + Dmat.setPreallocationNNZ(nnz) Dmat.setOption(PETSc.Mat.Option.NEW_NONZERO_ALLOCATION_ERR, True) assembler.arguments[0].data = Dmat.handle assembler() diff --git a/tests/firedrake/regression/test_fdm.py b/tests/firedrake/regression/test_fdm.py index 0bc370843a..7ec43252f6 100644 --- a/tests/firedrake/regression/test_fdm.py +++ b/tests/firedrake/regression/test_fdm.py @@ -1,4 +1,5 @@ import pytest +import numpy from firedrake import * from pyop2.utils import as_tuple from firedrake.petsc import DEFAULT_DIRECT_SOLVER @@ -339,7 +340,7 @@ def test_ipdg_direct_solver(fs): @pytest.mark.parallel @pytest.mark.parametrize("degree", range(1, 3)) @pytest.mark.parametrize("variant", ("spectral", "integral", "fdm")) -def test_tabulate_exterior_derivative(mesh, variant, degree): +def test_tabulate_gradient(mesh, variant, degree): from firedrake.preconditioners.fdm import tabulate_exterior_derivative tdim = mesh.topological_dimension() family = {1: "DG", 2: "RTCE", 3: "NCE"}[tdim] @@ -353,3 +354,26 @@ def test_tabulate_exterior_derivative(mesh, variant, degree): with u0.dat.vec as x, u1.dat.vec as y: D.mult(x, y) assert y.norm() < 1E-12 + + +# @pytest.mark.parallel(nprocs=2) +@pytest.mark.parametrize("mat_type", ("aij", "is")) +@pytest.mark.parametrize("degree", range(1, 3)) +@pytest.mark.parametrize("variant", ("spectral", "integral", "fdm")) +def test_tabulate_divergence(mesh, variant, degree, mat_type): + from firedrake.preconditioners.fdm import tabulate_exterior_derivative + tdim = mesh.topological_dimension() + family = {1: "CG", 2: "RTCF", 3: "NCF"}[tdim] + + V = FunctionSpace(mesh, family, degree, variant=variant) + Q = FunctionSpace(mesh, "DG", 0, variant=variant) + D = tabulate_exterior_derivative(V, Q, mat_type=mat_type, allow_repeated=True) + + # Dref = assemble(Interpolate(div(TrialFunction(V)), Q)).petscmat + Jdet = JacobianDeterminant(mesh) + s = abs(Jdet) / Jdet + Dref = assemble(inner(div(TrialFunction(V)) * s, TestFunction(Q))*dx).petscmat + alpha = -(-1)**tdim + Dref.axpy(-1*alpha, D.convert("aij")) + _, _, vals = Dref.getValuesCSR() + assert numpy.allclose(vals, 0) From fb3dda72b1023ac9d1e8604befc52bd36caa7512 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 1 Aug 2025 11:55:36 +0100 Subject: [PATCH 25/72] Fix tabulate_exterior_derivative in parallel --- firedrake/preconditioners/fdm.py | 33 +++++++++++++------------- tests/firedrake/regression/test_fdm.py | 29 ++++++++++++---------- 2 files changed, 33 insertions(+), 29 deletions(-) diff --git a/firedrake/preconditioners/fdm.py b/firedrake/preconditioners/fdm.py index 4ae835f59e..0adfef6649 100644 --- a/firedrake/preconditioners/fdm.py +++ b/firedrake/preconditioners/fdm.py @@ -1650,12 +1650,13 @@ def diff_blocks(tdim, formdegree, A00, A11, A10): def get_matis_lgmap(Vsub, allow_repeated): + """Construct the local to global mapping for MatIS assembly.""" lgmap = Vsub.dof_dset.lgmap if allow_repeated: indices = broken_function(Vsub, lgmap.indices).dat.data_ro else: indices = lgmap.indices.copy() - local_indices = numpy.arange(len(indices), dtype=PETSc.IntType) + local_indices = numpy.arange(indices.size, dtype=PETSc.IntType) cell_node_map = broken_function(Vsub, local_indices).dat.data_ro ghost = numpy.setdiff1d(local_indices, numpy.unique(cell_node_map), assume_unique=True) indices[ghost] = -1 @@ -1664,10 +1665,8 @@ def get_matis_lgmap(Vsub, allow_repeated): def tabulate_exterior_derivative(Vc, Vf, cbcs=[], fbcs=[], comm=None, mat_type="aij", allow_repeated=False): - """ - Tabulate exterior derivative: Vc -> Vf as an explicit sparse matrix. - Works for any tensor-product basis. These are the same matrices one needs for HypreAMS and friends. - """ + """Tabulate exterior derivative: Vc -> Vf as an explicit sparse matrix. + Works for any tensor-product basis. These are the same matrices one needs for HypreAMS and friends.""" if comm is None: comm = Vf.comm ec = Vc.finat_element @@ -1720,28 +1719,29 @@ def tabulate_exterior_derivative(Vc, Vf, cbcs=[], fbcs=[], comm=None, mat_type=" temp.destroy() eye.destroy() - sizes = tuple(V.dof_dset.layout_vec.getSizes() for V in (Vf, Vc)) - preallocator = PETSc.Mat().create(comm=comm) - preallocator.setType(PETSc.Mat.Type.PREALLOCATOR) - preallocator.setSizes(sizes) - preallocator.setUp() - spaces = (Vf, Vc) bcs = (fbcs, cbcs) if mat_type == "is": assert not any(bcs) lgmaps = tuple(get_matis_lgmap(V, allow_repeated) for V in spaces) - indices = tuple(mask_local_indices(V, V.dof_dset.lgmap, allow_repeated) for V in spaces) else: + allow_repeated = False lgmaps = tuple(V.local_to_global_map(bcs) for V, bcs in zip(spaces, bcs)) - indices = tuple(op2.Dat(V.dof_dset, lgmap.indices)(op2.READ, V.cell_node_map()) - for V, lgmap in zip(spaces, lgmaps)) - kernel = ElementKernel(Dhat, name="exterior_derivative").kernel() + indices_acc = tuple(mask_local_indices(V, V.dof_dset.lgmap, allow_repeated) for V in spaces) + + sizes = tuple(V.dof_dset.layout_vec.getSizes() for V in (Vf, Vc)) + preallocator = PETSc.Mat().create(comm=comm) + preallocator.setType(PETSc.Mat.Type.PREALLOCATOR) + preallocator.setSizes(sizes) + preallocator.setLGMap(*lgmaps) + preallocator.setUp() + + kernel = ElementKernel(Dhat, name="exterior_derivative").kernel(mat_type=mat_type) assembler = op2.ParLoop(kernel, Vc.mesh().cell_set, *(op2.PassthroughArg(op2.OpaqueType("Mat"), m.handle) for m in (preallocator, Dhat)), - *indices) + *indices_acc) assembler() preallocator.assemble() nnz = get_preallocation(preallocator, sizes[0][0]) @@ -1755,6 +1755,7 @@ def tabulate_exterior_derivative(Vc, Vf, cbcs=[], fbcs=[], comm=None, mat_type=" Dmat.setLGMap(*lgmaps) Dmat.setPreallocationNNZ(nnz) Dmat.setOption(PETSc.Mat.Option.NEW_NONZERO_ALLOCATION_ERR, True) + Dmat.setUp() assembler.arguments[0].data = Dmat.handle assembler() diff --git a/tests/firedrake/regression/test_fdm.py b/tests/firedrake/regression/test_fdm.py index 7ec43252f6..71f8fe8454 100644 --- a/tests/firedrake/regression/test_fdm.py +++ b/tests/firedrake/regression/test_fdm.py @@ -337,16 +337,17 @@ def test_ipdg_direct_solver(fs): assert norm(u_exact-uh, "H1") < 1.0E-8 -@pytest.mark.parallel +@pytest.mark.parallel(nprocs=2) +@pytest.mark.parametrize("mat_type", ("aij",)) @pytest.mark.parametrize("degree", range(1, 3)) @pytest.mark.parametrize("variant", ("spectral", "integral", "fdm")) -def test_tabulate_gradient(mesh, variant, degree): +def test_tabulate_gradient(mesh, variant, degree, mat_type): from firedrake.preconditioners.fdm import tabulate_exterior_derivative tdim = mesh.topological_dimension() family = {1: "DG", 2: "RTCE", 3: "NCE"}[tdim] + V0 = FunctionSpace(mesh, "Lagrange", degree, variant=variant) V1 = FunctionSpace(mesh, family, degree-(tdim == 1), variant=variant) - D = tabulate_exterior_derivative(V0, V1) u0 = Function(V0).interpolate(1) @@ -356,10 +357,10 @@ def test_tabulate_gradient(mesh, variant, degree): assert y.norm() < 1E-12 -# @pytest.mark.parallel(nprocs=2) -@pytest.mark.parametrize("mat_type", ("aij", "is")) -@pytest.mark.parametrize("degree", range(1, 3)) -@pytest.mark.parametrize("variant", ("spectral", "integral", "fdm")) +@pytest.mark.parallel(nprocs=2) +@pytest.mark.parametrize("mat_type", ("aij", "is",)) +@pytest.mark.parametrize("degree", (1, 2)) +@pytest.mark.parametrize("variant", ("integral",)) def test_tabulate_divergence(mesh, variant, degree, mat_type): from firedrake.preconditioners.fdm import tabulate_exterior_derivative tdim = mesh.topological_dimension() @@ -368,12 +369,14 @@ def test_tabulate_divergence(mesh, variant, degree, mat_type): V = FunctionSpace(mesh, family, degree, variant=variant) Q = FunctionSpace(mesh, "DG", 0, variant=variant) D = tabulate_exterior_derivative(V, Q, mat_type=mat_type, allow_repeated=True) - - # Dref = assemble(Interpolate(div(TrialFunction(V)), Q)).petscmat Jdet = JacobianDeterminant(mesh) - s = abs(Jdet) / Jdet - Dref = assemble(inner(div(TrialFunction(V)) * s, TestFunction(Q))*dx).petscmat - alpha = -(-1)**tdim - Dref.axpy(-1*alpha, D.convert("aij")) + alpha = (-1) ** (1 + tdim) + s = Function(Q).interpolate(alpha * abs(Jdet) / Jdet) + with s.dat.vec as svec: + D.diagonalScale(svec, None) + + Dref = assemble(inner(div(TrialFunction(V)), TestFunction(Q))*dx).petscmat + Dij = D if D.type.endswith("aij") else D.convert(Dref.type, PETSc.Mat()) + Dref.axpy(-1, Dij) _, _, vals = Dref.getValuesCSR() assert numpy.allclose(vals, 0) From 052bedd0cec08cd844f38f89d1a5556593a53baf Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 1 Aug 2025 12:20:49 +0100 Subject: [PATCH 26/72] BDDC get_divergence_mat --- firedrake/preconditioners/bddc.py | 52 +++++++++++++++++--------- tests/firedrake/regression/test_fdm.py | 2 + 2 files changed, 37 insertions(+), 17 deletions(-) diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index cee8fd359b..6c69a0d750 100644 --- a/firedrake/preconditioners/bddc.py +++ b/firedrake/preconditioners/bddc.py @@ -8,8 +8,7 @@ from firedrake.functionspace import FunctionSpace, VectorFunctionSpace, TensorFunctionSpace from firedrake.preconditioners.fdm import tabulate_exterior_derivative from firedrake.preconditioners.hiptmair import curl_to_grad -from firedrake import assemble -from ufl import curl, div, H1, H2, HCurl, HDiv, inner, dx +from ufl import curl, div, H1, H2, HCurl, HDiv, inner, dx, JacobianDeterminant from pyop2.utils import as_tuple import numpy @@ -95,15 +94,19 @@ def initialize(self, pc): tdim = V.mesh().topological_dimension() if tdim >= 2 and V.finat_element.formdegree == tdim-1: - # Should we use a callable like for hypre_ams? - B = appctx.get("divergence_mat", None) - if B is None: - from firedrake.assemble import assemble - d = {HCurl: curl, HDiv: div}[sobolev_space] - Q = V.reconstruct(family="DG", degree=degree-1) - b = inner(d(TrialFunction(V)), TestFunction(Q)) * dx(degree=2*(degree-1)) - B = assemble(b, mat_type="is") - bddcpc.setBDDCDivergenceMat(B.petscmat) + get_divergence = appctx.get("get_divergence_mat", get_divergence_mat) + qdegree = degree - 1 + if variant in ("fdm", "integral"): + qdegree = 0 + A, P = pc.getOperators() + allow_repeated = P.getISAllowRepeated() + divergence = get_divergence(V, degree=qdegree, mat_type="is", allow_repeated=allow_repeated) + try: + div_args, div_kwargs = divergence + except ValueError: + div_args = (divergence,) + div_kwargs = dict() + bddcpc.setBDDCDivergenceMat(*div_args, **div_kwargs) elif sobolev_space == HCurl: # Should we use a callable like for hypre_ams? gradient = appctx.get("discrete_gradient", None) @@ -143,9 +146,9 @@ def get_vertex_dofs(V): return vertex_dofs -def get_divergence_mat(V): +def get_divergence_mat(V, degree=None, mat_type="aij", allow_repeated=False): sobolev_space = V.ufl_element().sobolev_space - degree = max(as_tuple(V.ufl_element().degree())) + vdegree = max(as_tuple(V.ufl_element().degree())) d = {HCurl: curl, HDiv: div}[sobolev_space] if V.shape == (): make_function_space = FunctionSpace @@ -153,10 +156,25 @@ def get_divergence_mat(V): make_function_space = VectorFunctionSpace else: make_function_space = TensorFunctionSpace - Q = make_function_space(V.mesh(), "DG", degree-1) - b = inner(d(TrialFunction(V)), TestFunction(Q)) * dx(degree=2*(degree-1)) - B = assemble.assemble(b, mat_type="is") - return (B.petscmat,), dict() + + if degree is None: + degree = vdegree-1 + Q = make_function_space(V.mesh(), "DG", degree) + if False: + from firedrake import assemble + b = inner(d(TrialFunction(V)), TestFunction(Q)) * dx(degree=degree+vdegree-1) + B = assemble(b, mat_type=mat_type).petscmat + else: + B = tabulate_exterior_derivative(V, Q, mat_type=mat_type, allow_repeated=allow_repeated) + # Fix sign + tdim = V.mesh().topological_dimension() + alpha = (-1) ** (tdim-1) + Jdet = JacobianDeterminant(V.mesh()) + s = Function(Q).interpolate(alpha * abs(Jdet) / Jdet) + with s.dat.vec as svec: + B.diagonalScale(svec, None) + + return (B,), dict() def get_discrete_gradient(V): diff --git a/tests/firedrake/regression/test_fdm.py b/tests/firedrake/regression/test_fdm.py index 71f8fe8454..c2631867fb 100644 --- a/tests/firedrake/regression/test_fdm.py +++ b/tests/firedrake/regression/test_fdm.py @@ -369,6 +369,8 @@ def test_tabulate_divergence(mesh, variant, degree, mat_type): V = FunctionSpace(mesh, family, degree, variant=variant) Q = FunctionSpace(mesh, "DG", 0, variant=variant) D = tabulate_exterior_derivative(V, Q, mat_type=mat_type, allow_repeated=True) + + # Fix sign Jdet = JacobianDeterminant(mesh) alpha = (-1) ** (1 + tdim) s = Function(Q).interpolate(alpha * abs(Jdet) / Jdet) From 1ec8faa9f20c7076d1f2c82e0849216f55938397 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 1 Aug 2025 15:08:29 +0100 Subject: [PATCH 27/72] Fix high-order divergence tabulation --- firedrake/preconditioners/bddc.py | 24 +++++++++--------------- firedrake/preconditioners/fdm.py | 2 +- tests/firedrake/regression/test_fdm.py | 18 +++++++++--------- 3 files changed, 19 insertions(+), 25 deletions(-) diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index 6c69a0d750..bc6ac2cf2c 100644 --- a/firedrake/preconditioners/bddc.py +++ b/firedrake/preconditioners/bddc.py @@ -95,12 +95,9 @@ def initialize(self, pc): tdim = V.mesh().topological_dimension() if tdim >= 2 and V.finat_element.formdegree == tdim-1: get_divergence = appctx.get("get_divergence_mat", get_divergence_mat) - qdegree = degree - 1 - if variant in ("fdm", "integral"): - qdegree = 0 A, P = pc.getOperators() allow_repeated = P.getISAllowRepeated() - divergence = get_divergence(V, degree=qdegree, mat_type="is", allow_repeated=allow_repeated) + divergence = get_divergence(V, mat_type="is", allow_repeated=allow_repeated) try: div_args, div_kwargs = divergence except ValueError: @@ -146,9 +143,10 @@ def get_vertex_dofs(V): return vertex_dofs -def get_divergence_mat(V, degree=None, mat_type="aij", allow_repeated=False): +def get_divergence_mat(V, mat_type="aij", allow_repeated=False): + from firedrake import assemble sobolev_space = V.ufl_element().sobolev_space - vdegree = max(as_tuple(V.ufl_element().degree())) + degree = max(as_tuple(V.ufl_element().degree())) d = {HCurl: curl, HDiv: div}[sobolev_space] if V.shape == (): make_function_space = FunctionSpace @@ -157,20 +155,16 @@ def get_divergence_mat(V, degree=None, mat_type="aij", allow_repeated=False): else: make_function_space = TensorFunctionSpace - if degree is None: - degree = vdegree-1 - Q = make_function_space(V.mesh(), "DG", degree) + qdegree = degree-1 + Q = make_function_space(V.mesh(), "DG", 0, variant=f"integral({qdegree})") if False: - from firedrake import assemble - b = inner(d(TrialFunction(V)), TestFunction(Q)) * dx(degree=degree+vdegree-1) + b = inner(d(TrialFunction(V)), TestFunction(Q)) * dx(degree=qdegree) B = assemble(b, mat_type=mat_type).petscmat else: B = tabulate_exterior_derivative(V, Q, mat_type=mat_type, allow_repeated=allow_repeated) - # Fix sign - tdim = V.mesh().topological_dimension() - alpha = (-1) ** (tdim-1) + # Fix scale Jdet = JacobianDeterminant(V.mesh()) - s = Function(Q).interpolate(alpha * abs(Jdet) / Jdet) + s = assemble(inner(TrialFunction(Q)*(1/Jdet), TestFunction(Q))*dx(degree=0), diagonal=True) with s.dat.vec as svec: B.diagonalScale(svec, None) diff --git a/firedrake/preconditioners/fdm.py b/firedrake/preconditioners/fdm.py index 0adfef6649..8f4b37d4f8 100644 --- a/firedrake/preconditioners/fdm.py +++ b/firedrake/preconditioners/fdm.py @@ -1635,7 +1635,7 @@ def diff_blocks(tdim, formdegree, A00, A11, A10): A_blocks = [[A00.kron(A10)], [A10.kron(A00)]] elif formdegree == 1: A_blocks = [[A10.kron(A11), A11.kron(A10)]] - A_blocks[-1][-1].scale(-1) + A_blocks[0][0].scale(-1) elif tdim == 3: if formdegree == 0: A_blocks = [[kron3(A00, A00, A10)], [kron3(A00, A10, A00)], [kron3(A10, A00, A00)]] diff --git a/tests/firedrake/regression/test_fdm.py b/tests/firedrake/regression/test_fdm.py index c2631867fb..a5535dd4bc 100644 --- a/tests/firedrake/regression/test_fdm.py +++ b/tests/firedrake/regression/test_fdm.py @@ -338,8 +338,8 @@ def test_ipdg_direct_solver(fs): @pytest.mark.parallel(nprocs=2) -@pytest.mark.parametrize("mat_type", ("aij",)) -@pytest.mark.parametrize("degree", range(1, 3)) +@pytest.mark.parametrize("mat_type", ("aij", "is")) +@pytest.mark.parametrize("degree", (1, 4)) @pytest.mark.parametrize("variant", ("spectral", "integral", "fdm")) def test_tabulate_gradient(mesh, variant, degree, mat_type): from firedrake.preconditioners.fdm import tabulate_exterior_derivative @@ -358,27 +358,27 @@ def test_tabulate_gradient(mesh, variant, degree, mat_type): @pytest.mark.parallel(nprocs=2) -@pytest.mark.parametrize("mat_type", ("aij", "is",)) -@pytest.mark.parametrize("degree", (1, 2)) -@pytest.mark.parametrize("variant", ("integral",)) +@pytest.mark.parametrize("mat_type", ("aij", "is")) +@pytest.mark.parametrize("degree", (1, 4)) +@pytest.mark.parametrize("variant", ("spectral", "integral", "fdm")) def test_tabulate_divergence(mesh, variant, degree, mat_type): from firedrake.preconditioners.fdm import tabulate_exterior_derivative tdim = mesh.topological_dimension() family = {1: "CG", 2: "RTCF", 3: "NCF"}[tdim] V = FunctionSpace(mesh, family, degree, variant=variant) - Q = FunctionSpace(mesh, "DG", 0, variant=variant) + Q = FunctionSpace(mesh, "DG", 0, variant=f"integral({degree-1})") D = tabulate_exterior_derivative(V, Q, mat_type=mat_type, allow_repeated=True) - # Fix sign + # Fix scale Jdet = JacobianDeterminant(mesh) - alpha = (-1) ** (1 + tdim) - s = Function(Q).interpolate(alpha * abs(Jdet) / Jdet) + s = assemble(inner(TrialFunction(Q)*(1/Jdet), TestFunction(Q))*dx, diagonal=True) with s.dat.vec as svec: D.diagonalScale(svec, None) Dref = assemble(inner(div(TrialFunction(V)), TestFunction(Q))*dx).petscmat Dij = D if D.type.endswith("aij") else D.convert(Dref.type, PETSc.Mat()) + Dref.axpy(-1, Dij) _, _, vals = Dref.getValuesCSR() assert numpy.allclose(vals, 0) From 671261b2448125c8f438a6e8536eb2cddedd39a3 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 1 Aug 2025 16:28:46 +0100 Subject: [PATCH 28/72] Cleanup --- firedrake/preconditioners/bddc.py | 34 +++++++++--------------- firedrake/preconditioners/facet_split.py | 6 +++-- firedrake/preconditioners/fdm.py | 4 +-- 3 files changed, 18 insertions(+), 26 deletions(-) diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index bc6ac2cf2c..70e59febad 100644 --- a/firedrake/preconditioners/bddc.py +++ b/firedrake/preconditioners/bddc.py @@ -8,7 +8,7 @@ from firedrake.functionspace import FunctionSpace, VectorFunctionSpace, TensorFunctionSpace from firedrake.preconditioners.fdm import tabulate_exterior_derivative from firedrake.preconditioners.hiptmair import curl_to_grad -from ufl import curl, div, H1, H2, HCurl, HDiv, inner, dx, JacobianDeterminant +from ufl import H1, H2, HCurl, inner, dx, JacobianDeterminant from pyop2.utils import as_tuple import numpy @@ -145,28 +145,18 @@ def get_vertex_dofs(V): def get_divergence_mat(V, mat_type="aij", allow_repeated=False): from firedrake import assemble - sobolev_space = V.ufl_element().sobolev_space degree = max(as_tuple(V.ufl_element().degree())) - d = {HCurl: curl, HDiv: div}[sobolev_space] - if V.shape == (): - make_function_space = FunctionSpace - elif len(V.shape) == 1: - make_function_space = VectorFunctionSpace - else: - make_function_space = TensorFunctionSpace - - qdegree = degree-1 - Q = make_function_space(V.mesh(), "DG", 0, variant=f"integral({qdegree})") - if False: - b = inner(d(TrialFunction(V)), TestFunction(Q)) * dx(degree=qdegree) - B = assemble(b, mat_type=mat_type).petscmat - else: - B = tabulate_exterior_derivative(V, Q, mat_type=mat_type, allow_repeated=allow_repeated) - # Fix scale - Jdet = JacobianDeterminant(V.mesh()) - s = assemble(inner(TrialFunction(Q)*(1/Jdet), TestFunction(Q))*dx(degree=0), diagonal=True) - with s.dat.vec as svec: - B.diagonalScale(svec, None) + + Q = TensorFunctionSpace(V.mesh(), "DG", 0, variant=f"integral({degree-1})", shape=V.value_shape[:-1]) + # d = {HCurl: curl, HDiv: div}[V.ufl_element().sobolev_space] + # b = inner(d(TrialFunction(V)), TestFunction(Q)) * dx(degree=degree-1) + # B = assemble(b, mat_type=mat_type).petscmat + + B = tabulate_exterior_derivative(V, Q, mat_type=mat_type, allow_repeated=allow_repeated) + Jdet = JacobianDeterminant(V.mesh()) + s = assemble(inner(TrialFunction(Q)*(1/Jdet), TestFunction(Q))*dx(degree=0), diagonal=True) + with s.dat.vec as svec: + B.diagonalScale(svec, None) return (B,), dict() diff --git a/firedrake/preconditioners/facet_split.py b/firedrake/preconditioners/facet_split.py index ac785af650..722319631f 100644 --- a/firedrake/preconditioners/facet_split.py +++ b/firedrake/preconditioners/facet_split.py @@ -1,7 +1,7 @@ from functools import partial from mpi4py import MPI from pyop2 import op2, PermutedMap -from finat.ufl import RestrictedElement, MixedElement, TensorElement, VectorElement +from finat.ufl import BrokenElement, RestrictedElement, MixedElement, TensorElement, VectorElement from firedrake.petsc import PETSc from firedrake.preconditioners.base import PCBase from firedrake.bcs import restricted_function_space @@ -209,7 +209,9 @@ def restrict(ele, restriction_domain): if isinstance(ele, VectorElement): return type(ele)(restrict(ele._sub_element, restriction_domain), dim=ele.num_sub_elements) elif isinstance(ele, TensorElement): - return type(ele)(restrict(ele._sub_element, restriction_domain), shape=ele._shape, symmetry=ele._symmety) + return type(ele)(restrict(ele._sub_element, restriction_domain), shape=ele._shape, symmetry=ele._symmetry) + elif restriction_domain == "broken": + return BrokenElement(ele) else: return RestrictedElement(ele, restriction_domain) diff --git a/firedrake/preconditioners/fdm.py b/firedrake/preconditioners/fdm.py index 8f4b37d4f8..ce73b99e3f 100644 --- a/firedrake/preconditioners/fdm.py +++ b/firedrake/preconditioners/fdm.py @@ -8,7 +8,7 @@ evaluate_dual, get_permutation_to_nodal_elements, cache_generate_code) -from firedrake.preconditioners.facet_split import split_dofs, restricted_dofs +from firedrake.preconditioners.facet_split import restrict, restricted_dofs, split_dofs from firedrake.formmanipulation import ExtractSubBlock from firedrake.functionspace import FunctionSpace, MixedFunctionSpace from firedrake.function import Function @@ -41,7 +41,7 @@ def broken_function(V, val): - W = FunctionSpace(V.mesh(), finat.ufl.BrokenElement(V.ufl_element())) + W = FunctionSpace(V.mesh(), restrict(V.ufl_element(), "broken")) w = Function(W, dtype=val.dtype) v = Function(V, val=val) domain = "{[i]: 0 <= i < v.dofs}" From 6926ad1ab6ba078dc5c565b5ad9b1a9b426ddf6e Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 1 Aug 2025 17:12:41 +0100 Subject: [PATCH 29/72] cleanup --- firedrake/preconditioners/bddc.py | 27 +++++++++------------------ 1 file changed, 9 insertions(+), 18 deletions(-) diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index 70e59febad..7622ba3d60 100644 --- a/firedrake/preconditioners/bddc.py +++ b/firedrake/preconditioners/bddc.py @@ -8,7 +8,7 @@ from firedrake.functionspace import FunctionSpace, VectorFunctionSpace, TensorFunctionSpace from firedrake.preconditioners.fdm import tabulate_exterior_derivative from firedrake.preconditioners.hiptmair import curl_to_grad -from ufl import H1, H2, HCurl, inner, dx, JacobianDeterminant +from ufl import H1, H2, inner, dx, JacobianDeterminant from pyop2.utils import as_tuple import numpy @@ -94,9 +94,9 @@ def initialize(self, pc): tdim = V.mesh().topological_dimension() if tdim >= 2 and V.finat_element.formdegree == tdim-1: - get_divergence = appctx.get("get_divergence_mat", get_divergence_mat) A, P = pc.getOperators() allow_repeated = P.getISAllowRepeated() + get_divergence = appctx.get("get_divergence_mat", get_divergence_mat) divergence = get_divergence(V, mat_type="is", allow_repeated=allow_repeated) try: div_args, div_kwargs = divergence @@ -104,17 +104,12 @@ def initialize(self, pc): div_args = (divergence,) div_kwargs = dict() bddcpc.setBDDCDivergenceMat(*div_args, **div_kwargs) - elif sobolev_space == HCurl: - # Should we use a callable like for hypre_ams? - gradient = appctx.get("discrete_gradient", None) - if gradient is None: - from firedrake.preconditioners.fdm import tabulate_exterior_derivative - from firedrake.preconditioners.hiptmair import curl_to_grad - Q = V.reconstruct(element=curl_to_grad(V.ufl_element())) - gradient = tabulate_exterior_derivative(Q, V) - if variant == 'fdm': - corners = get_vertex_dofs(Q) - gradient.compose('_elements_corners', corners) + elif tdim >= 2 and V.finat_element.formdegree == 1: + get_gradient = appctx.get("get_discrete_gradient", get_discrete_gradient) + gradient = get_gradient(V) + try: + grad_args, grad_kwargs = gradient + except ValueError: grad_args = (gradient,) grad_kwargs = dict() bddcpc.setBDDCDiscreteGradient(*grad_args, **grad_kwargs) @@ -146,13 +141,9 @@ def get_vertex_dofs(V): def get_divergence_mat(V, mat_type="aij", allow_repeated=False): from firedrake import assemble degree = max(as_tuple(V.ufl_element().degree())) - Q = TensorFunctionSpace(V.mesh(), "DG", 0, variant=f"integral({degree-1})", shape=V.value_shape[:-1]) - # d = {HCurl: curl, HDiv: div}[V.ufl_element().sobolev_space] - # b = inner(d(TrialFunction(V)), TestFunction(Q)) * dx(degree=degree-1) - # B = assemble(b, mat_type=mat_type).petscmat - B = tabulate_exterior_derivative(V, Q, mat_type=mat_type, allow_repeated=allow_repeated) + Jdet = JacobianDeterminant(V.mesh()) s = assemble(inner(TrialFunction(Q)*(1/Jdet), TestFunction(Q))*dx(degree=0), diagonal=True) with s.dat.vec as svec: From f5e1651fe53bf833f81c8a43bbc2dc8b5b5f25ba Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 2 Aug 2025 17:49:10 +0100 Subject: [PATCH 30/72] Fix Q1 --- firedrake/preconditioners/fdm.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/firedrake/preconditioners/fdm.py b/firedrake/preconditioners/fdm.py index ce73b99e3f..ce2c0420b6 100644 --- a/firedrake/preconditioners/fdm.py +++ b/firedrake/preconditioners/fdm.py @@ -218,15 +218,15 @@ def allocate_matrix(self, Amat, V, J, bcs, fcp, pmat_type, use_static_condensati Vfacet = None Vbig = V ebig = V.ufl_element() - _, fdofs = split_dofs(V.finat_element) + idofs, fdofs = split_dofs(V.finat_element) elif len(ifacet) == 1: Vfacet = V[ifacet[0]] ebig, = set(unrestrict_element(Vsub.ufl_element()) for Vsub in V) Vbig = V.reconstruct(element=ebig) - if len(V) > 1: - dims = [Vsub.finat_element.space_dimension() for Vsub in V] - assert sum(dims) == Vbig.finat_element.space_dimension() + space_dim = Vbig.finat_element.space_dimension() + assert space_dim == sum(Vsub.finat_element.space_dimension() for Vsub in V) fdofs = restricted_dofs(Vfacet.finat_element, Vbig.finat_element) + idofs = numpy.setdiff1d(numpy.arange(space_dim, dtype=fdofs.dtype), fdofs) else: raise ValueError("Expecting at most one FunctionSpace restricted onto facets.") self.embedding_element = ebig @@ -245,7 +245,7 @@ def allocate_matrix(self, Amat, V, J, bcs, fcp, pmat_type, use_static_condensati # Dictionary with kernel to compute the Schur complement self.schur_kernel = {} - if V == Vbig and Vbig.finat_element.formdegree == 0: + if V == Vbig and Vbig.finat_element.formdegree == 0 and len(idofs) > 0 and pmat_type.endswith("aij"): # If we are in H(grad), we just pad with zeros on the statically-condensed pattern self.schur_kernel[V] = SchurComplementPattern elif Vfacet and use_static_condensation: From 9ff5cf6b042ce0c83fa92bb3bec5e9bd02578351 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 5 Aug 2025 12:25:03 +0100 Subject: [PATCH 31/72] Test tabulate_exterior_derivative --- tests/firedrake/regression/test_fdm.py | 48 ++++++++++++++++++++------ 1 file changed, 37 insertions(+), 11 deletions(-) diff --git a/tests/firedrake/regression/test_fdm.py b/tests/firedrake/regression/test_fdm.py index a5535dd4bc..5a54d53e94 100644 --- a/tests/firedrake/regression/test_fdm.py +++ b/tests/firedrake/regression/test_fdm.py @@ -338,7 +338,7 @@ def test_ipdg_direct_solver(fs): @pytest.mark.parallel(nprocs=2) -@pytest.mark.parametrize("mat_type", ("aij", "is")) +@pytest.mark.parametrize("mat_type", ("aij",)) @pytest.mark.parametrize("degree", (1, 4)) @pytest.mark.parametrize("variant", ("spectral", "integral", "fdm")) def test_tabulate_gradient(mesh, variant, degree, mat_type): @@ -348,13 +348,40 @@ def test_tabulate_gradient(mesh, variant, degree, mat_type): V0 = FunctionSpace(mesh, "Lagrange", degree, variant=variant) V1 = FunctionSpace(mesh, family, degree-(tdim == 1), variant=variant) - D = tabulate_exterior_derivative(V0, V1) + D = tabulate_exterior_derivative(V0, V1, mat_type=mat_type) + + M = assemble(inner(TrialFunction(V1), TestFunction(V1))*dx).petscmat + Dij = D if D.type.endswith("aij") else D.convert(M.type, PETSc.Mat()) + B = M.matMult(Dij) - u0 = Function(V0).interpolate(1) - u1 = Function(V1) - with u0.dat.vec as x, u1.dat.vec as y: - D.mult(x, y) - assert y.norm() < 1E-12 + Bref = assemble(inner(grad(TrialFunction(V0)), TestFunction(V1))*dx).petscmat + Bref.axpy(-1, B) + _, _, vals = Bref.getValuesCSR() + assert numpy.allclose(vals, 0) + + +@pytest.mark.parallel(nprocs=2) +@pytest.mark.parametrize("mat_type", ("aij",)) +@pytest.mark.parametrize("degree", (1, 4)) +@pytest.mark.parametrize("variant", ("spectral", "integral", "fdm")) +def test_tabulate_curl(mesh, variant, degree, mat_type): + from firedrake.preconditioners.fdm import tabulate_exterior_derivative + tdim = mesh.topological_dimension() + family1 = {1: "CG", 2: "CG", 3: "NCE"}[tdim] + family2 = {1: "DG", 2: "RTCF", 3: "NCF"}[tdim] + + V1 = FunctionSpace(mesh, family1, degree, variant=variant) + V2 = FunctionSpace(mesh, family2, degree-(tdim == 1), variant=variant) + D = tabulate_exterior_derivative(V1, V2, mat_type=mat_type) + + M = assemble((-1)**(tdim-1)*inner(TrialFunction(V2), TestFunction(V2))*dx).petscmat + Dij = D if D.type.endswith("aij") else D.convert(M.type, PETSc.Mat()) + B = M.matMult(Dij) + + Bref = assemble(inner(curl(TrialFunction(V1)), TestFunction(V2))*dx).petscmat + Bref.axpy(-1, B) + _, _, vals = Bref.getValuesCSR() + assert numpy.allclose(vals, 0) @pytest.mark.parallel(nprocs=2) @@ -372,13 +399,12 @@ def test_tabulate_divergence(mesh, variant, degree, mat_type): # Fix scale Jdet = JacobianDeterminant(mesh) - s = assemble(inner(TrialFunction(Q)*(1/Jdet), TestFunction(Q))*dx, diagonal=True) - with s.dat.vec as svec: - D.diagonalScale(svec, None) + M = assemble(inner(TrialFunction(Q)*(1/Jdet), TestFunction(Q))*dx, diagonal=True) + with M.dat.vec as mvec: + D.diagonalScale(mvec, None) Dref = assemble(inner(div(TrialFunction(V)), TestFunction(Q))*dx).petscmat Dij = D if D.type.endswith("aij") else D.convert(Dref.type, PETSc.Mat()) - Dref.axpy(-1, Dij) _, _, vals = Dref.getValuesCSR() assert numpy.allclose(vals, 0) From de9240b9db29deca32c8d71a9008aba27523d456 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 5 Aug 2025 12:25:20 +0100 Subject: [PATCH 32/72] H(curl) constraints for BDDC --- firedrake/preconditioners/bddc.py | 56 +++++++++++++++++++++++++------ 1 file changed, 46 insertions(+), 10 deletions(-) diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index 7622ba3d60..2836418064 100644 --- a/firedrake/preconditioners/bddc.py +++ b/firedrake/preconditioners/bddc.py @@ -104,7 +104,8 @@ def initialize(self, pc): div_args = (divergence,) div_kwargs = dict() bddcpc.setBDDCDivergenceMat(*div_args, **div_kwargs) - elif tdim >= 2 and V.finat_element.formdegree == 1: + + elif tdim >= 3 and V.finat_element.formdegree == 1: get_gradient = appctx.get("get_discrete_gradient", get_discrete_gradient) gradient = get_gradient(V) try: @@ -130,12 +131,19 @@ def applyTranspose(self, pc, x, y): self.pc.applyTranspose(x, y) -def get_vertex_dofs(V): - W = V.reconstruct(element=restrict(V.ufl_element(), "vertex")) +def get_dof_mapping(V, W): indices = get_restriction_indices(V, W) indices = V.dof_dset.lgmap.apply(indices) - vertex_dofs = PETSc.IS().createGeneral(indices, comm=V.comm) - return vertex_dofs + return PETSc.IS().createGeneral(indices, comm=V.comm) + + +def get_restricted_dofs(V, domain): + W = FunctionSpace(V.mesh(), restrict(V.ufl_element(), domain)) + return get_dof_mapping(V, W) + + +def get_low_order_dofs(V): + return get_dof_mapping(V, V.reconstruct(degree=1)) def get_divergence_mat(V, mat_type="aij", allow_repeated=False): @@ -153,13 +161,41 @@ def get_divergence_mat(V, mat_type="aij", allow_repeated=False): def get_discrete_gradient(V): - degree = max(as_tuple(V.ufl_element().degree())) - variant = V.ufl_element().variant() + from firedrake import VectorSpaceBasis + Q = FunctionSpace(V.mesh(), curl_to_grad(V.ufl_element())) + degree = max(as_tuple(Q.ufl_element().degree())) + variant = Q.ufl_element().variant() gradient = tabulate_exterior_derivative(Q, V) - if variant == 'fdm': - corners = get_vertex_dofs(Q) - gradient.compose('_elements_corners', corners) + + if variant == "fdm": + vdofs = get_restricted_dofs(Q, "vertex") + gradient.compose('_elements_corners', vdofs) + + # Constant moments along edges + wdofs = get_low_order_dofs(V) + v0 = Function(V) + with v0.dat.vec as y: + y.setValues(wdofs, numpy.ones((wdofs.getSizes()[0],), dtype=PETSc.RealType)) + + # Linear moments along edges + # Constant moments of gradients of H1 edge bubbles + p = Function(Q).interpolate(1) + q = Function(Q) + + edofs = get_restricted_dofs(Q, "edge") + edofs = edofs.difference(vdofs) + with q.dat.vec as qvec, p.dat.vec as pvec: + qvec.setValues(edofs, pvec.getValues(edofs)) + + v1 = Function(V) + with v1.dat.vec as y, q.dat.vec as x: + gradient.mult(x, y) + y.chop(1E-14) + + nsp = VectorSpaceBasis([v0, v1]) + gradient.compose("_constraints", nsp.nullspace()) + grad_args = (gradient,) grad_kwargs = {'order': degree} return grad_args, grad_kwargs From b8bcb0f4c6df213bf0c9f646d8c322dde63d76ad Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 6 Aug 2025 12:30:36 +0100 Subject: [PATCH 33/72] cleanup --- firedrake/preconditioners/bddc.py | 12 +++-- firedrake/preconditioners/fdm.py | 88 +++++++++++++++---------------- 2 files changed, 51 insertions(+), 49 deletions(-) diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index 2836418064..c26ec48858 100644 --- a/firedrake/preconditioners/bddc.py +++ b/firedrake/preconditioners/bddc.py @@ -63,7 +63,7 @@ def initialize(self, pc): # Handle boundary dofs ctx = get_appctx(dm) - bcs = tuple(ctx._problem.bcs) + bcs = tuple(ctx._problem.dirichlet_bcs()) if V.extruded: boundary_nodes = numpy.unique(numpy.concatenate(list(map(V.boundary_nodes, ("on_boundary", "top", "bottom"))))) else: @@ -156,27 +156,27 @@ def get_divergence_mat(V, mat_type="aij", allow_repeated=False): s = assemble(inner(TrialFunction(Q)*(1/Jdet), TestFunction(Q))*dx(degree=0), diagonal=True) with s.dat.vec as svec: B.diagonalScale(svec, None) - - return (B,), dict() + return (B,), {} def get_discrete_gradient(V): from firedrake import VectorSpaceBasis Q = FunctionSpace(V.mesh(), curl_to_grad(V.ufl_element())) - degree = max(as_tuple(Q.ufl_element().degree())) - variant = Q.ufl_element().variant() gradient = tabulate_exterior_derivative(Q, V) + variant = Q.ufl_element().variant() if variant == "fdm": vdofs = get_restricted_dofs(Q, "vertex") gradient.compose('_elements_corners', vdofs) + if False: # Constant moments along edges wdofs = get_low_order_dofs(V) v0 = Function(V) with v0.dat.vec as y: y.setValues(wdofs, numpy.ones((wdofs.getSizes()[0],), dtype=PETSc.RealType)) + y.assemble() # Linear moments along edges # Constant moments of gradients of H1 edge bubbles @@ -187,6 +187,7 @@ def get_discrete_gradient(V): edofs = edofs.difference(vdofs) with q.dat.vec as qvec, p.dat.vec as pvec: qvec.setValues(edofs, pvec.getValues(edofs)) + qvec.assemble() v1 = Function(V) with v1.dat.vec as y, q.dat.vec as x: @@ -196,6 +197,7 @@ def get_discrete_gradient(V): nsp = VectorSpaceBasis([v0, v1]) gradient.compose("_constraints", nsp.nullspace()) + degree = max(as_tuple(Q.ufl_element().degree())) grad_args = (gradient,) grad_kwargs = {'order': degree} return grad_args, grad_kwargs diff --git a/firedrake/preconditioners/fdm.py b/firedrake/preconditioners/fdm.py index ce2c0420b6..9599d1b628 100644 --- a/firedrake/preconditioners/fdm.py +++ b/firedrake/preconditioners/fdm.py @@ -40,33 +40,6 @@ __all__ = ("FDMPC", "PoissonFDMPC") -def broken_function(V, val): - W = FunctionSpace(V.mesh(), restrict(V.ufl_element(), "broken")) - w = Function(W, dtype=val.dtype) - v = Function(V, val=val) - domain = "{[i]: 0 <= i < v.dofs}" - instructions = """ - for i - w[i] = v[i] - end - """ - par_loop((domain, instructions), ufl.dx, {'w': (w, op2.WRITE), 'v': (v, op2.READ)}) - return w - - -def mask_local_indices(V, lgmap, repeated=False): - mask = lgmap.indices - if repeated: - w = broken_function(V, mask) - V = w.function_space() - mask = w.dat.data_ro_with_halos - indices = numpy.arange(len(mask), dtype=PETSc.IntType) - indices[mask == -1] = -1 - indices_dat = V.make_dat(val=indices) - indices_acc = indices_dat(op2.READ, V.cell_node_map()) - return indices_acc - - class FDMPC(PCBase): """ A preconditioner for tensor-product elements that changes the shape @@ -710,7 +683,7 @@ def insert_mode(self): def assembly_lgmaps(self): if self.mat_type != "is": return {Vsub: Vsub.dof_dset.lgmap for Vsub in self.V} - return {Vsub: get_matis_lgmap(Vsub, self.allow_repeated) for Vsub in self.V} + return {Vsub: unghosted_lgmap(Vsub, Vsub.dof_dset.lgmap, self.allow_repeated) for Vsub in self.V} def setup_block(self, Vrow, Vcol): # Preallocate the auxiliary sparse operator @@ -1649,15 +1622,42 @@ def diff_blocks(tdim, formdegree, A00, A11, A10): return A_blocks -def get_matis_lgmap(Vsub, allow_repeated): +def broken_function(V, val): + W = FunctionSpace(V.mesh(), restrict(V.ufl_element(), "broken")) + w = Function(W, dtype=val.dtype) + v = Function(V, val=val) + domain = "{[i]: 0 <= i < v.dofs}" + instructions = """ + for i + w[i] = v[i] + end + """ + par_loop((domain, instructions), ufl.dx, {'w': (w, op2.WRITE), 'v': (v, op2.READ)}) + return w + + +def mask_local_indices(V, lgmap, allow_repeated): + mask = lgmap.indices + if allow_repeated: + w = broken_function(V, mask) + V = w.function_space() + mask = w.dat.data_ro_with_halos + + indices = numpy.arange(mask.size, dtype=PETSc.IntType) + indices[mask == -1] = -1 + indices_dat = V.make_dat(val=indices) + indices_acc = indices_dat(op2.READ, V.cell_node_map()) + return indices_acc + + +def unghosted_lgmap(V, lgmap, allow_repeated): """Construct the local to global mapping for MatIS assembly.""" - lgmap = Vsub.dof_dset.lgmap if allow_repeated: - indices = broken_function(Vsub, lgmap.indices).dat.data_ro + indices = broken_function(V, lgmap.indices).dat.data_ro else: indices = lgmap.indices.copy() local_indices = numpy.arange(indices.size, dtype=PETSc.IntType) - cell_node_map = broken_function(Vsub, local_indices).dat.data_ro + cell_node_map = broken_function(V, local_indices).dat.data_ro ghost = numpy.setdiff1d(local_indices, numpy.unique(cell_node_map), assume_unique=True) indices[ghost] = -1 @@ -1669,6 +1669,7 @@ def tabulate_exterior_derivative(Vc, Vf, cbcs=[], fbcs=[], comm=None, mat_type=" Works for any tensor-product basis. These are the same matrices one needs for HypreAMS and friends.""" if comm is None: comm = Vf.comm + ec = Vc.finat_element ef = Vf.finat_element if ef.formdegree - ec.formdegree != 1: @@ -1719,28 +1720,27 @@ def tabulate_exterior_derivative(Vc, Vf, cbcs=[], fbcs=[], comm=None, mat_type=" temp.destroy() eye.destroy() + if mat_type != "is": + allow_repeated = False spaces = (Vf, Vc) bcs = (fbcs, cbcs) + lgmaps = tuple(V.local_to_global_map(bcs) for V, bcs in zip(spaces, bcs)) + indices_acc = tuple(mask_local_indices(V, lgmap, allow_repeated) for V, lgmap in zip(spaces, lgmaps)) if mat_type == "is": - assert not any(bcs) - lgmaps = tuple(get_matis_lgmap(V, allow_repeated) for V in spaces) - else: - allow_repeated = False - lgmaps = tuple(V.local_to_global_map(bcs) for V, bcs in zip(spaces, bcs)) + lgmaps = tuple(unghosted_lgmap(V, lgmap, allow_repeated) for V, lgmap in zip(spaces, lgmaps)) - indices_acc = tuple(mask_local_indices(V, V.dof_dset.lgmap, allow_repeated) for V in spaces) - - sizes = tuple(V.dof_dset.layout_vec.getSizes() for V in (Vf, Vc)) + sizes = tuple(V.dof_dset.layout_vec.getSizes() for V in spaces) preallocator = PETSc.Mat().create(comm=comm) preallocator.setType(PETSc.Mat.Type.PREALLOCATOR) preallocator.setSizes(sizes) preallocator.setLGMap(*lgmaps) + preallocator.setOption(PETSc.Mat.Option.IGNORE_ZERO_ENTRIES, False) preallocator.setUp() - kernel = ElementKernel(Dhat, name="exterior_derivative").kernel(mat_type=mat_type) - assembler = op2.ParLoop(kernel, + kernel = ElementKernel(Dhat, name="exterior_derivative") + assembler = op2.ParLoop(kernel.kernel(mat_type=mat_type), Vc.mesh().cell_set, - *(op2.PassthroughArg(op2.OpaqueType("Mat"), m.handle) for m in (preallocator, Dhat)), + *kernel.make_args(preallocator), *indices_acc) assembler() preallocator.assemble() @@ -1758,8 +1758,8 @@ def tabulate_exterior_derivative(Vc, Vf, cbcs=[], fbcs=[], comm=None, mat_type=" Dmat.setUp() assembler.arguments[0].data = Dmat.handle assembler() - Dmat.assemble() + Dhat.destroy() return Dmat From 6e216b4a20560843e3797291daad557205e3dfdf Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 6 Aug 2025 16:49:24 +0100 Subject: [PATCH 34/72] Refactor --- firedrake/preconditioners/fdm.py | 100 +++++++++++++++---------------- 1 file changed, 50 insertions(+), 50 deletions(-) diff --git a/firedrake/preconditioners/fdm.py b/firedrake/preconditioners/fdm.py index 9599d1b628..d473115fb0 100644 --- a/firedrake/preconditioners/fdm.py +++ b/firedrake/preconditioners/fdm.py @@ -686,50 +686,25 @@ def assembly_lgmaps(self): return {Vsub: unghosted_lgmap(Vsub, Vsub.dof_dset.lgmap, self.allow_repeated) for Vsub in self.V} def setup_block(self, Vrow, Vcol): - # Preallocate the auxiliary sparse operator + """Preallocate the auxiliary sparse operator.""" sizes = tuple(Vsub.dof_dset.layout_vec.getSizes() for Vsub in (Vrow, Vcol)) rmap = self.assembly_lgmaps[Vrow] cmap = self.assembly_lgmaps[Vcol] on_diag = Vrow == Vcol ptype = self.mat_type if on_diag else PETSc.Mat.Type.AIJ - preallocator = PETSc.Mat().create(comm=self.comm) - preallocator.setType(PETSc.Mat.Type.PREALLOCATOR) - preallocator.setSizes(sizes) - preallocator.setISAllowRepeated(self.allow_repeated) - preallocator.setLGMap(rmap, cmap) - preallocator.setOption(PETSc.Mat.Option.IGNORE_ZERO_ENTRIES, False) - if ptype.endswith("sbaij"): - preallocator.setOption(PETSc.Mat.Option.IGNORE_LOWER_TRIANGULAR, True) - preallocator.setUp() + preallocator = get_preallocator(self.comm, sizes, rmap, cmap, mat_type=ptype) self.set_values(preallocator, Vrow, Vcol) preallocator.assemble() - dnz, onz = get_preallocation(preallocator, sizes[0][0]) - if on_diag: - numpy.maximum(dnz, 1, out=dnz) + P = allocate_matrix(preallocator, ptype, on_diag=on_diag, allow_repeated=self.allow_repeated) preallocator.destroy() - P = PETSc.Mat().create(comm=self.comm) - P.setType(ptype) - P.setSizes(sizes) - P.setISAllowRepeated(self.allow_repeated) - P.setLGMap(rmap, cmap) - if on_diag and ptype == "is" and self.allow_repeated: - bsize = Vrow.finat_element.space_dimension() * Vrow.block_size + + if on_diag and P.type == "is" and self.allow_repeated: + bsize = Vrow.block_size * Vrow.finat_element.space_dimension() local_mat = P.getISLocalMat() nblocks = local_mat.getSize()[0] // bsize sizes = numpy.full((nblocks,), bsize, dtype=PETSc.IntType) local_mat.setVariableBlockSizes(sizes) - P.setPreallocationNNZ((dnz, onz)) - - if not (ptype.endswith("sbaij") or ptype == "is"): - P.setOption(PETSc.Mat.Option.UNUSED_NONZERO_LOCATION_ERR, True) - P.setOption(PETSc.Mat.Option.NEW_NONZERO_ALLOCATION_ERR, True) - P.setOption(PETSc.Mat.Option.STRUCTURALLY_SYMMETRIC, on_diag) - P.setOption(PETSc.Mat.Option.FORCE_DIAGONAL_ENTRIES, True) - P.setOption(PETSc.Mat.Option.KEEP_NONZERO_PATTERN, True) - if ptype.endswith("sbaij"): - P.setOption(PETSc.Mat.Option.IGNORE_LOWER_TRIANGULAR, True) - P.setUp() return P @PETSc.Log.EventDecorator("FDMSetValues") @@ -1623,6 +1598,7 @@ def diff_blocks(tdim, formdegree, A00, A11, A10): def broken_function(V, val): + """Return a Function(V, val=val) interpolated onto the broken space.""" W = FunctionSpace(V.mesh(), restrict(V.ufl_element(), "broken")) w = Function(W, dtype=val.dtype) v = Function(V, val=val) @@ -1637,6 +1613,7 @@ def broken_function(V, val): def mask_local_indices(V, lgmap, allow_repeated): + """Return a numpy array with the masked local indices.""" mask = lgmap.indices if allow_repeated: w = broken_function(V, mask) @@ -1660,10 +1637,48 @@ def unghosted_lgmap(V, lgmap, allow_repeated): cell_node_map = broken_function(V, local_indices).dat.data_ro ghost = numpy.setdiff1d(local_indices, numpy.unique(cell_node_map), assume_unique=True) indices[ghost] = -1 - return PETSc.LGMap().create(indices, bsize=lgmap.getBlockSize(), comm=lgmap.getComm()) +def get_preallocator(comm, sizes, rmap, cmap, mat_type=None): + """Set up a matrix preallocator.""" + preallocator = PETSc.Mat().create(comm=comm) + preallocator.setType(PETSc.Mat.Type.PREALLOCATOR) + preallocator.setSizes(sizes) + preallocator.setLGMap(rmap, cmap) + preallocator.setOption(PETSc.Mat.Option.IGNORE_ZERO_ENTRIES, False) + if mat_type is not None and mat_type.endswith("sbaij"): + preallocator.setOption(PETSc.Mat.Option.IGNORE_LOWER_TRIANGULAR, True) + preallocator.setUp() + return preallocator + + +def allocate_matrix(preallocator, mat_type, on_diag=False, allow_repeated=False): + """Set up a matrix from a preallocator.""" + sizes = preallocator.getSizes() + nnz = get_preallocation(preallocator, sizes[0][0]) + if on_diag: + numpy.maximum(nnz[0], 1, out=nnz[0]) + + A = PETSc.Mat().create(comm=preallocator.getComm()) + A.setType(mat_type) + A.setSizes(sizes) + A.setBlockSize(preallocator.getBlockSize()) + A.setISAllowRepeated(allow_repeated) + A.setLGMap(*preallocator.getLGMap()) + A.setPreallocationNNZ(nnz) + if mat_type.endswith("sbaij"): + A.setOption(PETSc.Mat.Option.IGNORE_LOWER_TRIANGULAR, True) + if not (mat_type.endswith("sbaij") or mat_type == "is"): + A.setOption(PETSc.Mat.Option.UNUSED_NONZERO_LOCATION_ERR, True) + A.setOption(PETSc.Mat.Option.NEW_NONZERO_ALLOCATION_ERR, True) + A.setOption(PETSc.Mat.Option.STRUCTURALLY_SYMMETRIC, on_diag) + A.setOption(PETSc.Mat.Option.FORCE_DIAGONAL_ENTRIES, on_diag) + A.setOption(PETSc.Mat.Option.KEEP_NONZERO_PATTERN, True) + A.setUp() + return A + + def tabulate_exterior_derivative(Vc, Vf, cbcs=[], fbcs=[], comm=None, mat_type="aij", allow_repeated=False): """Tabulate exterior derivative: Vc -> Vf as an explicit sparse matrix. Works for any tensor-product basis. These are the same matrices one needs for HypreAMS and friends.""" @@ -1730,12 +1745,7 @@ def tabulate_exterior_derivative(Vc, Vf, cbcs=[], fbcs=[], comm=None, mat_type=" lgmaps = tuple(unghosted_lgmap(V, lgmap, allow_repeated) for V, lgmap in zip(spaces, lgmaps)) sizes = tuple(V.dof_dset.layout_vec.getSizes() for V in spaces) - preallocator = PETSc.Mat().create(comm=comm) - preallocator.setType(PETSc.Mat.Type.PREALLOCATOR) - preallocator.setSizes(sizes) - preallocator.setLGMap(*lgmaps) - preallocator.setOption(PETSc.Mat.Option.IGNORE_ZERO_ENTRIES, False) - preallocator.setUp() + preallocator = get_preallocator(comm, sizes, *lgmaps) kernel = ElementKernel(Dhat, name="exterior_derivative") assembler = op2.ParLoop(kernel.kernel(mat_type=mat_type), @@ -1744,22 +1754,12 @@ def tabulate_exterior_derivative(Vc, Vf, cbcs=[], fbcs=[], comm=None, mat_type=" *indices_acc) assembler() preallocator.assemble() - nnz = get_preallocation(preallocator, sizes[0][0]) - preallocator.destroy() - Dmat = PETSc.Mat().create(comm=comm) - Dmat.setType(mat_type) - Dmat.setSizes(sizes) - Dmat.setBlockSize(Vf.block_size) - Dmat.setISAllowRepeated(allow_repeated) - Dmat.setLGMap(*lgmaps) - Dmat.setPreallocationNNZ(nnz) - Dmat.setOption(PETSc.Mat.Option.NEW_NONZERO_ALLOCATION_ERR, True) - Dmat.setUp() + Dmat = allocate_matrix(preallocator, mat_type, allow_repeated=allow_repeated) assembler.arguments[0].data = Dmat.handle + preallocator.destroy() assembler() Dmat.assemble() - Dhat.destroy() return Dmat From 3e5bac316dd528314a377ab14e68a5f1600b5747 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 7 Aug 2025 15:58:07 +0100 Subject: [PATCH 35/72] FDMPC: Support other variants --- firedrake/preconditioners/bddc.py | 45 +++---------------------- firedrake/preconditioners/fdm.py | 5 ++- tests/firedrake/regression/test_bddc.py | 4 +-- tests/firedrake/regression/test_fdm.py | 9 ++--- 4 files changed, 13 insertions(+), 50 deletions(-) diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index c26ec48858..57717b3c3b 100644 --- a/firedrake/preconditioners/bddc.py +++ b/firedrake/preconditioners/bddc.py @@ -58,7 +58,7 @@ def initialize(self, pc): # Do not use CSR of local matrix to define dofs connectivity unless requested # Using the CSR only makes sense for H1/H2 problems is_h1h2 = sobolev_space in [H1, H2] - if "pc_bddc_use_local_mat_graph" not in opts and (not is_h1h2 or variant == "fdm"): + if "pc_bddc_use_local_mat_graph" not in opts and (not is_h1h2 or variant in {"fdm", "demkowicz", "demkowiczmass"}): opts["pc_bddc_use_local_mat_graph"] = False # Handle boundary dofs @@ -131,21 +131,13 @@ def applyTranspose(self, pc, x, y): self.pc.applyTranspose(x, y) -def get_dof_mapping(V, W): +def get_restricted_dofs(V, domain): + W = FunctionSpace(V.mesh(), restrict(V.ufl_element(), domain)) indices = get_restriction_indices(V, W) indices = V.dof_dset.lgmap.apply(indices) return PETSc.IS().createGeneral(indices, comm=V.comm) -def get_restricted_dofs(V, domain): - W = FunctionSpace(V.mesh(), restrict(V.ufl_element(), domain)) - return get_dof_mapping(V, W) - - -def get_low_order_dofs(V): - return get_dof_mapping(V, V.reconstruct(degree=1)) - - def get_divergence_mat(V, mat_type="aij", allow_repeated=False): from firedrake import assemble degree = max(as_tuple(V.ufl_element().degree())) @@ -160,43 +152,14 @@ def get_divergence_mat(V, mat_type="aij", allow_repeated=False): def get_discrete_gradient(V): - from firedrake import VectorSpaceBasis - Q = FunctionSpace(V.mesh(), curl_to_grad(V.ufl_element())) gradient = tabulate_exterior_derivative(Q, V) variant = Q.ufl_element().variant() - if variant == "fdm": + if variant in {"fdm", "demkowicz", "demkowiczmass"}: vdofs = get_restricted_dofs(Q, "vertex") gradient.compose('_elements_corners', vdofs) - if False: - # Constant moments along edges - wdofs = get_low_order_dofs(V) - v0 = Function(V) - with v0.dat.vec as y: - y.setValues(wdofs, numpy.ones((wdofs.getSizes()[0],), dtype=PETSc.RealType)) - y.assemble() - - # Linear moments along edges - # Constant moments of gradients of H1 edge bubbles - p = Function(Q).interpolate(1) - q = Function(Q) - - edofs = get_restricted_dofs(Q, "edge") - edofs = edofs.difference(vdofs) - with q.dat.vec as qvec, p.dat.vec as pvec: - qvec.setValues(edofs, pvec.getValues(edofs)) - qvec.assemble() - - v1 = Function(V) - with v1.dat.vec as y, q.dat.vec as x: - gradient.mult(x, y) - y.chop(1E-14) - - nsp = VectorSpaceBasis([v0, v1]) - gradient.compose("_constraints", nsp.nullspace()) - degree = max(as_tuple(Q.ufl_element().degree())) grad_args = (gradient,) grad_kwargs = {'order': degree} diff --git a/firedrake/preconditioners/fdm.py b/firedrake/preconditioners/fdm.py index d473115fb0..3c640a1916 100644 --- a/firedrake/preconditioners/fdm.py +++ b/firedrake/preconditioners/fdm.py @@ -107,7 +107,10 @@ def initialize(self, pc): # Transform the problem into the space with FDM shape functions V = J.arguments()[-1].function_space() - V_fdm = V.reconstruct(variant=self._variant) + if self._variant == "fdm" and V.ufl_element().variant() in {"fdm", "demkowicz", "demkowiczmass"}: + V_fdm = V + else: + V_fdm = V.reconstruct(variant=self._variant) if V == V_fdm: J_fdm, bcs_fdm = (J, bcs) else: diff --git a/tests/firedrake/regression/test_bddc.py b/tests/firedrake/regression/test_bddc.py index b22f113ea1..183d6bf3a8 100644 --- a/tests/firedrake/regression/test_bddc.py +++ b/tests/firedrake/regression/test_bddc.py @@ -138,10 +138,10 @@ def mesh(request): @pytest.mark.parametrize("degree", range(1, 3)) @pytest.mark.parametrize("variant", ("spectral", "fdm")) def test_vertex_dofs(mesh, variant, degree): - from firedrake.preconditioners.bddc import get_vertex_dofs + from firedrake.preconditioners.bddc import get_restricted_dofs P1 = FunctionSpace(mesh, "Lagrange", 1, variant=variant) V0 = FunctionSpace(mesh, "Lagrange", degree, variant=variant) - v = get_vertex_dofs(V0) + v = get_restricted_dofs(V0, "vertex") assert v.getSizes() == P1.dof_dset.layout_vec.getSizes() diff --git a/tests/firedrake/regression/test_fdm.py b/tests/firedrake/regression/test_fdm.py index 5a54d53e94..50aa3e7df3 100644 --- a/tests/firedrake/regression/test_fdm.py +++ b/tests/firedrake/regression/test_fdm.py @@ -339,8 +339,7 @@ def test_ipdg_direct_solver(fs): @pytest.mark.parallel(nprocs=2) @pytest.mark.parametrize("mat_type", ("aij",)) -@pytest.mark.parametrize("degree", (1, 4)) -@pytest.mark.parametrize("variant", ("spectral", "integral", "fdm")) +@pytest.mark.parametrize("variant,degree", [("spectral", 1), ("spectral", 4), ("integral", 4), ("fdm", 4)]) def test_tabulate_gradient(mesh, variant, degree, mat_type): from firedrake.preconditioners.fdm import tabulate_exterior_derivative tdim = mesh.topological_dimension() @@ -362,8 +361,7 @@ def test_tabulate_gradient(mesh, variant, degree, mat_type): @pytest.mark.parallel(nprocs=2) @pytest.mark.parametrize("mat_type", ("aij",)) -@pytest.mark.parametrize("degree", (1, 4)) -@pytest.mark.parametrize("variant", ("spectral", "integral", "fdm")) +@pytest.mark.parametrize("variant,degree", [("spectral", 1), ("spectral", 4), ("integral", 4), ("fdm", 4)]) def test_tabulate_curl(mesh, variant, degree, mat_type): from firedrake.preconditioners.fdm import tabulate_exterior_derivative tdim = mesh.topological_dimension() @@ -386,8 +384,7 @@ def test_tabulate_curl(mesh, variant, degree, mat_type): @pytest.mark.parallel(nprocs=2) @pytest.mark.parametrize("mat_type", ("aij", "is")) -@pytest.mark.parametrize("degree", (1, 4)) -@pytest.mark.parametrize("variant", ("spectral", "integral", "fdm")) +@pytest.mark.parametrize("variant,degree", [("spectral", 1), ("spectral", 4), ("integral", 4), ("fdm", 4)]) def test_tabulate_divergence(mesh, variant, degree, mat_type): from firedrake.preconditioners.fdm import tabulate_exterior_derivative tdim = mesh.topological_dimension() From e8be078c71abb5258e11b3203540f59844098151 Mon Sep 17 00:00:00 2001 From: Stefano Zampini Date: Fri, 8 Aug 2025 10:16:59 +0200 Subject: [PATCH 36/72] bddc: allow view of preconditioning matrix --- firedrake/preconditioners/bddc.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index 57717b3c3b..d2ed5b1b03 100644 --- a/firedrake/preconditioners/bddc.py +++ b/firedrake/preconditioners/bddc.py @@ -54,6 +54,9 @@ def initialize(self, pc): bddcpc.setOperators(*pc.getOperators()) bddcpc.setType(PETSc.PC.Type.BDDC) + # Allow viewing preconditioning matrix for debugging purposes + P.viewFromOptions('-view_mat', bddcpc) + opts = PETSc.Options(bddcpc.getOptionsPrefix()) # Do not use CSR of local matrix to define dofs connectivity unless requested # Using the CSR only makes sense for H1/H2 problems From c98ea74beda6173d814029deb105ced99ad3f450 Mon Sep 17 00:00:00 2001 From: Stefano Zampini Date: Tue, 19 Aug 2025 11:28:51 +0300 Subject: [PATCH 37/72] bddc: attach constants of the H(grad) space to discrete gradient --- firedrake/preconditioners/bddc.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index d2ed5b1b03..24a4c7afc2 100644 --- a/firedrake/preconditioners/bddc.py +++ b/firedrake/preconditioners/bddc.py @@ -55,7 +55,7 @@ def initialize(self, pc): bddcpc.setType(PETSc.PC.Type.BDDC) # Allow viewing preconditioning matrix for debugging purposes - P.viewFromOptions('-view_mat', bddcpc) + P.viewFromOptions('-view_pmat', bddcpc) opts = PETSc.Options(bddcpc.getOptionsPrefix()) # Do not use CSR of local matrix to define dofs connectivity unless requested @@ -97,7 +97,6 @@ def initialize(self, pc): tdim = V.mesh().topological_dimension() if tdim >= 2 and V.finat_element.formdegree == tdim-1: - A, P = pc.getOperators() allow_repeated = P.getISAllowRepeated() get_divergence = appctx.get("get_divergence_mat", get_divergence_mat) divergence = get_divergence(V, mat_type="is", allow_repeated=allow_repeated) @@ -141,7 +140,7 @@ def get_restricted_dofs(V, domain): return PETSc.IS().createGeneral(indices, comm=V.comm) -def get_divergence_mat(V, mat_type="aij", allow_repeated=False): +def get_divergence_mat(V, mat_type="is", allow_repeated=False): from firedrake import assemble degree = max(as_tuple(V.ufl_element().degree())) Q = TensorFunctionSpace(V.mesh(), "DG", 0, variant=f"integral({degree-1})", shape=V.value_shape[:-1]) @@ -155,9 +154,14 @@ def get_divergence_mat(V, mat_type="aij", allow_repeated=False): def get_discrete_gradient(V): + from firedrake import Constant + from firedrake.nullspace import VectorSpaceBasis + Q = FunctionSpace(V.mesh(), curl_to_grad(V.ufl_element())) gradient = tabulate_exterior_derivative(Q, V) - + nsp = VectorSpaceBasis([Function(Q).interpolate(Constant(1))]) + nsp.orthonormalize() + gradient.setNearNullSpace(nsp.nullspace()) variant = Q.ufl_element().variant() if variant in {"fdm", "demkowicz", "demkowiczmass"}: vdofs = get_restricted_dofs(Q, "vertex") From 4312274d37751e0c2a86f70c0b9a3b4f08c6b956 Mon Sep 17 00:00:00 2001 From: Stefano Zampini Date: Thu, 21 Aug 2025 18:15:49 +0300 Subject: [PATCH 38/72] ImplicitMatrixContext: allow for an arbitrary submatrix extraction --- firedrake/matrix_free/operators.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/firedrake/matrix_free/operators.py b/firedrake/matrix_free/operators.py index 18a767557d..12a5611945 100644 --- a/firedrake/matrix_free/operators.py +++ b/firedrake/matrix_free/operators.py @@ -377,11 +377,20 @@ def createSubMatrix(self, mat, row_is, col_is, target=None): row_ises = self._y.function_space().dof_dset.field_ises col_ises = self._x.function_space().dof_dset.field_ises - row_inds = find_sub_block(row_is, row_ises, comm=self.comm) - if row_is == col_is and row_ises == col_ises: - col_inds = row_inds - else: - col_inds = find_sub_block(col_is, col_ises, comm=self.comm) + try: + row_inds = find_sub_block(row_is, row_ises, comm=self.comm) + if row_is == col_is and row_ises == col_ises: + col_inds = row_inds + else: + col_inds = find_sub_block(col_is, col_ises, comm=self.comm) + except LookupError: + # No match for sub_block + # use default PETSc implementation via MATSHELL + popmethod = self.createSubMatrix + self.createSubMatrix = None + submat = mat.createSubMatrix(row_is, col_is) + self.createSubMatrix = popmethod + return submat splitter = ExtractSubBlock() asub = splitter.split(self.a, From 2616dc08a892f20b0bc503f03f113942d3ac7ac1 Mon Sep 17 00:00:00 2001 From: Stefano Zampini Date: Tue, 2 Sep 2025 18:17:46 +0300 Subject: [PATCH 39/72] bddc nedelec: attach constants as nullspace --- firedrake/preconditioners/bddc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index 24a4c7afc2..552437710c 100644 --- a/firedrake/preconditioners/bddc.py +++ b/firedrake/preconditioners/bddc.py @@ -161,7 +161,7 @@ def get_discrete_gradient(V): gradient = tabulate_exterior_derivative(Q, V) nsp = VectorSpaceBasis([Function(Q).interpolate(Constant(1))]) nsp.orthonormalize() - gradient.setNearNullSpace(nsp.nullspace()) + gradient.setNullSpace(nsp.nullspace()) variant = Q.ufl_element().variant() if variant in {"fdm", "demkowicz", "demkowiczmass"}: vdofs = get_restricted_dofs(Q, "vertex") From 0b3d6e755255e2464aa1b8d283106b87cd4477ec Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 8 Sep 2025 12:40:33 +0100 Subject: [PATCH 40/72] Update .github/workflows/core.yml --- .github/workflows/core.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/.github/workflows/core.yml b/.github/workflows/core.yml index cbbe56cadc..995cbf71e6 100644 --- a/.github/workflows/core.yml +++ b/.github/workflows/core.yml @@ -221,7 +221,6 @@ jobs: --extra-index-url https://download.pytorch.org/whl/cpu \ "$(echo ./firedrake-repo/dist/firedrake-*.tar.gz)[ci,docs]" - pip install -I "firedrake-fiat @ git+https://github.com/firedrakeproject/fiat.git@pbrubeck/fix/permutations" firedrake-clean pip list From b16c0f66ef016564e8be59689e46c176cf920dff Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 8 Sep 2025 15:38:12 +0100 Subject: [PATCH 41/72] Remove unrelated test --- tests/firedrake/regression/test_assemble.py | 16 ---------------- 1 file changed, 16 deletions(-) diff --git a/tests/firedrake/regression/test_assemble.py b/tests/firedrake/regression/test_assemble.py index 10cffcc71f..43d823d883 100644 --- a/tests/firedrake/regression/test_assemble.py +++ b/tests/firedrake/regression/test_assemble.py @@ -401,19 +401,3 @@ def test_assemble_tensor_empty_shape(mesh): v = Function(V).assign(1) expected = assemble(inner(v, v)*dx) assert np.allclose(result, expected) - - -@pytest.mark.parallel(nprocs=2) -def test_assemble_integral_quads(): - no_overlap = {"overlap_type": (DistributedMeshOverlapType.NONE, 0)} - mesh0 = UnitSquareMesh(3, 1, quadrilateral=True, distribution_parameters=no_overlap) - mesh1 = UnitSquareMesh(3, 1, quadrilateral=True) - - degree = 3 - variant = "integral" - V0 = FunctionSpace(mesh0, "Lagrange", degree, variant=variant) - V1 = FunctionSpace(mesh1, "Lagrange", degree, variant=variant) - c0 = assemble(conj(TestFunction(V0))*dx, bcs=DirichletBC(V0, 0, "on_boundary")) - c1 = assemble(conj(TestFunction(V1))*dx, bcs=DirichletBC(V1, 0, "on_boundary")) - - assert np.linalg.norm(c0.dat.data) == np.linalg.norm(c1.dat.data) From 6a1938bb0f2ac8565d1b1be458bb009da8fb11c3 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 24 Sep 2025 09:14:04 +0100 Subject: [PATCH 42/72] Update pyop2/types/mat.py --- pyop2/types/mat.py | 1 - 1 file changed, 1 deletion(-) diff --git a/pyop2/types/mat.py b/pyop2/types/mat.py index ec54e15414..4326198453 100644 --- a/pyop2/types/mat.py +++ b/pyop2/types/mat.py @@ -621,7 +621,6 @@ def _init_monolithic(self): rlgmap = rset.unblocked_lgmap clgmap = cset.unblocked_lgmap if self.mat_type == "is": - # FIXME monolithic lgmaps rlgmap, clgmap = self.sparsity._lgmaps create = mat.createIS else: From 508c4bbb5a4321939ef165b1c206abecbc57459f Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 25 Sep 2025 11:06:12 +0100 Subject: [PATCH 43/72] Fix mixed IS bcs --- firedrake/assemble.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index abbfa21e46..10bd6f7447 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -1508,6 +1508,8 @@ def _apply_bc(self, tensor, bc, u=None): # Set diagonal entries on bc nodes to 1 if the current # block is on the matrix diagonal and its index matches the # index of the function space the bc is defined on. + if op2tensor.handle.getType() == "is": + op2tensor.handle.assemble() op2tensor[index, index].set_local_diagonal_entries(bc.nodes, idx=component, diag_val=self.weight) # Handle off-diagonal block involving real function space. # "lgmaps" is correctly constructed in _matrix_arg, but From e3c1dd7c952f01a169c4ef7014e84aebf38e6691 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 29 Sep 2025 11:36:53 +0100 Subject: [PATCH 44/72] cleanup --- firedrake/assemble.py | 58 ------------------------------------------- pyop2/types/mat.py | 26 +++++++++++++++---- 2 files changed, 21 insertions(+), 63 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 69b113c5fc..5664c263c5 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -1420,18 +1420,6 @@ def _make_sparsity(test, trial, mat_type, sub_mat_type, maps_and_regions): except SparsityFormatError: raise ValueError("Monolithic matrix assembly not supported for systems " "with R-space blocks") - - # TODO reconstruct dof_dset with the unghosted lgmap - if mat_type == "is": - rmap = unghosted_lgmap(sparsity._dsets[0].lgmap, test.function_space()) - cmap = unghosted_lgmap(sparsity._dsets[1].lgmap, trial.function_space()) - sparsity._lgmaps = (rmap, cmap) - elif mat_type == "nest" and sub_mat_type == "is": - for i, j in numpy.ndindex(sparsity.shape): - block = sparsity[i, j] - rmap = unghosted_lgmap(block._dsets[0].lgmap, test.function_space()[i]) - cmap = unghosted_lgmap(block._dsets[1].lgmap, trial.function_space()[j]) - block._lgmaps = (rmap, cmap) return sparsity def _make_maps_and_regions(self): @@ -2244,49 +2232,3 @@ def index_function_spaces(form, indices): return tuple(a.ufl_function_space()[i] for i, a in zip(indices, form.arguments())) else: raise AssertionError - - -def masked_lgmap(lgmap, mask, block=True): - if block: - indices = lgmap.block_indices.copy() - bsize = lgmap.getBlockSize() - else: - indices = lgmap.indices.copy() - bsize = 1 - - if len(mask) > 0: - indices[mask] = -1 - return PETSc.LGMap().create(indices=indices, bsize=bsize, comm=lgmap.comm) - - -def unghosted_lgmap(lgmap, V, block=True): - if block: - ndofs = lgmap.getBlockIndices().size - else: - ndofs = lgmap.getIndices().size - mask = numpy.arange(ndofs, dtype=PETSc.IntType) - - mesh = V._mesh - mesh_dm = mesh.topology_dm - start, end = mesh_dm.getHeightStratum(0) - for i, W in enumerate(V): - iset = V.dof_dset.local_ises[i] - W_local_indices = iset.indices - bsize = 1 if block else iset.getBlockSize() - section = W.dm.getDefaultSection() - for seed in range(start, end): - # Do not loop over ghost cells - if mesh_dm.getLabelValue("pyop2_ghost", seed) != -1: - continue - closure, _ = mesh_dm.getTransitiveClosure(seed, useCone=True) - for p in closure: - dof = section.getDof(p) - if dof <= 0: - continue - off = section.getOffset(p) - # Local indices within W - W_indices = slice(bsize * off, bsize * (off + dof)) - mask[W_local_indices[W_indices]] = -1 - - mask = mask[mask > -1] - return masked_lgmap(lgmap, mask, block=block) diff --git a/pyop2/types/mat.py b/pyop2/types/mat.py index 4326198453..79e6b2c349 100644 --- a/pyop2/types/mat.py +++ b/pyop2/types/mat.py @@ -58,7 +58,6 @@ def __init__(self, dsets, maps_and_regions, name=None, nest=None, block_sparse=N if self._initialized: return self._dsets = dsets - self._lgmaps = None self._maps_and_regions = maps_and_regions self._block_sparse = block_sparse self._diagonal_block = diagonal_block @@ -342,6 +341,16 @@ def masked_lgmap(lgmap, mask, block=True): return PETSc.LGMap().create(indices=indices, bsize=bsize, comm=lgmap.comm) +def unghosted_lgmap(dset, node_maps): + lgmap = dset.lgmap + ndofs = lgmap.getBlockIndices().size + mask = np.arange(ndofs, dtype=PETSc.IntType) + for local_ises, node_map in zip(dset.local_ises, node_maps): + mask[local_ises.indices[node_map.values]] = -1 + mask = mask[mask > -1] + return masked_lgmap(lgmap, mask) + + class AbstractMat(DataCarrier, abc.ABC): r"""OP2 matrix data. A ``Mat`` is defined on a sparsity pattern and holds a value for each element in the :class:`Sparsity`. @@ -618,12 +627,17 @@ def _init_dense(self): def _init_monolithic(self): mat = PETSc.Mat() rset, cset = self.sparsity.dsets - rlgmap = rset.unblocked_lgmap - clgmap = cset.unblocked_lgmap if self.mat_type == "is": - rlgmap, clgmap = self.sparsity._lgmaps + rmaps = [None for _ in range(len(rset.local_ises))] + cmaps = [None for _ in range(len(cset.local_ises))] + for (i, j) in self.sparsity._maps_and_regions: + rmaps[i], cmaps[j], _ = tuple(self.sparsity._maps_and_regions[(i, j)])[0] + rlgmap = unghosted_lgmap(rset, rmaps) + clgmap = unghosted_lgmap(cset, cmaps) create = mat.createIS else: + rlgmap = rset.unblocked_lgmap + clgmap = cset.unblocked_lgmap create = mat.createAIJ size = ((self.nrows, None), (self.ncols, None)) create(size, bsize=1, comm=self.comm) @@ -695,7 +709,9 @@ def _init_block(self): rdim, cdim = self.dims[0][0] if self.mat_type == "is": - row_lg, col_lg = self.sparsity._lgmaps + rmap, cmap, _ = tuple(self.sparsity._maps_and_regions[(0, 0)])[0] + row_lg = unghosted_lgmap(rset, [rmap]) + col_lg = unghosted_lgmap(cset, [cmap]) block_sparse = False create = mat.createIS elif rdim == cdim and rdim > 1 and self.sparsity._block_sparse: From 508c5406389262d6c8161454adf57940ebd1e059 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 29 Sep 2025 11:38:48 +0100 Subject: [PATCH 45/72] remove negative index workaround --- pyop2/types/mat.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/pyop2/types/mat.py b/pyop2/types/mat.py index 79e6b2c349..5dafb91082 100644 --- a/pyop2/types/mat.py +++ b/pyop2/types/mat.py @@ -852,10 +852,6 @@ def set_local_diagonal_entries(self, rows, diag_val=1.0, idx=None): rows = rows.reshape(-1, 1) if self.handle.type == "is": self.handle.assemble() - # PETSc does not properly handle local dofs that map - # to a negative global index - rmap, _ = self.handle.getLGMap() - rows = rows[rmap.apply(rows) > -1] self.handle.zeroRowsColumnsLocal(rows, diag_val) else: self.change_assembly_state(Mat.INSERT_VALUES) @@ -976,10 +972,6 @@ def set_local_diagonal_entries(self, rows, diag_val=1.0, idx=None): rows = rows.reshape(-1, 1) if self.handle.type == "is": self.handle.assemble() - # PETSc does not properly handle local dofs that map - # to a negative global index - rmap, _ = self.handle.getLGMap() - rows = rows[rmap.apply(rows) > -1] self.handle.zeroRowsColumnsLocal(rows, diag_val) else: self.change_assembly_state(Mat.INSERT_VALUES) From 60d26d0c396285af2ad1e702402cdc4b5c4a70f7 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 29 Sep 2025 13:05:22 +0100 Subject: [PATCH 46/72] Harder test with MixedFunctionSpace(Vector, Scalar) --- pyop2/types/mat.py | 31 +++++++++++++++------ tests/firedrake/regression/test_assemble.py | 20 ++++++++----- 2 files changed, 35 insertions(+), 16 deletions(-) diff --git a/pyop2/types/mat.py b/pyop2/types/mat.py index 5dafb91082..7f522c6a25 100644 --- a/pyop2/types/mat.py +++ b/pyop2/types/mat.py @@ -342,13 +342,25 @@ def masked_lgmap(lgmap, mask, block=True): def unghosted_lgmap(dset, node_maps): - lgmap = dset.lgmap - ndofs = lgmap.getBlockIndices().size - mask = np.arange(ndofs, dtype=PETSc.IntType) - for local_ises, node_map in zip(dset.local_ises, node_maps): - mask[local_ises.indices[node_map.values]] = -1 - mask = mask[mask > -1] - return masked_lgmap(lgmap, mask) + if len(node_maps) == 1: + cmap, = node_maps + mask = np.setdiff1d( + cmap.values_with_halo[cmap.iterset.size:], + cmap.values[:cmap.iterset.size], + ) + else: + mask_pieces = [] + for iset, cmap in zip(dset.local_ises, node_maps): + to_mask = np.setdiff1d( + cmap.values_with_halo[cmap.iterset.size:], + cmap.values[:cmap.iterset.size] + ) + bs = iset.block_size + if bs > 1: + to_mask = np.concatenate([i + bs * to_mask for i in range(bs)]) + mask_pieces.append(iset.indices[to_mask]) + mask = np.concatenate(mask_pieces) + return masked_lgmap(dset.lgmap, mask) class AbstractMat(DataCarrier, abc.ABC): @@ -630,8 +642,9 @@ def _init_monolithic(self): if self.mat_type == "is": rmaps = [None for _ in range(len(rset.local_ises))] cmaps = [None for _ in range(len(cset.local_ises))] - for (i, j) in self.sparsity._maps_and_regions: - rmaps[i], cmaps[j], _ = tuple(self.sparsity._maps_and_regions[(i, j)])[0] + for (i, j), maps_and_regions in self.sparsity._maps_and_regions.items(): + for item in maps_and_regions: + rmaps[i], cmaps[j], _ = item rlgmap = unghosted_lgmap(rset, rmaps) clgmap = unghosted_lgmap(cset, cmaps) create = mat.createIS diff --git a/tests/firedrake/regression/test_assemble.py b/tests/firedrake/regression/test_assemble.py index 43d823d883..7cb5128d6e 100644 --- a/tests/firedrake/regression/test_assemble.py +++ b/tests/firedrake/regression/test_assemble.py @@ -119,22 +119,28 @@ def test_mat_nest_real_block_assembler_correctly_reuses_tensor(mesh): @pytest.mark.parametrize("shape,mat_type", [("scalar", "is"), ("vector", "is"), ("mixed", "is"), ("mixed", "nest")]) @pytest.mark.parametrize("dirichlet_bcs", [False, True]) def test_assemble_matis(mesh, shape, mat_type, dirichlet_bcs): - if shape == "vector": + if shape == "scalar": + V = FunctionSpace(mesh, "CG", 1) + elif shape == "vector": + V = VectorFunctionSpace(mesh, "CG", 1, dim=3) + elif shape == "mixed": V = VectorFunctionSpace(mesh, "CG", 1) + Q = FunctionSpace(mesh, "CG", 1) + V = V * Q else: - V = FunctionSpace(mesh, "CG", 1) - if shape == "mixed": - V = V * V - if V.value_size == 1: + raise ValueError(f"Unrecognized shape {shape}.") + + if V.value_size == 1 or shape == "mixed": A = 1 else: - A = as_matrix([[2, -1], [-1, 2]]) + A = as_matrix([[2, -1, 0], [-1, 2, -1], [0, -1, 2]]) u = TrialFunction(V) v = TestFunction(V) a = inner(A * grad(u), grad(v))*dx if dirichlet_bcs: - bcs = [DirichletBC(V.sub(i), 0, (i % 4+1, (i+2) % 4+1)) for i in range(V.value_size)] + subspaces = [V.sub(i).sub(j) for i in range(len(V)) for j in range(len(V.sub(i)))] + bcs = [DirichletBC(subspaces[i], 0, (i % 4+1, (i+2) % 4+1)) for i in range(len(subspaces))] else: bcs = None From a5d5bbe81f4c2f4d312b2fd535192dac2d7a8a6e Mon Sep 17 00:00:00 2001 From: Stefano Zampini Date: Thu, 2 Oct 2025 10:57:11 +0300 Subject: [PATCH 47/72] Fix BC for MATIS block --- firedrake/assemble.py | 10 ++++++---- firedrake/functionspaceimpl.py | 26 ++++++++++++++++++-------- 2 files changed, 24 insertions(+), 12 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 5664c263c5..cce1efbbbc 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -2024,16 +2024,18 @@ def collect_lgmaps(self): row_bcs, col_bcs = self._filter_bcs(i, j) # the tensor is already indexed rlgmap, clgmap = self._tensor.local_to_global_maps - rlgmap = self.test_function_space[i].local_to_global_map(row_bcs, rlgmap) - clgmap = self.trial_function_space[j].local_to_global_map(col_bcs, clgmap) + ismatis = self._tensor.handle.getType() == 'is' + rlgmap = self.test_function_space[i].local_to_global_map(row_bcs, (rlgmap, ismatis)) + clgmap = self.trial_function_space[j].local_to_global_map(col_bcs, (clgmap, ismatis)) return ((rlgmap, clgmap),) else: lgmaps = [] for i, j in self.get_indicess(): row_bcs, col_bcs = self._filter_bcs(i, j) rlgmap, clgmap = self._tensor[i, j].local_to_global_maps - rlgmap = self.test_function_space[i].local_to_global_map(row_bcs, rlgmap) - clgmap = self.trial_function_space[j].local_to_global_map(col_bcs, clgmap) + ismatis = self._tensor[i, j].handle.getType() == 'is' + rlgmap = self.test_function_space[i].local_to_global_map(row_bcs, (rlgmap, ismatis)) + clgmap = self.trial_function_space[j].local_to_global_map(col_bcs, (clgmap, ismatis)) lgmaps.append((rlgmap, clgmap)) return tuple(lgmaps) else: diff --git a/firedrake/functionspaceimpl.py b/firedrake/functionspaceimpl.py index 4dda3266e1..8faa66d82a 100644 --- a/firedrake/functionspaceimpl.py +++ b/firedrake/functionspaceimpl.py @@ -837,7 +837,9 @@ def local_to_global_map(self, bcs, lgmap=None): # this space has been recursively split out from [e.g. inside # fieldsplit] if bcs is None or len(bcs) == 0: - return lgmap or self.dof_dset.lgmap + if lgmap and lgmap[0]: + return lgmap[0] + return self.dof_dset.lgmap for bc in bcs: fs = bc.function_space() while fs.component is not None and fs.parent is not None: @@ -846,7 +848,7 @@ def local_to_global_map(self, bcs, lgmap=None): raise RuntimeError("DirichletBC defined on a different FunctionSpace!") unblocked = any(bc.function_space().component is not None for bc in bcs) - if lgmap is None: + if lgmap is None or lgmap[0] is None: lgmap = self.dof_dset.lgmap if unblocked: indices = lgmap.indices.copy() @@ -856,10 +858,16 @@ def local_to_global_map(self, bcs, lgmap=None): bsize = lgmap.getBlockSize() assert bsize == self.block_size else: - # MatBlock case, LGMap is already unrolled. - indices = lgmap.block_indices.copy() - bsize = lgmap.getBlockSize() - unblocked = True + # MatBlock case, the LGMap is implementation dependent + lgmap, ismatis = lgmap + if ismatis: + bsize = 1 + indices = lgmap.indices.copy() + unblocked = False + else: + indices = lgmap.block_indices.copy() + bsize = lgmap.getBlockSize() + unblocked = True nodes = [] for bc in bcs: if bc.function_space().component is not None: @@ -970,7 +978,9 @@ def __hash__(self): self.boundary_set)) def local_to_global_map(self, bcs, lgmap=None): - return lgmap or self.dof_dset.lgmap + if lgmap and lgmap[0]: + return lgmap[0] + return self.dof_dset.lgmap def collapse(self): return type(self)(self.function_space.collapse(), boundary_set=self.boundary_set) @@ -1173,7 +1183,7 @@ def exterior_facet_node_map(self): function space nodes.""" return op2.MixedMap(s.exterior_facet_node_map() for s in self) - def local_to_global_map(self, bcs): + def local_to_global_map(self, bcs, lgmap=None): r"""Return a map from process local dof numbering to global dof numbering. If BCs is provided, mask out those dofs which match the BC nodes.""" From fc9ecb2f95f2a275e873da0d6b4b476444f4a99b Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 2 Oct 2025 21:04:35 +0100 Subject: [PATCH 48/72] add mat_type kawrg to local_to_global_map --- firedrake/assemble.py | 12 ++++++------ firedrake/functionspaceimpl.py | 21 ++++++++------------- 2 files changed, 14 insertions(+), 19 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index cce1efbbbc..db8634fbcd 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -2024,18 +2024,18 @@ def collect_lgmaps(self): row_bcs, col_bcs = self._filter_bcs(i, j) # the tensor is already indexed rlgmap, clgmap = self._tensor.local_to_global_maps - ismatis = self._tensor.handle.getType() == 'is' - rlgmap = self.test_function_space[i].local_to_global_map(row_bcs, (rlgmap, ismatis)) - clgmap = self.trial_function_space[j].local_to_global_map(col_bcs, (clgmap, ismatis)) + mat_type = self._tensor.handle.getType() + rlgmap = self.test_function_space[i].local_to_global_map(row_bcs, rlgmap, mat_type=mat_type) + clgmap = self.trial_function_space[j].local_to_global_map(col_bcs, clgmap, mat_type=mat_type) return ((rlgmap, clgmap),) else: lgmaps = [] for i, j in self.get_indicess(): row_bcs, col_bcs = self._filter_bcs(i, j) rlgmap, clgmap = self._tensor[i, j].local_to_global_maps - ismatis = self._tensor[i, j].handle.getType() == 'is' - rlgmap = self.test_function_space[i].local_to_global_map(row_bcs, (rlgmap, ismatis)) - clgmap = self.trial_function_space[j].local_to_global_map(col_bcs, (clgmap, ismatis)) + mat_type = self._tensor[i, j].handle.getType() + rlgmap = self.test_function_space[i].local_to_global_map(row_bcs, rlgmap, mat_type=mat_type) + clgmap = self.trial_function_space[j].local_to_global_map(col_bcs, clgmap, mat_type=mat_type) lgmaps.append((rlgmap, clgmap)) return tuple(lgmaps) else: diff --git a/firedrake/functionspaceimpl.py b/firedrake/functionspaceimpl.py index 8faa66d82a..17aac10075 100644 --- a/firedrake/functionspaceimpl.py +++ b/firedrake/functionspaceimpl.py @@ -828,7 +828,7 @@ def boundary_nodes(self, sub_domain): return self._shared_data.boundary_nodes(self, sub_domain) @PETSc.Log.EventDecorator() - def local_to_global_map(self, bcs, lgmap=None): + def local_to_global_map(self, bcs, lgmap=None, mat_type=None): r"""Return a map from process local dof numbering to global dof numbering. If BCs is provided, mask out those dofs which match the BC nodes.""" @@ -837,9 +837,7 @@ def local_to_global_map(self, bcs, lgmap=None): # this space has been recursively split out from [e.g. inside # fieldsplit] if bcs is None or len(bcs) == 0: - if lgmap and lgmap[0]: - return lgmap[0] - return self.dof_dset.lgmap + return lgmap or self.dof_dset.lgmap for bc in bcs: fs = bc.function_space() while fs.component is not None and fs.parent is not None: @@ -848,7 +846,7 @@ def local_to_global_map(self, bcs, lgmap=None): raise RuntimeError("DirichletBC defined on a different FunctionSpace!") unblocked = any(bc.function_space().component is not None for bc in bcs) - if lgmap is None or lgmap[0] is None: + if lgmap is None: lgmap = self.dof_dset.lgmap if unblocked: indices = lgmap.indices.copy() @@ -859,8 +857,7 @@ def local_to_global_map(self, bcs, lgmap=None): assert bsize == self.block_size else: # MatBlock case, the LGMap is implementation dependent - lgmap, ismatis = lgmap - if ismatis: + if mat_type == "is": bsize = 1 indices = lgmap.indices.copy() unblocked = False @@ -977,10 +974,8 @@ def __hash__(self): return hash((self.mesh(), self.dof_dset, self.ufl_element(), self.boundary_set)) - def local_to_global_map(self, bcs, lgmap=None): - if lgmap and lgmap[0]: - return lgmap[0] - return self.dof_dset.lgmap + def local_to_global_map(self, bcs, lgmap=None, mat_type=None): + return lgmap or self.dof_dset.lgmap def collapse(self): return type(self)(self.function_space.collapse(), boundary_set=self.boundary_set) @@ -1183,7 +1178,7 @@ def exterior_facet_node_map(self): function space nodes.""" return op2.MixedMap(s.exterior_facet_node_map() for s in self) - def local_to_global_map(self, bcs, lgmap=None): + def local_to_global_map(self, bcs, lgmap=None, mat_type=None): r"""Return a map from process local dof numbering to global dof numbering. If BCs is provided, mask out those dofs which match the BC nodes.""" @@ -1429,7 +1424,7 @@ def top_nodes(self): ":class:`RealFunctionSpace` objects have no bottom nodes." return None - def local_to_global_map(self, bcs, lgmap=None): + def local_to_global_map(self, bcs, lgmap=None, mat_type=None): assert len(bcs) == 0 return None From 320a2737a4b51f138ee3fd93dcc751ab365ae999 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 2 Oct 2025 22:20:03 +0100 Subject: [PATCH 49/72] DROP BEFORE MERGE --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 8b03f6bcc7..21036a2afc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,7 +20,7 @@ dependencies = [ "mpi4py>3; python_version >= '3.13'", "mpi4py; python_version < '3.13'", # TODO RELEASE: use releases - "fenics-ufl @ git+https://github.com/FEniCS/ufl.git", + "fenics-ufl @ git+https://github.com/FEniCS/ufl.git@main", "firedrake-fiat @ git+https://github.com/firedrakeproject/fiat.git", "h5py>3.12.1", "libsupermesh", From 4cb84040eb3a19d149d87ea1bad9bc84bac19f00 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 15 Oct 2025 18:13:07 +0100 Subject: [PATCH 50/72] Fix unghosted_lgmap on ExtrudedMesh --- pyop2/types/mat.py | 43 +++++++++++++++++---- tests/firedrake/regression/test_assemble.py | 2 +- 2 files changed, 36 insertions(+), 9 deletions(-) diff --git a/pyop2/types/mat.py b/pyop2/types/mat.py index 7f522c6a25..b4c37c91ef 100644 --- a/pyop2/types/mat.py +++ b/pyop2/types/mat.py @@ -341,20 +341,47 @@ def masked_lgmap(lgmap, mask, block=True): return PETSc.LGMap().create(indices=indices, bsize=bsize, comm=lgmap.comm) +def mask_ghost_cells(cell_node_map): + """Return the local indices of the nodes that belong to ghost cells.""" + own_cells = cell_node_map.iterset.size + owned = cell_node_map.values[:own_cells] + ghost = cell_node_map.values_with_halo[own_cells:] + offset = cell_node_map.offset + if offset is None or ghost.size == 0: + # Non-extruded case + mask = np.setdiff1d(ghost, owned) + elif cell_node_map.iterset.constant_layers: + # Extruded case + mask_pieces = [] + owned = owned.copy() + ghost = ghost.copy() + quotient = cell_node_map.offset_quotient + layers = cell_node_map.iterset.layers + for i in range(layers-1): + if quotient is not None and i == layers-2: + # Periodic extruded case + owned -= quotient + ghost -= quotient + mask_pieces.append(np.setdiff1d(ghost, owned)) + owned += offset + ghost += offset + mask = np.concatenate(mask_pieces) + else: + raise NotImplementedError("MatIS does not support variable extrusion with overlap.") + return mask + + def unghosted_lgmap(dset, node_maps): + """Return a local-to-global map where the nodes on ghost cell are masked out.""" if len(node_maps) == 1: + # Non-mixed case cmap, = node_maps - mask = np.setdiff1d( - cmap.values_with_halo[cmap.iterset.size:], - cmap.values[:cmap.iterset.size], - ) + mask = mask_ghost_cells(cmap) else: + # Mixed case mask_pieces = [] for iset, cmap in zip(dset.local_ises, node_maps): - to_mask = np.setdiff1d( - cmap.values_with_halo[cmap.iterset.size:], - cmap.values[:cmap.iterset.size] - ) + to_mask = mask_ghost_cells(cmap) bs = iset.block_size if bs > 1: to_mask = np.concatenate([i + bs * to_mask for i in range(bs)]) diff --git a/tests/firedrake/regression/test_assemble.py b/tests/firedrake/regression/test_assemble.py index 7cb5128d6e..e1625ef0df 100644 --- a/tests/firedrake/regression/test_assemble.py +++ b/tests/firedrake/regression/test_assemble.py @@ -130,7 +130,7 @@ def test_assemble_matis(mesh, shape, mat_type, dirichlet_bcs): else: raise ValueError(f"Unrecognized shape {shape}.") - if V.value_size == 1 or shape == "mixed": + if V.value_size == 1: A = 1 else: A = as_matrix([[2, -1, 0], [-1, 2, -1], [0, -1, 2]]) From cf1bd88a417a253f94bfdfbed05da7dec09483e9 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 17 Oct 2025 15:28:55 +0100 Subject: [PATCH 51/72] Fix test --- tests/firedrake/regression/test_assemble.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/tests/firedrake/regression/test_assemble.py b/tests/firedrake/regression/test_assemble.py index e1625ef0df..c19f36bb38 100644 --- a/tests/firedrake/regression/test_assemble.py +++ b/tests/firedrake/regression/test_assemble.py @@ -139,11 +139,20 @@ def test_assemble_matis(mesh, shape, mat_type, dirichlet_bcs): v = TestFunction(V) a = inner(A * grad(u), grad(v))*dx if dirichlet_bcs: - subspaces = [V.sub(i).sub(j) for i in range(len(V)) for j in range(len(V.sub(i)))] - bcs = [DirichletBC(subspaces[i], 0, (i % 4+1, (i+2) % 4+1)) for i in range(len(subspaces))] + subspaces = [V] if len(V) == 1 else [V.sub(i) for i in range(len(V))] + components = [] + for i, Vi in enumerate(subspaces): + if Vi.block_size == 1: + components.append(Vi) + else: + components.extend(Vi.sub(j) for j in range(Vi.block_size)) + + assert len(components) == V.value_size + bcs = [DirichletBC(components[i], 0, (i % 4+1, (i+2) % 4+1)) for i in range(len(components))] else: bcs = None + aij_ref = assemble(a, bcs=bcs, mat_type="aij").petscmat ais = assemble(a, bcs=bcs, mat_type=mat_type, sub_mat_type="is").petscmat aij = PETSc.Mat() From 74a52ec7975a5b1ba4cb2e110654c006e6073ef5 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 22 Oct 2025 09:32:10 +0100 Subject: [PATCH 52/72] Drop unrelated change --- pyop2/types/mat.py | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/pyop2/types/mat.py b/pyop2/types/mat.py index b4c37c91ef..9da14e5d91 100644 --- a/pyop2/types/mat.py +++ b/pyop2/types/mat.py @@ -852,21 +852,14 @@ def zero(self): self.handle.zeroEntries() @mpi.collective - def zero_rows(self, rows, diag_val=1.0, idx=None): + def zero_rows(self, rows, diag_val=1.0): """Zeroes the specified rows of the matrix, with the exception of the diagonal entry, which is set to diag_val. May be used for applying strong boundary conditions. :param rows: a :class:`Subset` or an iterable""" - rows = rows.indices if isinstance(rows, Subset) else rows - rows = np.asarray(rows, dtype=dtypes.IntType) - rbs, _ = self.dims[0][0] - if rbs > 1: - if idx is not None: - rows = rbs * rows + idx - else: - rows = np.dstack([rbs*rows + i for i in range(rbs)]).flatten() self.assemble() + rows = rows.indices if isinstance(rows, Subset) else rows self.handle.zeroRowsLocal(rows, diag_val) def _flush_assembly(self): From 4c50572c85d09b86dd89a90e9e783d9e74bcfaca Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 22 Oct 2025 09:34:05 +0100 Subject: [PATCH 53/72] Update pyop2/types/mat.py --- pyop2/types/mat.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyop2/types/mat.py b/pyop2/types/mat.py index 9da14e5d91..b6539b1ed7 100644 --- a/pyop2/types/mat.py +++ b/pyop2/types/mat.py @@ -667,8 +667,8 @@ def _init_monolithic(self): mat = PETSc.Mat() rset, cset = self.sparsity.dsets if self.mat_type == "is": - rmaps = [None for _ in range(len(rset.local_ises))] - cmaps = [None for _ in range(len(cset.local_ises))] + rmaps = [None for _ in rset.local_ises] + cmaps = [None for _ in cset.local_ises] for (i, j), maps_and_regions in self.sparsity._maps_and_regions.items(): for item in maps_and_regions: rmaps[i], cmaps[j], _ = item From cb83725724ac72b075699e8266608c41e6171d8c Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 27 Oct 2025 13:48:59 +0000 Subject: [PATCH 54/72] Apply suggestions from code review --- firedrake/assemble.py | 1 + firedrake/matrix_free/operators.py | 2 +- pyop2/types/mat.py | 2 ++ 3 files changed, 4 insertions(+), 1 deletion(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 6fac7e17a4..fab06e4b37 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -1460,6 +1460,7 @@ def _apply_bc(self, tensor, bc, u=None): # block is on the matrix diagonal and its index matches the # index of the function space the bc is defined on. if op2tensor.handle.getType() == "is": + # Flag the entire matrix as assembled before indexing the diagonal block op2tensor.handle.assemble() op2tensor[index, index].set_local_diagonal_entries(bc.nodes, idx=component, diag_val=self.weight) # Handle off-diagonal block involving real function space. diff --git a/firedrake/matrix_free/operators.py b/firedrake/matrix_free/operators.py index 12a5611945..0c1070b1b3 100644 --- a/firedrake/matrix_free/operators.py +++ b/firedrake/matrix_free/operators.py @@ -384,7 +384,7 @@ def createSubMatrix(self, mat, row_is, col_is, target=None): else: col_inds = find_sub_block(col_is, col_ises, comm=self.comm) except LookupError: - # No match for sub_block + # Attemping to extract a submatrix that does not match with a subfield # use default PETSc implementation via MATSHELL popmethod = self.createSubMatrix self.createSubMatrix = None diff --git a/pyop2/types/mat.py b/pyop2/types/mat.py index b6539b1ed7..3dc747fbe8 100644 --- a/pyop2/types/mat.py +++ b/pyop2/types/mat.py @@ -696,6 +696,8 @@ def _init_monolithic(self): # We completely fill the allocated matrix when zeroing the # entries, so raise an error if we "missed" one. if self.mat_type != "is": + # The local matrix will have fewer nonzeros than the one prescribed + # in the global sparsity pattern mat.setOption(mat.Option.UNUSED_NONZERO_LOCATION_ERR, True) mat.setOption(mat.Option.IGNORE_OFF_PROC_ENTRIES, False) mat.setOption(mat.Option.NEW_NONZERO_ALLOCATION_ERR, True) From 46e3354c43d270612796c9fd35d1a346bc2deb77 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 27 Oct 2025 13:53:20 +0000 Subject: [PATCH 55/72] Update firedrake/matrix_free/operators.py --- firedrake/matrix_free/operators.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/firedrake/matrix_free/operators.py b/firedrake/matrix_free/operators.py index 0c1070b1b3..25b2a9704e 100644 --- a/firedrake/matrix_free/operators.py +++ b/firedrake/matrix_free/operators.py @@ -385,7 +385,7 @@ def createSubMatrix(self, mat, row_is, col_is, target=None): col_inds = find_sub_block(col_is, col_ises, comm=self.comm) except LookupError: # Attemping to extract a submatrix that does not match with a subfield - # use default PETSc implementation via MATSHELL + # use default PETSc implementation (MatCreateSubMatrixVirtual) via MATSHELL popmethod = self.createSubMatrix self.createSubMatrix = None submat = mat.createSubMatrix(row_is, col_is) From a0e01cdb8bf62db86c2ac94225751c26c745026a Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 27 Oct 2025 16:40:19 +0000 Subject: [PATCH 56/72] topological_dimension --- tests/firedrake/regression/test_bddc.py | 2 +- tests/firedrake/regression/test_fdm.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/firedrake/regression/test_bddc.py b/tests/firedrake/regression/test_bddc.py index d8025f24b6..bec55df92f 100644 --- a/tests/firedrake/regression/test_bddc.py +++ b/tests/firedrake/regression/test_bddc.py @@ -151,7 +151,7 @@ def test_vertex_dofs(mesh, variant, degree): def test_bddc_fdm(mesh, family, degree): variant = "fdm" bcs = True - tdim = mesh.topological_dimension() + tdim = mesh.topological_dimension expected = 6 if tdim == 2 else 11 assert solve_riesz_map(mesh, family, degree, variant, bcs) <= expected diff --git a/tests/firedrake/regression/test_fdm.py b/tests/firedrake/regression/test_fdm.py index f35a0b4d78..872b5eeb9b 100644 --- a/tests/firedrake/regression/test_fdm.py +++ b/tests/firedrake/regression/test_fdm.py @@ -342,7 +342,7 @@ def test_ipdg_direct_solver(fs): @pytest.mark.parametrize("variant,degree", [("spectral", 1), ("spectral", 4), ("integral", 4), ("fdm", 4)]) def test_tabulate_gradient(mesh, variant, degree, mat_type): from firedrake.preconditioners.fdm import tabulate_exterior_derivative - tdim = mesh.topological_dimension() + tdim = mesh.topological_dimension family = {1: "DG", 2: "RTCE", 3: "NCE"}[tdim] V0 = FunctionSpace(mesh, "Lagrange", degree, variant=variant) @@ -364,7 +364,7 @@ def test_tabulate_gradient(mesh, variant, degree, mat_type): @pytest.mark.parametrize("variant,degree", [("spectral", 1), ("spectral", 4), ("integral", 4), ("fdm", 4)]) def test_tabulate_curl(mesh, variant, degree, mat_type): from firedrake.preconditioners.fdm import tabulate_exterior_derivative - tdim = mesh.topological_dimension() + tdim = mesh.topological_dimension family1 = {1: "CG", 2: "CG", 3: "NCE"}[tdim] family2 = {1: "DG", 2: "RTCF", 3: "NCF"}[tdim] @@ -387,7 +387,7 @@ def test_tabulate_curl(mesh, variant, degree, mat_type): @pytest.mark.parametrize("variant,degree", [("spectral", 1), ("spectral", 4), ("integral", 4), ("fdm", 4)]) def test_tabulate_divergence(mesh, variant, degree, mat_type): from firedrake.preconditioners.fdm import tabulate_exterior_derivative - tdim = mesh.topological_dimension() + tdim = mesh.topological_dimension family = {1: "CG", 2: "RTCF", 3: "NCF"}[tdim] V = FunctionSpace(mesh, family, degree, variant=variant) From 53b22e16083050e205b57afaaa440ae33bbddbbf Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 29 Oct 2025 17:01:46 +0000 Subject: [PATCH 57/72] Apply suggestions from code review Co-authored-by: Josh Hope-Collins --- firedrake/matrix_free/operators.py | 4 ++-- tests/firedrake/regression/test_assemble.py | 1 - tests/firedrake/regression/test_bddc.py | 1 + 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/firedrake/matrix_free/operators.py b/firedrake/matrix_free/operators.py index 25b2a9704e..a91380980f 100644 --- a/firedrake/matrix_free/operators.py +++ b/firedrake/matrix_free/operators.py @@ -384,8 +384,8 @@ def createSubMatrix(self, mat, row_is, col_is, target=None): else: col_inds = find_sub_block(col_is, col_ises, comm=self.comm) except LookupError: - # Attemping to extract a submatrix that does not match with a subfield - # use default PETSc implementation (MatCreateSubMatrixVirtual) via MATSHELL + # Attemping to extract a submatrix that does not match with a subfield. + # Use default PETSc implementation (MatCreateSubMatrixVirtual) via MATSHELL instead. popmethod = self.createSubMatrix self.createSubMatrix = None submat = mat.createSubMatrix(row_is, col_is) diff --git a/tests/firedrake/regression/test_assemble.py b/tests/firedrake/regression/test_assemble.py index c19f36bb38..9972715746 100644 --- a/tests/firedrake/regression/test_assemble.py +++ b/tests/firedrake/regression/test_assemble.py @@ -180,7 +180,6 @@ def test_assemble_matis(mesh, shape, mat_type, dirichlet_bcs): assert ais.type == "is" ais.convert("aij", aij) - aij_ref = assemble(a, bcs=bcs, mat_type="aij").petscmat aij_ref.axpy(-1, aij) ind, iptr, values = aij_ref.getValuesCSR() assert np.allclose(values, 0) diff --git a/tests/firedrake/regression/test_bddc.py b/tests/firedrake/regression/test_bddc.py index bec55df92f..c4f98d998f 100644 --- a/tests/firedrake/regression/test_bddc.py +++ b/tests/firedrake/regression/test_bddc.py @@ -138,6 +138,7 @@ def mesh(request): @pytest.mark.parametrize("degree", range(1, 3)) @pytest.mark.parametrize("variant", ("spectral", "fdm")) def test_vertex_dofs(mesh, variant, degree): + """Check that we extract the right number of vertex dofs from a high order Lagrange space.""" from firedrake.preconditioners.bddc import get_restricted_dofs P1 = FunctionSpace(mesh, "Lagrange", 1, variant=variant) V0 = FunctionSpace(mesh, "Lagrange", degree, variant=variant) From 09dec977eb7970926f548c31cd569d8b9595a6cd Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 30 Oct 2025 09:09:50 +0000 Subject: [PATCH 58/72] Test BDDC for H(curl) and H(div) --- tests/firedrake/regression/test_bddc.py | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/tests/firedrake/regression/test_bddc.py b/tests/firedrake/regression/test_bddc.py index c4f98d998f..a4fb514bc2 100644 --- a/tests/firedrake/regression/test_bddc.py +++ b/tests/firedrake/regression/test_bddc.py @@ -147,8 +147,7 @@ def test_vertex_dofs(mesh, variant, degree): @pytest.mark.parallel -@pytest.mark.parametrize("degree", (4,)) -@pytest.mark.parametrize("family", "Q") +@pytest.mark.parametrize("family,degree", [("Q", 4)]) def test_bddc_fdm(mesh, family, degree): variant = "fdm" bcs = True @@ -158,12 +157,23 @@ def test_bddc_fdm(mesh, family, degree): @pytest.mark.parallel -@pytest.mark.parametrize("degree", (4,)) -@pytest.mark.parametrize("family", "Q") +@pytest.mark.parametrize("family,degree", [("Q", 4)]) @pytest.mark.parametrize("vector", (False, True), ids=("scalar", "vector")) -def test_bddc_aij(mesh, family, degree, vector): +def test_bddc_aij_quad(mesh, family, degree, vector): variant = None bcs = True tdim = mesh.topological_dimension expected = 7 if tdim == 2 else 11 assert solve_riesz_map(mesh, family, degree, variant, bcs, vector=vector) <= expected + + +@pytest.mark.parallel +@pytest.mark.parametrize("family,degree", [("CG", 3), ("N1curl", 3), ("N1div", 3)]) +def test_bddc_aij_simplex(family, degree): + nx = 4 + mesh = UnitCubeMesh(nx, nx, nx) + variant = None + bcs = True + tdim = mesh.topological_dimension + expected = {"CG": 13, "N1curl": 14, "N1div": 12}[family] + assert solve_riesz_map(mesh, family, degree, variant, bcs) <= expected From 8939fc4541cb9be9a551c53dee0bf4044359e4aa Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 30 Oct 2025 09:11:18 +0000 Subject: [PATCH 59/72] Update firedrake/preconditioners/bddc.py --- firedrake/preconditioners/bddc.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index 0d394257e8..1c0797f4b2 100644 --- a/firedrake/preconditioners/bddc.py +++ b/firedrake/preconditioners/bddc.py @@ -54,9 +54,6 @@ def initialize(self, pc): bddcpc.setOperators(*pc.getOperators()) bddcpc.setType(PETSc.PC.Type.BDDC) - # Allow viewing preconditioning matrix for debugging purposes - P.viewFromOptions('-view_pmat', bddcpc) - opts = PETSc.Options(bddcpc.getOptionsPrefix()) # Do not use CSR of local matrix to define dofs connectivity unless requested # Using the CSR only makes sense for H1/H2 problems From 31c80d8c4c954798e8fc5bfb3ff3c236f674a6b4 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 30 Oct 2025 10:04:25 +0000 Subject: [PATCH 60/72] docs --- firedrake/functionspaceimpl.py | 28 +++++++++++++++++++++++++--- 1 file changed, 25 insertions(+), 3 deletions(-) diff --git a/firedrake/functionspaceimpl.py b/firedrake/functionspaceimpl.py index 05b0218aa8..ead1866dd7 100644 --- a/firedrake/functionspaceimpl.py +++ b/firedrake/functionspaceimpl.py @@ -855,7 +855,28 @@ def boundary_nodes(self, sub_domain): def local_to_global_map(self, bcs, lgmap=None, mat_type=None): r"""Return a map from process local dof numbering to global dof numbering. - If BCs is provided, mask out those dofs which match the BC nodes.""" + Parameters + ---------- + bcs: [firedrake.bcs.BCBase] + If provided, mask out those dofs which match the BC nodes. + lgmap: PETSc.LGMap + The base local-to-global map, which might be partially masked. + mat_type: str + The matrix assembly type. This is required as different matrix types + handle the LGMap differently for MixedFunctionSpace. + + Note + ---- + For a :func:`.VectorFunctionSpace` or :func:`.TensorFunctionSpace` the returned + LGMap will be the scalar one, unless the bcs are imposed on a particular component. + For a :class:`MixedFunctionSpace` the returned LGMap is unblocked, + unless mat_type == "is". + + Returns + ------- + PETSc.LGMap + A local-to-global map with masked BC dofs. + """ # Caching these things is too complicated, since it depends # not just on the bcs, but also the parent space, and anything # this space has been recursively split out from [e.g. inside @@ -881,13 +902,14 @@ def local_to_global_map(self, bcs, lgmap=None, mat_type=None): assert bsize == self.block_size else: # MatBlock case, the LGMap is implementation dependent + bsize = lgmap.getBlockSize() + assert bsize == self.block_size if mat_type == "is": - bsize = 1 indices = lgmap.indices.copy() unblocked = False else: + # LGMap is already unrolled indices = lgmap.block_indices.copy() - bsize = lgmap.getBlockSize() unblocked = True nodes = [] for bc in bcs: From 7c31da6c755caa175b251c71d6e8e4a0b281cad6 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 30 Oct 2025 11:15:26 +0000 Subject: [PATCH 61/72] implement is_lagrange --- firedrake/preconditioners/bddc.py | 28 ++++++++++++++++++++++--- tests/firedrake/regression/test_bddc.py | 4 +--- 2 files changed, 26 insertions(+), 6 deletions(-) diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index 1c0797f4b2..2ab9ebb6c3 100644 --- a/firedrake/preconditioners/bddc.py +++ b/firedrake/preconditioners/bddc.py @@ -10,6 +10,7 @@ from firedrake.preconditioners.hiptmair import curl_to_grad from ufl import H1, H2, inner, dx, JacobianDeterminant from pyop2.utils import as_tuple +import gem import numpy __all__ = ("BDDCPC",) @@ -58,7 +59,7 @@ def initialize(self, pc): # Do not use CSR of local matrix to define dofs connectivity unless requested # Using the CSR only makes sense for H1/H2 problems is_h1h2 = sobolev_space in [H1, H2] - if "pc_bddc_use_local_mat_graph" not in opts and (not is_h1h2 or variant in {"fdm", "demkowicz", "demkowiczmass"}): + if "pc_bddc_use_local_mat_graph" not in opts and (not is_h1h2 or not is_lagrange(V.finat_element)): opts["pc_bddc_use_local_mat_graph"] = False # Handle boundary dofs @@ -159,8 +160,7 @@ def get_discrete_gradient(V): nsp = VectorSpaceBasis([Function(Q).interpolate(Constant(1))]) nsp.orthonormalize() gradient.setNullSpace(nsp.nullspace()) - variant = Q.ufl_element().variant() - if variant in {"fdm", "demkowicz", "demkowiczmass"}: + if not is_lagrange(Q.finat_element): vdofs = get_restricted_dofs(Q, "vertex") gradient.compose('_elements_corners', vdofs) @@ -168,3 +168,25 @@ def get_discrete_gradient(V): grad_args = (gradient,) grad_kwargs = {'order': degree} return grad_args, grad_kwargs + + +def is_lagrange(finat_element): + """Returns whether finat_element.dual_basis consists only of point evaluation dofs.""" + try: + Q, ps = finat_element.dual_basis + except NotImplementedError: + return False + # Inspect the weight matrix + # Lagrange elements have gem.Delta as the only terminal nodes + children = [Q] + while children: + nodes = [] + for c in children: + if isinstance(c, gem.Delta): + pass + elif isinstance(c, gem.gem.Terminal): + return False + else: + nodes.extend(c.children) + children = nodes + return True diff --git a/tests/firedrake/regression/test_bddc.py b/tests/firedrake/regression/test_bddc.py index a4fb514bc2..60cf1f11ea 100644 --- a/tests/firedrake/regression/test_bddc.py +++ b/tests/firedrake/regression/test_bddc.py @@ -118,8 +118,7 @@ def solve_riesz_map(mesh, family, degree, variant, bcs, condense=False, vector=F sp = solver_parameters(condense, variant=variant) solver = LinearVariationalSolver(problem, near_nullspace=nsp, - solver_parameters=sp, - options_prefix="") + solver_parameters=sp) solver.solve() return solver.snes.getLinearSolveIterations() @@ -174,6 +173,5 @@ def test_bddc_aij_simplex(family, degree): mesh = UnitCubeMesh(nx, nx, nx) variant = None bcs = True - tdim = mesh.topological_dimension expected = {"CG": 13, "N1curl": 14, "N1div": 12}[family] assert solve_riesz_map(mesh, family, degree, variant, bcs) <= expected From 20b49dcde0f986b24a027edf9ea33d480e81dd0b Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 30 Oct 2025 18:51:11 +0000 Subject: [PATCH 62/72] remove unrelated change --- firedrake/matrix_free/operators.py | 19 +++++-------------- 1 file changed, 5 insertions(+), 14 deletions(-) diff --git a/firedrake/matrix_free/operators.py b/firedrake/matrix_free/operators.py index a91380980f..18a767557d 100644 --- a/firedrake/matrix_free/operators.py +++ b/firedrake/matrix_free/operators.py @@ -377,20 +377,11 @@ def createSubMatrix(self, mat, row_is, col_is, target=None): row_ises = self._y.function_space().dof_dset.field_ises col_ises = self._x.function_space().dof_dset.field_ises - try: - row_inds = find_sub_block(row_is, row_ises, comm=self.comm) - if row_is == col_is and row_ises == col_ises: - col_inds = row_inds - else: - col_inds = find_sub_block(col_is, col_ises, comm=self.comm) - except LookupError: - # Attemping to extract a submatrix that does not match with a subfield. - # Use default PETSc implementation (MatCreateSubMatrixVirtual) via MATSHELL instead. - popmethod = self.createSubMatrix - self.createSubMatrix = None - submat = mat.createSubMatrix(row_is, col_is) - self.createSubMatrix = popmethod - return submat + row_inds = find_sub_block(row_is, row_ises, comm=self.comm) + if row_is == col_is and row_ises == col_ises: + col_inds = row_inds + else: + col_inds = find_sub_block(col_is, col_ises, comm=self.comm) splitter = ExtractSubBlock() asub = splitter.split(self.a, From f356c0f54f379781356c499bdcf43c3c28a70b5f Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 31 Oct 2025 14:00:09 +0000 Subject: [PATCH 63/72] test for correctness --- tests/firedrake/regression/test_bddc.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/firedrake/regression/test_bddc.py b/tests/firedrake/regression/test_bddc.py index 60cf1f11ea..53c41e5b84 100644 --- a/tests/firedrake/regression/test_bddc.py +++ b/tests/firedrake/regression/test_bddc.py @@ -120,6 +120,8 @@ def solve_riesz_map(mesh, family, degree, variant, bcs, condense=False, vector=F solver = LinearVariationalSolver(problem, near_nullspace=nsp, solver_parameters=sp) solver.solve() + uerr = Function(V).assign(uh - u_exact) + assert (assemble(a(uerr, uerr)) / assemble(a(u_exact, u_exact))) ** 0.5 < 1E-8 return solver.snes.getLinearSolveIterations() From dde83a81c48a351a4cfd0dda96caa2059bcdd705 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 31 Oct 2025 14:42:52 +0000 Subject: [PATCH 64/72] test mesh independence --- tests/firedrake/regression/test_bddc.py | 43 +++++++++++++------------ 1 file changed, 22 insertions(+), 21 deletions(-) diff --git a/tests/firedrake/regression/test_bddc.py b/tests/firedrake/regression/test_bddc.py index 53c41e5b84..2d51d55b55 100644 --- a/tests/firedrake/regression/test_bddc.py +++ b/tests/firedrake/regression/test_bddc.py @@ -6,8 +6,7 @@ def bddc_params(): chol = { "pc_type": "cholesky", - "pc_factor_mat_solver_type": "petsc", - "pc_factor_mat_ordering_type": "natural", + "pc_factor_mat_solver_type": DEFAULT_DIRECT_SOLVER, } sp = { "mat_type": "is", @@ -15,7 +14,7 @@ def bddc_params(): "pc_python_type": "firedrake.BDDCPC", "bddc_pc_bddc_neumann": chol, "bddc_pc_bddc_dirichlet": chol, - "bddc_pc_bddc_coarse": DEFAULT_DIRECT_SOLVER, + "bddc_pc_bddc_coarse": chol, } return sp @@ -126,21 +125,23 @@ def solve_riesz_map(mesh, family, degree, variant, bcs, condense=False, vector=F @pytest.fixture(params=(2, 3), ids=("square", "cube")) -def mesh(request): +def mh(request): dim = request.param - nx = 4 - msh = UnitSquareMesh(nx, nx, quadrilateral=True) + nx = 3 + base = UnitSquareMesh(nx, nx, quadrilateral=True) + mh = MeshHierarchy(base, 1) if dim == 3: - msh = ExtrudedMesh(msh, nx) - return msh + mh = ExtrudedMeshHierarchy(mh, height=1, base_layer=nx) + return mh @pytest.mark.parallel @pytest.mark.parametrize("degree", range(1, 3)) @pytest.mark.parametrize("variant", ("spectral", "fdm")) -def test_vertex_dofs(mesh, variant, degree): +def test_vertex_dofs(mh, variant, degree): """Check that we extract the right number of vertex dofs from a high order Lagrange space.""" from firedrake.preconditioners.bddc import get_restricted_dofs + mesh = mh[-1] P1 = FunctionSpace(mesh, "Lagrange", 1, variant=variant) V0 = FunctionSpace(mesh, "Lagrange", degree, variant=variant) v = get_restricted_dofs(V0, "vertex") @@ -149,31 +150,31 @@ def test_vertex_dofs(mesh, variant, degree): @pytest.mark.parallel @pytest.mark.parametrize("family,degree", [("Q", 4)]) -def test_bddc_fdm(mesh, family, degree): +def test_bddc_fdm(mh, family, degree): + """Test h-independence of condition number by measuring iteration counts""" variant = "fdm" bcs = True - tdim = mesh.topological_dimension - expected = 6 if tdim == 2 else 11 - assert solve_riesz_map(mesh, family, degree, variant, bcs) <= expected + its = [solve_riesz_map(m, family, degree, variant, bcs) for m in mh] + assert (np.diff(its) <= 1).all() @pytest.mark.parallel @pytest.mark.parametrize("family,degree", [("Q", 4)]) @pytest.mark.parametrize("vector", (False, True), ids=("scalar", "vector")) -def test_bddc_aij_quad(mesh, family, degree, vector): +def test_bddc_aij_quad(mh, family, degree, vector): + """Test h-independence of condition number by measuring iteration counts""" variant = None bcs = True - tdim = mesh.topological_dimension - expected = 7 if tdim == 2 else 11 - assert solve_riesz_map(mesh, family, degree, variant, bcs, vector=vector) <= expected + its = [solve_riesz_map(m, family, degree, variant, bcs, vector=vector) for m in mh] + assert (np.diff(its) <= 1).all() @pytest.mark.parallel @pytest.mark.parametrize("family,degree", [("CG", 3), ("N1curl", 3), ("N1div", 3)]) def test_bddc_aij_simplex(family, degree): - nx = 4 - mesh = UnitCubeMesh(nx, nx, nx) + """Test h-independence of condition number by measuring iteration counts""" variant = None bcs = True - expected = {"CG": 13, "N1curl": 14, "N1div": 12}[family] - assert solve_riesz_map(mesh, family, degree, variant, bcs) <= expected + meshes = [UnitCubeMesh(nx, nx, nx) for nx in (3, 6)] + its = [solve_riesz_map(m, family, degree, variant, bcs) for m in meshes] + assert (np.diff(its) <= 1).all() From a87093e4284a44034dc58ca2134ec2e6c6b05619 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 31 Oct 2025 15:01:38 +0000 Subject: [PATCH 65/72] cleanup --- tests/firedrake/regression/test_bddc.py | 30 +++++++++++++++---------- 1 file changed, 18 insertions(+), 12 deletions(-) diff --git a/tests/firedrake/regression/test_bddc.py b/tests/firedrake/regression/test_bddc.py index 2d51d55b55..6c9d25ed4b 100644 --- a/tests/firedrake/regression/test_bddc.py +++ b/tests/firedrake/regression/test_bddc.py @@ -3,6 +3,11 @@ from firedrake.petsc import DEFAULT_DIRECT_SOLVER +@pytest.fixture +def rg(): + return RandomGenerator(PCG64(seed=123456789)) + + def bddc_params(): chol = { "pc_type": "cholesky", @@ -21,7 +26,7 @@ def bddc_params(): def solver_parameters(static_condensation=False, variant=None): rtol = 1E-8 - atol = 1E-12 + atol = 0 sp_bddc = bddc_params() if variant != "fdm": sp = sp_bddc @@ -56,6 +61,7 @@ def solver_parameters(static_condensation=False, variant=None): sp.update({ "ksp_type": "cg", "ksp_norm_type": "natural", + "ksp_converged_reason": None, "ksp_monitor": None, "ksp_rtol": rtol, "ksp_atol": atol, @@ -65,7 +71,7 @@ def solver_parameters(static_condensation=False, variant=None): return sp -def solve_riesz_map(mesh, family, degree, variant, bcs, condense=False, vector=False): +def solve_riesz_map(rg, mesh, family, degree, variant, bcs, condense=False, vector=False): dirichlet_ids = [] if bcs: dirichlet_ids = ["on_boundary"] @@ -91,11 +97,10 @@ def solve_riesz_map(mesh, family, degree, variant, bcs, condense=False, vector=F formdegree = V.finat_element.formdegree if formdegree == 0: - a = inner(d(u), d(v)) * dx(degree=2*degree) + a = inner(d(u), d(v)) * dx else: - a = (inner(u, v) + inner(d(u), d(v))) * dx(degree=2*degree) + a = (inner(u, v) + inner(d(u), d(v))) * dx - rg = RandomGenerator(PCG64(seed=123456789)) u_exact = rg.uniform(V, -1, 1) L = ufl.replace(a, {u: u_exact}) bcs = [DirichletBC(V, u_exact, sub) for sub in dirichlet_ids] @@ -150,31 +155,32 @@ def test_vertex_dofs(mh, variant, degree): @pytest.mark.parallel @pytest.mark.parametrize("family,degree", [("Q", 4)]) -def test_bddc_fdm(mh, family, degree): +@pytest.mark.parametrize("condense", (False, True)) +def test_bddc_fdm(rg, mh, family, degree, condense): """Test h-independence of condition number by measuring iteration counts""" variant = "fdm" bcs = True - its = [solve_riesz_map(m, family, degree, variant, bcs) for m in mh] - assert (np.diff(its) <= 1).all() + its = [solve_riesz_map(rg, m, family, degree, variant, bcs, condense=condense) for m in mh] + assert (np.diff(its) <= 2).all() @pytest.mark.parallel @pytest.mark.parametrize("family,degree", [("Q", 4)]) @pytest.mark.parametrize("vector", (False, True), ids=("scalar", "vector")) -def test_bddc_aij_quad(mh, family, degree, vector): +def test_bddc_aij_quad(rg, mh, family, degree, vector): """Test h-independence of condition number by measuring iteration counts""" variant = None bcs = True - its = [solve_riesz_map(m, family, degree, variant, bcs, vector=vector) for m in mh] + its = [solve_riesz_map(rg, m, family, degree, variant, bcs, vector=vector) for m in mh] assert (np.diff(its) <= 1).all() @pytest.mark.parallel @pytest.mark.parametrize("family,degree", [("CG", 3), ("N1curl", 3), ("N1div", 3)]) -def test_bddc_aij_simplex(family, degree): +def test_bddc_aij_simplex(rg, family, degree): """Test h-independence of condition number by measuring iteration counts""" variant = None bcs = True meshes = [UnitCubeMesh(nx, nx, nx) for nx in (3, 6)] - its = [solve_riesz_map(m, family, degree, variant, bcs) for m in meshes] + its = [solve_riesz_map(rg, m, family, degree, variant, bcs) for m in meshes] assert (np.diff(its) <= 1).all() From b10e62a158a33478cc99df74b5714f63951aebec Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 31 Oct 2025 15:10:19 +0000 Subject: [PATCH 66/72] cleanup --- tests/firedrake/regression/test_bddc.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/firedrake/regression/test_bddc.py b/tests/firedrake/regression/test_bddc.py index 6c9d25ed4b..5430cc1b2c 100644 --- a/tests/firedrake/regression/test_bddc.py +++ b/tests/firedrake/regression/test_bddc.py @@ -124,8 +124,9 @@ def solve_riesz_map(rg, mesh, family, degree, variant, bcs, condense=False, vect solver = LinearVariationalSolver(problem, near_nullspace=nsp, solver_parameters=sp) solver.solve() + rtol = sp.get("ksp_rtol", 1E-8) uerr = Function(V).assign(uh - u_exact) - assert (assemble(a(uerr, uerr)) / assemble(a(u_exact, u_exact))) ** 0.5 < 1E-8 + assert (assemble(a(uerr, uerr)) / assemble(a(u_exact, u_exact))) ** 0.5 < rtol return solver.snes.getLinearSolveIterations() From a2d8f303cff34f52389cfd8b3d52bd7445561cc9 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 31 Oct 2025 16:45:09 +0000 Subject: [PATCH 67/72] Apply suggestions from code review --- tests/firedrake/regression/test_bddc.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/firedrake/regression/test_bddc.py b/tests/firedrake/regression/test_bddc.py index 5430cc1b2c..c2ad4c2a65 100644 --- a/tests/firedrake/regression/test_bddc.py +++ b/tests/firedrake/regression/test_bddc.py @@ -173,7 +173,7 @@ def test_bddc_aij_quad(rg, mh, family, degree, vector): variant = None bcs = True its = [solve_riesz_map(rg, m, family, degree, variant, bcs, vector=vector) for m in mh] - assert (np.diff(its) <= 1).all() + assert (np.diff(its) <= 2).all() @pytest.mark.parallel @@ -184,4 +184,4 @@ def test_bddc_aij_simplex(rg, family, degree): bcs = True meshes = [UnitCubeMesh(nx, nx, nx) for nx in (3, 6)] its = [solve_riesz_map(rg, m, family, degree, variant, bcs) for m in meshes] - assert (np.diff(its) <= 1).all() + assert (np.diff(its) <= 2).all() From 73cdcc2e1419101df12de0637e15732209eec416 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 3 Nov 2025 11:36:59 +0000 Subject: [PATCH 68/72] add more tests --- firedrake/preconditioners/bddc.py | 7 ++- tests/firedrake/regression/test_bddc.py | 57 ++++++++++++++----------- 2 files changed, 39 insertions(+), 25 deletions(-) diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index 2ab9ebb6c3..91e7246787 100644 --- a/firedrake/preconditioners/bddc.py +++ b/firedrake/preconditioners/bddc.py @@ -157,7 +157,12 @@ def get_discrete_gradient(V): Q = FunctionSpace(V.mesh(), curl_to_grad(V.ufl_element())) gradient = tabulate_exterior_derivative(Q, V) - nsp = VectorSpaceBasis([Function(Q).interpolate(Constant(1))]) + basis = Function(Q) + try: + basis.interpolate(Constant(1)) + except NotImplementedError: + basis.project(Constant(1)) + nsp = VectorSpaceBasis([basis]) nsp.orthonormalize() gradient.setNullSpace(nsp.nullspace()) if not is_lagrange(Q.finat_element): diff --git a/tests/firedrake/regression/test_bddc.py b/tests/firedrake/regression/test_bddc.py index c2ad4c2a65..2b7e0502dd 100644 --- a/tests/firedrake/regression/test_bddc.py +++ b/tests/firedrake/regression/test_bddc.py @@ -24,14 +24,14 @@ def bddc_params(): return sp -def solver_parameters(static_condensation=False, variant=None): - rtol = 1E-8 - atol = 0 +def solver_parameters(cellwise=False, condense=False, variant=None, rtol=1E-10, atol=0): sp_bddc = bddc_params() - if variant != "fdm": + if not cellwise: + assert not condense sp = sp_bddc - elif static_condensation: + elif condense: + assert variant == "fdm" sp = { "pc_type": "python", "pc_python_type": "firedrake.FacetSplitPC", @@ -46,10 +46,12 @@ def solver_parameters(static_condensation=False, variant=None): "facet_fdm_pc_fieldsplit_diag_use_amat": False, "facet_fdm_pc_fieldsplit_off_diag_use_amat": False, "facet_fdm_fieldsplit_ksp_type": "preonly", - "facet_fdm_fieldsplit_0_pc_type": "jacobi", + "facet_fdm_fieldsplit_0_pc_type": "bjacobi", + "facet_fdm_fieldsplit_0_pc_type_sub_pc_type": "icc", "facet_fdm_fieldsplit_1": sp_bddc, } else: + assert variant == "fdm" sp = { "pc_type": "python", "pc_python_type": "firedrake.FDMPC", @@ -62,7 +64,6 @@ def solver_parameters(static_condensation=False, variant=None): "ksp_type": "cg", "ksp_norm_type": "natural", "ksp_converged_reason": None, - "ksp_monitor": None, "ksp_rtol": rtol, "ksp_atol": atol, }) @@ -71,7 +72,9 @@ def solver_parameters(static_condensation=False, variant=None): return sp -def solve_riesz_map(rg, mesh, family, degree, variant, bcs, condense=False, vector=False): +def solve_riesz_map(rg, mesh, family, degree, variant, bcs, cellwise=False, condense=False, vector=False): + """Solve the riesz map for a random manufactured solution and return the + square root of the estimated condition number.""" dirichlet_ids = [] if bcs: dirichlet_ids = ["on_boundary"] @@ -120,20 +123,25 @@ def solve_riesz_map(rg, mesh, family, degree, variant, bcs, condense=False, vect uh = Function(V, name="solution") problem = LinearVariationalProblem(a, L, uh, bcs=bcs) - sp = solver_parameters(condense, variant=variant) + rtol = 1E-8 + sp = solver_parameters(cellwise=cellwise, condense=condense, variant=variant, rtol=rtol) + sp.setdefault("ksp_view_singularvalues", None) solver = LinearVariationalSolver(problem, near_nullspace=nsp, solver_parameters=sp) solver.solve() - rtol = sp.get("ksp_rtol", 1E-8) uerr = Function(V).assign(uh - u_exact) assert (assemble(a(uerr, uerr)) / assemble(a(u_exact, u_exact))) ** 0.5 < rtol - return solver.snes.getLinearSolveIterations() + + ew = solver.snes.ksp.computeEigenvalues() + assert min(ew) >= 1.0 + kappa = max(abs(ew)) / min(abs(ew)) + return kappa ** 0.5 @pytest.fixture(params=(2, 3), ids=("square", "cube")) def mh(request): dim = request.param - nx = 3 + nx = 4 base = UnitSquareMesh(nx, nx, quadrilateral=True) mh = MeshHierarchy(base, 1) if dim == 3: @@ -154,34 +162,35 @@ def test_vertex_dofs(mh, variant, degree): assert v.getSizes() == P1.dof_dset.layout_vec.getSizes() -@pytest.mark.parallel -@pytest.mark.parametrize("family,degree", [("Q", 4)]) +@pytest.mark.parallel([1, 3]) +@pytest.mark.parametrize("family,degree", [("Q", 4), ("E", 3), ("F", 3)]) @pytest.mark.parametrize("condense", (False, True)) -def test_bddc_fdm(rg, mh, family, degree, condense): +def test_bddc_cellwise_fdm(rg, mh, family, degree, condense): """Test h-independence of condition number by measuring iteration counts""" variant = "fdm" bcs = True - its = [solve_riesz_map(rg, m, family, degree, variant, bcs, condense=condense) for m in mh] - assert (np.diff(its) <= 2).all() + sqrt_kappa = [solve_riesz_map(rg, m, family, degree, variant, bcs, cellwise=True, condense=condense) for m in mh] + assert (np.diff(sqrt_kappa) <= 0.1).all(), str(sqrt_kappa) @pytest.mark.parallel @pytest.mark.parametrize("family,degree", [("Q", 4)]) @pytest.mark.parametrize("vector", (False, True), ids=("scalar", "vector")) def test_bddc_aij_quad(rg, mh, family, degree, vector): - """Test h-independence of condition number by measuring iteration counts""" + """Test h-dependence of condition number by measuring iteration counts""" variant = None bcs = True - its = [solve_riesz_map(rg, m, family, degree, variant, bcs, vector=vector) for m in mh] - assert (np.diff(its) <= 2).all() + sqrt_kappa = [solve_riesz_map(rg, m, family, degree, variant, bcs, vector=vector) for m in mh] + assert (np.diff(sqrt_kappa) <= 0.5).all(), str(sqrt_kappa) @pytest.mark.parallel @pytest.mark.parametrize("family,degree", [("CG", 3), ("N1curl", 3), ("N1div", 3)]) def test_bddc_aij_simplex(rg, family, degree): - """Test h-independence of condition number by measuring iteration counts""" + """Test h-dependence of condition number by measuring iteration counts""" variant = None bcs = True - meshes = [UnitCubeMesh(nx, nx, nx) for nx in (3, 6)] - its = [solve_riesz_map(rg, m, family, degree, variant, bcs) for m in meshes] - assert (np.diff(its) <= 2).all() + base = UnitCubeMesh(2, 2, 2) + meshes = MeshHierarchy(base, 2) + sqrt_kappa = [solve_riesz_map(rg, m, family, degree, variant, bcs) for m in meshes] + assert (np.diff(sqrt_kappa) <= 0.5).all(), str(sqrt_kappa) From 480057003ffe21ed34853aaac4bd19a48b08de62 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 3 Nov 2025 15:44:05 +0000 Subject: [PATCH 69/72] fix complex --- firedrake/preconditioners/fdm.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/firedrake/preconditioners/fdm.py b/firedrake/preconditioners/fdm.py index 258a4367d5..3b43c989f2 100644 --- a/firedrake/preconditioners/fdm.py +++ b/firedrake/preconditioners/fdm.py @@ -473,7 +473,7 @@ def assemble_coefficients(self, J, fcp, block_diagonal=False): dku = ufl.div(u) if sobolev == ufl.HDiv else ufl.curl(u) eps = expand_derivatives(ufl.diff(ufl.replace(expand_derivatives(dku), {ufl.grad(u): du}), du)) if sobolev == ufl.HDiv: - map_grad = lambda p: ufl.outer(p, eps/tdim) + map_grad = lambda p: ufl.conj(ufl.outer(p, eps/tdim)) elif len(eps.ufl_shape) == 3: map_grad = lambda p: ufl.dot(p, eps/2) else: @@ -2227,7 +2227,7 @@ def assemble_coefficients(self, J, fcp): if not isinstance(alpha, ufl.constantvalue.Zero): Q = FunctionSpace(mesh, finat.ufl.TensorElement(DG, shape=alpha.ufl_shape)) tensor = coefficients.setdefault("alpha", Function(Q.dual())) - assembly_callables.append(partial(get_assembler(ufl.inner(TestFunction(Q), alpha)*dx, form_compiler_parameters=fcp).assemble, tensor=tensor)) + assembly_callables.append(partial(get_assembler(ufl.inner(alpha, TestFunction(Q))*dx, form_compiler_parameters=fcp).assemble, tensor=tensor)) # get zero-th order coefficent ref_val = [ufl.variable(t) for t in args_J] @@ -2248,7 +2248,7 @@ def assemble_coefficients(self, J, fcp): beta = ufl.diag_vector(beta) Q = FunctionSpace(mesh, finat.ufl.TensorElement(DG, shape=beta.ufl_shape) if beta.ufl_shape else DG) tensor = coefficients.setdefault("beta", Function(Q.dual())) - assembly_callables.append(partial(get_assembler(ufl.inner(TestFunction(Q), beta)*dx, form_compiler_parameters=fcp).assemble, tensor=tensor)) + assembly_callables.append(partial(get_assembler(ufl.inner(beta, TestFunction(Q))*dx, form_compiler_parameters=fcp).assemble, tensor=tensor)) family = "CG" if tdim == 1 else "DGT" degree = 1 if tdim == 1 else 0 @@ -2270,11 +2270,11 @@ def assemble_coefficients(self, J, fcp): Q = FunctionSpace(mesh, finat.ufl.TensorElement(DGT, shape=G.ufl_shape)) tensor = coefficients.setdefault("Gq_facet", Function(Q.dual())) - assembly_callables.append(partial(get_assembler(ifacet_inner(TestFunction(Q), G), form_compiler_parameters=fcp).assemble, tensor=tensor)) + assembly_callables.append(partial(get_assembler(ifacet_inner(G, TestFunction(Q)), form_compiler_parameters=fcp).assemble, tensor=tensor)) PT = Piola.T Q = FunctionSpace(mesh, finat.ufl.TensorElement(DGT, shape=PT.ufl_shape)) tensor = coefficients.setdefault("PT_facet", Function(Q.dual())) - assembly_callables.append(partial(get_assembler(ifacet_inner(TestFunction(Q), PT), form_compiler_parameters=fcp).assemble, tensor=tensor)) + assembly_callables.append(partial(get_assembler(ifacet_inner(PT, TestFunction(Q)), form_compiler_parameters=fcp).assemble, tensor=tensor)) # make DGT functions with BC flags shape = V.ufl_element().reference_value_shape @@ -2294,7 +2294,7 @@ def assemble_coefficients(self, J, fcp): if beta.ufl_shape: beta = ufl.diag_vector(beta) ds_ext = ufl.Measure(itype, domain=mesh, subdomain_id=it.subdomain_id(), metadata=md) - forms.append(ufl.inner(test, beta)*ds_ext) + forms.append(ufl.inner(beta, test)*ds_ext) tensor = coefficients.setdefault("bcflags", Function(Q.dual())) if len(forms): From c88b5e4993b79891fadce25c11408ea4b5f12d6c Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 3 Nov 2025 15:52:05 +0000 Subject: [PATCH 70/72] Extruded periodic --- firedrake/preconditioners/bddc.py | 3 ++- tests/firedrake/regression/test_bddc.py | 1 + 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index 91e7246787..24e47320f2 100644 --- a/firedrake/preconditioners/bddc.py +++ b/firedrake/preconditioners/bddc.py @@ -65,7 +65,8 @@ def initialize(self, pc): # Handle boundary dofs ctx = get_appctx(dm) bcs = tuple(ctx._problem.dirichlet_bcs()) - if V.extruded: + mesh = V.mesh().unique() + if mesh.extruded and not mesh.extruded_periodic: boundary_nodes = numpy.unique(numpy.concatenate(list(map(V.boundary_nodes, ("on_boundary", "top", "bottom"))))) else: boundary_nodes = V.boundary_nodes("on_boundary") diff --git a/tests/firedrake/regression/test_bddc.py b/tests/firedrake/regression/test_bddc.py index 2b7e0502dd..40a4b422cf 100644 --- a/tests/firedrake/regression/test_bddc.py +++ b/tests/firedrake/regression/test_bddc.py @@ -62,6 +62,7 @@ def solver_parameters(cellwise=False, condense=False, variant=None, rtol=1E-10, sp.update({ "ksp_type": "cg", + "ksp_max_it": 20, "ksp_norm_type": "natural", "ksp_converged_reason": None, "ksp_rtol": rtol, From 4a37082cda6ed1c447a2ece62e5c7e46f53b6b07 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 4 Nov 2025 08:55:12 +0000 Subject: [PATCH 71/72] Apply suggestions from code review --- tests/firedrake/regression/test_bddc.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/firedrake/regression/test_bddc.py b/tests/firedrake/regression/test_bddc.py index 40a4b422cf..c111c07181 100644 --- a/tests/firedrake/regression/test_bddc.py +++ b/tests/firedrake/regression/test_bddc.py @@ -40,7 +40,7 @@ def solver_parameters(cellwise=False, condense=False, variant=None, rtol=1E-10, "facet_fdm_static_condensation": True, "facet_fdm_pc_use_amat": False, "facet_fdm_mat_type": "is", - "facet_fdm_mat_is_allow_repeated": True, + "facet_fdm_mat_is_allow_repeated": cellwise, "facet_fdm_pc_type": "fieldsplit", "facet_fdm_pc_fieldsplit_type": "symmetric_multiplicative", "facet_fdm_pc_fieldsplit_diag_use_amat": False, @@ -56,7 +56,7 @@ def solver_parameters(cellwise=False, condense=False, variant=None, rtol=1E-10, "pc_type": "python", "pc_python_type": "firedrake.FDMPC", "fdm_pc_use_amat": False, - "fdm_mat_is_allow_repeated": True, + "fdm_mat_is_allow_repeated": cellwise, "fdm": sp_bddc, } From 0ce9ad22b5f5dddfb5c1685efaa8c6713df6c15e Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 4 Nov 2025 08:57:33 +0000 Subject: [PATCH 72/72] Apply suggestions from code review --- pyop2/types/mat.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyop2/types/mat.py b/pyop2/types/mat.py index 338fc1f874..47e1cf8732 100644 --- a/pyop2/types/mat.py +++ b/pyop2/types/mat.py @@ -372,7 +372,7 @@ def mask_ghost_cells(cell_node_map): def unghosted_lgmap(dset, node_maps): - """Return a local-to-global map where the nodes on ghost cell are masked out.""" + """Return a local-to-global map where the nodes on ghost cells are masked out.""" if len(node_maps) == 1: # Non-mixed case cmap, = node_maps