Merge pull request #4656 from firedrakeproject/pbrubeck/merge-release-into-main

DO NOT SQUASH merge release into main

Author: pbrubeck

docs/source/conf.py

@@ -172,6 +172,8 @@
     r'https://www.hilton.com/en/hotels/leehnhn-hilton-leeds-city/',
     r'https://www.radissonhotels.com/en-us/hotels/park-plaza-leeds',
     r'https://www.radissonhotels.com/en-us/hotels/radisson-blu-leeds'
+    r'https://www.radissonhotels.com/en-us/hotels/radisson-blu-leeds',
+    r'https://all.accor.com/hotel/*',
 ]
 linkcheck_timeout = 30
 
@@ -413,10 +415,8 @@
 
 intersphinx_mapping = {
     'pyop2': ('https://op2.github.io/PyOP2', None),
-    'ufl': ('https://fenics.readthedocs.io/projects/ufl/en/latest/', None),
     'ufl': ('https://docs.fenicsproject.org/ufl/main/', None),
-    'FIAT': ('https://fenics.readthedocs.io/projects/fiat/en/latest/', None),
-    'FInAT': ('https://finat.github.io/FInAT/', None),
+    'FIAT': ('https://firedrakeproject.org/fiat', None),
     'petsctools': ('https://firedrakeproject.org/petsctools/', None),
     'mpi4py': ('https://mpi4py.readthedocs.io/en/stable/', None),
     'h5py': ('http://docs.h5py.org/en/latest/', None),

docs/source/install.rst

@@ -590,7 +590,7 @@ should be followed:
 
 #. Install Firedrake in editable mode without build isolation::
 
-   $ pip install --no-build-isolation --no-binary h5py --editable './firedrake[check]'
+   $ pip install --no-build-isolation --no-binary h5py --editable './firedrake[check,docs]'
 
 
 Editing subpackages

firedrake/adjoint_utils/blocks/solving.py

@@ -163,6 +163,13 @@ def _should_compute_boundary_adjoint(self, relevant_dependencies):
     def adj_sol(self):
         return self.adj_state
 
+    @adj_sol.setter
+    def adj_sol(self, value):
+        if self.adj_state is None:
+            self.adj_state = value.copy(deepcopy=True)
+        else:
+            self.adj_state.assign(value)
+
     def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
         fwd_block_variable = self.get_outputs()[0]
         u = fwd_block_variable.output
@@ -187,7 +194,7 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
         adj_sol, adj_sol_bdy = self._assemble_and_solve_adj_eq(
             dFdu_form, dJdu, compute_bdy
         )
-        self.adj_state = adj_sol
+        self.adj_sol = adj_sol
         if self.adj_cb is not None:
             self.adj_cb(adj_sol)
         if self.adj_bdy_cb is not None and compute_bdy:
@@ -408,7 +415,7 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs,
             firedrake.derivative(dFdu_form, fwd_block_variable.saved_output,
                                  tlm_output))
 
-        adj_sol = self.adj_state
+        adj_sol = self.adj_sol
         if adj_sol is None:
             raise RuntimeError("Hessian computation was run before adjoint.")
         bdy = self._should_compute_boundary_adjoint(relevant_dependencies)
@@ -726,7 +733,7 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
             relevant_dependencies
         )
         adj_sol, adj_sol_bdy = self._adjoint_solve(adj_inputs[0], compute_bdy)
-        self.adj_state = adj_sol
+        self.adj_sol = adj_sol
         if self.adj_cb is not None:
             self.adj_cb(adj_sol)
         if self.adj_bdy_cb is not None and compute_bdy:

firedrake/adjoint_utils/variational_solver.py

@@ -27,7 +27,7 @@ def wrapper(self, *args, **kwargs):
                     # Try again without expanding derivatives,
                     # as dFdu might have been simplied to an empty Form
                     self._ad_adj_F = adjoint(dFdu, derivatives_expanded=True)
-            except (TypeError, NotImplementedError):
+            except (ValueError, TypeError, NotImplementedError):
                 self._ad_adj_F = None
             self._ad_kwargs = {'Jp': self.Jp, 'form_compiler_parameters': self.form_compiler_parameters, 'is_linear': self.is_linear}
             self._ad_count_map = {}

firedrake/assemble.py

@@ -60,17 +60,23 @@ def assemble(expr, *args, **kwargs):
         `ufl.classes.Measure` in the form. For example, if a
         ``quadrature_degree`` of 4 is specified in this argument, but a degree of
         3 is requested in the measure, the latter will be used.
-    mat_type : str
+    mat_type : str | None
         String indicating how a 2-form (matrix) should be
         assembled -- either as a monolithic matrix (``"aij"`` or ``"baij"``),
