Improved rigor for marker_z and line_z

Author: Ickaser

ext/UnitfulExt.jl

@@ -22,7 +22,7 @@ Main recipe
     if axisletter === :z && get(plotattributes, :seriestype, :nothing) ∈ clims_types
         u = get(plotattributes, :zunit, _unit(eltype(x)))
         ustripattribute!(plotattributes, :clims, u)
-        append_unit_if_needed!(plotattributes, :colorbar_title, u)
+        append_cbar_unit_if_needed!(plotattributes, u)
     end
     fixaxis!(plotattributes, x, axisletter)
 end
@@ -34,7 +34,7 @@ function fixaxis!(attr, x, axisletter)
     axis = Symbol(axisletter, :axis)       # xaxis, yaxis, zaxis
     # if the subplot already exists with data, use that unit
     sp = get(attr, :subplot, 1)
-    if sp ≤ length(attr[:plot_object])
+    if sp ≤ length(attr[:plot_object]) && attr[:plot_object].n > 0
         spu = getaxisunit(attr[:plot_object][sp][axis])
         if !isnothing(spu)
             u = spu
@@ -51,9 +51,9 @@ function fixaxis!(attr, x, axisletter)
         ustripattribute!(attr, :fillrange, u)
     end
     fixaspectratio!(attr, u, axisletter)
-    fixmarkercolor!(attr)
+    fixseriescolor!(attr, :marker_z)
+    fixseriescolor!(attr, :line_z)
     fixmarkersize!(attr)
-    fixlinecolor!(attr)
     _ustrip.(u, x)  # strip the unit
 end
 
@@ -62,7 +62,7 @@ end
     u = get(plotattributes, :zunit, _unit(eltype(z)))
     ustripattribute!(plotattributes, :clims, u)
     z = fixaxis!(plotattributes, z, :z)
-    append_unit_if_needed!(plotattributes, :colorbar_title, u)
+    append_cbar_unit_if_needed!(plotattributes, u)
     x, y, z
 end
 
@@ -154,15 +154,30 @@ function fixaspectratio!(attr, u, axisletter)
 end
 
 # Markers / lines
-function fixmarkercolor!(attr)
-    u = ustripattribute!(attr, :marker_z)
-    ustripattribute!(attr, :clims, u)
-    u == NoUnits || append_unit_if_needed!(attr, :colorbar_title, u)
-end
-function fixlinecolor!(attr)
-    u = ustripattribute!(attr, :line_z)
+function fixseriescolor!(attr, key)
+    sp = get(attr, :subplot, 1)
+    # Precedence to user-passed zunit
+    if haskey(attr, :zunit) 
+        u = attr[:zunit]
+        ustripattribute!(attr, key, u)
+    # Then to an existing subplot's colorbar title
+    elseif sp ≤ length(attr[:plot_object]) && attr[:plot_object].n > 0
+        cbar_title = get(attr[:plot_object][sp], :colorbar_title, nothing)
+        spu = (cbar_title isa UnitfulString ? cbar_title.unit : nothing)
+        if !isnothing(spu)
+            u = spu
+            ustripattribute!(attr, key, u)
+        else
+            u = ustripattribute!(attr, key)
+        end
+    # Otherwise, get from the attribute
+    else
+        u = ustripattribute!(attr, key)
+    end
     ustripattribute!(attr, :clims, u)
-    u == NoUnits || append_unit_if_needed!(attr, :colorbar_title, u)
+    # fixseriescolor! is called for each axis, so after the first pass,
+    # u will be NoUnits and we don't want to append unit again
+    u == NoUnits || append_cbar_unit_if_needed!(attr, u)
 end
 fixmarkersize!(attr) = ustripattribute!(attr, :markersize)
 
@@ -192,6 +207,8 @@ end
 
 #=======================================
 Label string containing unit information
+Used only for colorbars, etc., which don't 
+have a bettter place for storing units
 =======================================#
 
 const APS = Plots.AbstractProtectedString
@@ -202,42 +219,45 @@ end
 
 #=====================================
 Append unit to labels when appropriate
-This is needed for colorbars, mostly,
-since axes have their own handling
+This is needed for colorbars, etc., since axes have
+distinct unit handling
 =====================================#
 
-append_unit_if_needed!(attr, key, u) =
-    append_unit_if_needed!(attr, key, get(attr, key, nothing), u)
+append_cbar_unit_if_needed!(attr, u) =
+    append_cbar_unit_if_needed!(attr, get(attr, :colorbar_title, nothing), u)
 # dispatch on the type of `label`
-append_unit_if_needed!(attr, key, label::Plots.ProtectedString, u) = nothing
-append_unit_if_needed!(attr, key, label::UnitfulString, u) = nothing
-function append_unit_if_needed!(attr, key, label::Nothing, u)
-    attr[key] = if attr[:plot_object].backend == Plots.PGFPlotsXBackend()
+append_cbar_unit_if_needed!(attr, label::UnitfulString, u) = nothing
+function append_cbar_unit_if_needed!(attr, label::Nothing, u)
+    unitformat = get(attr, Symbol(:z, :unitformat), :round)
+    if unitformat ∈ [:nounit, :none, false, nothing]
+        return attr[:colorbar_title] = UnitfulString("", u)
+    end
+    attr[:colorbar_title] = if Plots.backend_name() === :pgfplotsx
         UnitfulString(LaTeXString(latexify(u)), u)
     else
         UnitfulString(string(u), u)
     end
 end
