From 05d20e23289a352d87fb6225a06726d343b28293 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 18 Jun 2026 21:44:37 +0100 Subject: [PATCH 1/7] FormProduct assembly into LRC matrix --- firedrake/assemble.py | 58 +++++++++++++++++++ .../regression/test_assemble_baseform.py | 57 ++++++++++++++++++ 2 files changed, 115 insertions(+) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 6dd7c4f03f..74de5c4e88 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -400,6 +400,55 @@ def _as_pyop2_type(tensor, indices=None): assert indices is None return tensor + @staticmethod + def _vector_columns_to_dense(functions): + if not functions: + raise ValueError("Cannot build an LRC factor matrix with no columns") + + first, *_ = functions + with first.dat.vec_ro as vec: + comm = vec.comm + local_size = vec.getLocalSize() + global_size = vec.getSize() + + dense = PETSc.Mat().createDense( + size=((local_size, global_size), (len(functions), len(functions))), comm=comm + ) + dense.setUp() + values = dense.getDenseArray() + for i, function in enumerate(functions): + with function.dat.vec_ro as vec: + if vec.getLocalSize() != local_size or vec.getSize() != global_size: + raise ValueError("LRC factor vectors do not share the same layout") + values[:, i] = vec.getArray(readonly=True) + dense.assemble() + return dense + + def _assemble_form_product_lrc(self, expr, tensor, bcs, assembled_factors): + if self._mat_type != "lrc": + raise ValueError("FormProduct assembly requires mat_type='lrc'") + if tensor is not None: + raise NotImplementedError("Assembly of FormProduct into an existing tensor is not supported") + if bcs: + raise NotImplementedError("Boundary conditions on LRC FormProduct assembly are not supported") + if len(expr.factors()) != 2: + raise NotImplementedError("LRC FormProduct assembly currently supports exactly two factors") + if len(expr.arguments()) != 2: + raise ValueError("LRC FormProduct assembly requires aggregate rank 2") + if any(len(factor.arguments()) != 1 for factor in expr.factors()): + raise ValueError("LRC FormProduct assembly requires rank-1 factors") + if len(assembled_factors) != 2: + raise TypeError("Not enough operands for FormProduct") + if not all(isinstance(factor, (firedrake.Cofunction, firedrake.Function)) for factor in assembled_factors): + raise TypeError("LRC FormProduct factors must assemble to Functions or Cofunctions") + + U = self._vector_columns_to_dense((assembled_factors[0],)) + V = self._vector_columns_to_dense((assembled_factors[1],)) + petscmat = PETSc.Mat().createLRC(None, U, None, V) + petscmat.assemble() + return Matrix(expr, petscmat, bcs=bcs, options_prefix=self._options_prefix, + fc_params=self._form_compiler_params) + def assemble(self, tensor=None, current_state=None): """Assemble the form. @@ -467,6 +516,8 @@ def base_form_assembly_visitor(self, expr, tensor, bcs, *args): else: raise AssertionError return assembler.assemble(tensor=tensor) + elif isinstance(expr, ufl.FormProduct): + return self._assemble_form_product_lrc(expr, tensor, bcs, args) elif isinstance(expr, ufl.Adjoint): if len(args) != 1: raise TypeError("Not enough operands for Adjoint") @@ -683,12 +734,19 @@ def base_form_preorder_traversal(expr, visitor, visited={}): def reconstruct_node_from_operands(expr, operands): if isinstance(expr, (ufl.Adjoint, ufl.Action)): return expr._ufl_expr_reconstruct_(*operands) + elif isinstance(expr, ufl.FormProduct): + return ufl.FormProduct(*operands) if operands else expr elif isinstance(expr, ufl.FormSum): return ufl.FormSum(*[(op, w) for op, w in zip(operands, expr.weights())]) return expr @staticmethod def base_form_operands(expr): + if isinstance(expr, ufl.FormProduct): + if (len(expr.factors()) == 2 and len(expr.arguments()) == 2 + and all(len(factor.arguments()) == 1 for factor in expr.factors())): + return expr.ufl_operands + return [] if isinstance(expr, (ufl.FormSum, ufl.Adjoint, ufl.Action)): return expr.ufl_operands if isinstance(expr, ufl.Form): diff --git a/tests/firedrake/regression/test_assemble_baseform.py b/tests/firedrake/regression/test_assemble_baseform.py index 8f76b70cb0..60a5578527 100644 --- a/tests/firedrake/regression/test_assemble_baseform.py +++ b/tests/firedrake/regression/test_assemble_baseform.py @@ -382,3 +382,60 @@ def test_assemble_baseform_return_tensor_if_given(): assert b0 is not b1 assert b1 is tensor + + +def test_assemble_formproduct_lrc(mesh): + V = FunctionSpace(mesh, "CG", 1) + v = TestFunction(V) + x = SpatialCoordinate(mesh)[0] + f = Function(V).interpolate(1 + x) + g = Function(V).interpolate(2 - x) + probe = Function(V).interpolate(x) + + row_form = inner(f, v) * dx + col_form = inner(g, v) * dx + product = ufl.FormProduct(row_form, col_form) + + A = assemble(product, mat_type="lrc") + assert A.petscmat.getType() == "lrc" + + row = assemble(row_form) + col = assemble(col_form) + with probe.dat.vec_ro as probe_vec, row.dat.vec_ro as row_vec, col.dat.vec_ro as col_vec: + result_vec = row_vec.duplicate() + A.petscmat.mult(probe_vec, result_vec) + expected = row_vec.getArray(readonly=True) * col_vec.dot(probe_vec) + assert np.allclose(result_vec.getArray(readonly=True), expected, rtol=1e-13, atol=1e-13) + + +def test_assemble_formproduct_requires_lrc_mat_type(mesh): + V = FunctionSpace(mesh, "CG", 1) + v = TestFunction(V) + form = v * dx + product = ufl.FormProduct(form, form) + + with pytest.raises(ValueError, match="mat_type='lrc'"): + assemble(product) + + +def test_assemble_formproduct_lrc_rejects_non_rank_one_factor(mesh): + V = FunctionSpace(mesh, "CG", 1) + v = TestFunction(V) + u = TrialFunction(V) + linear = v * dx + bilinear = inner(u, v) * dx + product = ufl.FormProduct(linear, bilinear) + + with pytest.raises(ValueError, match="aggregate rank 2"): + assemble(product, mat_type="lrc") + + +def test_assemble_formproduct_lrc_rejects_bcs(mesh): + V = FunctionSpace(mesh, "CG", 1) + v = TestFunction(V) + form = v * dx + product = ufl.FormProduct(form, form) + bc = DirichletBC(V, 0, "on_boundary") + + with pytest.raises(NotImplementedError, match="Boundary conditions"): + assemble(product, mat_type="lrc", bcs=bc) From c36454768979d0e50cd361aff9d22c1b9e0f08b0 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 18 Jun 2026 21:47:43 +0100 Subject: [PATCH 2/7] DROP BEFORE MERGE --- .github/actions/install/action.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/actions/install/action.yml b/.github/actions/install/action.yml index 8ba3dbbe57..75ac62ef83 100644 --- a/.github/actions/install/action.yml +++ b/.github/actions/install/action.yml @@ -152,6 +152,7 @@ runs: --extra-index-url https://download.pytorch.org/whl/cpu \ "./firedrake-repo[${{ inputs.deps }}]" + pip install -v --no-deps --ignore-installed git+https://github.com/firedrakeproject/ufl.git@pbrubeck/form-product firedrake-clean pip list From 92c17c808dc6aec3f5010c93f72cef09a47507f5 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 19 Jun 2026 05:45:37 +0100 Subject: [PATCH 3/7] Support scalar factors --- firedrake/assemble.py | 80 +++++++++++++++---- .../regression/test_assemble_baseform.py | 67 ++++++++++++++++ 2 files changed, 131 insertions(+), 16 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 74de5c4e88..586fce115d 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -400,6 +400,24 @@ def _as_pyop2_type(tensor, indices=None): assert indices is None return tensor + @staticmethod + def _as_scalar_value(value): + if isinstance(value, ufl.ZeroBaseForm): + value = 0.0 + elif isinstance(value, ufl.constantvalue.Zero): + value = 0.0 + elif isinstance(value, ufl.constantvalue.ScalarValue): + value = value.value() + elif isinstance(value, (firedrake.Constant, firedrake.Function)): + value = value.dat.data_ro + + if isinstance(value, numpy.ndarray): + # Assert singleton ndarray + value = value.item() + if not isinstance(value, numbers.Complex): + raise ValueError("Expecting a scalar expression") + return value + @staticmethod def _vector_columns_to_dense(functions): if not functions: @@ -424,7 +442,7 @@ def _vector_columns_to_dense(functions): dense.assemble() return dense - def _assemble_form_product_lrc(self, expr, tensor, bcs, assembled_factors): + def _assemble_form_product_lrc(self, expr, tensor, bcs, assembled_factors, scale=1): if self._mat_type != "lrc": raise ValueError("FormProduct assembly requires mat_type='lrc'") if tensor is not None: @@ -443,12 +461,51 @@ def _assemble_form_product_lrc(self, expr, tensor, bcs, assembled_factors): raise TypeError("LRC FormProduct factors must assemble to Functions or Cofunctions") U = self._vector_columns_to_dense((assembled_factors[0],)) + U.scale(scale) V = self._vector_columns_to_dense((assembled_factors[1],)) petscmat = PETSc.Mat().createLRC(None, U, None, V) petscmat.assemble() return Matrix(expr, petscmat, bcs=bcs, options_prefix=self._options_prefix, fc_params=self._form_compiler_params) + def _assemble_form_product(self, expr, tensor, bcs, assembled_factors): + ranks = tuple(len(factor.arguments()) for factor in expr.factors()) + if len(assembled_factors) != len(expr.factors()): + return self._assemble_form_product_lrc(expr, tensor, bcs, assembled_factors) + if sum(ranks) > 2 or not any(rank == 0 for rank in ranks): + return self._assemble_form_product_lrc(expr, tensor, bcs, assembled_factors) + + scalar_weight = PETSc.ScalarType(1) + higher_rank_factors = [] + assembled_higher_rank_factors = [] + for factor, rank, assembled_factor in zip(expr.factors(), ranks, assembled_factors): + if rank == 0: + scalar_weight *= self._as_scalar_value(assembled_factor) + else: + higher_rank_factors.append(factor) + assembled_higher_rank_factors.append(assembled_factor) + + if not higher_rank_factors: + return tensor.assign(scalar_weight) if tensor else scalar_weight + + if len(higher_rank_factors) == 1: + assembled_higher_rank_form = assembled_higher_rank_factors[0] + elif len(higher_rank_factors) == 2 and all(len(factor.arguments()) == 1 + for factor in higher_rank_factors): + higher_rank_form = ufl.FormProduct(*higher_rank_factors) + if tensor is not None: + return self._assemble_form_product_lrc( + higher_rank_form, tensor, bcs, assembled_higher_rank_factors) + assembled_higher_rank_form = self._assemble_form_product_lrc( + higher_rank_form, None, bcs, assembled_higher_rank_factors, scale=scalar_weight) + return Matrix(expr, assembled_higher_rank_form.petscmat, bcs=bcs, + options_prefix=self._options_prefix, fc_params=self._form_compiler_params) + else: + raise ValueError("FormProduct preprocessing requires remaining aggregate rank <= 2") + + weighted_form = ufl.FormSum((assembled_higher_rank_form, scalar_weight)) + return self.base_form_assembly_visitor(weighted_form, tensor, bcs, assembled_higher_rank_form) + def assemble(self, tensor=None, current_state=None): """Assemble the form. @@ -517,7 +574,7 @@ def base_form_assembly_visitor(self, expr, tensor, bcs, *args): raise AssertionError return assembler.assemble(tensor=tensor) elif isinstance(expr, ufl.FormProduct): - return self._assemble_form_product_lrc(expr, tensor, bcs, args) + return self._assemble_form_product(expr, tensor, bcs, args) elif isinstance(expr, ufl.Adjoint): if len(args) != 1: raise TypeError("Not enough operands for Adjoint") @@ -576,19 +633,7 @@ def base_form_assembly_visitor(self, expr, tensor, bcs, *args): # Assemble weights weights = [] for w in expr.weights(): - if isinstance(w, ufl.constantvalue.Zero): - w = 0.0 - elif isinstance(w, ufl.constantvalue.ScalarValue): - w = w.value() - elif isinstance(w, (firedrake.Constant, firedrake.Function)): - w = w.dat.data_ro - - if isinstance(w, numpy.ndarray): - # Assert singleton ndarray - w = w.item() - if not isinstance(w, numbers.Complex): - raise ValueError("Expecting a scalar weight expression") - weights.append(w) + weights.append(self._as_scalar_value(w)) # Scalar FormSum if all(isinstance(op, numbers.Complex) for op in args): @@ -743,8 +788,11 @@ def reconstruct_node_from_operands(expr, operands): @staticmethod def base_form_operands(expr): if isinstance(expr, ufl.FormProduct): + ranks = tuple(len(factor.arguments()) for factor in expr.factors()) if (len(expr.factors()) == 2 and len(expr.arguments()) == 2 - and all(len(factor.arguments()) == 1 for factor in expr.factors())): + and all(rank == 1 for rank in ranks)): + return expr.ufl_operands + if sum(ranks) <= 2 and any(rank == 0 for rank in ranks): return expr.ufl_operands return [] if isinstance(expr, (ufl.FormSum, ufl.Adjoint, ufl.Action)): diff --git a/tests/firedrake/regression/test_assemble_baseform.py b/tests/firedrake/regression/test_assemble_baseform.py index 60a5578527..48c34fa853 100644 --- a/tests/firedrake/regression/test_assemble_baseform.py +++ b/tests/firedrake/regression/test_assemble_baseform.py @@ -408,6 +408,73 @@ def test_assemble_formproduct_lrc(mesh): assert np.allclose(result_vec.getArray(readonly=True), expected, rtol=1e-13, atol=1e-13) +def test_assemble_formproduct_lrc_with_scalar_factors(mesh): + V = FunctionSpace(mesh, "CG", 1) + v = TestFunction(V) + x = SpatialCoordinate(mesh)[0] + f = Function(V).interpolate(1 + x) + g = Function(V).interpolate(2 - x) + probe = Function(V).interpolate(x) + + scalar_a = Constant(2.0) * dx(domain=mesh) + scalar_b = Constant(3.0) * dx(domain=mesh) + row_form = inner(f, v) * dx + col_form = inner(g, v) * dx + product = ufl.FormProduct(scalar_a, row_form, scalar_b, col_form) + + A = assemble(product, mat_type="lrc") + assert A.petscmat.getType() == "lrc" + + scale = assemble(scalar_a) * assemble(scalar_b) + row = assemble(row_form) + col = assemble(col_form) + with probe.dat.vec_ro as probe_vec, row.dat.vec_ro as row_vec, col.dat.vec_ro as col_vec: + result_vec = row_vec.duplicate() + A.petscmat.mult(probe_vec, result_vec) + expected = scale * row_vec.getArray(readonly=True) * col_vec.dot(probe_vec) + assert np.allclose(result_vec.getArray(readonly=True), expected, rtol=1e-13, atol=1e-13) + + +def test_assemble_formproduct_rank_two_with_scalar_factors(mesh): + V = FunctionSpace(mesh, "CG", 1) + v = TestFunction(V) + u = TrialFunction(V) + scalar_a = Constant(2.0) * dx(domain=mesh) + scalar_b = Constant(3.0) * dx(domain=mesh) + bilinear = inner(u, v) * dx + product = ufl.FormProduct(scalar_a, bilinear, scalar_b) + + A = assemble(product) + expected = assemble(assemble(scalar_a) * assemble(scalar_b) * bilinear) + assert np.allclose(A.petscmat[:, :], expected.petscmat[:, :], rtol=1e-13, atol=1e-13) + + +def test_assemble_formproduct_rank_one_with_scalar_factors(mesh): + V = FunctionSpace(mesh, "CG", 1) + v = TestFunction(V) + x = SpatialCoordinate(mesh)[0] + f = Function(V).interpolate(1 + x) + scalar_a = Constant(2.0) * dx(domain=mesh) + scalar_b = Constant(3.0) * dx(domain=mesh) + linear = inner(f, v) * dx + product = ufl.FormProduct(scalar_a, linear, scalar_b) + + b = assemble(product) + expected = assemble(assemble(scalar_a) * assemble(scalar_b) * linear) + assert np.allclose(b.dat.data_ro, expected.dat.data_ro, rtol=1e-13, atol=1e-13) + + +def test_assemble_formproduct_rank_zero_factors(mesh): + scalar_a = Constant(2.0) * dx(domain=mesh) + scalar_b = Constant(3.0) * dx(domain=mesh) + scalar_c = Constant(4.0) * dx(domain=mesh) + product = ufl.FormProduct(scalar_a, scalar_b, scalar_c) + + actual = assemble(product) + expected = assemble(scalar_a) * assemble(scalar_b) * assemble(scalar_c) + assert np.allclose(actual, expected, rtol=1e-13, atol=1e-13) + + def test_assemble_formproduct_requires_lrc_mat_type(mesh): V = FunctionSpace(mesh, "CG", 1) v = TestFunction(V) From 3d1ef13e41915d081e20c9acc05665af1de61633 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 19 Jun 2026 10:40:43 +0100 Subject: [PATCH 4/7] Assembly general LRC matrices --- firedrake/assemble.py | 116 ++++++++++++++- .../regression/test_assemble_baseform.py | 140 ++++++++++++++++++ 2 files changed, 249 insertions(+), 7 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 586fce115d..b0b4294b31 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -442,6 +442,33 @@ def _vector_columns_to_dense(functions): dense.assemble() return dense + @staticmethod + def _lrc_weights_to_vec(weights): + c = PETSc.Vec().createSeq(len(weights), comm=PETSc.COMM_SELF) + c.setValues(range(len(weights)), weights) + c.assemble() + return c + + @staticmethod + def _is_lrc_form_product(expr): + if not isinstance(expr, ufl.FormProduct): + return False + + ranks = tuple(len(factor.arguments()) for factor in expr.factors()) + return (len(expr.arguments()) == 2 + and ranks.count(1) == 2 + and all(rank in (0, 1) for rank in ranks)) + + def _assemble_lrc_matrix(self, expr, base_matrix, row_factors, col_factors, weights, bcs): + U = self._vector_columns_to_dense(row_factors) + V = self._vector_columns_to_dense(col_factors) + c = self._lrc_weights_to_vec(weights) + A = base_matrix.petscmat if base_matrix is not None else None + petscmat = PETSc.Mat().createLRC(A, U, c, V) + petscmat.assemble() + return Matrix(expr, petscmat, bcs=bcs, options_prefix=self._options_prefix, + fc_params=self._form_compiler_params) + def _assemble_form_product_lrc(self, expr, tensor, bcs, assembled_factors, scale=1): if self._mat_type != "lrc": raise ValueError("FormProduct assembly requires mat_type='lrc'") @@ -460,13 +487,83 @@ def _assemble_form_product_lrc(self, expr, tensor, bcs, assembled_factors, scale if not all(isinstance(factor, (firedrake.Cofunction, firedrake.Function)) for factor in assembled_factors): raise TypeError("LRC FormProduct factors must assemble to Functions or Cofunctions") - U = self._vector_columns_to_dense((assembled_factors[0],)) - U.scale(scale) - V = self._vector_columns_to_dense((assembled_factors[1],)) - petscmat = PETSc.Mat().createLRC(None, U, None, V) - petscmat.assemble() - return Matrix(expr, petscmat, bcs=bcs, options_prefix=self._options_prefix, - fc_params=self._form_compiler_params) + return self._assemble_lrc_matrix( + expr, None, (assembled_factors[0],), (assembled_factors[1],), (scale,), bcs + ) + + def _normalise_lrc_form_product(self, expr, weight): + scalar_weight = PETSc.ScalarType(weight) + rank_one_factors = [] + for factor in expr.factors(): + rank = len(factor.arguments()) + if rank == 0: + assembled_factor = firedrake.assemble( + factor, form_compiler_parameters=self._form_compiler_params + ) + scalar_weight *= self._as_scalar_value(assembled_factor) + else: + rank_one_factors.append(factor) + + if len(rank_one_factors) != 2: + raise ValueError("LRC FormProduct assembly requires exactly two rank-1 factors") + return rank_one_factors[0], rank_one_factors[1], scalar_weight + + def _assemble_lrc_factor(self, factor): + assembled = firedrake.assemble( + factor, form_compiler_parameters=self._form_compiler_params + ) + if not isinstance(assembled, (firedrake.Cofunction, firedrake.Function)): + raise TypeError("LRC FormProduct factors must assemble to Functions or Cofunctions") + return assembled + + def _assemble_form_sum_lrc(self, expr, tensor, bcs): + if self._mat_type != "lrc" or not isinstance(expr, ufl.FormSum): + return None + + weights = [self._as_scalar_value(w) for w in expr.weights()] + base_terms = [] + lrc_terms = [] + for component, weight in zip(expr.components(), weights): + if self._is_lrc_form_product(component): + lrc_terms.append((component, weight)) + else: + base_terms.append((component, weight)) + + if not lrc_terms: + return None + if tensor is not None: + raise NotImplementedError("Assembly of FormSum with LRC terms into an existing tensor is not supported") + if bcs: + raise NotImplementedError("Boundary conditions on LRC FormSum assembly are not supported") + + base_matrix = None + if base_terms: + base_form = ufl.FormSum(*base_terms) + base_mat_type = self._sub_mat_type + base_matrix = firedrake.assemble( + base_form, + bcs=bcs, + form_compiler_parameters=self._form_compiler_params, + mat_type=base_mat_type, + sub_mat_type=self._sub_mat_type, + options_prefix=self._options_prefix, + weight=self._weight, + allocation_integral_types=self.allocation_integral_types, + is_base_form_preprocessed=True, + ) + if not isinstance(base_matrix, MatrixBase): + raise TypeError("LRC base form must assemble to a matrix") + + row_factors = [] + col_factors = [] + lrc_weights = [] + for product, weight in lrc_terms: + row_factor, col_factor, lrc_weight = self._normalise_lrc_form_product(product, weight) + row_factors.append(self._assemble_lrc_factor(row_factor)) + col_factors.append(self._assemble_lrc_factor(col_factor)) + lrc_weights.append(lrc_weight) + + return self._assemble_lrc_matrix(expr, base_matrix, row_factors, col_factors, lrc_weights, bcs) def _assemble_form_product(self, expr, tensor, bcs, assembled_factors): ranks = tuple(len(factor.arguments()) for factor in expr.factors()) @@ -528,6 +625,11 @@ def assemble(self, tensor=None, current_state=None): in a post-order fashion and evaluating the nodes on the fly. """ + if isinstance(self._form, ufl.FormSum) and len(self._form.arguments()) == 2: + lrc_result = self._assemble_form_sum_lrc(self._form, tensor, self._bcs) + if lrc_result is not None: + return lrc_result + def visitor(e, *operands): t = tensor if e is self._form else None # Deal with 2-form bcs inside the visitor diff --git a/tests/firedrake/regression/test_assemble_baseform.py b/tests/firedrake/regression/test_assemble_baseform.py index 48c34fa853..3d40967b8c 100644 --- a/tests/firedrake/regression/test_assemble_baseform.py +++ b/tests/firedrake/regression/test_assemble_baseform.py @@ -435,6 +435,146 @@ def test_assemble_formproduct_lrc_with_scalar_factors(mesh): assert np.allclose(result_vec.getArray(readonly=True), expected, rtol=1e-13, atol=1e-13) +def _assert_lrc_action(A, probe, terms, base_matrix=None): + with probe.dat.vec_ro as probe_vec: + actual_vec = probe_vec.duplicate() + A.petscmat.mult(probe_vec, actual_vec) + expected = np.zeros_like(actual_vec.getArray(readonly=True)) + + if base_matrix is not None: + base_vec = probe_vec.duplicate() + base_matrix.petscmat.mult(probe_vec, base_vec) + expected += base_vec.getArray(readonly=True) + + for weight, row, col in terms: + with row.dat.vec_ro as row_vec, col.dat.vec_ro as col_vec: + expected += weight * row_vec.getArray(readonly=True) * col_vec.dot(probe_vec) + + assert np.allclose(actual_vec.getArray(readonly=True), expected, rtol=1e-13, atol=1e-13) + + +def test_assemble_formsum_lrc_with_base_matrix(mesh): + V = FunctionSpace(mesh, "CG", 1) + v = TestFunction(V) + u = TrialFunction(V) + x, y = SpatialCoordinate(mesh) + f1 = Function(V).interpolate(1 + x) + g1 = Function(V).interpolate(2 - x) + f2 = Function(V).interpolate(1 + y) + g2 = Function(V).interpolate(2 - y) + probe = Function(V).interpolate(x + y) + + base_form = inner(u, v) * dx + row_form1 = inner(f1, v) * dx + col_form1 = inner(g1, v) * dx + row_form2 = inner(f2, v) * dx + col_form2 = inner(g2, v) * dx + product1 = ufl.FormProduct(row_form1, col_form1) + product2 = ufl.FormProduct(row_form2, col_form2) + formsum = ufl.FormSum((base_form, 0.75), (product1, 1.5), (product2, -0.25)) + + A = assemble(formsum, mat_type="lrc", sub_mat_type="aij") + assert A.petscmat.getType() == "lrc" + + base_matrix = assemble(0.75 * base_form, mat_type="aij") + terms = ((1.5, assemble(row_form1), assemble(col_form1)), + (-0.25, assemble(row_form2), assemble(col_form2))) + _assert_lrc_action(A, probe, terms, base_matrix=base_matrix) + + +def test_assemble_formsum_lrc_base_uses_sub_mat_type(mesh): + V = VectorFunctionSpace(mesh, "CG", 1) + v = TestFunction(V) + u = TrialFunction(V) + x, y = SpatialCoordinate(mesh) + f = Function(V).interpolate(as_vector((1 + x, 2 + y))) + g = Function(V).interpolate(as_vector((2 - x, 1 - y))) + + base_form = inner(u, v) * dx + product = ufl.FormProduct(inner(f, v) * dx, inner(g, v) * dx) + formsum = ufl.FormSum((base_form, 1), (product, 1)) + + A = assemble(formsum, mat_type="lrc", sub_mat_type="baij") + base, _, _, _ = A.petscmat.getLRCMats() + assert base.getType().endswith("baij") + + +def test_assemble_formsum_lrc_without_base_matrix(mesh): + V = FunctionSpace(mesh, "CG", 1) + v = TestFunction(V) + x, y = SpatialCoordinate(mesh) + f1 = Function(V).interpolate(1 + x) + g1 = Function(V).interpolate(2 - x) + f2 = Function(V).interpolate(1 + y) + g2 = Function(V).interpolate(2 - y) + probe = Function(V).interpolate(x - y) + + row_form1 = inner(f1, v) * dx + col_form1 = inner(g1, v) * dx + row_form2 = inner(f2, v) * dx + col_form2 = inner(g2, v) * dx + product1 = ufl.FormProduct(row_form1, col_form1) + product2 = ufl.FormProduct(row_form2, col_form2) + formsum = ufl.FormSum((product1, 1.5), (product2, -0.25)) + + A = assemble(formsum, mat_type="lrc") + assert A.petscmat.getType() == "lrc" + + terms = ((1.5, assemble(row_form1), assemble(col_form1)), + (-0.25, assemble(row_form2), assemble(col_form2))) + _assert_lrc_action(A, probe, terms) + + +def test_assemble_formsum_lrc_with_scalar_product_factors(mesh): + V = FunctionSpace(mesh, "CG", 1) + v = TestFunction(V) + u = TrialFunction(V) + x = SpatialCoordinate(mesh)[0] + f = Function(V).interpolate(1 + x) + g = Function(V).interpolate(2 - x) + probe = Function(V).interpolate(x) + + base_form = inner(u, v) * dx + scalar_a = Constant(2.0) * dx(domain=mesh) + scalar_b = Constant(3.0) * dx(domain=mesh) + row_form = inner(f, v) * dx + col_form = inner(g, v) * dx + product = ufl.FormProduct(scalar_a, row_form, scalar_b, col_form) + formsum = ufl.FormSum((base_form, 1), (product, 0.5)) + + A = assemble(formsum, mat_type="lrc", sub_mat_type="aij") + base_matrix = assemble(base_form, mat_type="aij") + scale = 0.5 * assemble(scalar_a) * assemble(scalar_b) + terms = ((scale, assemble(row_form), assemble(col_form)),) + _assert_lrc_action(A, probe, terms, base_matrix=base_matrix) + + +def test_assemble_formsum_lrc_rejects_tensor(mesh): + V = FunctionSpace(mesh, "CG", 1) + v = TestFunction(V) + u = TrialFunction(V) + form = inner(u, v) * dx + product = ufl.FormProduct(v * dx, v * dx) + formsum = ufl.FormSum((form, 1), (product, 1)) + tensor = assemble(form) + + with pytest.raises(NotImplementedError, match="existing tensor"): + assemble(formsum, mat_type="lrc", tensor=tensor) + + +def test_assemble_formsum_lrc_rejects_bcs(mesh): + V = FunctionSpace(mesh, "CG", 1) + v = TestFunction(V) + u = TrialFunction(V) + form = inner(u, v) * dx + product = ufl.FormProduct(v * dx, v * dx) + formsum = ufl.FormSum((form, 1), (product, 1)) + bc = DirichletBC(V, 0, "on_boundary") + + with pytest.raises(NotImplementedError, match="Boundary conditions on LRC FormSum"): + assemble(formsum, mat_type="lrc", bcs=bc) + + def test_assemble_formproduct_rank_two_with_scalar_factors(mesh): V = FunctionSpace(mesh, "CG", 1) v = TestFunction(V) From 5ce2a8985d6eb7631c3c23234ee420cbd43f1425 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 19 Jun 2026 15:45:08 +0100 Subject: [PATCH 5/7] Support DirichletBC --- firedrake/assemble.py | 74 ++++++++++++++----- .../regression/test_assemble_baseform.py | 44 ++++++++--- 2 files changed, 88 insertions(+), 30 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index b0b4294b31..1458c65e31 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -459,14 +459,44 @@ def _is_lrc_form_product(expr): and ranks.count(1) == 2 and all(rank in (0, 1) for rank in ranks)) - def _assemble_lrc_matrix(self, expr, base_matrix, row_factors, col_factors, weights, bcs): + @staticmethod + def _bc_matches_space(bc, function_space): + if not isinstance(bc, DirichletBC): + return False + + if ufl.duals.is_dual(function_space): + function_space = function_space.dual() + + V = bc.function_space() + while True: + bc_space = V.dual() if ufl.duals.is_dual(V) else V + if bc_space == function_space: + return True + if V.parent is None: + return False + V = V.parent + + @staticmethod + def _check_lrc_bcs(bcs): + if any(not isinstance(bc, DirichletBC) for bc in bcs): + raise NotImplementedError("Only DirichletBC objects are supported on LRC assembly") + + def _lrc_factor_bcs(self, factor, bcs): + if not bcs: + return () + + arg, = factor.arguments() + function_space = arg.function_space() + return tuple(bc for bc in bcs if self._bc_matches_space(bc, function_space)) + + def _assemble_lrc_matrix(self, expr, base_matrix, row_factors, col_factors, weights): U = self._vector_columns_to_dense(row_factors) V = self._vector_columns_to_dense(col_factors) c = self._lrc_weights_to_vec(weights) A = base_matrix.petscmat if base_matrix is not None else None petscmat = PETSc.Mat().createLRC(A, U, c, V) petscmat.assemble() - return Matrix(expr, petscmat, bcs=bcs, options_prefix=self._options_prefix, + return Matrix(expr, petscmat, bcs=(), options_prefix=self._options_prefix, fc_params=self._form_compiler_params) def _assemble_form_product_lrc(self, expr, tensor, bcs, assembled_factors, scale=1): @@ -474,21 +504,25 @@ def _assemble_form_product_lrc(self, expr, tensor, bcs, assembled_factors, scale raise ValueError("FormProduct assembly requires mat_type='lrc'") if tensor is not None: raise NotImplementedError("Assembly of FormProduct into an existing tensor is not supported") - if bcs: - raise NotImplementedError("Boundary conditions on LRC FormProduct assembly are not supported") + self._check_lrc_bcs(bcs) if len(expr.factors()) != 2: raise NotImplementedError("LRC FormProduct assembly currently supports exactly two factors") if len(expr.arguments()) != 2: raise ValueError("LRC FormProduct assembly requires aggregate rank 2") if any(len(factor.arguments()) != 1 for factor in expr.factors()): raise ValueError("LRC FormProduct assembly requires rank-1 factors") - if len(assembled_factors) != 2: - raise TypeError("Not enough operands for FormProduct") - if not all(isinstance(factor, (firedrake.Cofunction, firedrake.Function)) for factor in assembled_factors): - raise TypeError("LRC FormProduct factors must assemble to Functions or Cofunctions") + if bcs: + row_factor = self._assemble_lrc_factor(expr.factors()[0], bcs) + col_factor = self._assemble_lrc_factor(expr.factors()[1], bcs) + else: + if len(assembled_factors) != 2: + raise TypeError("Not enough operands for FormProduct") + if not all(isinstance(factor, (firedrake.Cofunction, firedrake.Function)) for factor in assembled_factors): + raise TypeError("LRC FormProduct factors must assemble to Functions or Cofunctions") + row_factor, col_factor = assembled_factors return self._assemble_lrc_matrix( - expr, None, (assembled_factors[0],), (assembled_factors[1],), (scale,), bcs + expr, None, (row_factor,), (col_factor,), (scale,) ) def _normalise_lrc_form_product(self, expr, weight): @@ -508,9 +542,12 @@ def _normalise_lrc_form_product(self, expr, weight): raise ValueError("LRC FormProduct assembly requires exactly two rank-1 factors") return rank_one_factors[0], rank_one_factors[1], scalar_weight - def _assemble_lrc_factor(self, factor): + def _assemble_lrc_factor(self, factor, bcs): assembled = firedrake.assemble( - factor, form_compiler_parameters=self._form_compiler_params + factor, + bcs=self._lrc_factor_bcs(factor, bcs), + form_compiler_parameters=self._form_compiler_params, + is_base_form_preprocessed=True, ) if not isinstance(assembled, (firedrake.Cofunction, firedrake.Function)): raise TypeError("LRC FormProduct factors must assemble to Functions or Cofunctions") @@ -533,8 +570,7 @@ def _assemble_form_sum_lrc(self, expr, tensor, bcs): return None if tensor is not None: raise NotImplementedError("Assembly of FormSum with LRC terms into an existing tensor is not supported") - if bcs: - raise NotImplementedError("Boundary conditions on LRC FormSum assembly are not supported") + self._check_lrc_bcs(bcs) base_matrix = None if base_terms: @@ -557,13 +593,13 @@ def _assemble_form_sum_lrc(self, expr, tensor, bcs): row_factors = [] col_factors = [] lrc_weights = [] - for product, weight in lrc_terms: - row_factor, col_factor, lrc_weight = self._normalise_lrc_form_product(product, weight) - row_factors.append(self._assemble_lrc_factor(row_factor)) - col_factors.append(self._assemble_lrc_factor(col_factor)) + for prod, weight in lrc_terms: + row_factor, col_factor, lrc_weight = self._normalise_lrc_form_product(prod, weight) + row_factors.append(self._assemble_lrc_factor(row_factor, bcs)) + col_factors.append(self._assemble_lrc_factor(col_factor, bcs)) lrc_weights.append(lrc_weight) - return self._assemble_lrc_matrix(expr, base_matrix, row_factors, col_factors, lrc_weights, bcs) + return self._assemble_lrc_matrix(expr, base_matrix, row_factors, col_factors, lrc_weights) def _assemble_form_product(self, expr, tensor, bcs, assembled_factors): ranks = tuple(len(factor.arguments()) for factor in expr.factors()) @@ -595,7 +631,7 @@ def _assemble_form_product(self, expr, tensor, bcs, assembled_factors): higher_rank_form, tensor, bcs, assembled_higher_rank_factors) assembled_higher_rank_form = self._assemble_form_product_lrc( higher_rank_form, None, bcs, assembled_higher_rank_factors, scale=scalar_weight) - return Matrix(expr, assembled_higher_rank_form.petscmat, bcs=bcs, + return Matrix(expr, assembled_higher_rank_form.petscmat, bcs=(), options_prefix=self._options_prefix, fc_params=self._form_compiler_params) else: raise ValueError("FormProduct preprocessing requires remaining aggregate rank <= 2") diff --git a/tests/firedrake/regression/test_assemble_baseform.py b/tests/firedrake/regression/test_assemble_baseform.py index 3d40967b8c..ef29dfa96c 100644 --- a/tests/firedrake/regression/test_assemble_baseform.py +++ b/tests/firedrake/regression/test_assemble_baseform.py @@ -562,17 +562,29 @@ def test_assemble_formsum_lrc_rejects_tensor(mesh): assemble(formsum, mat_type="lrc", tensor=tensor) -def test_assemble_formsum_lrc_rejects_bcs(mesh): +def test_assemble_formsum_lrc_with_bcs(mesh): V = FunctionSpace(mesh, "CG", 1) v = TestFunction(V) u = TrialFunction(V) - form = inner(u, v) * dx - product = ufl.FormProduct(v * dx, v * dx) - formsum = ufl.FormSum((form, 1), (product, 1)) + x, y = SpatialCoordinate(mesh) + f = Function(V).interpolate(1 + x) + g = Function(V).interpolate(2 - y) + probe = Function(V).interpolate(1 + x + y) + + base_form = inner(u, v) * dx + row_form = inner(f, v) * dx + col_form = inner(g, v) * dx + product = ufl.FormProduct(row_form, col_form) + formsum = ufl.FormSum((base_form, 1), (product, 0.5)) bc = DirichletBC(V, 0, "on_boundary") - with pytest.raises(NotImplementedError, match="Boundary conditions on LRC FormSum"): - assemble(formsum, mat_type="lrc", bcs=bc) + A = assemble(formsum, mat_type="lrc", sub_mat_type="aij", bcs=bc) + assert A.petscmat.getType() == "lrc" + assert not A.has_bcs + + base_matrix = assemble(base_form, mat_type="aij", bcs=bc) + terms = ((0.5, assemble(row_form, bcs=bc), assemble(col_form, bcs=bc)),) + _assert_lrc_action(A, probe, terms, base_matrix=base_matrix) def test_assemble_formproduct_rank_two_with_scalar_factors(mesh): @@ -637,12 +649,22 @@ def test_assemble_formproduct_lrc_rejects_non_rank_one_factor(mesh): assemble(product, mat_type="lrc") -def test_assemble_formproduct_lrc_rejects_bcs(mesh): +def test_assemble_formproduct_lrc_with_bcs(mesh): V = FunctionSpace(mesh, "CG", 1) v = TestFunction(V) - form = v * dx - product = ufl.FormProduct(form, form) + x, y = SpatialCoordinate(mesh) + f = Function(V).interpolate(1 + x) + g = Function(V).interpolate(2 - y) + probe = Function(V).interpolate(1 + x + y) + + row_form = inner(f, v) * dx + col_form = inner(g, v) * dx + product = ufl.FormProduct(row_form, col_form) bc = DirichletBC(V, 0, "on_boundary") - with pytest.raises(NotImplementedError, match="Boundary conditions"): - assemble(product, mat_type="lrc", bcs=bc) + A = assemble(product, mat_type="lrc", bcs=bc) + assert A.petscmat.getType() == "lrc" + assert not A.has_bcs + + terms = ((1, assemble(row_form, bcs=bc), assemble(col_form, bcs=bc)),) + _assert_lrc_action(A, probe, terms) From 2d879df5a254a837753399d9f6df48dbbe794fbd Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 19 Jun 2026 21:18:49 +0100 Subject: [PATCH 6/7] _LRCDescriptor --- firedrake/assemble.py | 183 ++++++++++++++++++++++-------------------- 1 file changed, 94 insertions(+), 89 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 1458c65e31..ea79f9c382 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -319,6 +319,31 @@ def assemble(self, tensor=None, current_state=None): """ +class _LRCDescriptor: + """Unassembled low-rank matrix contribution. + + The descriptor lets BaseForm traversal carry low-rank columns through + parent nodes such as FormSum before a single MATLRC is created. + """ + + def __init__(self, expr, base_terms=(), row_factors=(), col_factors=(), weights=()): + self.expr = expr + self.base_terms = tuple(base_terms) + self.row_factors = tuple(row_factors) + self.col_factors = tuple(col_factors) + self.weights = tuple(weights) + + def scaled(self, weight): + weight = PETSc.ScalarType(weight) + return _LRCDescriptor( + self.expr, + base_terms=tuple((matrix, weight * w) for matrix, w in self.base_terms), + row_factors=self.row_factors, + col_factors=self.col_factors, + weights=tuple(weight * w for w in self.weights), + ) + + class BaseFormAssembler(AbstractFormAssembler): """Base form assembler. @@ -449,16 +474,6 @@ def _lrc_weights_to_vec(weights): c.assemble() return c - @staticmethod - def _is_lrc_form_product(expr): - if not isinstance(expr, ufl.FormProduct): - return False - - ranks = tuple(len(factor.arguments()) for factor in expr.factors()) - return (len(expr.arguments()) == 2 - and ranks.count(1) == 2 - and all(rank in (0, 1) for rank in ranks)) - @staticmethod def _bc_matches_space(bc, function_space): if not isinstance(bc, DirichletBC): @@ -489,17 +504,36 @@ def _lrc_factor_bcs(self, factor, bcs): function_space = arg.function_space() return tuple(bc for bc in bcs if self._bc_matches_space(bc, function_space)) - def _assemble_lrc_matrix(self, expr, base_matrix, row_factors, col_factors, weights): - U = self._vector_columns_to_dense(row_factors) - V = self._vector_columns_to_dense(col_factors) - c = self._lrc_weights_to_vec(weights) - A = base_matrix.petscmat if base_matrix is not None else None + def _lrc_base_petscmat(self, base_terms): + if not base_terms: + return None + + if len(base_terms) == 1: + matrix, weight = base_terms[0] + if weight == 1: + return matrix.petscmat + + result = PETSc.Mat() + for i, (matrix, weight) in enumerate(base_terms): + if i == 0: + matrix.petscmat.copy(result=result) + result.scale(weight) + else: + result.axpy(weight, matrix.petscmat) + result.assemble() + return result + + def _assemble_lrc_descriptor(self, descriptor): + U = self._vector_columns_to_dense(descriptor.row_factors) + V = self._vector_columns_to_dense(descriptor.col_factors) + c = self._lrc_weights_to_vec(descriptor.weights) + A = self._lrc_base_petscmat(descriptor.base_terms) petscmat = PETSc.Mat().createLRC(A, U, c, V) petscmat.assemble() - return Matrix(expr, petscmat, bcs=(), options_prefix=self._options_prefix, + return Matrix(descriptor.expr, petscmat, bcs=(), options_prefix=self._options_prefix, fc_params=self._form_compiler_params) - def _assemble_form_product_lrc(self, expr, tensor, bcs, assembled_factors, scale=1): + def _form_product_lrc_descriptor(self, expr, tensor, bcs, assembled_factors, scale=1): if self._mat_type != "lrc": raise ValueError("FormProduct assembly requires mat_type='lrc'") if tensor is not None: @@ -521,26 +555,8 @@ def _assemble_form_product_lrc(self, expr, tensor, bcs, assembled_factors, scale raise TypeError("LRC FormProduct factors must assemble to Functions or Cofunctions") row_factor, col_factor = assembled_factors - return self._assemble_lrc_matrix( - expr, None, (row_factor,), (col_factor,), (scale,) - ) - - def _normalise_lrc_form_product(self, expr, weight): - scalar_weight = PETSc.ScalarType(weight) - rank_one_factors = [] - for factor in expr.factors(): - rank = len(factor.arguments()) - if rank == 0: - assembled_factor = firedrake.assemble( - factor, form_compiler_parameters=self._form_compiler_params - ) - scalar_weight *= self._as_scalar_value(assembled_factor) - else: - rank_one_factors.append(factor) - - if len(rank_one_factors) != 2: - raise ValueError("LRC FormProduct assembly requires exactly two rank-1 factors") - return rank_one_factors[0], rank_one_factors[1], scalar_weight + return _LRCDescriptor(expr, row_factors=(row_factor,), col_factors=(col_factor,), + weights=(PETSc.ScalarType(scale),)) def _assemble_lrc_factor(self, factor, bcs): assembled = firedrake.assemble( @@ -553,60 +569,41 @@ def _assemble_lrc_factor(self, factor, bcs): raise TypeError("LRC FormProduct factors must assemble to Functions or Cofunctions") return assembled - def _assemble_form_sum_lrc(self, expr, tensor, bcs): - if self._mat_type != "lrc" or not isinstance(expr, ufl.FormSum): + def _form_sum_lrc_descriptor(self, expr, tensor, bcs, assembled_components, weights): + if not any(isinstance(component, _LRCDescriptor) for component in assembled_components): return None - weights = [self._as_scalar_value(w) for w in expr.weights()] - base_terms = [] - lrc_terms = [] - for component, weight in zip(expr.components(), weights): - if self._is_lrc_form_product(component): - lrc_terms.append((component, weight)) - else: - base_terms.append((component, weight)) - - if not lrc_terms: - return None + if self._mat_type != "lrc": + raise ValueError("FormSum with FormProduct terms requires mat_type='lrc'") if tensor is not None: raise NotImplementedError("Assembly of FormSum with LRC terms into an existing tensor is not supported") self._check_lrc_bcs(bcs) - base_matrix = None - if base_terms: - base_form = ufl.FormSum(*base_terms) - base_mat_type = self._sub_mat_type - base_matrix = firedrake.assemble( - base_form, - bcs=bcs, - form_compiler_parameters=self._form_compiler_params, - mat_type=base_mat_type, - sub_mat_type=self._sub_mat_type, - options_prefix=self._options_prefix, - weight=self._weight, - allocation_integral_types=self.allocation_integral_types, - is_base_form_preprocessed=True, - ) - if not isinstance(base_matrix, MatrixBase): - raise TypeError("LRC base form must assemble to a matrix") - + base_terms = [] row_factors = [] col_factors = [] lrc_weights = [] - for prod, weight in lrc_terms: - row_factor, col_factor, lrc_weight = self._normalise_lrc_form_product(prod, weight) - row_factors.append(self._assemble_lrc_factor(row_factor, bcs)) - col_factors.append(self._assemble_lrc_factor(col_factor, bcs)) - lrc_weights.append(lrc_weight) + for component, weight in zip(assembled_components, weights): + if isinstance(component, _LRCDescriptor): + descriptor = component.scaled(weight) + base_terms.extend(descriptor.base_terms) + row_factors.extend(descriptor.row_factors) + col_factors.extend(descriptor.col_factors) + lrc_weights.extend(descriptor.weights) + elif isinstance(component, MatrixBase): + base_terms.append((component, PETSc.ScalarType(weight))) + else: + raise TypeError("Mismatching FormSum shapes") - return self._assemble_lrc_matrix(expr, base_matrix, row_factors, col_factors, lrc_weights) + return _LRCDescriptor(expr, base_terms=base_terms, row_factors=row_factors, + col_factors=col_factors, weights=lrc_weights) def _assemble_form_product(self, expr, tensor, bcs, assembled_factors): ranks = tuple(len(factor.arguments()) for factor in expr.factors()) if len(assembled_factors) != len(expr.factors()): - return self._assemble_form_product_lrc(expr, tensor, bcs, assembled_factors) + return self._form_product_lrc_descriptor(expr, tensor, bcs, assembled_factors) if sum(ranks) > 2 or not any(rank == 0 for rank in ranks): - return self._assemble_form_product_lrc(expr, tensor, bcs, assembled_factors) + return self._form_product_lrc_descriptor(expr, tensor, bcs, assembled_factors) scalar_weight = PETSc.ScalarType(1) higher_rank_factors = [] @@ -626,13 +623,13 @@ def _assemble_form_product(self, expr, tensor, bcs, assembled_factors): elif len(higher_rank_factors) == 2 and all(len(factor.arguments()) == 1 for factor in higher_rank_factors): higher_rank_form = ufl.FormProduct(*higher_rank_factors) - if tensor is not None: - return self._assemble_form_product_lrc( - higher_rank_form, tensor, bcs, assembled_higher_rank_factors) - assembled_higher_rank_form = self._assemble_form_product_lrc( - higher_rank_form, None, bcs, assembled_higher_rank_factors, scale=scalar_weight) - return Matrix(expr, assembled_higher_rank_form.petscmat, bcs=(), - options_prefix=self._options_prefix, fc_params=self._form_compiler_params) + descriptor = self._form_product_lrc_descriptor( + higher_rank_form, tensor, bcs, assembled_higher_rank_factors, + scale=scalar_weight) + return _LRCDescriptor(expr, + row_factors=descriptor.row_factors, + col_factors=descriptor.col_factors, + weights=descriptor.weights) else: raise ValueError("FormProduct preprocessing requires remaining aggregate rank <= 2") @@ -661,11 +658,6 @@ def assemble(self, tensor=None, current_state=None): in a post-order fashion and evaluating the nodes on the fly. """ - if isinstance(self._form, ufl.FormSum) and len(self._form.arguments()) == 2: - lrc_result = self._assemble_form_sum_lrc(self._form, tensor, self._bcs) - if lrc_result is not None: - return lrc_result - def visitor(e, *operands): t = tensor if e is self._form else None # Deal with 2-form bcs inside the visitor @@ -675,6 +667,8 @@ def visitor(e, *operands): # DAG assembly: traverse the DAG in a post-order fashion and evaluate the node on the fly. visited = {} result = BaseFormAssembler.base_form_postorder_traversal(self._form, visitor, visited) + if isinstance(result, _LRCDescriptor): + result = self._assemble_lrc_descriptor(result) # Deal with 1-form bcs outside the visitor rank = len(self._form.arguments()) @@ -704,8 +698,9 @@ def base_form_assembly_visitor(self, expr, tensor, bcs, *args): assembler = OneFormAssembler(form, form_compiler_parameters=self._form_compiler_params, zero_bc_nodes=self._zero_bc_nodes, diagonal=self._diagonal, weight=self._weight) elif rank == 2: + mat_type = self._sub_mat_type if self._mat_type == "lrc" else self._mat_type assembler = TwoFormAssembler(form, bcs=bcs, form_compiler_parameters=self._form_compiler_params, - mat_type=self._mat_type, sub_mat_type=self._sub_mat_type, + mat_type=mat_type, sub_mat_type=self._sub_mat_type, options_prefix=self._options_prefix, appctx=self._appctx, weight=self._weight, allocation_integral_types=self.allocation_integral_types) else: @@ -717,6 +712,8 @@ def base_form_assembly_visitor(self, expr, tensor, bcs, *args): if len(args) != 1: raise TypeError("Not enough operands for Adjoint") mat, = args + if isinstance(mat, _LRCDescriptor): + mat = self._assemble_lrc_descriptor(mat) result = tensor.petscmat if tensor else PETSc.Mat() # Out-of-place Hermitian transpose mat.petscmat.hermitianTranspose(out=result) @@ -728,6 +725,10 @@ def base_form_assembly_visitor(self, expr, tensor, bcs, *args): if len(args) != 2: raise TypeError("Not enough operands for Action") lhs, rhs = args + if isinstance(lhs, _LRCDescriptor): + lhs = self._assemble_lrc_descriptor(lhs) + if isinstance(rhs, _LRCDescriptor): + rhs = self._assemble_lrc_descriptor(rhs) if isinstance(lhs, MatrixBase): if isinstance(rhs, (firedrake.Cofunction, firedrake.Function)): petsc_mat = lhs.petscmat @@ -773,6 +774,10 @@ def base_form_assembly_visitor(self, expr, tensor, bcs, *args): for w in expr.weights(): weights.append(self._as_scalar_value(w)) + lrc_descriptor = self._form_sum_lrc_descriptor(expr, tensor, bcs, args, weights) + if lrc_descriptor is not None: + return lrc_descriptor + # Scalar FormSum if all(isinstance(op, numbers.Complex) for op in args): result = numpy.dot(weights, args) From 13c5a92c6d6f9acf417848c302fc7685a8528238 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 19 Jun 2026 22:15:47 +0100 Subject: [PATCH 7/7] cleanup --- firedrake/assemble.py | 222 +++++++++--------- .../regression/test_assemble_baseform.py | 4 +- 2 files changed, 108 insertions(+), 118 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index ea79f9c382..4c1d44c233 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -326,12 +326,13 @@ class _LRCDescriptor: parent nodes such as FormSum before a single MATLRC is created. """ - def __init__(self, expr, base_terms=(), row_factors=(), col_factors=(), weights=()): + def __init__(self, expr, base_terms=(), row_factors=(), col_factors=(), weights=(), bcs=None): self.expr = expr self.base_terms = tuple(base_terms) self.row_factors = tuple(row_factors) self.col_factors = tuple(col_factors) self.weights = tuple(weights) + self.bcs = bcs def scaled(self, weight): weight = PETSc.ScalarType(weight) @@ -343,6 +344,97 @@ def scaled(self, weight): weights=tuple(weight * w for w in self.weights), ) + @cached_property + def petscmat(self): + U = self._vector_columns_to_dense(self.row_factors, bcs=self.bcs) + V = self._vector_columns_to_dense(self.col_factors, bcs=self.bcs) + c = self._lrc_weights_to_vec(self.weights) + A = self._lrc_base_petscmat(self.base_terms) + petscmat = PETSc.Mat().createLRC(A, U, c, V) + petscmat.assemble() + return petscmat + + @staticmethod + def _vector_columns_to_dense(functions, bcs=None): + if not functions: + raise ValueError("Cannot build an LRC factor matrix with no columns") + + first, *_ = functions + with first.dat.vec_ro as vec: + comm = vec.comm + local_size = vec.getLocalSize() + global_size = vec.getSize() + + dense = PETSc.Mat().createDense( + size=((local_size, global_size), (len(functions), len(functions))), comm=comm + ) + dense.setUp() + values = dense.getDenseArray() + for i, function in enumerate(functions): + function = _LRCDescriptor._apply_bcs(function, bcs) + with function.dat.vec_ro as vec: + if vec.getLocalSize() != local_size or vec.getSize() != global_size: + raise ValueError("LRC factor vectors do not share the same layout") + values[:, i] = vec.getArray(readonly=True) + dense.assemble() + return dense + + @staticmethod + def _lrc_weights_to_vec(weights): + c = PETSc.Vec().createSeq(len(weights), comm=PETSc.COMM_SELF) + c.setValues(range(len(weights)), weights) + c.assemble() + return c + + @staticmethod + def _bc_matches_space(bc, function_space): + V = bc.function_space() + while True: + if V == function_space or V.dual() == function_space: + return True + if V.parent is None: + return False + V = V.parent + + @staticmethod + def _lrc_factor_bcs(factor, bcs): + if not bcs: + return () + + arg, = factor.arguments() + function_space = arg.function_space() + return tuple(bc for bc in bcs if _LRCDescriptor._bc_matches_space(bc, function_space)) + + @staticmethod + def _lrc_base_petscmat(base_terms): + if not base_terms: + return None + + if len(base_terms) == 1: + matrix, weight = base_terms[0] + if weight == 1: + return matrix.petscmat + + result = PETSc.Mat() + for i, (matrix, weight) in enumerate(base_terms): + if i == 0: + matrix.petscmat.copy(result=result) + result.scale(weight) + else: + result.axpy(weight, matrix.petscmat) + result.assemble() + return result + + @staticmethod + def _apply_bcs(factor, bcs): + bcs = _LRCDescriptor._lrc_factor_bcs(factor, bcs) + if not bcs: + return factor + factor = factor.copy(deepcopy=True) + for bc in bcs: + bc.zero(factor) + return factor + class BaseFormAssembler(AbstractFormAssembler): """Base form assembler. @@ -443,131 +535,24 @@ def _as_scalar_value(value): raise ValueError("Expecting a scalar expression") return value - @staticmethod - def _vector_columns_to_dense(functions): - if not functions: - raise ValueError("Cannot build an LRC factor matrix with no columns") - - first, *_ = functions - with first.dat.vec_ro as vec: - comm = vec.comm - local_size = vec.getLocalSize() - global_size = vec.getSize() - - dense = PETSc.Mat().createDense( - size=((local_size, global_size), (len(functions), len(functions))), comm=comm - ) - dense.setUp() - values = dense.getDenseArray() - for i, function in enumerate(functions): - with function.dat.vec_ro as vec: - if vec.getLocalSize() != local_size or vec.getSize() != global_size: - raise ValueError("LRC factor vectors do not share the same layout") - values[:, i] = vec.getArray(readonly=True) - dense.assemble() - return dense - - @staticmethod - def _lrc_weights_to_vec(weights): - c = PETSc.Vec().createSeq(len(weights), comm=PETSc.COMM_SELF) - c.setValues(range(len(weights)), weights) - c.assemble() - return c - - @staticmethod - def _bc_matches_space(bc, function_space): - if not isinstance(bc, DirichletBC): - return False - - if ufl.duals.is_dual(function_space): - function_space = function_space.dual() - - V = bc.function_space() - while True: - bc_space = V.dual() if ufl.duals.is_dual(V) else V - if bc_space == function_space: - return True - if V.parent is None: - return False - V = V.parent - - @staticmethod - def _check_lrc_bcs(bcs): - if any(not isinstance(bc, DirichletBC) for bc in bcs): - raise NotImplementedError("Only DirichletBC objects are supported on LRC assembly") - - def _lrc_factor_bcs(self, factor, bcs): - if not bcs: - return () - - arg, = factor.arguments() - function_space = arg.function_space() - return tuple(bc for bc in bcs if self._bc_matches_space(bc, function_space)) - - def _lrc_base_petscmat(self, base_terms): - if not base_terms: - return None - - if len(base_terms) == 1: - matrix, weight = base_terms[0] - if weight == 1: - return matrix.petscmat - - result = PETSc.Mat() - for i, (matrix, weight) in enumerate(base_terms): - if i == 0: - matrix.petscmat.copy(result=result) - result.scale(weight) - else: - result.axpy(weight, matrix.petscmat) - result.assemble() - return result - - def _assemble_lrc_descriptor(self, descriptor): - U = self._vector_columns_to_dense(descriptor.row_factors) - V = self._vector_columns_to_dense(descriptor.col_factors) - c = self._lrc_weights_to_vec(descriptor.weights) - A = self._lrc_base_petscmat(descriptor.base_terms) - petscmat = PETSc.Mat().createLRC(A, U, c, V) - petscmat.assemble() - return Matrix(descriptor.expr, petscmat, bcs=(), options_prefix=self._options_prefix, - fc_params=self._form_compiler_params) - def _form_product_lrc_descriptor(self, expr, tensor, bcs, assembled_factors, scale=1): if self._mat_type != "lrc": raise ValueError("FormProduct assembly requires mat_type='lrc'") if tensor is not None: raise NotImplementedError("Assembly of FormProduct into an existing tensor is not supported") - self._check_lrc_bcs(bcs) if len(expr.factors()) != 2: raise NotImplementedError("LRC FormProduct assembly currently supports exactly two factors") if len(expr.arguments()) != 2: raise ValueError("LRC FormProduct assembly requires aggregate rank 2") if any(len(factor.arguments()) != 1 for factor in expr.factors()): raise ValueError("LRC FormProduct assembly requires rank-1 factors") - if bcs: - row_factor = self._assemble_lrc_factor(expr.factors()[0], bcs) - col_factor = self._assemble_lrc_factor(expr.factors()[1], bcs) - else: - if len(assembled_factors) != 2: - raise TypeError("Not enough operands for FormProduct") - if not all(isinstance(factor, (firedrake.Cofunction, firedrake.Function)) for factor in assembled_factors): - raise TypeError("LRC FormProduct factors must assemble to Functions or Cofunctions") - row_factor, col_factor = assembled_factors - - return _LRCDescriptor(expr, row_factors=(row_factor,), col_factors=(col_factor,), - weights=(PETSc.ScalarType(scale),)) - - def _assemble_lrc_factor(self, factor, bcs): - assembled = firedrake.assemble( - factor, - bcs=self._lrc_factor_bcs(factor, bcs), - form_compiler_parameters=self._form_compiler_params, - is_base_form_preprocessed=True, - ) - if not isinstance(assembled, (firedrake.Cofunction, firedrake.Function)): + if len(assembled_factors) != 2: + raise TypeError("Not enough operands for FormProduct") + if not all(isinstance(factor, (firedrake.Cofunction, firedrake.Function)) for factor in assembled_factors): raise TypeError("LRC FormProduct factors must assemble to Functions or Cofunctions") - return assembled + + return _LRCDescriptor(expr, row_factors=assembled_factors[:1], col_factors=assembled_factors[1:], + weights=(PETSc.ScalarType(scale),), bcs=bcs) def _form_sum_lrc_descriptor(self, expr, tensor, bcs, assembled_components, weights): if not any(isinstance(component, _LRCDescriptor) for component in assembled_components): @@ -577,7 +562,6 @@ def _form_sum_lrc_descriptor(self, expr, tensor, bcs, assembled_components, weig raise ValueError("FormSum with FormProduct terms requires mat_type='lrc'") if tensor is not None: raise NotImplementedError("Assembly of FormSum with LRC terms into an existing tensor is not supported") - self._check_lrc_bcs(bcs) base_terms = [] row_factors = [] @@ -598,6 +582,11 @@ def _form_sum_lrc_descriptor(self, expr, tensor, bcs, assembled_components, weig return _LRCDescriptor(expr, base_terms=base_terms, row_factors=row_factors, col_factors=col_factors, weights=lrc_weights) + def _assemble_lrc_descriptor(self, descriptor): + return Matrix(descriptor.expr, descriptor.petscmat, bcs=descriptor.bcs, + options_prefix=self._options_prefix, + fc_params=self._form_compiler_params) + def _assemble_form_product(self, expr, tensor, bcs, assembled_factors): ranks = tuple(len(factor.arguments()) for factor in expr.factors()) if len(assembled_factors) != len(expr.factors()): @@ -667,8 +656,9 @@ def visitor(e, *operands): # DAG assembly: traverse the DAG in a post-order fashion and evaluate the node on the fly. visited = {} result = BaseFormAssembler.base_form_postorder_traversal(self._form, visitor, visited) + if isinstance(result, _LRCDescriptor): - result = self._assemble_lrc_descriptor(result) + return self._assemble_lrc_descriptor(result) # Deal with 1-form bcs outside the visitor rank = len(self._form.arguments()) diff --git a/tests/firedrake/regression/test_assemble_baseform.py b/tests/firedrake/regression/test_assemble_baseform.py index ef29dfa96c..17aad98279 100644 --- a/tests/firedrake/regression/test_assemble_baseform.py +++ b/tests/firedrake/regression/test_assemble_baseform.py @@ -580,7 +580,7 @@ def test_assemble_formsum_lrc_with_bcs(mesh): A = assemble(formsum, mat_type="lrc", sub_mat_type="aij", bcs=bc) assert A.petscmat.getType() == "lrc" - assert not A.has_bcs + assert A.has_bcs base_matrix = assemble(base_form, mat_type="aij", bcs=bc) terms = ((0.5, assemble(row_form, bcs=bc), assemble(col_form, bcs=bc)),) @@ -664,7 +664,7 @@ def test_assemble_formproduct_lrc_with_bcs(mesh): A = assemble(product, mat_type="lrc", bcs=bc) assert A.petscmat.getType() == "lrc" - assert not A.has_bcs + assert A.has_bcs terms = ((1, assemble(row_form, bcs=bc), assemble(col_form, bcs=bc)),) _assert_lrc_action(A, probe, terms)