Colorbar enhancements (JuliaPlots#3346)

* CompatHelper: bump compat for "Showoff" to "1.0"

* fix series-segments for empty series

* fix wireframe on pyplot

* colorbar redesign

* minimal working version

* reduce code duplication for colorbar ticks

* fix aspect_ratio in GR with legend=:outertopright

* fix GR test failure

* new release [skip ci]

* colorbar scale supported

* Added weights example to ? histogram

Helps to clarify the use of weights (which differs from StatsBase functions)

* Update precompile_*.jl file

* minor version bump [skip ci]

* working prototype

* fixed formatting, added colorbar docs

* colorbar redesign

* minimal working version

* reduce code duplication for colorbar ticks

* fix GR test failure

* colorbar scale supported

* working prototype

* fixed formatting, added colorbar docs

Co-authored-by: Daniel Schwabeneder <daschw@disroot.org>

src/Plots.jl

@@ -193,6 +193,7 @@ const _plotly_min_js_filename = "plotly-1.57.1.min.js"
 
 include("types.jl")
 include("utils.jl")
+include("colorbars.jl")
 include("axes.jl")
 include("args.jl")
 include("components.jl")

src/axes.jl

@@ -137,11 +137,10 @@ const _label_func_tex = Dict{Symbol,Function}(
 labelfunc_tex(scale::Symbol) = get(_label_func_tex, scale, convert_sci_unicode)
 
 
-function optimal_ticks_and_labels(sp::Subplot, axis::Axis, ticks = nothing)
-    amin, amax = axis_limits(sp, axis[:letter])
+function optimal_ticks_and_labels(ticks, alims, scale, formatter)
+    amin, amax = alims
 
     # scale the limits
-    scale = axis[:scale]
     sf = RecipesPipeline.scale_func(scale)
 
     # If the axis input was a Date or DateTime use a special logic to find
@@ -152,7 +151,7 @@ function optimal_ticks_and_labels(sp::Subplot, axis::Axis, ticks = nothing)
     # rather than on the input format
     # TODO: maybe: non-trivial scale (:ln, :log2, :log10) for date/datetime
     if ticks === nothing && scale == :identity
-        if axis[:formatter] == RecipesPipeline.dateformatter
+        if formatter == RecipesPipeline.dateformatter
             # optimize_datetime_ticks returns ticks and labels(!) based on
             # integers/floats corresponding to the DateTime type. Thus, the axes
             # limits, which resulted from converting the Date type to integers,
@@ -163,7 +162,7 @@ function optimal_ticks_and_labels(sp::Subplot, axis::Axis, ticks = nothing)
                 k_min = 2, k_max = 4)
             # Now the ticks are converted back to floats corresponding to Dates.
             return ticks / 864e5, labels
-        elseif axis[:formatter] == RecipesPipeline.datetimeformatter
+        elseif formatter == RecipesPipeline.datetimeformatter
             return optimize_datetime_ticks(amin, amax; k_min = 2, k_max = 4)
         end
     end
@@ -187,14 +186,13 @@ function optimal_ticks_and_labels(sp::Subplot, axis::Axis, ticks = nothing)
             # chosen  ticks is not too much bigger than amin - amax:
             strict_span = false,
         )
-        axis[:lims] = map(RecipesPipeline.inverse_scale_func(scale), (viewmin, viewmax))
+        # axis[:lims] = map(RecipesPipeline.inverse_scale_func(scale), (viewmin, viewmax))
     else
         scaled_ticks = map(sf, (filter(t -> amin <= t <= amax, ticks)))
     end
     unscaled_ticks = map(RecipesPipeline.inverse_scale_func(scale), scaled_ticks)
 
     labels = if any(isfinite, unscaled_ticks)
-        formatter = axis[:formatter]
         if formatter in (:auto, :plain, :scientific, :engineering)
             map(labelfunc(scale, backend()), Showoff.showoff(scaled_ticks, formatter))
         elseif formatter == :latex
