Is making assign only work for weighted sums too restrictive?

[connorjward]

I have recently encountered two occasions where our new policy of making assign only work for weighted sums seems a bit clunky:

• The line p += assemble(dt * inner(nabla_grad(v), nabla_grad(phi))*dx) / assemble(v*dx) for the linear wave equation demo now fails without interpolate (@JDBetteridge found this).
• @nbouziani wanted to run A.assign(ufl.conj(B)) which similarly fails without interpolate.

With the way I have rewritten assign to use numpy operations, it would be quite easy to modify the tree visitor to add support for these operations (division and conj) since numpy has analogues for them. In general assign would then work for any expression where all the 'bits' are in the same function space (and be fast since we'd be using numpy). Is this worth doing? I confess I never quite understood the "assign used to let you do things that weren't mathematically legal" argument so please tell me if this is a bad idea.

I have another related question. AIUI assign is pretty much "assemble some expression into a Function and then use it to modify an existing Function". In other words I think the following are equivalent:

• u = Function(V).assign(expr) and u = assemble(expr)
• u += expr and u.dat.data += assemble(expr).dat.data.

With that in mind does it therefore make sense for the assign code I have written to actually belong in assemble.py?

[dham]

The first one of these is a solve masquerading as an assign as a hack. Currently that would definitely break the adjoint. Arguably we should find a way to specify that the matrix in a solver is diagonal.

The second one is something we didn't think of because we didn't think of the complex case. real, imag and conj should be allowed. Currently we don't know how to differentiate those anyway so that doesn't break anything.

The last point about assemble is bang on. I think assemble actually means "do the numerical computation described by this symbolic expression" and assign falls exactly into that.

[connorjward]

If the adjoint for assign in these cases is undefined then how come wrapping the expression in an interpolate fixes things? This makes me think that assign is just (faster) syntactic sugar for interpolate.

Could we not somehow use the same adjoint code path for both assign and interpolate?

[dham]

assign is a special case of interpolate, so in some sense this is just a special case of interpolate. However, adjoint interpolation currently always involves assembling a matrix. Conversely, the divide by diagonal inverse is used in explicit timestepping, which is usually associated with short, fast time steps. This is therefore likely to be a performance-critical operation. We therefore need to be careful about deciding things are "just" sugar.

[connorjward]

Could we not do something like the following:

    def _ad_annotate_assign(assign):
        @wraps(assign)
        def wrapper(self, expr, *args, **kwargs):
            if is_weighted_sum_of_coefficients(expr):
                # use the current adjoint assign approach
            else:
                # fall back to using the same approach as interpolate
                return annotate_interpolate(self, expr, *args, **kwargs)
        return wrapper

Then we can use the "fast path" approach for assign in the forwards direction in all cases but also ensure a valid adjoint.

[dham]

Yes, I guess that's OK. We can sort out the solve case subsequently when someone complains about the cost.

[connorjward]

OK. I'll create an issue and tackle this at some point. Thanks for your help!