Translation of parametrized diffusion equation from Fenics to Fenicsx

Hello there,
I am translating some Fenics code from GitHub - MoGeist/diffusion_PPDE to fenicsx. In this code I try to solve the parametrized diffusion equation n times. This equation looks as follows in general: -\nabla (a(x,y) \nabla u) = f in [0,1]^2 and 0 on the boundary. Here a(x,y)= a_0 + sum_{j=1}^p y_j sin(x_j). In this case, x,y \in R^p and y_j are uniformly drawn from [0,1]. Thus a(x,y) is a sum of trigonometric functions, which is generated in the following way:

def trigonometric_basis(n, p, mu=1, sigma=1):
	'''Generate list of expressions and array of coefficients in polynomial basis (T1)

	n: 		number of samples
	p: 		number of basis fcts.
	mu: 	shift
	sigma: 	decay rate
	'''	

	num_terms = p
	print(n,p)
	coeff = np.random.uniform(0, 1., (n, num_terms))
	expr_list = []
	
	template_a = "+(1+{}*{}*(1+{}))"
	template_sin = "sin({}*3.14159*x[0])*sin({}*3.14159*x[1])"	

	for i in range(n):
		temp = str(mu)  
		for j in range(1,num_terms+1):
			sin = template_sin.format(floor((j+2)/2), ceil((j+2)/2))
			a = template_a.format(j**(sigma), str(coeff[i,j-1]), sin) 
			temp += a 
		expr_list.append(temp)
	print(coeff)
	return expr_list, coeff	

In fenics, this was now solvable with:

def solver(expr):
	'''Solver for parametric diffusion equation'''
	
	u = Function(V) 
	v = TestFunction(V)

	a = Expression(expr, degree=params['inter_deg']) 
	F = dot(a*grad(u), grad(v))*dx - f*v*dx

	solve(F == 0, u, bc)
	
	return np.array(u.vector().get_local())

Since a(x,y) is generated with a for loop, I am not sure how to translate this into the form mentioned on Equivalent for expression in dolfinx since here the functions are typed by hand. A minimal example of how this would look like is as follows:

import dolfinx as dnx
import ufl
import numpy as np
from petsc4py.PETSc import ScalarType
from dolfinx.fem.petsc import LinearProblem
from petsc4py import PETSc
import multiprocessing as mp
from mpi4py import MPI
from math import floor, ceil


default_parameters = {
    'num_samples': 1,
    'basis_fct': 'mixed_trigonometric',
    'mu': 1, # shift hyper-parameter 
    'mode': 'x', # only required for cookies    
    'sigma': 1, # only required for trigonometric polynomials
    'num_basis': 2, # actual number of basis functions may vary
    'finite_element': 'P', # see http://femtable.org/
    'inter_deg': 1, # degree of interpolation
    'mesh_size': 10, # number of edges in the mesh (number of nodes: edges + 1)
    'right_hand_side': '20+10*x[0]-5*x[1]', 
    'folder_name': 'some_dataset', # folder where data is saved
    'part': 1,
    'part_max': 1 
}

# create mesh and function space
msh = dnx.mesh.create_unit_square(MPI.COMM_WORLD, 10, 10)
x = ufl.SpatialCoordinate(msh)
f = 20+10*x[0]-5*x[1]
V = dnx.fem.functionspace(msh, ('P', 1))

facets = dnx.mesh.locate_entities_boundary( msh,
    dim=1,
    marker=lambda x: np.isclose(x[0], 0.0)
    | np.isclose(x[0], 1.0)
    | np.isclose(x[1], 0.0)
    | np.isclose(x[1], 1.0),
)
dofs = dnx.fem.locate_dofs_topological(V=V, entity_dim=1, entities=facets)
bc = dnx.fem.dirichletbc(value=ScalarType(0), dofs=dofs, V=V)

bcs = [bc]

def trigonometric_basis(n, p, mu=1, sigma=1):
    '''Generate list of expressions and array of coefficients in polynomial basis (T1)'''
    num_terms = p
    coeff = np.random.uniform(0, 1., (n, num_terms))
    expr_list = []
    
    template_a = "+(1+{}*{}*(1+{}))"
    template_sin = "sin({}*pi*x[0])*sin({}*pi*x[1])"
    
    for i in range(n):
        temp = str(mu)
        for j in range(1, num_terms+1):
            sin = template_sin.format(floor((j+2)/2), ceil((j+2)/2))
            a = template_a.format(j**(sigma), str(coeff[i, j-1]), sin)
            temp += a
        expr_list.append(temp)
    return expr_list, coeff

expr_list, coeff = trigonometric_basis(5,2,1)
for expr in expr_list:
    uh = dnx.fem.Function(V) 
    v = ufl.TestFunction(V)
    a(expr) = ...
    a_bilinear = a(expr)*ufl.dot(ufl.grad(u), ufl.grad(v)) * ufl.dx
    L = f * v * ufl.dx
    problem = LinearProblem(a_bilinear, L, bcs=[bc], petsc_options={"ksp_type": "preonly", "pc_type": "lu"})
    uh = problem.solve()

Could someone please show me how to write this? I am new in fenicsx and I can’t find any similar problems. It might be an easy question but I have no idea how to do this, any help would be greatly appreciated!

Best,
Daan

I won’t try to put my answer into your code, but hopefully it should be helpful enough for you to do it yourself.

You may want to create a python function

def generate_a(parameter):
    def _(x):
       return parameter[0] + parameter[1] * np.sin(x[0]) + parameter[2] * np.sin(x[1]) 

    return _

Say that V is a FE space suitable for representing the coefficient a, then you can interpolate it as

a_interp = dolfinx.fem.Function(V)
define a bilinear form that uses a_interp

parameter_values = [[1, 2, 3], [4, 5, 6]]
for p in parameter_values:  # for instance, the first iteration will consider the parameter values [1,2,3]
   a_interp.interpolate(generate_a(p))
   solve your problem that uses a_interp