Compute complex argument of Function

Hello FENICsX users.

I’m trying to perform a fairly basic algebraic operation - obtain the argument of a complex scalar valued dolfinx Function. In order to do that, I tried to use an Expression as follows :

from mpi4py.MPI import COMM_WORLD as comm
from petsc4py import PETSc as pet
import dolfinx as dfx
import ufl

domain = dfx.mesh.create_unit_square(comm, 8, 8, dfx.mesh.CellType.quadrilateral)
FE = ufl.FiniteElement("CG",domain.ufl_cell(),1)
FS = dfx.fem.FunctionSpace(domain,FE)
u = dfx.fem.Function(FS)
u.interpolate(lambda x: x[0]+dfx.fem.Constant(domain, pet.ScalarType(1j))*x[1])
expr=dfx.fem.Expression(ufl.atan_2(ufl.imag(u),ufl.real(u)),
						FS.element.interpolation_points())

This fails with RuntimeError: Not supported in current scalar mode. I’m running this in complex mode of course inside a docker container. Does anyone have a workaround ?

I guess I’m going to answer my own question. It’s fairly straightforward to handle the arrays directly. Still unclear why the related UFL functions fail.

u,v = dfx.fem.Function(FS),dfx.fem.Function(FS)
u.interpolate(lambda x: x[0]+1j*x[1])
v.x.array[:]=np.arctan2(u.x.array.imag,u.x.array.real)
v.x.scatter_forward()

To me this seems related to the fact that atan is mapped to:
catan: https://github.com/FEniCS/ffcx/blob/c9831436449eeae702acf3855efc90a5327fd4d9/ffcx/codegeneration/C/ufl_to_cnodes.py#L89
while atan2 doesn’t have a complex counter-part:
atan2, atan2f, atan2l - cppreference.com