@@ -221,51 +219,52 @@ end
 # return (continuous_values, discrete_values) for the ticks on this axis
 function get_ticks(sp::Subplot, axis::Axis; update = true)
     if update || !haskey(axis.plotattributes, :optimized_ticks)
+        dvals = axis[:discrete_values]
         ticks = _transform_ticks(axis[:ticks])
-        if ticks in (:none, nothing, false)
-            axis.plotattributes[:optimized_ticks] = nothing
+        axis.plotattributes[:optimized_ticks] = if ticks isa Symbol && ticks !== :none &&
+            ispolar(sp) && axis[:letter] === :x && !isempty(dvals)
+            collect(0:pi/4:7pi/4), string.(0:45:315)
         else
-            # treat :native ticks as :auto
-            ticks = ticks == :native ? :auto : ticks
-
-            dvals = axis[:discrete_values]
-            cv, dv = if typeof(ticks) <: Symbol
-                if !isempty(dvals)
-                    # discrete ticks...
-                    n = length(dvals)
-                    rng = if ticks == :auto && n > 15
-                        Δ = ceil(Int, n / 10)
-                        Δ:Δ:n
-                    else # if ticks == :all
-                        1:n
-                    end
-                    axis[:continuous_values][rng], dvals[rng]
-                elseif ispolar(axis.sps[1]) && axis[:letter] == :x
-                    #force theta axis to be full circle
-                    (collect(0:pi/4:7pi/4), string.(0:45:315))
-                else
-                    # compute optimal ticks and labels
-                    optimal_ticks_and_labels(sp, axis)
-                end
-            elseif typeof(ticks) <: Union{AVec, Int}
-                if !isempty(dvals) && typeof(ticks) <: Int
-                    rng = Int[round(Int,i) for i in range(1, stop=length(dvals), length=ticks)]
-                    axis[:continuous_values][rng], dvals[rng]
-                else
-                    # override ticks, but get the labels
-                    optimal_ticks_and_labels(sp, axis, ticks)
-                end
-            elseif typeof(ticks) <: NTuple{2, Any}
-                # assuming we're passed (ticks, labels)
-                ticks
-            else
-                error("Unknown ticks type in get_ticks: $(typeof(ticks))")
-            end
-            axis.plotattributes[:optimized_ticks] = (cv, dv)
+            cvals = axis[:continuous_values]
+            alims = axis_limits(sp, axis[:letter])
+            scale = axis[:scale]
+            formatter = axis[:formatter]
+            get_ticks(ticks, cvals, dvals, alims, scale, formatter)
         end
     end
-    axis.plotattributes[:optimized_ticks]
+    return axis.plotattributes[:optimized_ticks]
+end
+
+function get_ticks(ticks::Symbol, cvals::T, dvals, args...) where T
+    if ticks === :none
+        return T[], String[]
+    elseif !isempty(dvals)
+        n = length(dvals)
+        if ticks === :all || n < 16
+            return cvals, string.(dvals)
+        else
+            Δ = ceil(Int, n / 10)
+            rng = Δ:Δ:n
+            return cvals[rng], string.(dvals[rng])
+        end
+    else
+        return optimal_ticks_and_labels(nothing, args...)
+    end
+end
+get_ticks(ticks::AVec, cvals, dvals, args...) = optimal_ticks_and_labels(ticks, args...)
+function get_ticks(ticks::Int, dvals, cvals, args...)
+    if !isempty(dvals)
+        rng = round.(Int, range(1, stop=length(dvals), length=ticks))
+        cvals[rng], string.(dvals[rng])
+    else
+        optimal_ticks_and_labels(ticks, args...)
+    end
 end
