How to Interpolate Special Functions

Hi,

I have a similar question to this particular post. I’m wondering how to interpolate special functions into dolfinx.
Unlike other posts I’ve seen however I’m wondering how to make them dependent on the values of other fem.Function objects.

Let me show a motivating example. For ease of understanding I’ve copied most of the diffusion example and implemented a trivial function as a test case.

import scipy.constants as con
import ufl
from dolfinx import fem, io, mesh
from dolfinx.fem.petsc import (
    apply_lifting,
    assemble_matrix,
    assemble_vector,
    create_vector,
    set_bc,
)
from mpi4py import MPI
from petsc4py import PETSc

# Define temporal parameters
t = 0  # Start time
T = 1.0  # Final time
num_steps = 50
dt = T / num_steps  # time step size

# Define mesh
nx, ny = 50, 50
domain = mesh.create_rectangle(
    MPI.COMM_WORLD,
    [np.array([-2, -2]), np.array([2, 2])],
    [nx, ny],
    mesh.CellType.triangle,
)
V = fem.functionspace(domain, ('Lagrange', 1))


def ElectronHeatCapacity(T, Ne):
    """Imagine a more complex funciton here. In reality there are fermi-dirac integrals ect.
    For a specific example look at npplus fermi code.
    """

    return 3 / 2 * con.k


def initial_condition(x, a=5):
    return np.exp(-a * (x[0] ** 2 + x[1] ** 2))


u_n = fem.Function(V)
u_n.name = 'u_n'
u_n.interpolate(initial_condition)
cp = fem.Function(V)
cp.name = 'cp'

ne = fem.Function(V)
ne.name = 'ne'
ne.interpolate(lambda x: np.ones_like(x[1]) * 4e21)

# what I'm currently doing.
cp.x.array[:] = ElectronHeatCapacity(uh.x.array, ne.x.array)


# Create boundary condition
fdim = domain.topology.dim - 1
boundary_facets = mesh.locate_entities_boundary(
    domain, fdim, lambda x: np.full(x.shape[1], True, dtype=bool)
)
bc = fem.dirichletbc(
    PETSc.ScalarType(0), fem.locate_dofs_topological(V, fdim, boundary_facets), V
)

xdmf = io.XDMFFile(domain.comm, 'diffusion.xdmf', 'w')
xdmf.write_mesh(domain)

# Define solution variable, and interpolate initial solution for visualization in Paraview
uh = fem.Function(V)
uh.name = 'uh'
uh.interpolate(initial_condition)
xdmf.write_function(uh, t)


u, v = ufl.TrialFunction(V), ufl.TestFunction(V)
f = fem.Constant(domain, PETSc.ScalarType(0))
a = cp * u * v * ufl.dx + dt * ufl.dot(ufl.grad(u), ufl.grad(v)) * ufl.dx
L = (u_n + dt * f) * v * ufl.dx


bilinear_form = fem.form(a)
linear_form = fem.form(L)

A = assemble_matrix(bilinear_form, bcs=[bc])
A.assemble()
b = create_vector(linear_form)

solver = PETSc.KSP().create(domain.comm)
solver.setOperators(A)
solver.setType(PETSc.KSP.Type.PREONLY)
solver.getPC().setType(PETSc.PC.Type.LU)

for i in range(num_steps):
    t += dt

    # Update the right hand side reusing the initial vector
    with b.localForm() as loc_b:
        loc_b.set(0)
    assemble_vector(b, linear_form)

    # Apply Dirichlet boundary condition to the vector
    apply_lifting(b, [bilinear_form], [[bc]])
    b.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
    set_bc(b, [bc])

    # Solve linear problem
    solver.solve(b, uh.vector)
    uh.x.scatter_forward()

    ## update the heat capacity
    cp.x.array[:] = ElectronHeatCapacity(uh.x.array, ne.x.array)

    # Update solution at previous time step (u_n)
    u_n.x.array[:] = uh.x.array

    # Write solution to file
    xdmf.write_function(uh, t)
    # Update plot
xdmf.close()

The only real change I’ve implemented here is the Cp specific heat fem.Function. The cp depends on the Ne the electron density and the current temperature. I have a function I call that gets that takes in those variables and return the specific heat. (Right now that function is trivial but imagine I’m calling some fermi-dirac integral solver function)

My current approach is to just directly operate of the DOF values. I’m curious if this I the best approach? Would it scale properly when using multiple MPI processes for example? Also are there any advantages to using quadrature elements for the internal values instead?

Thanks for any and all suggestions.

