added limited support for logarithmic units

jw3126 · May 1, 2022 · b3e5ac3 · b3e5ac3
1 parent 7741be7
commit b3e5ac3
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 16 deletions.
diff --git a/src/UnitfulRecipes.jl b/src/UnitfulRecipes.jl
@@ -1,7 +1,7 @@
 module UnitfulRecipes
 
 using RecipesBase
-using Unitful: Quantity, unit, ustrip, Unitful, dimension, Units
+using Unitful: Quantity, unit, ustrip, Unitful, dimension, Units, LogScaled, logunit, MixedUnits, Level, Gain
 export @P_str
 
 const clims_types = (:contour, :contourf, :heatmap, :surface)
@@ -10,11 +10,12 @@ const clims_types = (:contour, :contourf, :heatmap, :surface)
 Main recipe
 ==========#
 
-@recipe function f(::Type{T}, x::T) where T <: AbstractArray{<:Union{Missing,<:Quantity}}
+@recipe function f(::Type{T}, x::T) where T <: AbstractArray{<:Union{Missing,<:Quantity, <:LogScaled}}
     axisletter = plotattributes[:letter]   # x, y, or z
     if (axisletter == :z) &&
         get(plotattributes, :seriestype, :nothing) ∈ clims_types
-        u = get(plotattributes, :zunit, unit(eltype(x)))
+        # u = get(plotattributes, :zunit, unit(eltype(x)))
+        u = get(plotattributes, :zunit, eltype(x) <: LogScaled ? logunit(eltype(x)) : unit(eltype(x)))
         ustripattribute!(plotattributes, :clims, u)
         append_unit_if_needed!(plotattributes, :colorbar_title, u)
     end
@@ -30,7 +31,7 @@ function fixaxis!(attr, x, axisletter)
     axisunit = Symbol(axisletter, :unit)   # xunit, yunit, zunit
     axis = Symbol(axisletter, :axis)       # xaxis, yaxis, zaxis
     # Get the unit
-    u = pop!(attr, axisunit, unit(eltype(x)))
+    u = pop!(attr, axisunit, eltype(x) <: LogScaled ? logunit(eltype(x)) : unit(eltype(x)))
     # If the subplot already exists with data, get its unit
     sp = get(attr, :subplot, 1)
     if sp ≤ length(attr[:plot_object]) && attr[:plot_object].n > 0
@@ -54,22 +55,28 @@ function fixaxis!(attr, x, axisletter)
     fixmarkersize!(attr)
     fixlinecolor!(attr)
     # Strip the unit
-    ustrip.(u, x)
+    if eltype(x) <: Level
+        ustrip(x)
+    elseif eltype(x) <: Gain
+        getfield.(x, :val)
+    else
+        ustrip.(u, x)
+    end
 end
 
 # Recipe for (x::AVec, y::AVec, z::Surface) types
 const AVec = AbstractVector
 const AMat{T} = AbstractArray{T,2} where T
 @recipe function f(x::AVec, y::AVec, z::AMat{T}) where T <: Quantity
-    u = get(plotattributes, :zunit, unit(eltype(z)))
+    u = get(plotattributes, :zunit, eltype(z) <: LogScaled ? logunit(eltype(z)) : unit(eltype(z)))
     ustripattribute!(plotattributes, :clims, u)
     z = fixaxis!(plotattributes, z, :z)
     append_unit_if_needed!(plotattributes, :colorbar_title, u)
     x, y, z
 end
 
 # Recipe for vectors of vectors
-@recipe function f(::Type{T}, x::T) where T <: AbstractVector{<:AbstractVector{<:Union{Missing,<:Quantity}}}
+@recipe function f(::Type{T}, x::T) where T <: AbstractVector{<:AbstractVector{<:Union{Missing,<:Quantity, <:LogScaled}}}
     axisletter = plotattributes[:letter]   # x, y, or z
     [fixaxis!(plotattributes, x, axisletter) for x in x]
 end
@@ -81,19 +88,19 @@ end
 end
 
 # Recipes for functions
-@recipe function f(f::Function, x::T) where T <: AVec{<:Union{Missing,<:Quantity}}
+@recipe function f(f::Function, x::T) where T <: AVec{<:Union{Missing,<:Quantity, <:LogScaled}}
     x, f.(x)
 end