+get_ticks(ticks::NTuple{2, Any}, args...) = ticks
+get_ticks(::Nothing, cvals::T, args...) where T = T[], String[]
+get_ticks(ticks::Bool, args...) =
+    ticks ? get_ticks(:auto, args...) : get_ticks(nothing, args...)
+get_ticks(::T, args...) where T = error("Unknown ticks type in get_ticks: $T")
 
 _transform_ticks(ticks) = ticks
 _transform_ticks(ticks::AbstractArray{T}) where T <: Dates.TimeType = Dates.value.(ticks)

src/backends.jl

@@ -531,6 +531,10 @@ const _pyplot_attr = merge_with_base_supported([
     :guidefontfamily, :guidefontsize, :guidefontcolor,
     :grid, :gridalpha, :gridstyle, :gridlinewidth,
     :legend, :legendtitle, :colorbar, :colorbar_title, :colorbar_entry,
+    :colorbar_ticks, :colorbar_tickfontfamily, :colorbar_tickfontsize,
+    :colorbar_tickfonthalign, :colorbar_tickfontvalign,
+    :colorbar_tickfontrotation, :colorbar_tickfontcolor,
+    :colorbar_scale,
     :marker_z, :line_z, :fill_z,
     :levels,
     :ribbon, :quiver, :arrow,

src/backends/gr.jl

@@ -743,7 +743,7 @@ function _update_min_padding!(sp::Subplot{GRBackend})
         xticks, yticks, zticks = get_ticks(sp, xaxis), get_ticks(sp, yaxis), get_ticks(sp, zaxis)
         # Add margin for x and y ticks
         h = 0mm
-        if !(xticks in (nothing, false, :none))
+        if !isempty(first(xticks))
             gr_set_font(
                 tickfont(xaxis),
                 halign = (:left, :hcenter, :right)[sign(xaxis[:rotation]) + 2],
@@ -754,7 +754,7 @@ function _update_min_padding!(sp::Subplot{GRBackend})
             l = 0.01 + last(gr_get_ticks_size(xticks, xaxis[:rotation]))
             h = max(h, 1mm + get_size(sp)[2] * l * px)
         end
-        if !(yticks in (nothing, false, :none))
+        if !isempty(first(yticks))
             gr_set_font(
                 tickfont(yaxis),
                 halign = (:left, :hcenter, :right)[sign(yaxis[:rotation]) + 2],
@@ -774,7 +774,7 @@ function _update_min_padding!(sp::Subplot{GRBackend})
             end
         end
 
-        if !(zticks in (nothing, false, :none))
+        if !isempty(first(zticks))
             gr_set_font(
                 tickfont(zaxis),
                 halign = (zaxis[:mirror] ? :left : :right),
@@ -825,7 +825,7 @@ function _update_min_padding!(sp::Subplot{GRBackend})
     else
         # Add margin for x and y ticks
         xticks, yticks = get_ticks(sp, sp[:xaxis]), get_ticks(sp, sp[:yaxis])
-        if !(xticks in (nothing, false, :none))
+        if !isempty(first(xticks))
             gr_set_tickfont(sp, :x)
             l = 0.01 + last(gr_get_ticks_size(xticks, sp[:xaxis][:rotation]))
             h = 1mm + get_size(sp)[2] * l * px
@@ -835,7 +835,7 @@ function _update_min_padding!(sp::Subplot{GRBackend})
                 bottompad += h
             end
         end
-        if !(yticks in (nothing, false, :none))
+        if !isempty(first(yticks))
             gr_set_tickfont(sp, :y)
             l = 0.01 + first(gr_get_ticks_size(yticks, sp[:yaxis][:rotation]))
             w = 1mm + get_size(sp)[1] * l * px

src/backends/pyplot.jl

@@ -807,31 +807,37 @@ function py_compute_axis_minval(sp::Subplot, axis::Axis)
     minval
 end
 
-function py_set_scale(ax, sp::Subplot, axis::Axis)
-    scale = axis[:scale]
-    letter = axis[:letter]
+function py_set_scale(ax, sp::Subplot, scale::Symbol, letter::Symbol)
     scale in supported_scales() || return @warn("Unhandled scale value in pyplot: $scale")
     func = getproperty(ax, Symbol("set_", letter, "scale"))
     if PyPlot.version ≥ v"3.3" # https://matplotlib.org/3.3.0/api/api_changes.html
-        letter = Symbol("")
+        pyletter = Symbol("")
+    else
+        pyletter = letter
     end
     kw = KW()
     arg = if scale == :identity
         "linear"
     else
-        kw[Symbol(:base,letter)] = if scale == :ln
+        kw[Symbol(:base, pyletter)] = if scale == :ln
             ℯ
         elseif scale == :log2
             2
         elseif scale == :log10
             10
         end
-        kw[Symbol(:linthresh,letter)] = NaNMath.max(1e-16, py_compute_axis_minval(sp, axis))
+        axis = sp[Symbol(letter, :axis)]
+        kw[Symbol(:linthresh, pyletter)] = NaNMath.max(1e-16, py_compute_axis_minval(sp, axis))
         "symlog"
     end
     func(arg; kw...)
 end
 
+function py_set_scale(ax, sp::Subplot, axis::Axis)
+    scale = axis[:scale]
+    letter = axis[:letter]
+    py_set_scale(ax, sp, scale, letter)
+end
 
 function py_set_axis_colors(sp, ax, a::Axis)
     for (loc, spine) in ax.spines
@@ -972,24 +978,31 @@ function _before_layout_calcs(plt::Plot{PyPlotBackend})
 
             end
 
-            cb."set_label"(sp[:colorbar_title],size=py_thickness_scale(plt, sp[:yaxis][:guidefontsize]),family=sp[:yaxis][:guidefontfamily], color = py_color(sp[:yaxis][:guidefontcolor]))
+            cb."set_label"(sp[:colorbar_title],size=py_thickness_scale(plt, sp[:colorbar_titlefontsize]),family=sp[:colorbar_titlefontfamily], color = py_color(sp[:colorbar_titlefontcolor]))
 
             # cb."formatter".set_useOffset(false)   # This for some reason does not work, must be a pyplot bug, instead this is a workaround:
             cb."formatter".set_powerlimits((-Inf, Inf))
             cb."update_ticks"()
 
+            env = "\\mathregular"  # matches the outer fonts https://matplotlib.org/tutorials/text/mathtext.html
+            ticks = get_colorbar_ticks(sp)
+
             if sp[:colorbar] in (:top, :bottom)
                 axis = sp[:xaxis]  # colorbar inherits from x axis
                 cbar_axis = cb."ax"."xaxis"
+                ticks_letter=:x
             else
                 axis = sp[:yaxis]  # colorbar inherits from y axis
                 cbar_axis = cb."ax"."yaxis"
+                ticks_letter=:y
             end
+            py_set_scale(cb.ax, sp, sp[:colorbar_scale], ticks_letter)
+            sp[:colorbar_ticks] == :native ? nothing : py_set_ticks(cb.ax, ticks, ticks_letter, env)
 
             for lab in cbar_axis."get_ticklabels"()
-                  lab."set_fontsize"(py_thickness_scale(plt, axis[:tickfontsize]))
-                  lab."set_family"(axis[:tickfontfamily])
-                  lab."set_color"(py_color(axis[:tickfontcolor]))
+                  lab."set_fontsize"(py_thickness_scale(plt, sp[:colorbar_tickfontsize]))
+                  lab."set_family"(sp[:colorbar_tickfontfamily])
+                  lab."set_color"(py_color(sp[:colorbar_tickfontcolor]))
             end
 
             # Adjust thickness of the cbar ticks

src/colorbars.jl

@@ -0,0 +1,90 @@
+# These functions return an operator for use in `get_clims(::Seres, op)`
+process_clims(lims::Tuple{<:Number,<:Number}) = (zlims -> ifelse.(isfinite.(lims), lims, zlims)) ∘ ignorenan_extrema
+process_clims(s::Union{Symbol,Nothing,Missing}) = ignorenan_extrema
+# don't specialize on ::Function otherwise python functions won't work
+process_clims(f) = f
+
+function get_clims(sp::Subplot, op=process_clims(sp[:clims]))
+    zmin, zmax = Inf, -Inf
+    for series in series_list(sp)
+        if series[:colorbar_entry]
+            zmin, zmax = _update_clims(zmin, zmax, get_clims(series, op)...)
+        end
+    end
+    return zmin <= zmax ? (zmin, zmax) : (NaN, NaN)
+end
+
+function get_clims(sp::Subplot, series::Series, op=process_clims(sp[:clims]))
+    zmin, zmax = if series[:colorbar_entry]
+        get_clims(sp, op)
+    else
+        get_clims(series, op)
+    end
+    return zmin <= zmax ? (zmin, zmax) : (NaN, NaN)
+end
+
+"""
+    get_clims(::Series, op=Plots.ignorenan_extrema)
+
+Finds the limits for the colorbar by taking the "z-values" for the series and passing them into `op`,
+which must return the tuple `(zmin, zmax)`. The default op is the extrema of the finite
+values of the input.
+"""
+function get_clims(series::Series, op=ignorenan_extrema)
+    zmin, zmax = Inf, -Inf
+    z_colored_series = (:contour, :contour3d, :heatmap, :histogram2d, :surface, :hexbin)
+    for vals in (series[:seriestype] in z_colored_series ? series[:z] : nothing, series[:line_z], series[:marker_z], series[:fill_z])
+        if (typeof(vals) <: AbstractSurface) && (eltype(vals.surf) <: Union{Missing, Real})
+            zmin, zmax = _update_clims(zmin, zmax, op(vals.surf)...)
+        elseif (vals !== nothing) && (eltype(vals) <: Union{Missing, Real})
+            zmin, zmax = _update_clims(zmin, zmax, op(vals)...)
+        end
+    end
+    return zmin <= zmax ? (zmin, zmax) : (NaN, NaN)
+end
+
+_update_clims(zmin, zmax, emin, emax) = NaNMath.min(zmin, emin), NaNMath.max(zmax, emax)
+
+@enum ColorbarStyle cbar_gradient cbar_fill cbar_lines
+
+function colorbar_style(series::Series)
+    colorbar_entry = series[:colorbar_entry]
+    if !(colorbar_entry isa Bool)
+        @warn "Non-boolean colorbar_entry ignored."
+        colorbar_entry = true
+    end
+
+    if !colorbar_entry
+        nothing
+    elseif isfilledcontour(series)
+        cbar_fill
+    elseif iscontour(series)
+        cbar_lines
+    elseif series[:seriestype] ∈ (:heatmap,:surface) ||
+            any(series[z] !== nothing for z ∈ [:marker_z,:line_z,:fill_z])
+        cbar_gradient
+    else
+        nothing
+    end
+end
+
+hascolorbar(series::Series) = colorbar_style(series) !== nothing
+hascolorbar(sp::Subplot) = sp[:colorbar] != :none && any(hascolorbar(s) for s in series_list(sp))
+
+function get_colorbar_ticks(sp::Subplot; update = true)
+    if update || !haskey(sp.attr, :colorbar_optimized_ticks)
+        ticks = _transform_ticks(sp[:colorbar_ticks])
+        cvals = sp[:colorbar_continuous_values]
+        dvals = sp[:colorbar_discrete_values]
+        clims = get_clims(sp)
+        scale = sp[:colorbar_scale]
+        formatter = sp[:colorbar_formatter]
+        sp.attr[:colorbar_optimized_ticks] =
+            get_ticks(ticks, cvals, dvals, clims, scale, formatter)
+    end
+    return sp.attr[:colorbar_optimized_ticks]
+end
+
+function _update_subplot_colorbars(sp::Subplot)
+    # Dynamic callback from the pipeline if needed
+end