-function append_unit_if_needed!(attr, key, label::S, u) where {S<:AbstractString}
-    isempty(label) && return attr[key] = UnitfulString(label, u)
-    if attr[:plot_object].backend == Plots.PGFPlotsXBackend()
-        attr[key] = UnitfulString(
+function append_cbar_unit_if_needed!(attr, label::S, u) where {S<:AbstractString}
+    isempty(label) && return attr[:colorbar_title] = UnitfulString(label, u)
+    attr[:colorbar_title] = if Plots.backend_name() ≡ :pgfplotsx
+        UnitfulString(
             LaTeXString(
                 Plots.format_unit_label(
                     label,
                     latexify(u),
-                    get(attr, Symbol(get(attr, :letter, ""), :unitformat), :round),
+                    get(attr, :zunitformat, :round),
                 ),
             ),
             u,
         )
     else
-        attr[key] = UnitfulString(
+        UnitfulString(
             S(
                 Plots.format_unit_label(
                     label,
                     u,
-                    get(attr, Symbol(get(attr, :letter, ""), :unitformat), :round),
+                    get(attr, :zunitformat, :round),
                 ),
             ),
             u,
@@ -246,9 +266,8 @@ function append_unit_if_needed!(attr, key, label::S, u) where {S<:AbstractString
 end
 
 getaxisunit(::Nothing) = nothing
-getaxisunit(s::UnitfulString) = s.unit
-getaxisunit(a::Axis) = getaxisunit(a[:unit])
 getaxisunit(u) = u
+getaxisunit(a::Axis) = getaxisunit(a[:unit])
 
 #==============
 Fix annotations
@@ -280,8 +299,10 @@ function Plots.locate_annotation(
     rel::NTuple{N,<:MissingOrQuantity},
     label,
 ) where {N}
-    units = getaxisunit(sp.attr[:xaxis], sp.attr[:yaxis], sp.attr[:zaxis])
-    Plots.locate_annotation(sp, _ustrip.(zip(units, rel)), label)
+    units = getaxisunit(sp.attr[:xaxis]),
+    getaxisunit(sp.attr[:yaxis]),
+    getaxisunit(sp.attr[:zaxis])
+    PlotsBase.locate_annotation(sp, _ustrip.(zip(units, rel)), label)
 end
 
 #==================#

test/test_unitful.jl

@@ -5,6 +5,7 @@ using Unitful: m, cm, s, DimensionError
 xguide(pl, idx = length(pl.subplots)) = Plots.get_guide(pl.subplots[idx].attr[:xaxis])
 yguide(pl, idx = length(pl.subplots)) = Plots.get_guide(pl.subplots[idx].attr[:yaxis])
 zguide(pl, idx = length(pl.subplots)) = Plots.get_guide(pl.subplots[idx].attr[:zaxis])
+ctitle(pl, idx = length(pl.subplots)) = pl.subplots[idx].attr[:colorbar_title]
 xseries(pl, idx = length(pl.series_list)) = pl.series_list[idx].plotattributes[:x]
 yseries(pl, idx = length(pl.series_list)) = pl.series_list[idx].plotattributes[:y]
 zseries(pl, idx = length(pl.series_list)) = pl.series_list[idx].plotattributes[:z]
@@ -269,7 +270,13 @@ end
         x, y = rand(10) * us[1], rand(10) * us[2]
         @test scatter(x, y) isa Plots.Plot
         @test scatter(x, y, markersize = x) isa Plots.Plot
-        @test scatter(x, y, line_z = x) isa Plots.Plot
+
+        @test scatter(x, y, marker_z = x) isa Plots.Plot
+        if us[1] != us[2] && us[1] != 1 && us[2] != 1 # Non-matching dimensions
+            @test_throws DimensionError scatter!(x, y, marker_z = y) 
+        else # One is dimensionless, or have same dimensions
+            @test scatter!(x, y, marker_z = y) isa PlotsBase.Plot #
+        end
     end
 
     @testset "contour(x::$(us[1]), y::$(us[2]))" for us in collect(
@@ -286,6 +293,18 @@ end
         @test plot(y, label = P"meters") isa Plots.Plot
     end
 
+    @testset "colorbar title" begin
+
+        x, y = (1:0.01:2) * m, (1:0.02:2) * s
+        z = x' ./ y
+        pl = contour(x, y, z) 
+        @test ctitle(pl) ∈ ["m s^-1", "m s⁻¹"]
+        pl = contourf(x, y, z, zunit = u"km/hr")
+        @test ctitle(pl) ∈ ["km hr^-1", "km hr⁻¹"]
+        pl = heatmap(x, y, z, zunit = u"cm/s", zunitformat = :square, colorbar_title = "v")
+        @test ctitle(pl) ∈ ["v [cm s^-1]", "v [cm s⁻¹]"]
+    end
+
     @testset "twinx (#4750)" begin
         y = rand(10) * u"m"
         pl = plot(y; xlabel = "check", ylabel = "hello")