Error taking derivatives with mixed function space

[danshapero]

When I run the code below, I get the error Expecting scalar coefficient in this branch. I believe this is a genuine bug but can't rule out user error. This is the same problem as in #1723 so I tried Ivan's trick of passing the form through ufl.algorithms.expand_derivatives first, but it didn't improve matters.

import firedrake
from firedrake import sqrt, inner, max_value, Constant, sym, grad, tr, dx, ds

nx, ny = 32, 32
Lx, Ly = 20.0, 20.0
mesh = firedrake.RectangleMesh(nx, ny, Lx, Ly, diagonal="crossed")
x = firedrake.SpatialCoordinate(mesh)

Q = firedrake.FunctionSpace(mesh, "CG", 1)
V = firedrake.VectorFunctionSpace(mesh, "CG", 1)
Z = Q * V

y = Constant((Lx / 2, Ly / 2))
R = Constant(Lx / 4)
r = sqrt(inner(x - y, x - y))

b_0 = Constant(0.1)
expr = b_0 / (1 + ((r - R) / (0.5 * R))**2)
b = firedrake.interpolate(expr, Q)

def ε(u):
    return sym(grad(u))

g = Constant(9.8)
C = Constant(0.01)
a_0 = Constant(1e-3)
ρ = Constant(1420)
μ = Constant(2)

a = a_0 * firedrake.exp(-(r / R)**2)

def equation(z):
    Z = z.function_space()
    ϕ, v = firedrake.TestFunctions(Z)
    h, u = firedrake.split(z)

    viscosity = h * μ * (inner(ε(u), ε(v)) + tr(ε(u)) * tr(ε(v))) * dx
    s = h + b
    gravity = -ρ * g * h * inner(grad(s), v) * dx
    friction = C * inner(u, v) * dx
    momentum_eq = viscosity + friction - gravity

    cell_fluxes = -h * inner(u, grad(ϕ)) * dx
    n = firedrake.FacetNormal(Z.mesh())
    u_n = max_value(0, inner(u, n))
    boundary_fluxes = h * u_n * ϕ * ds
    sources = a * ϕ * dx
    mass_eq = sources - boundary_fluxes - cell_fluxes

    return momentum_eq + mass_eq

def conserved_variables(z):
    h, u = firedrake.split(z)
    return firedrake.as_vector((h, 0, 0))

z = firedrake.Function(Z)
F = equation(z)
z_n = z.copy(deepcopy=True)
dF = firedrake.derivative(F, z, z_n - z)

w = firedrake.TestFunction(Z)
Q = inner(conserved_variables(z), w) * dx
dQ = firedrake.derivative(Q, z, z_n - z)

dt = firedrake.Constant(0.1)
form = dQ - dt * dF - dt * F
problem = firedrake.NonlinearVariationalProblem(form, z_n)

Changing momentum_eq = friction without the viscosity or gravity terms makes the error go away, but I'm not sure how much that helps diagnose the problem.

[wence-]

Although you can put "anything" of the right shape into the third slot of derivative, it seems like the differentiation handler for GateauxDerivative doesn't actually handle all cases correctly. Basically, it's only expecting a few things.

Workaround:

tz = TrialFunction(Z)
dF = firedrake.derivative(F, z, tz)
...
dQ = firedrake.deriative(Q, z, tz)
form = firedrake.replace(dQ - dt * DF - dt * F, {tz: z_n - z})

Which at least doesn't produce an error.

[wence-]

dF = firedrake.derivative(F, z, tz)
...
dQ = firedrake.deriative(Q, z, tz)
form = firedrake.replace(dQ - dt * DF - dt * F, {tz: z_n - z})

Better:

dF = firedrake.derivative(F, z)
dQ = firedrake.derivative(Q, z)
form = action(dQ - dt * dF - dt * F, z_n - z)

[danshapero]

Both of those fixes work, thanks for the suggestions Lawrence!