Using Python functions in FEniCS variational formulation

Dear Pavan,

Thank you for for your shared interest in the question. The code runs to completion but the output is nonsensical. Something seems to be getting “stuck”. I cleaned up the code a bit and reposted it below with 1) the original formulation of the chemical potential using internal ufl operators and 2) using a new formulation of the chemical potential that uses an external operator; you just need to uncomment/comment the appropriate sections. If you run and visualize the output using the two approaches, you will see that the concentration field differs dramatically.

import os
import numpy as np

from mpi4py import MPI
from petsc4py import PETSc

import ufl
from basix.ufl import element, mixed_element, quadrature_element
from dolfinx import default_real_type, log, plot
from dolfinx.fem import Function, functionspace, form, assemble_vector, assemble_matrix
from dolfinx.fem.petsc import NonlinearProblem
from dolfinx.io import XDMFFile
from dolfinx.mesh import CellType, create_unit_interval
from dolfinx.nls.petsc import NewtonSolver
from dolfinx_external_operator import (
    FEMExternalOperator,
    evaluate_external_operators,
    evaluate_operands,
    replace_external_operators)

from ufl import grad, inner, derivative, Measure

lmbda = 1.0e-02  # surface parameter
dt = 1.0e-06  # time step
theta = 0.5  # time stepping family, e.g. theta=1 -> backward Euler, theta=0.5 -> Crank-Nicholson
N= 1000
msh = create_unit_interval(MPI.COMM_WORLD, N)

P1 = element("Lagrange", msh.basix_cell(), 1)

ME = functionspace(msh, mixed_element([P1, P1]))
SE = functionspace(msh, P1)

quadrature_degree = 1
QE = quadrature_element(msh.topology.cell_name(), degree=quadrature_degree, value_shape=())
Q = functionspace(msh, QE)
dx = Measure("dx", metadata={"quadrature_scheme": "default", "quadrature_degree": quadrature_degree})

me = ufl.TrialFunction(ME)
(c_hat, mu_hat) = ufl.TestFunctions(ME)

u = Function(ME)  # current solution
u0 = Function(ME)  # solution from previous converged step

rng = np.random.default_rng(42)
u.sub(0).interpolate(lambda x: 0.63 + 0.02 * (0.5 - rng.random(x.shape[1])))
u.x.scatter_forward()

(c, mu) = ufl.split(u)
(c0, mu0) = ufl.split(u0)

mu_mid = (1.0 - theta) * mu0 + theta * mu

################################################################################################
######## ORIGINAL IMPLEMENTATION OF THE CHEMICAL POTENTIAL USING UFL OPERATORS #################
################################################################################################
"""
mu_original = 200 * c * (1 - c)**2 - 200 * c**2 * (1 - c)
F0 = inner(c, c_hat) * dx - inner(c0, c_hat) * dx + dt * inner(grad(mu_mid), grad(c_hat)) * dx
F1 = inner(mu, mu_hat) * dx - inner(mu_original, mu_hat) * dx - lmbda * inner(grad(c), grad(mu_hat)) * dx
F = F0 + F1

problem = NonlinearProblem(F, u)


"""
################################################################################################
######## ATTEMPTED IMPLEMENTATION OF THE CHEMICAL POTENTIAL USING EXTERNAL OPERATOR ############
################################################################################################


mu_external = FEMExternalOperator(c, function_space=Q)

F0 = inner(c, c_hat) * dx - inner(c0, c_hat) * dx + dt * inner(grad(mu_mid), grad(c_hat)) * dx
F1 = inner(mu, mu_hat) * dx - inner(mu_external, mu_hat) * dx - lmbda * inner(grad(c), grad(mu_hat)) * dx
F = F0 + F1

def mu_external_impl(c):
    output = 200 * c * (1 - c)**2 - 200 * c**2 * (1 - c)
    return(output.reshape(-1))

def dmudc_external_impl(c):
    output = 200 * (1 - c)**2 - 800 * c * (1 - c) + 200 * c**2
    return(output.reshape(-1))

def mu_external_function(derivatives):
    if derivatives == (0,): 
        return(mu_external_impl)
    elif derivatives == (1,): 
        return(dmudc_external_impl)
    else:
        return(NotImplementedError)
    
mu_external.external_function = mu_external_function

J = derivative(F, u, me)
J_expanded = ufl.algorithms.expand_derivatives(J)
F_replaced, F_external_operators = replace_external_operators(F)
J_replaced, J_external_operators = replace_external_operators(J_expanded)
evaluated_operands = evaluate_operands(F_external_operators)
_ = evaluate_external_operators(F_external_operators, evaluated_operands)
_ = evaluate_external_operators(J_external_operators, evaluated_operands)

F_compiled = form(F_replaced)
J_compiled = form(J_replaced)

problem = NonlinearProblem(F_compiled, u, J=J_compiled)

################################################################################################
################################################################################################
################################################################################################

solver = NewtonSolver(MPI.COMM_WORLD, problem)
solver.convergence_criterion = "incremental"
solver.rtol = np.sqrt(np.finfo(default_real_type).eps) * 1e-2

ksp = solver.krylov_solver
opts = PETSc.Options()  # type: ignore
option_prefix = ksp.getOptionsPrefix()
opts[f"{option_prefix}ksp_type"] = "preonly"
opts[f"{option_prefix}pc_type"] = "lu"
sys = PETSc.Sys()  # type: ignore
# For factorisation prefer MUMPS, then superlu_dist, then default.
if sys.hasExternalPackage("mumps"):
    opts[f"{option_prefix}pc_factor_mat_solver_type"] = "mumps"
elif sys.hasExternalPackage("superlu_dist"):
    opts[f"{option_prefix}pc_factor_mat_solver_type"] = "superlu_dist"
ksp.setFromOptions()

file = XDMFFile(MPI.COMM_WORLD, "demo_ch/output.xdmf", "w")
file.write_mesh(msh)

# Step in time
t = 0.0
T = 100 * dt

u0.x.array[:] = u.x.array
while t < T:
    t += dt
    r = solver.solve(u)
    
    print(f"Step {int(t / dt)}: num iterations: {r[0]}")
    file.write_function(u.sub(0), t)
    
    u0.x.array[:] = u.x.array[:]

file.close()

Screenshots of the output running either option of the code.