-        a block matrix (``"nest"``), or left as a `matrix.ImplicitMatrix` giving
+        a block matrix (``"nest"``), or left as a :class:`firedrake.matrix.ImplicitMatrix` giving
         matrix-free actions (``'matfree'``). If not supplied, the default value in
-        ``parameters["default_matrix_type"]`` is used.  BAIJ differs
-        from AIJ in that only the block sparsity rather than the dof
+        ``parameters["default_matrix_type"]`` is used.  ``"baij"``` differs
+        from ``"aij"`` in that only the block sparsity rather than the DoF
         sparsity is constructed.  This can result in some memory
         savings, but does not work with all PETSc preconditioners.
-        BAIJ matrices only make sense for non-mixed matrices.
-    sub_mat_type : str
+        ``"baij"`` matrices only make sense for non-mixed matrices with arguments
+        on a :func:`firedrake.functionspace.VectorFunctionSpace`.
+
+        NOTE
+        ----
+        For the assembly of a 0-form or 1-form arising from the action of a 2-form,
+        the default matrix type is ``"matfree"``.
+    sub_mat_type : str | None
         String indicating the matrix type to
         use *inside* a nested block matrix.  Only makes sense if
         ``mat_type`` is ``nest``.  May be one of ``"aij"`` or ``"baij"``.  If
@@ -154,7 +160,10 @@ def get_assembler(form, *args, **kwargs):
     is_base_form_preprocessed = kwargs.pop('is_base_form_preprocessed', False)
     fc_params = kwargs.get('form_compiler_parameters', None)
     if isinstance(form, ufl.form.BaseForm) and not is_base_form_preprocessed:
-        mat_type = kwargs.get('mat_type', None)
+        # If not assembling a matrix, internal BaseForm nodes are matfree by default
+        # Otherwise, the default matrix type is firedrake.parameters["default_matrix_type"]
+        default_mat_type = "matfree" if len(form.arguments()) < 2 else None
+        mat_type = kwargs.get('mat_type', default_mat_type)
         # Preprocess the DAG and restructure the DAG
         # Only pre-process `form` once beforehand to avoid pre-processing for each assembly call
         form = BaseFormAssembler.preprocess_base_form(form, mat_type=mat_type, form_compiler_parameters=fc_params)

firedrake/interpolation.py

@@ -930,8 +930,8 @@ def callable():
         return callable
     else:
         loops = []
-
-        if access == op2.INC:
+        # Initialise to zero if needed
+        if access is op2.INC:
             loops.append(tensor.zero)
 
         # Arguments in the operand are allowed to be from a MixedFunctionSpace
@@ -957,7 +957,7 @@ def callable():
         for indices, sub_expr in expressions.items():
             sub_tensor = tensor[indices[0]] if rank == 1 else tensor
             loops.extend(_interpolator(sub_tensor, sub_expr, subset, access, bcs=bcs))
-
+        # Apply bcs
         if bcs and rank == 1:
             loops.extend(partial(bc.apply, f) for bc in bcs)
 
@@ -1038,32 +1038,36 @@ def _interpolator(tensor, expr, subset, access, bcs=None):
     parameters = {}
     parameters['scalar_type'] = utils.ScalarType
 
-    callables = ()
+    copyin = ()
+    copyout = ()
 
     # For the matfree adjoint 1-form and the 0-form, the cellwise kernel will add multiple
     # contributions from the facet DOFs of the dual argument.
     # The incoming Cofunction needs to be weighted by the reciprocal of the DOF multiplicity.
     needs_weight = isinstance(dual_arg, ufl.Cofunction) and not to_element.is_dg()
     if needs_weight:
-        # Compute the reciprocal of the DOF multiplicity
+        # Create a buffer for the weighted Cofunction
         W = dual_arg.function_space()
+        v = firedrake.Function(W)
+        expr = expr._ufl_expr_reconstruct_(operand, v=v)
+        copyin += (partial(dual_arg.dat.copy, v.dat),)
+
+        # Compute the reciprocal of the DOF multiplicity
+        wdat = W.make_dat()
+        m_ = get_interp_node_map(source_mesh, target_mesh, W)
         wsize = W.finat_element.space_dimension() * W.block_size
         kernel_code = f"""
         void multiplicity(PetscScalar *restrict w) {{
             for (PetscInt i=0; i<{wsize}; i++) w[i] += 1;
         }}"""
-        kernel = op2.Kernel(kernel_code, "multiplicity", requires_zeroed_output_arguments=False)
-        weight = firedrake.Function(W)
-        m_ = get_interp_node_map(source_mesh, target_mesh, W)
-        op2.par_loop(kernel, cell_set, weight.dat(op2.INC, m_))
-        with weight.dat.vec as w:
+        kernel = op2.Kernel(kernel_code, "multiplicity")
+        op2.par_loop(kernel, cell_set, wdat(op2.INC, m_))
+        with wdat.vec as w:
             w.reciprocal()
 
-        # Create a buffer for the weighted Cofunction and a callable to apply the weight
-        v = firedrake.Function(W)
-        expr = expr._ufl_expr_reconstruct_(operand, v=v)
-        with weight.dat.vec_ro as w, dual_arg.dat.vec_ro as x, v.dat.vec_wo as y:
-            callables += (partial(y.pointwiseMult, x, w),)
+        # Create a callable to apply the weight
+        with wdat.vec_ro as w, v.dat.vec as y:
+            copyin += (partial(y.pointwiseMult, y, w),)
 
     # We need to pass both the ufl element and the finat element
     # because the finat elements might not have the right mapping
@@ -1079,7 +1083,7 @@ def _interpolator(tensor, expr, subset, access, bcs=None):
     coefficient_numbers = kernel.coefficient_numbers
     needs_external_coords = kernel.needs_external_coords
     name = kernel.name
-    kernel = op2.Kernel(ast, name, requires_zeroed_output_arguments=True,
+    kernel = op2.Kernel(ast, name, requires_zeroed_output_arguments=(access is not op2.INC),
                         flop_count=kernel.flop_count, events=(kernel.event,))
 
     parloop_args = [kernel, cell_set]
@@ -1092,17 +1096,12 @@ def _interpolator(tensor, expr, subset, access, bcs=None):
         output = tensor
         tensor = op2.Dat(tensor.dataset)
         if access is not op2.WRITE:
-            copyin = (partial(output.copy, tensor), )
-        else:
-            copyin = ()
-        copyout = (partial(tensor.copy, output), )
-    else:
-        copyin = ()
-        copyout = ()
+            copyin += (partial(output.copy, tensor), )
+        copyout += (partial(tensor.copy, output), )
     if isinstance(tensor, op2.Global):
         parloop_args.append(tensor(access))
     elif isinstance(tensor, op2.Dat):
-        V_dest = arguments[-1].function_space() if isinstance(dual_arg, ufl.Cofunction) else V
+        V_dest = arguments[-1].function_space()
         m_ = get_interp_node_map(source_mesh, target_mesh, V_dest)
         parloop_args.append(tensor(access, m_))
     else:
@@ -1162,11 +1161,10 @@ def _interpolator(tensor, expr, subset, access, bcs=None):
                 parloop_args.append(target_ref_coords.dat(op2.READ, m_))
 
     parloop = op2.ParLoop(*parloop_args)
-    parloop_compute_callable = parloop.compute
     if isinstance(tensor, op2.Mat):
-        return parloop_compute_callable, tensor.assemble
+        return parloop, tensor.assemble
     else:
-        return copyin + callables + (parloop_compute_callable, ) + copyout
+        return copyin + (parloop, ) + copyout
 
 
 def get_interp_node_map(source_mesh, target_mesh, fs):

firedrake/mg/kernels.py

@@ -30,6 +30,8 @@
 from tsfc.driver import TSFCIntegralDataInfo
 from tsfc.kernel_interface.common import lower_integral_type
 from tsfc.parameters import default_parameters
+from tsfc.ufl_utils import apply_mapping, simplify_abs
+
 from finat.element_factory import create_element
 from finat.quadrature import make_quadrature
 from firedrake.pointquery_utils import dX_norm_square, X_isub_dX, init_X, inside_check, is_affine, celldist_l1_c_expr
@@ -113,28 +115,34 @@ def compile_element(expression, dual_space=None, parameters=None,
     # # Collect required coefficients
 
     try:
+        # Forward interpolation: expression has a coefficient
         arg, = extract_coefficients(expression)
         argument_multiindices = ()
         coefficient = True
-        if expression.ufl_shape:
-            tensor_indices = tuple(gem.Index() for s in expression.ufl_shape)
-        else:
-            tensor_indices = ()
     except ValueError:
+        # Adjoint interpolation: expression has an argument
         arg, = extract_arguments(expression)
         finat_elem = create_element(arg.ufl_element())
-        argument_multiindices = (finat_elem.get_indices(), )
-        argument_multiindex, = argument_multiindices
-        value_shape = finat_elem.value_shape
+        argument_multiindex = finat_elem.get_indices()
+        argument_multiindices = (argument_multiindex, )
+        coefficient = False
+
+    # Map into reference values
+    domain = extract_unique_domain(expression)
+    expression = apply_mapping(expression, arg.ufl_element(), domain)
+    value_shape = expression.ufl_shape
+
+    # Get indices for the output tensor
+    if coefficient:
+        tensor_indices = tuple(gem.Index() for s in value_shape)
+    else:
         if value_shape:
             tensor_indices = argument_multiindex[-len(value_shape):]
         else:
             tensor_indices = ()
-        coefficient = False
 
     # Replace coordinates (if any)
     builder = firedrake_interface.KernelBuilderBase(scalar_type=ScalarType)
-    domain = extract_unique_domain(expression)
     builder._domain_integral_type_map = {domain: "cell"}
     # Translate to GEM
     cell = domain.ufl_cell()
@@ -151,7 +159,7 @@ def compile_element(expression, dual_space=None, parameters=None,
     context = tsfc.fem.GemPointContext(**config)
 
     # Abs-simplification
-    expression = tsfc.ufl_utils.simplify_abs(expression, complex_mode)
+    expression = simplify_abs(expression, complex_mode)
 
     # Translate UFL -> GEM
     if coefficient:
@@ -177,7 +185,6 @@ def compile_element(expression, dual_space=None, parameters=None,
         return_variable = gem.Indexed(gem.Variable('R', finat_elem.index_shape), argument_multiindex)
         result = gem.Indexed(result, tensor_indices)
         if dual_space:
-            value_shape = dual_space.value_shape
             if value_shape:
                 var = gem.Indexed(gem.Variable("b", value_shape), tensor_indices)
                 b_arg = [lp.GlobalArg("b", dtype=ScalarType, shape=value_shape)]
@@ -219,7 +226,7 @@ def prolong_kernel(expression):
         assert hierarchy._meshes[int(idx)].cell_set._extruded
     V = expression.function_space()
     key = (("prolong",)
-           + V.value_shape
+           + (V.block_size,)
            + entity_dofs_key(V.finat_element.complex.get_topology())
            + entity_dofs_key(V.finat_element.entity_dofs())
            + entity_dofs_key(coordinates.function_space().finat_element.entity_dofs()))
@@ -283,7 +290,7 @@ def prolong_kernel(expression):
                "evaluate": eval_code,
                "spacedim": element.cell.get_spatial_dimension(),
                "ncandidate": hierarchy.fine_to_coarse_cells[levelf].shape[1],
-               "Rdim": V.value_size,
+               "Rdim": V.block_size,
                "inside_cell": inside_check(element.cell, eps=1e-8, X="Xref"),
                "celldist_l1_c_expr": celldist_l1_c_expr(element.cell, X="Xref"),
                "Xc_cell_inc": coords_element.space_dimension(),
@@ -301,7 +308,7 @@ def restrict_kernel(Vf, Vc):
     if Vf.extruded:
         assert Vc.extruded
     key = (("restrict",)
-           + Vf.value_shape
+           + (Vf.block_size,)
            + entity_dofs_key(Vf.finat_element.complex.get_topology())
            + entity_dofs_key(Vc.finat_element.complex.get_topology())
            + entity_dofs_key(Vf.finat_element.entity_dofs())
@@ -389,7 +396,7 @@ def inject_kernel(Vf, Vc):
     else:
         level_ratio = 1
     key = (("inject", level_ratio)
-           + Vf.value_shape
+           + (Vf.block_size,)
            + entity_dofs_key(Vc.finat_element.complex.get_topology())
            + entity_dofs_key(Vf.finat_element.complex.get_topology())
            + entity_dofs_key(Vc.finat_element.entity_dofs())
@@ -464,7 +471,7 @@ def inject_kernel(Vf, Vc):
             "celldist_l1_c_expr": celldist_l1_c_expr(Vc.finat_element.cell, X="Xref"),
             "tdim": Vc.mesh().topological_dimension(),
             "ncandidate": ncandidate,
-            "Rdim": numpy.prod(Vf.value_shape),
+            "Rdim": Vf.block_size,
             "Xf_cell_inc": coords_element.space_dimension(),
             "f_cell_inc": Vf_element.space_dimension()
         }
@@ -608,7 +615,8 @@ def dg_injection_kernel(Vf, Vc, ncell):
     # Coarse basis function evaluated at fine quadrature points
     phi_c = fem.fiat_to_ufl(Vce.point_evaluation(0, X_a, (Vce.cell.get_dimension(), 0)), 0)
 
-    tensor_indices = tuple(gem.Index(extent=d) for d in f.ufl_shape)
+    index_shape = f.ufl_element().reference_value_shape
+    tensor_indices = tuple(gem.Index(extent=d) for d in index_shape)
 
     phi_c = gem.Indexed(phi_c, argument_multiindex + tensor_indices)
     fexpr = gem.Indexed(fexpr, tensor_indices)