michelangelo@hp:~/Ongoing.Projects/NGVU2/mytests/xdmf$ python3 fun_in_xdmf.py 
Traceback (most recent call last):
  File "/home/michelangelo/Ongoing.Projects/NGVU2/mytests/xdmf/fun_in_xdmf.py", line 65, in <module>
    xdmf.write_function(uh * 5.0, 0)
  File "/usr/lib/petsc/lib/python3/dist-packages/dolfinx/io/utils.py", line 235, in write_function
    super().write_function(getattr(u, "_cpp_object", u), t, mesh_xpath)
TypeError: write_function(): incompatible function arguments. The following argument types are supported:
    1. (self: dolfinx.cpp.io.XDMFFile, u: dolfinx.cpp.fem.Function_float32, t: float, mesh_xpath: str = "/Xdmf/Domain/Grid[@GridType='Uniform'][1]") -> None
    2. (self: dolfinx.cpp.io.XDMFFile, u: dolfinx.cpp.fem.Function_float64, t: float, mesh_xpath: str = "/Xdmf/Domain/Grid[@GridType='Uniform'][1]") -> None
    3. (self: dolfinx.cpp.io.XDMFFile, u: dolfinx.cpp.fem.Function_complex64, t: float, mesh_xpath: str = "/Xdmf/Domain/Grid[@GridType='Uniform'][1]") -> None
    4. (self: dolfinx.cpp.io.XDMFFile, u: dolfinx.cpp.fem.Function_complex128, t: float, mesh_xpath: str = "/Xdmf/Domain/Grid[@GridType='Uniform'][1]") -> None

Invoked with: <dolfinx.io.utils.XDMFFile object at 0x7f23ea04c3b0>, Product(FloatValue(5.0), Coefficient(FunctionSpace(Mesh(blocked element (Basix element (P, quadrilateral, 1, gll_warped, unset, False), (2,)), 0), Basix element (P, quadrilateral, 1, gll_warped, unset, False)), 1)), 0, "/Xdmf/Domain/Grid[@GridType='Uniform'][1]"

from mpi4py import MPI
from petsc4py import PETSc
from petsc4py.PETSc import ScalarType

import ufl
import dolfinx
import numpy as np
import pyvista

comm = MPI.COMM_WORLD
pyvista.OFF_SCREEN = False
plot_sol = True


# Define a Python function that returns the constant value
class ConstantFun:
    def __init__(self):
        self.myval = 0

    def __eval__(self, x):
        self.values = np.full(x.shape[1], self.myval)  # Assuming x is a 2D array where x.shape[1] gives the number of points
        return self.values

# Create the domain (mesh) and its connectivity
nx, ny = 8, 8
cell_type = dolfinx.mesh.CellType.quadrilateral
domain = dolfinx.mesh.create_unit_square(comm, nx, ny, cell_type)
tdim = domain.topology.dim
fdim = tdim - 1
domain.topology.create_connectivity(fdim, tdim)

# FunctionSpace, Trial, and Test Functions
V = dolfinx.fem.FunctionSpace(domain, ("Lagrange", 1))
u = ufl.TrialFunction(V)
v = ufl.TestFunction(V)

# Dirichlet Boundary Conditions
uh = dolfinx.fem.Function(V)
# uh.interpolate(lambda x: 1 + x[0]**2 + 2 * x[1]**2)
const_expr = ConstantFun()
const_expr.myval = 40.0
uh.interpolate(const_expr.__eval__)  # Interpolate using the Python function

# Create Boundary conditions
boundary_facets = dolfinx.mesh.exterior_facet_indices(domain.topology)
boundary_dofs = dolfinx.fem.locate_dofs_topological(V, fdim, boundary_facets)
bc = dolfinx.fem.dirichletbc(uh, boundary_dofs)

# Source term
f = dolfinx.fem.Constant(domain, dolfinx.default_scalar_type(-6))

# Bilinear and Linear forms
a = ufl.dot(ufl.grad(u), ufl.grad(v)) * ufl.dx
L = f * v * ufl.dx

from dolfinx.fem.petsc import LinearProblem
problem = LinearProblem(
    a, L, bcs=[bc],
    petsc_options={"ksp_type": "preonly", "pc_type": "lu"}
    )
uh = problem.solve()

xdmf = dolfinx.io.XDMFFile(MPI.COMM_WORLD, 'uh.xdmf', "w")
xdmf.write_mesh(domain)
xdmf.write_function(uh * 5.0, 0)

print(pyvista.global_theme.jupyter_backend)

if plot_sol is True:
    u_topology, u_cell_types, u_geometry = dolfinx.plot.vtk_mesh(V)

    u_grid = pyvista.UnstructuredGrid(u_topology, u_cell_types, u_geometry)
    u_grid.point_data["u"] = uh.x.array.real
    u_grid.set_active_scalars("u")
    u_plotter = pyvista.Plotter()
    u_plotter.add_mesh(u_grid, show_edges=True)
    u_plotter.view_xy()
    if not pyvista.OFF_SCREEN:
        u_plotter.show()