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:
http://docs.enthought.com/mayavi/mayavi/auto/example_atomic_orbital.html

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)
            plot.isorange[]
        else
            nlevels * ((vrange[2] - vrange[1]) / N) # TODO calculate this
        end
        @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
            end
            RGBAf(Colors.color(c), line ? alpha : 0.0)
        end
    end

    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)
end

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
end

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)
fig
1 Like