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 diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 6dd7c4f03f..4c1d44c233 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -319,6 +319,123 @@ 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=(), 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) + 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), + ) + + @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. @@ -400,6 +517,114 @@ 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 + + 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") + 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") + + 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): + 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") + + base_terms = [] + row_factors = [] + col_factors = [] + lrc_weights = [] + 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 _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()): + 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._form_product_lrc_descriptor(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) + 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") + + 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. @@ -432,6 +657,9 @@ def visitor(e, *operands): visited = {} result = BaseFormAssembler.base_form_postorder_traversal(self._form, visitor, visited) + if isinstance(result, _LRCDescriptor): + return self._assemble_lrc_descriptor(result) + # Deal with 1-form bcs outside the visitor rank = len(self._form.arguments()) if rank == 1 and not isinstance(result, ufl.ZeroBaseForm): @@ -460,17 +688,22 @@ 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: raise AssertionError return assembler.assemble(tensor=tensor) + elif isinstance(expr, ufl.FormProduct): + 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") 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) @@ -482,6 +715,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 @@ -525,19 +762,11 @@ 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)) + + 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): @@ -683,12 +912,22 @@ 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): + ranks = tuple(len(factor.arguments()) for factor in expr.factors()) + if (len(expr.factors()) == 2 and len(expr.arguments()) == 2 + 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)): 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..17aad98279 100644 --- a/tests/firedrake/regression/test_assemble_baseform.py +++ b/tests/firedrake/regression/test_assemble_baseform.py @@ -382,3 +382,289 @@ 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_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 _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_with_bcs(mesh): + V = FunctionSpace(mesh, "CG", 1) + v = TestFunction(V) + u = TrialFunction(V) + 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") + + A = assemble(formsum, mat_type="lrc", sub_mat_type="aij", bcs=bc) + assert A.petscmat.getType() == "lrc" + 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)),) + _assert_lrc_action(A, probe, terms, base_matrix=base_matrix) + + +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) + 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_with_bcs(mesh): + V = FunctionSpace(mesh, "CG", 1) + v = TestFunction(V) + 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") + + A = assemble(product, mat_type="lrc", bcs=bc) + assert A.petscmat.getType() == "lrc" + assert A.has_bcs + + terms = ((1, assemble(row_form, bcs=bc), assemble(col_form, bcs=bc)),) + _assert_lrc_action(A, probe, terms)