Makie 3d line drawing performance

RecipeDemo.jl

module RecipeDemo

export IsotropicGaussian, IsotropicGMM, IsotropicMultiGMM, nonrecipe!

using StaticArrays

struct IsotropicGaussian{N,T}
    μ::SVector{N,T}
    σ::T
    ϕ::T
end
IsotropicGaussian{N}(μ::AbstractVector{T},σ::Real,ϕ::Real) where {N,T<:Real} = IsotropicGaussian{N,T}(SVector{N,T}(μ...),σ,ϕ)

struct IsotropicGMM{N,T}
    gaussians::Vector{IsotropicGaussian{N,T}}
end

struct IsotropicMultiGMM{N,T,K}
    gmms::Dict{K, IsotropicGMM{N,T}}
end



import MakieCore: plot!
using MakieCore: @recipe, lines!, mesh!, Theme
using GeometryBasics: GeometryBasics, Sphere
using Colors: RGB

const HALFWAY_RADIUS = sqrt(3) / 2
const EQUAL_VOL_CONST = 3*√π/4

const DEFAULT_COLORS = [            # CUD colors: https://jfly.uni-koeln.de/color/#assign
    RGB(0/255,   114/255, 178/255), # blue
    RGB(230/255, 159/255, 0/255  ), # orange
    RGB(0/255,   158/255, 115/255), # green
    RGB(204/255, 121/255, 167/255), # reddish purple
    RGB(86/255,  180/255, 233/255), # sky blue
    RGB(213/255, 94/255,  0/255  ), # vermillion
    RGB(240/255, 228/255, 66/255 ), # yellow
]

# use cached calculations for positions of points on a circle
const θs = range(0, 2π, length=32)        # use 32 points (arbitrary)
const cosθs = Float32[cos(θ) for θ in θs]
const sinθs = Float32[sin(θ) for θ in θs]

equal_volume_radius(σ, ϕ) = (EQUAL_VOL_CONST*abs(ϕ))^(1/3) * σ

const counter = Ref(0)

function flat_circle!(f, pos, r, dim::Int; kwargs...)
    if dim == 3
        xs = [r * cosθ + pos[1] for cosθ in cosθs]
        ys = [r * sinθ + pos[2] for sinθ in sinθs]
        zs = fill(pos[3], length(xs))
    elseif dim == 2
        xs = [r * cosθ + pos[1] for cosθ in cosθs]
        ys = fill(pos[2], length(xs))
        zs = [r * sinθ + pos[3] for sinθ in sinθs]
    elseif dim == 1
        ys = [r * cosθ + pos[2] for cosθ in cosθs]
        zs = [r * sinθ + pos[3] for sinθ in sinθs]
        xs = fill(pos[1], length(ys))
    end
    counter[] += 1
    lines!(f, xs,ys,zs; kwargs...)
end

function wire_sphere!(f, pos, r; kwargs...)
    for dim in 1:3
        flat_circle!(f, pos, r, dim; kwargs...)
        halfwaypos = Float32[0,0,0]
        halfwaypos[dim] = r / 2;
        flat_circle!(f, pos .- halfwaypos, r * HALFWAY_RADIUS, dim; kwargs...)
        flat_circle!(f, pos .+ halfwaypos, r * HALFWAY_RADIUS, dim; kwargs...)
    end
end

function solid_sphere!(f, pos, r; kwargs...)
    mesh!(f, Sphere(GeometryBasics.Point{3}(pos...), r); kwargs...)
end

@recipe(GaussianDisplay, g) do scene
    Theme(
        display = :wire,
        color = DEFAULT_COLORS[1],
    )
end

function plot!(gd::GaussianDisplay{<:NTuple{<:Any, <:IsotropicGaussian}})
    gauss = [gd[i][] for i=1:length(gd)]
    disp = gd[:display][]
    color = gd[:color][]
    label = gd[:label][]
    if disp === :wire
        for g in gauss
            wire_sphere!(gd, g.μ, g.σ; color=color, label)
        end
    elseif disp === :solid
        for g in gauss
            solid_sphere!(gd, g.μ, g.σ; color=color, label)
        end
    else
        throw(ArgumentError("Unrecognized display option: `$disp`"))
    end
    return gd
end

@recipe(GMMDisplay, g) do scene
    Theme(
        display = :wire,
        palette = DEFAULT_COLORS,
        color = nothing,
        label = "",
    )
end

function plot!(gd::GMMDisplay{<:NTuple{<:Any,<:IsotropicGMM}})
    gmms = [gd[i][] for i=1:length(gd)]
    len = length(gmms)
    disp = gd[:display][]
    color = gd[:color][]
    palette = gd[:palette][]
    label = gd[:label][]
    for (i,gmm) in enumerate(gmms)
        col = isnothing(color) ? palette[(i-1) % len + 1] : color
        gaussiandisplay!(gd, gmm.gaussians...; display=disp, color=col, label)
    end
    return gd
end

function plot!(gd::GMMDisplay{<:NTuple{<:Any,<:IsotropicMultiGMM{N,T,K}}}) where {N,T,K}
    mgmms = [gd[i][] for i=1:length(gd)]
    disp = gd[:display][]
    color = gd[:color][]
    palette = gd[:palette][]
    allkeys = Set{K}()
    for mgmm in mgmms
        allkeys = allkeys ∪ keys(mgmm.gmms)
    end
    len = length(allkeys)
    for (i,k) in enumerate(allkeys)
        col = isnothing(color) ? palette[(i-1) % len + 1] : color
        for mgmm in mgmms
            haskey(mgmm.gmms, k) && gmmdisplay!(gd, mgmm.gmms[k]; display=disp, color=col, palette=palette, label=string(k))
        end
    end
    return gd
end

function nonrecipe!(ax, mgmm)
    len = length(mgmm.gmms)
    for (i, gmm) in enumerate(values(mgmm.gmms))
        color = DEFAULT_COLORS[(i-1) % len + 1]
        for g in gmm.gaussians
            wire_sphere!(ax, g.μ, g.σ; color)
        end
    end
end

end   # module