-@recipe function f(x::T, f::Function) where T <: AVec{<:Union{Missing,<:Quantity}}
+@recipe function f(x::T, f::Function) where T <: AVec{<:Union{Missing,<:Quantity, <:LogScaled}}
     x, f.(x)
 end
-@recipe function f(x::T, y::AVec, f::Function) where T <: AVec{<:Union{Missing,<:Quantity}}
+@recipe function f(x::T, y::AVec, f::Function) where T <: AVec{<:Union{Missing,<:Quantity, <:LogScaled}}
     x, y, f.(x',y)
 end
-@recipe function f(x::AVec, y::T, f::Function) where T <: AVec{<:Union{Missing,<:Quantity}}
+@recipe function f(x::AVec, y::T, f::Function) where T <: AVec{<:Union{Missing,<:Quantity, <:LogScaled}}
     x, y, f.(x',y)
 end
-@recipe function f(x::T1, y::T2, f::Function) where {T1<:AVec{<:Union{Missing,<:Quantity}}, T2<:AVec{<:Union{Missing,<:Quantity}}}
+@recipe function f(x::T1, y::T2, f::Function) where {T1<:AVec{<:Union{Missing,<:Quantity, <:LogScaled}}, T2<:AVec{<:Union{Missing,<:Quantity, <:LogScaled}}}
     x, y, f.(x',y)
 end
 @recipe function f(f::Function, u::Units)
@@ -138,8 +145,16 @@ fixlinecolor!(attr) = ustripattribute!(attr, :line_z)
 function ustripattribute!(attr, key)
     if haskey(attr, key)
         v = attr[key]
-        u = unit(eltype(v))
-        attr[key] = ustrip.(u, v)
+        if eltype(v) <: Level
+            u = logunit(eltype(v))
+            attr[key] = ustrip(v)
+        elseif eltype(v) <: Gain
+            u = logunit(eltype(v))
+            attr[key] = v.val
+        else
+            u = unit(eltype(v))
+            attr[key] = ustrip.(u, v)
+        end
         return u
     else
         return Unitful.NoUnits
@@ -149,7 +164,11 @@ end
 function ustripattribute!(attr, key, u)
     if haskey(attr, key)
         v = attr[key]
-        if eltype(v) <: Quantity
+        if eltype(v) <: Level
+            attr[key] = ustrip(v)
+        elseif eltype(v) <: Gain
+            attr[key] = v.val
+        elseif eltype(v) <: Quantity
             attr[key] = ustrip.(u, v)
         end
     end
@@ -198,7 +217,7 @@ end
 Append unit to labels when appropriate
 =====================================#
 
-function append_unit_if_needed!(attr, key, u::Unitful.Units)
+function append_unit_if_needed!(attr, key, u::Union{Unitful.Units, Unitful.MixedUnits})
     label = get(attr, key, nothing)
     append_unit_if_needed!(attr, key, label, u)
 end

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -299,3 +299,25 @@ end
     plot!(plt, (1:3)m)
     @test yguide(plt) == "m"
 end
+
+@testset "LogScaled plots" begin
+    x, y = randn(3)*u"dBm", randn(3)*u"dB"
+
+    @testset "no keyword argument" begin
+        @test xguide(plot(x,y)) == "dBm"
+        @test xseries(plot(x,y)) ≈ ustrip.(x)
+        @test yguide(plot(x,y)) == "dB"
+        @test yseries(plot(x,y)) ≈ ustrip.(y)
+    end
+
+    @testset "labels" begin
+        @test xguide(plot(x, y, xlabel= "hello")) == "hello (dBm)"
+        @test xguide(plot(x, y, xlabel=P"hello")) == "hello"
+        @test yguide(plot(x, y, ylabel= "hello")) == "hello (dB)"
+        @test yguide(plot(x, y, ylabel=P"hello")) == "hello"
+        @test xguide(plot(x, y, xlabel= "hello", ylabel= "hello")) == "hello (dBm)"
+        @test xguide(plot(x, y, xlabel=P"hello", ylabel=P"hello")) == "hello"
+        @test yguide(plot(x, y, xlabel= "hello", ylabel= "hello")) == "hello (dB)"
+        @test yguide(plot(x, y, xlabel=P"hello", ylabel=P"hello")) == "hello"
+    end
+end