Merge pull request #50 from gustaphe/function

Plot function against unit

Project.toml

@@ -1,7 +1,7 @@
 name = "UnitfulRecipes"
 uuid = "42071c24-d89e-48dd-8a24-8a12d9b8861f"
 authors = ["Benoit Pasquier", "Jan Weidner"]
-version = "1.1.1"
+version = "1.2.0"
 
 [deps]
 RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"

docs/lit/examples/1_Examples.jl

@@ -150,3 +150,16 @@ x = (1.0:0.1:10) * GeV/c
 y = @. (2 + sin(x / (GeV/c))) * 0.4GeV/c^2 # a sine to make it pretty
 yerror = 10.9MeV/c^2 * exp.(randn(length(x))) # some noise for pretty again
 plot(x, y; yerror, title="My unitful data with yerror bars", lab="")
+
+# ## Functions
+#
+# In order to plot a unitful function on a unitful axis, supply as a second argument a
+# vector of unitful sample points, or the unit for the independent axis:
+
+model(x) = 1u"V"*exp(-((x-0.5u"s")/0.7u"s")^2)
+t = randn(10)u"s" # Sample points
+U = model.(t) + randn(10)u"dV" .|> u"V" # Noisy acquicisions
+plot(t, U; xlabel="t", ylabel="U", st=:scatter, label="Samples")
+plot!(model, t; st=:scatter, label="Noise removed")
+plot!(model, u"s"; label="True function")
+

src/UnitfulRecipes.jl

@@ -1,7 +1,7 @@
 module UnitfulRecipes
 
 using RecipesBase
-using Unitful: Quantity, unit, ustrip, Unitful, dimension
+using Unitful: Quantity, unit, ustrip, Unitful, dimension, Units
 export @P_str
 
 #==========
@@ -75,7 +75,30 @@ end
 @recipe function f(x::T1, y::T2, f::Function) where {T1<:AVec{<:Union{Missing,<:Quantity}}, T2<:AVec{<:Union{Missing,<:Quantity}}}
     x, y, f.(x',y)
 end
+@recipe function f(f::Function, u::Units)
+    uf = UnitFunction(f, [u])
+    recipedata = RecipesBase.apply_recipe(plotattributes, uf)
+    (_, xmin, xmax) = recipedata[1].args
+    return f, xmin*u, xmax*u
+end
+
+"""
+```julia
+UnitFunction
+```
+A function, bundled with the assumed units of each of its inputs.
 
+```julia
+f(x, y) = x^2 + y
+uf = UnitFunction(f, u"m", u"m^2")
+uf(3, 2) == f(3u"m", 2u"m"^2) == 7u"m^2"
+```
+"""
+struct UnitFunction <: Function
+    f::Function
+    u::Vector{Units}
+end
+(f::UnitFunction)(args...) = f.f((args .* f.u)...)
 
 #===============
 Attribute fixing

test/runtests.jl

@@ -114,6 +114,17 @@ end
         g(x,y) = x*y*m # If the unit comes from the function only then it throws
         @test_throws DimensionError plot(x, y, g) isa Plots.Plot
     end
+    @testset "plot(f, u)" begin
+        f(x) = x^2
+        pl = plot(x*m, f.(x*m))
+        @test plot!(pl, f, m) isa Plots.Plot
+        @test_throws DimensionError plot!(pl, f, s) isa Plots.Plot
+        pl = plot(f, m)
+        @test xguide(pl) == string(m)
+        @test yguide(pl) == string(m^2)
+        f(x) = exp(x/(3m))
+        @test plot(f, u"m") isa Plots.Plot
+    end
 end
 
 @testset "Moar plots" begin