How to plot isosurfaces of complex scalar fields?

I would like to plot isosurfaces of a complex-valued scalar field. This is commonly used to visualize atomic and molecular orbitals, where the isosurfaces are taken at some specified magnitudes of the underlying scalar field, and the phase is mapped to a colour using a cyclic colour map (e.g. like the one used in [ANN] DomainColoring(Toy).jl – Smooth Complex Plotting).

Here is an example of what it could look like, using Mayavi:

How could I accomplish this using (GL)Makie? I tried playing around with extending plot!(::Contour) to work with for complex volumes, but that plot type seems to assume a single colour per isosurface:

using GLMakie
import Makie: elconvert, float32type, el32convert, plot!, nan_extrema, to_levels
using Colors

float32type(x::AbstractArray{Complex{T}}) where T = Complex{float32type(T)}
el32convert(x::AbstractArray{<:Complex}) = elconvert(float32type(x), x)

function plot!(plot::Contour{<: Tuple{X, Y, Z, Vol}}) where {X, Y, Z, Vol<:AbstractArray{<:Complex,3}}
    @show eltype(Vol)
    x, y, z, volume = plot[1:4]
    @extract plot (colormap, levels, linewidth, alpha)
    avolume = lift(a -> abs.(a), volume)
    ϕvolume = lift(a -> angle.(a), volume)
    valuerange = lift(nan_extrema, plot, avolume)
    @show valuerange
    cliprange = replace_automatic!(()-> valuerange, plot, :colorrange)
    cmap = lift(plot, colormap, levels, alpha, cliprange, valuerange) do _cmap, l, alpha, cliprange, vrange
        levels = to_levels(l, vrange)
        nlevels = length(levels)
        N = 50 * nlevels

        iso_eps = if haskey(plot, :isorange)
            nlevels * ((vrange[2] - vrange[1]) / N) # TODO calculate this
        @show iso_eps
        cmap = to_colormap(_cmap)
        v_interval = cliprange[1] .. cliprange[2]
        # resample colormap and make the empty area between iso surfaces transparent
        map(1:N) do i
            i01 = (i-1) / (N - 1)
            c = Makie.interpolated_getindex(cmap, i01)
            # c = RGBA(1.0f0, 0.0f0, 0.0f0, 1.0f0)

            isoval = vrange[1] + (i01 * (vrange[2] - vrange[1]))
            line = reduce(levels, init = false) do v0, level
                isoval in v_interval || return false
                v0 || abs(level - isoval) <= iso_eps
            RGBAf(Colors.color(c), line ? alpha : 0.0)

    attr = copy(Attributes(plot))

    attr[:colorrange] = cliprange
    attr[:colormap] = cmap
    attr[:algorithm] = 7
    pop!(attr, :levels)
    pop!(attr, :alpha) # don't apply alpha 2 times

    # unused attributes
    pop!(attr, :labels)
    pop!(attr, :labelfont)
    pop!(attr, :labelsize)
    pop!(attr, :labelcolor)
    pop!(attr, :labelformatter)
    volume!(plot, attr, x, y, z, avolume)

function meshgrid(x, y, z)
    nx = length(x)
    ny = length(y)
    nz = length(z)

    X = repeat(x, outer=(1, ny, nz))
    Y = repeat(y', outer=(nx, 1, nz))
    Z = repeat(permutedims(z[:,:,:], (3,2,1)), outer=(nx, ny, 1))

    X, Y, Z

x = range(-1, stop=1, length=201)
y = range(-1, stop=1, length=201)
z = range(-1, stop=1, length=201)

X,Y,Z = meshgrid(x, y, z)

R² = X.^2 + Y.^2 + Z.^2
R = .√(R²)
Θ = atan.(.√(X.^2 + Y.^2), Z)
Φ = atan.(Y,X)

F = exp.(-R) .* exp.(im .* Φ)

fig = Figure()
ax = Axis3(fig[1,1], aspect=:equal, perspectiveness=1)
contour!(ax, x, y, z, F, levels=5)
1 Like