BUG: Cofunction.sub() assignment of a base form is not working

[Ig-dolci]

Steps to reproduce:

mesh = UnitSquareMesh(2, 2)
BDM = FunctionSpace(mesh, "BDM", 1)
DG = FunctionSpace(mesh, "DG", 0)
W = BDM * DG
sigma, u = TrialFunctions(W)
tau, v = TestFunctions(W)
c = Cofunction(W.dual())
c.sub(1).assign(Constant(1.0) * v * dx)

Failure:

File "/Users/ddolci/tes_fire_install/firedrake/src/firedrake/firedrake/cofunction.py", line 204, in assign
    return self.assign(assembled_expr, subset=subset)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "petsc4py/PETSc/Log.pyx", line 188, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "petsc4py/PETSc/Log.pyx", line 189, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "<decorator-gen-28>", line 2, in assign
  File "/Users/ddolci/tes_fire_install/firedrake/src/firedrake/firedrake/utils.py", line 82, in wrapper
    return f(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^
  File "/Users/ddolci/tes_fire_install/firedrake/src/firedrake/firedrake/cofunction.py", line 203, in assign
    assembled_expr = firedrake.assemble(expr)
                     ^^^^^^^^^^^^^^^^^^^^^^^^
  File "petsc4py/PETSc/Log.pyx", line 188, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "petsc4py/PETSc/Log.pyx", line 189, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "/Users/ddolci/tes_fire_install/firedrake/src/firedrake/firedrake/adjoint_utils/assembly.py", line 28, in wrapper
    form = BaseFormAssembler.preprocess_base_form(form)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ddolci/tes_fire_install/firedrake/src/firedrake/firedrake/assemble.py", line 820, in preprocess_base_form
    expr = BaseFormAssembler.expand_derivatives_form(expr, form_compiler_parameters)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ddolci/tes_fire_install/firedrake/src/firedrake/firedrake/assemble.py", line 864, in expand_derivatives_form
    return ufl.algorithms.ad.expand_derivatives(form)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ddolci/tes_fire_install/firedrake/src/ufl/ufl/algorithms/ad.py", line 31, in expand_derivatives
    form = apply_algebra_lowering(form)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ddolci/tes_fire_install/firedrake/src/ufl/ufl/algorithms/apply_algebra_lowering.py", line 146, in apply_algebra_lowering
    return map_integrand_dags(LowerCompoundAlgebra(), expr)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ddolci/tes_fire_install/firedrake/src/ufl/ufl/algorithms/map_integrands.py", line 73, in map_integrand_dags
    return map_integrands(lambda expr: map_expr_dag(function, expr, compress),
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ddolci/tes_fire_install/firedrake/src/ufl/ufl/algorithms/map_integrands.py", line 66, in map_integrands
    return function(integrand)
           ^^^^^^^^^^^^^^^^^^^
  File "/Users/ddolci/tes_fire_install/firedrake/src/ufl/ufl/algorithms/map_integrands.py", line 73, in <lambda>
    return map_integrands(lambda expr: map_expr_dag(function, expr, compress),
                                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ddolci/tes_fire_install/firedrake/src/ufl/ufl/corealg/map_dag.py", line 34, in map_expr_dag
    result, = map_expr_dags(function, [expression], compress=compress,
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ddolci/tes_fire_install/firedrake/src/ufl/ufl/corealg/map_dag.py", line 91, in map_expr_dags
    for v in traversal(expression):
  File "/Users/ddolci/tes_fire_install/firedrake/src/ufl/ufl/corealg/traversal.py", line 87, in unique_post_traversal
    visited.add(expr)
  File "/Users/ddolci/tes_fire_install/firedrake/src/ufl/ufl/coefficient.py", line 149, in __hash__
    hash(self._ufl_function_space),
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ddolci/tes_fire_install/firedrake/src/firedrake/firedrake/functionspaceimpl.py", line 309, in __hash__
    return hash((self.mesh(), self.topological))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ddolci/tes_fire_install/firedrake/src/ufl/ufl/core/ufl_type.py", line 38, in __hash__
    return hash(self._ufl_hash_data_())
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ddolci/tes_fire_install/firedrake/src/FInAT/finat/ufl/finiteelementbase.py", line 87, in __hash__
    return hash(self._ufl_hash_data_())
                ^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ddolci/tes_fire_install/firedrake/src/FInAT/finat/ufl/finiteelementbase.py", line 79, in _ufl_hash_data_
    return repr(self)
           ^^^^^^^^^^
RecursionError: maximum recursion depth exceeded

[Ig-dolci]

The issue arises from the assemble of a form in a Mixed Function Space. Consider the code:

from firedrake import *

mesh = UnitSquareMesh(2, 2)

BDM = FunctionSpace(mesh, "BDM", 1)
DG = FunctionSpace(mesh, "DG", 0)

# Create a Mixed Function Space 
W = BDM * DG

# Define test and trial functions for DG
v0 = TestFunction(DG)
sigma, u = TrialFunctions(W)
tau, v = TestFunctions(W)

# Assemble the form
aux = assemble(Constant(1.0) * v * dx)

aux = assemble(Constant(1.0) * v * dx) is returning a Cofunction in the Mixed Function Space BDM x DG. Therefore, when we try:

c = Cofunction(W.dual())
c.sub(1).assign(aux)

This assignment fails because the function spaces of aux and c.sub(1) do not match, i.e., aux.function_space() != c.sub(1).function_space().

[Ig-dolci]

That is not a bug. I was just doing the math wrong assignment.