With your current example, it is unclear to me why you couldn’t just use dolfinx.fem.Expression to define the expression you would like to interpolate, similar to: dolfinx/python/demo/demo_elasticity.py at 7c1812ead4ab7e88c1e848e12cfa79c77de56acb · FEniCS/dolfinx · GitHub

As your example does not include the

it is quite hard for me to imagine how to wrap a general interface around it.

There are three ways of doing interpolation in DOLFINx:

  1. Interpolate a dolfinx.fem.Function into another dolfinx.fem.Function
  2. Use a dolfinx.fem.Expression to define a ufl-expression that should be interpolated into a function
  3. Work directly on the spatial coordinates (physical coordinates) of the mesh sending in a lambda function taking in a 2 dimensional numpy array (x) and returning a function of those coordinates.

Yeah let me explain the function I’m using a bit more in actuality then.
I’ve added it below. As a quick explanation this is calculating the heat capacity of electrons in a plasma. It depends on the complete fermi integral ( see the function fd12).

Here is an example of what I need to solve.
F_j(u)=\int_0^\infty \frac{x^j}{e^{(x-u)}+1}dx

The npplus functions fd12 and such work similar to the special functions in Scipy. You input a numpy array and output a solution in a numpy array of the same shape.

I’m thinking option 3 is the one that makes the most sense, but I’m not sure how to implement it when I need values of two other fem.Function

def h32(x):
    return 4 / (3 * np.sqrt(np.pi)) * fd32(x)


def h12(x):
    return 2 / np.sqrt(np.pi) * fd12(x)


def hm12(x):
    return 1 / np.sqrt(np.pi) * fdm12(x)


def fermi_energy(ne: NDArray) -> NDArray:
    """
    Calculate the Fermi energy in eV for a given electron density
    ne = electron density in cm^-3

    Returns the Fermi energy in eV
    """
    ne_m = ne * 1e6
    return (
        (con.h * con.c) ** 2
        / (8 * con.m_e * con.c**2)
        * (3 / np.pi) ** (2 / 3)
        * ne_m ** (2 / 3)
        / con.eV
    )


def muF(T: NDArray[np.float64], Ne: NDArray[np.float64]) -> NDArray:
    """
    Calculate the chemical potential in eV for a given temperature and electron density
    Ne = electron density in cm^-3

    T = temperature in K

    Returns the chemical potential in eV
    """

    ef: NDArray[np.float64] = fermi_energy(ne=Ne)
    ret_val: NDArray[np.float64] = ef * (
        1
        - 1 / 3 * (np.pi * K_boltzman_ev * T / (2 * ef)) ** 2
        - np.pi**2 / 80 * (K_boltzman_ev * T / (ef)) ** 4
    )
    return ret_val


def beta(T):
    return 1 / (K_boltzman_ev * T)

def cp_thermal_physics(
    T: ArrayLike | float, Ne: ArrayLike | float
) -> ArrayLike | float:
    """
    Outputs electron heat capacity in J/mol

    Parameters
    ----------
    ne: number of electrons per cm^-3
    T: Temperature in deg K

    Returns
    ----------
    Cp: Electron Heat Capacity in J/mol-K
    """
    beta_ev = beta(T)
    mu_calc = muF(T, Ne=Ne)
    x = mu_calc * beta_ev
    lmd = np.exp(x)
    mask = np.isclose(lmd, 0)
    notmask = np.logical_not(mask)
    cp = np.zeros(shape=x.shape, dtype=float)
    rh32 = h32(x[notmask])
    rh12 = h12(x[notmask])
    rhm12 = hm12(x[notmask])

    cp[notmask] = 3 / 2 * K_boltzman_ev * (5 / 2 * rh32 / rh12 - 3 / 2 * rh12 / rhm12)
    cp[mask] = 3 / 2 * K_boltzman_ev

    return cp * con.eV * con.Avogadro

Is u here a function from dolfin? How does it depend on x?
x in your integral seems one dimensional, while u(x) would have to take in a 2D mesh.

X is just an integration variable (not having to do anything with like positions of the mesh or anything) sorry for not making that clear.

In this case we can take U as the value of a dolfin function at any given point. The idea is to take the U from the dolfin function run it through the function and then output a new transformed dolfin function that is used in further calculations.

I think you want to use @Andrash’s ExternalOperator for this.
See for instance: https://www.youtube.com/watch?v=y8goeapqfsw&ab_channel=FEniCSProject
and
GitHub - a-latyshev/dolfinx-external-operator: Extension of DOLFINx implementing the concept of external operator