triangularheatmap userplot recipe for plotting data defined on a honeycomb lattice

triangular_heatmap.jl

using Plots; backend()

# get a upper/lower triangle shape for the heatmap on an underlying honeycomb lattice
function _triangle(center::NTuple{2}, hexlength::Real, triangle_orientation::Symbol = :lower)
    @assert triangle_orientation in (:upper, :lower)
    sgn = triangle_orientation == :upper ? 1 : -1
    Plots.Shape(
        [center .+ edge for edge in
            [
                (0, sgn * hexlength),
                (-sqrt(3) * hexlength / 2, -sgn * hexlength / 2),
                (sqrt(3) * hexlength / 2, -sgn * hexlength / 2),
                (0, sgn * hexlength),
            ]
        ]
    )
end

@userplot TriangularHeatmap

# """
#     `triangularheatmap(positions::Vector{NTuple{2}}, values::Vector, hexlength::Real = 1.0, triangle_orientation::Vector{Symbol} = [iseven(i) ? :upper : :lower for i in eachindex(positions)])`
#
# `triangularheatmap!(positions::Vector{NTuple{2}}, values::Vector, hexlength::Real = 1.0, triangle_orientation::Vector{Symbol} = [iseven(i) ? :upper : :lower for i in eachindex(positions)])`
#
# Plot a heatmap with triangular shapes.
#
# Arguments:
# * `positions` Positions of the triangle centers (hexagon corners of the underlying honeycomb lattice)
# * `values` Function values to be plotted
# * `hexlength` edge length of the underlying honeycomb lattice
# * `triangle_orientation` Vector of triangle orientations (`:upper` or `:lower`)
# """

@recipe function f(thp::TriangularHeatmap)
    # get positions and data values
    positions, values = thp.args[1:2]
    # get the hexagon length
    hexlength = length(thp.args) > 2 ? thp.args[3] : 1.0
    # get the orientation of each triangle
    triangle_orientation = length(thp.args) > 3 ? thp.args[4] : [iseven(i) ? :upper : :lower for i in eachindex(positions)]

    # dirty hack to get a colorbar - plot invisible heatmap with
    # data extrema
    emin, emax = extrema(values)
    @series begin
       seriestype := :heatmap
       primary := false
       fill(positions[1][1], 2), fill(positions[1][2], 2), [emin emax; emax emin]
    end

    # get color gradient set by user
    usercol = get(d, :seriescolor, :inferno)
    mycgrad = cgrad(isa(usercol, Symbol) ? usercol : :inferno)
    # map data values to [0, 1] for indexing the gradient
    mycolors = [mycgrad[(v - emin) / (emax - emin)] for v in values]

    color --> mycolors
    linecolor --> mycolors
    label --> ""
    aspect_ratio --> :equal
    # we assume here that we need alternating lower and upper triangles
    # this may be modified according to the ordering of the positions
    [_triangle(positions[i], hexlength, triangle_orientation[i]) for i in eachindex(positions)]
end

# ------------------------------------------------------------------------------
# test the recipe

# dimensions
m, n = 20, 41
hexlength = 1.0

# get positions of corners of a honeycomb lattice with edge length `hexlength`
positions = [(i * sqrt(3) * hexlength / 2, 3 / 2 * j * hexlength + (iseven(i + j) ? hexlength / 2 : 0)) for i in 1:n for j in 1:m]

# show positions
scatter(positions, aspect_ratio = :equal, label = "")

# data (function) values at the positions
data = [i * j for i in 1:n for j in 1:m]

# orientation of the corresponding triangles
tri_or = [iseven(i + j) ? :lower : :upper for i in 1:n for j in 1:m]

# plot the triangular heatmap
triangularheatmap(positions, data, hexlength, tri_or)


# ------------------------------------------------------------------------------
# different gradients
plot([triangularheatmap(positions, data, hexlength, tri_or, color = grad) for grad in cgradients()]...)