Integrating FEniCS with PyMC3

Questions v3
1

@ivan I’ve read your article on including PDEs in PyMC3 (Including partial differential equations in your PyMC3 model) and have attempted to follow it with a more complicated example modeling wildfire spread. When I put it all together the sampler fails to sample at all, being stuck at 0 samples indefinitely. I know there are a lot of reasons this could be, but I’ve spent hours trying to figure this out so I’m hoping people can shed light on individual aspects of my code to help me understand what I’m doing right and wrong. Can you (or anyone reading this) help me understand if my integration with FEniCS is done correctly? Or if you have any other suggestions for why this isn’t working? I’m trying to recover the best parameters to be used in the PDE to get the best IoU metrics with data assimilation shown in the code below.

Here is my PDE function:

def solve_pde(k, A, B, C, Cs):
    # Define parameters.
    a,b = 0,10
    x1,y1 = a,a
    x2,y2 = b,b
    Nx = 200
    Ny = 200
    tf = 3
    Nt = tf*10
    t,dt = np.linspace(0,tf,Nt+1,retstep=True)
    dt = fa.Constant(dt)
    Ta = fa.Constant(20.) # Ambient Temperature.
    v = fenics.as_vector([fa.Constant(0.), fa.Constant(0.)]) # Wind vector.

    # Define function space.
    mesh = fa.RectangleMesh(fenics.Point(x1, y1), fenics.Point(x2, y2), Nx, Ny, 'crossed')
    V = fenics.FunctionSpace(mesh, 'CG', 1)

    # Define initial values.
    T_0 = fenics.Expression('1200/cosh(pow(x[0]-b,2)+pow(x[1]-b,2))+Ta', degree=2, b=4, Ta=Ta)
    T_n = fa.interpolate(T_0, V)
    funcstring = "(pow(x[0]-5,2)+pow(x[1]-5,2) <= 1.0)? 1.0 : 0.0;" # Fuel Circle
    S_0 = Expression(funcstring, degree=2)
    S_n = fa.interpolate(S_0, V)
    
    # Define Trial and Test Functions.
    T_n1 = fenics.TrialFunction(V)
    S_n1 = fenics.TrialFunction(V)
    q = fenics.TestFunction(V)
    p = fenics.TestFunction(V)

    # Define Boundary Conditions: BC = dot(grad(T_n),N) where N is the unit normal vector.  
    BC = fa.Expression("0.", degree=1) # Pure Neumann BCs.

    # Define variational forms.
    T_diff = T_n-Ta
    T_diff = T_diff if T_diff!=0. else T_diff+.1
    linearized = exp(-B/T_diff) + B*exp(-B/T_diff)/T_diff**2*dt # Adding .1 to avoid dividing by zero
    F_T = (T_n1-T_n)*q*dx \
            + dt*(k*dot(grad(T_n1),grad(q)) + dot(v,grad(T_n1))*q - A*(S_n*linearized*q - C*(T_n1-Ta)*q))*dx \
            + (k*BC*q)*ds
    F_S = (S_n1-S_n)*p*dx + dt*Cs*S_n*linearized*p*dx

    # Separate into bilinear and linear form.
    a_T, L_T = lhs(F_T), rhs(F_T)
    a_S, L_S = lhs(F_S), rhs(F_S)
    
    # Step in time.
    W = VectorFunctionSpace(mesh, 'CG', 1, dim=11)
    w = Function(W) # Function to store all perimeters. 
    j = 0
    T = fa.Function(V)
    S = fa.Function(V)
    for i in range(len(t)):  
        # Solve for T, update solution.
        fa.solve(a_T==L_T, T)
        T_n.assign(T)
        # Solve for S, update solution
        fa.solve(a_S==L_S, S)
        S_n.assign(S)

        # Save every 3rd result.
        skip = 3
        if i%skip==0:
            wa = w.vector().get_local()
            newT = T_n.vector().get_local()
            wa[j*len(newT):(j+1)*len(newT)] = newT
            w.vector().set_local(wa)
            j+=1
    return w

From here, I do the following:

template = (fa.Constant(0.0), fa.Constant(0.0), fa.Constant(0.0), fa.Constant(0.0), fa.Constant(0.0))
theano_fem_solver = create_fem_theano_op(template)(solve_pde)

# Get solution after perturbing parameters.
k_noise = fa.Constant(1.0 + np.random.normal(scale=.2))
A_noise = fa.Constant(10000. + np.random.normal(scale=1000))
B_noise = fa.Constant(300. + np.random.normal(scale=30))
C_noise = fa.Constant((1 + np.random.normal(scale=.2))/10000)
Cs_noise = fa.Constant(3. + np.random.normal(scale=.5))
Ts_noise = solve_pde(k_noise, A_noise, B_noise, C_noise, Cs_noise)
Ts_observed = to_numpy(Ts_noise).reshape(11,-1)

My custom likelihood is:

def tt_log_beta_pdf(x,a,b):
    log_beta = tt.gammaln(a) + tt.gammaln(b) - tt.gammaln(a+b)
    return (a-1)*tt.log(x) + (b-1)*tt.log(1-x) - log_beta

def wildfire_likelihood(T_preds, T_obs, burn=600):
    '''
    Calculates IOU score between observed and predicted 
    fire perimeters and maps it to a likelihood.  
    '''
    n_per, n_data = T_obs.shape
    burn_pred = T_preds >= burn
    burn_obs = T_obs >= burn
    logps = tt.zeros(n_per)
    for p in range(n_per):
        pred = burn_pred[p*n_data:(p+1)*n_data]
        obs = burn_obs[p].flatten()
        
        intersection = pred*obs
        union = pred + obs
        if tt.gt(union.sum(),0):
            if tt.gt(intersection.sum(),0):
                IoU = intersection.sum()/union.sum()
            else:
                IoU = 0.001
        else:
            IoU = 0.999
        logp = tt_log_beta_pdf(IoU,5,IoU) 
        logps = tt.set_subtensor(logps[p], logp)
    return logps.sum()

Followed by

with pm.Model() as model:
    # Set prior distributions.
    k = pm.Lognormal("k", mu=np.log(1), sigma=4, shape=(1,)) # Collision frequency
    A = pm.Lognormal("A", mu=np.log(10000.), sigma=1000, shape=(1,)) # Pre-exponential factor (rate)
    B = pm.Lognormal("B", mu=np.log(300.), sigma=100, shape=(1,)) # Proportionality Coefficient/Activation Energy
    C = pm.Lognormal("C", mu=np.log(1/10000), sigma=1/10000, shape=(1,)) # Scaled Heat Transfer Coefficient
    Cs = pm.Lognormal("Cs", mu=np.log(3), sigma=3, shape=(1,)) # Fuel Relative Disappearance Rate
    
    # Run the model.
    T_preds = theano_fem_solver(k, A, B, C, Cs)
    
    # Calculate the likelihood.
    likelihood = pm.Potential('likelihood', wildfire_likelihood(T_preds, Ts_observed))
    
    # MCMC:
    trace = pm.sample(tune=5, draws=5)

Any help is greatly appreciated!

2

cc @aseyboldt