Avoiding assembly in a non-linear mixed finite element reaction-diffusion problem

I am trying to solve a reaction-diffusion system between two chemicals with concentrations a and b. The diffusion co-efficient of b is much larger than that of a so it can be treated as spatially uniform i.e. b ∈ \mathbb{R}. There is also conservation of mass. The reaction-diffusion can be written as

\frac{∂a}{∂t} = D\nabla^2a + f(a, b)

The conservation of mass is an integral constraint.

b∫\mathrm{d}Ω + ∫a\,\mathrm{d}Ω = M,\ \mathrm{a\ constant}

f(a, b) is a non-linear function in a. So I linearized it using a^{k+1} = a^k + \delta a and b^{k+1} = b^k + \delta b as

f(a^{k+1}, b^{k+1}) ≈ f(a^k, b^k) + δa\,f_a(a^k, b^k) + δb\,f_b(a^k, b^k)

where f_a = ∂f/∂a and f_b = ∂f/∂b. I will use Newton’s method to solve for δa and δb at every time step.

With δa and δb as my trial functions and u and r as the corresponding test functions from the mixed P1×\mathbb{R} space I can write down the bilinear form for this problem as

ΔtD\int∇δa\cdot∇u\,\mathrm{d}x + \intδa\,u\,\mathrm{d}x - Δt∫f_a^k\,δa\,u\,\mathrm{d}x - Δt∫f_b^k\,δb\,u\,\mathrm{d}x - ∫(δa + δb)\,r\,\mathrm{d}x.

The linear form is

Δt∫f^k\,u\,\mathrm{d}x - Δt\,D∫∇a^k⋅∇u\,\mathrm{d}x - \int(a^k - a^n)u\,\mathrm{d}x - (α - b_k - a_k)r\,\mathrm{d}x.

Here a^n represents the solution at the previous time step. f^k, f_a^k and f_b^k are all evaluated at (a^k, b^k). I have denoted the total average concentration as \alpha = M/\int\mathrm{d}\Omega, a constant.

Now I am trying to avoid assembly as this is a time-dependent problem based on the discussion in Time-Dependent Problems. So I want to write the terms in the bilinear and linear forms using mass and stiffness matrices. For example, I can write

ΔtD\int∇δa\cdot∇u\,\mathrm{d}x = ΔtD\int\sum_j∇\hat{\phi}_i\cdot∇\phi_j\,δa_j\,\mathrm{d}x = ΔtDK_{ij}\delta a_j.

where δa = \sum_j \phi_jδa_j and u = \hat{\phi}_i where \phi_j represents the basis functions and K_{ij} = \int∇\hat{\phi}_i\cdot∇\phi_j is the stiffness matrix.

Question

How to write the following term in terms of the basis functions (mass matrix)?

Δt∫f_a^k\,δa\,u\,\mathrm{d}x

I have attached my working code for the linearized problem. I want to modify it so that I can avoid assembly in the time loop.

import numpy as np
import dolfin as df

mesh = df.RectangleMesh(df.Point(-10, -10), df.Point(10, 10), 101, 101)

P1 = df.FiniteElement('Lagrange', mesh.ufl_cell(), 1)
R = df.FiniteElement('Real', mesh.ufl_cell(), 0)
ME = df.FunctionSpace(mesh, P1*R)

AB = df.TrialFunction(ME)
UV = df.TestFunction(ME)
δa, δb = df.split(AB)
u, r = df.split(UV)

AnBn = df.Function(ME)
a_n, _ = df.split(AnBn)
AkBk = df.Function(ME)
a_k, b_k = df.split(AkBk)

# Some constants for my simulation
Δt = df.Constant(0.01)
D = df.Constant(0.5)
k = df.Constant(0.216)
K = df.Constant(3.81)
δ = df.Constant(4.45)
γ = df.Constant(3.57)
α = df.Constant(10.0)
vr = df.Constant(1.0)

# The reaction function and its derivatives wrt a, b
def f(a, b):
    return b*(k + γ*a**2/(K**2 + a**2)) - δ*a

def fa(a, b):
    return 2*a*b*γ*K**2/(K**2 + a**2)**2 - δ

def fb(a, b):
    return k + γ*a**2/(K**2 + a**2)

# The forms
bilinear = Δt*D*df.inner(df.grad(δa), df.grad(u))*df.dx + δa*u*df.dx -\
            Δt*fa(a_k, b_k)*δa*u*df.dx - Δt*fb(a_k, b_k)*δb*u*df.dx - (δa + δb)*r*df.dx

linear = Δt*f(a_k, b_k)*u*df.dx - Δt*D*df.inner(df.grad(a_k), df.grad(u))*df.dx -\
        (a_k - a_n)*u*df.dx - (α - b_k - a_k)*r*df.dx

AB = df.Function(ME)

# Newton's method solver parameters
ω = 1.0
ε = 1.0
tol = 1.0e-12
maxiter = 25

with df.XDMFFile('Test.xdmf') as out:
    out.parameters['rewrite_function_mesh'] = False
    out.parameters['functions_share_mesh'] = True
    
    # Apply initial condition
    cpp_code = """
    [&](){
        // r is the radial distance
        // s controls sharpness of the step function, choose based on mesh size
        // R is the radius of the circular base region

        double r = sqrt((x[0] - p)*(x[0] - p) + (x[1] - q)*(x[1] - q));
        if(r > 1.5*R){
            return am;
        }
        else{
            double exp1 = exp(R*s);
            double exp2 = exp(r*s);
            return (ap*exp1 + am*exp2)/(exp1 + exp2);
        }
    }()
    """
    ic = df.Expression(cpp_code, ap=5.0, am=0.5, p=2, q=3, R=2, s=5.0, degree=5)
    a0 = df.interpolate(ic, ME.sub(0).collapse())
    b0 = df.interpolate(df.Expression('9.354', degree=0), ME.sub(1).collapse())
    df.assign(AnBn, [a0, b0])
    df.assign(AkBk, [a0, b0])

    # Store initial state to file
    a0.rename('Concentration_A', 'a Function')
    out.write(a0, 0.0)

    # Start the time loop
    starttime = float(Δt)
    endtime = 2.00
    for t in np.arange(starttime, endtime + float(Δt), float(Δt)):
        print('time =', t)

        # The Newton's Method loop
        iternum = 0
        ε = 1.0
        while ε > tol and iternum < maxiter:
            iternum += 1
            Ax, bx = df.assemble_system(bilinear, linear)
            df.solve(Ax, AB.vector(), bx)
            ε = np.linalg.norm(AB.vector()[:], ord=np.Inf)
            print('\tnorm:', ε)
            AkBk.vector().axpy(ω, AB.vector())

        AnBn.assign(AkBk)

        # Store solution to file
        a_t = AnBn.split()[0]
        a_t.rename('Concentration_A', 'a Function')
        out.write(a_t, t)