Allow for Vector of parameters and use Turing as dependency (#22)

Author: theogf

Project.toml

@@ -1,23 +1,23 @@
 name = "Turkie"
 uuid = "8156cc02-0533-41cd-9345-13411ebe105f"
 authors = ["Theo Galy-Fajou <[email protected]> and contributors"]
-version = "0.1.1"
+version = "0.1.2"
 
 [deps]
 AbstractPlotting = "537997a7-5e4e-5d89-9595-2241ea00577e"
 ColorSchemes = "35d6a980-a343-548e-a6ea-1d62b119f2f4"
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
-DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8"
 KernelDensity = "5ab0869b-81aa-558d-bb23-cbf5423bbe9b"
 OnlineStats = "a15396b6-48d5-5d58-9928-6d29437db91e"
+Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
 
 [compat]
 AbstractPlotting = "0.15"
 ColorSchemes = "3.10"
 Colors = "0.12"
-DynamicPPL = "0.9, 0.10"
 KernelDensity = "0.5, 0.6"
 OnlineStats = "1.5"
+Turing = "0.15"
 julia = "1.4"
 
 [extras]

README.md

@@ -58,15 +58,18 @@ chain = sample(m, NUTS(0.65), 300; callback = cb) # Sample and plot at the same
 If you want to show only some variables you can give a `Dict` to `TurkieCallback` :
 
 ```julia
-cb = TurkieCallback(Dict(:m0 => [:trace, :mean],
-                        :s => [:autocov, :var]))
+cb = TurkieCallback(
+            (m0 = [:trace, :mean], s = [:autocov, :var])
+          )
 
 ```
 
 You can also directly pass `OnlineStats` object : 
 ```julia
 using OnlineStats
-cb = TurkieCallback(Dict(:v => [Mean(), AutoCov(20)]))
+cb = TurkieCallback(
+            (v = [Mean(), AutoCov(20)],)
+          )
 ```
 
 If you want to record the video do

docs/src/index.md

@@ -66,11 +66,13 @@ While sampling the callback object `cb` will be called and the statistics will b
 ## Tuning the quantities
 
 Of course the default is not always desirable.
-You can chose what variables and what quantities are shown by giving a `Dict` to `TurkieCallback` instead of a model.
+You can chose what variables and what quantities are shown by giving a `NamedTuple` to `TurkieCallback` instead of a model.
 For example,
 ```julia
-cb = TurkieCallback(Dict(:v => [:trace, :mean],
-                        :s => [:autocov, :var]))
+cb = TurkieCallback(
+                (v = [:trace, :mean],
+                s = [:autocov, :var])
+            )
 ```
 will only show the trace and the sample mean of `v` and the auto-covariance and variance of `s`.
 Pairs should be of the type `{Symbol,AbstractVector}`.

src/Turkie.jl

@@ -8,7 +8,7 @@ using AbstractPlotting.MakieLayout # Layouting tool
 using Colors, ColorSchemes # Colors tools
 using KernelDensity # To be able to give a KDE
 using OnlineStats # Estimators
-using DynamicPPL: VarInfo, Model
+using Turing: DynamicPPL.VarInfo, DynamicPPL.Model, Inference._params_to_array
 
 export TurkieCallback
 
@@ -23,19 +23,31 @@ name(::Val{T}) where {T} = string(T)
 name(s::OnlineStat) = string(nameof(typeof(s)))
 
 """
-    TurkieCallback(model::DynamicPPL.Model, plots::Series/AbstractVector; window=1000, kwargs...)
-
+    TurkieCallback(args...; kwargs....)
+    
+## Arguments
+- Option 1 : 
+    `model::DynamicPPL.Model, plots::Series/AbstractVector=[:histkde, Mean(Float32), Variance(Float32), AutoCov(20, Float32)]`
+
+For each of the variables of the given model each `plot` from `plots` will be plotted
+Multidimensional variable will be automatically have indices added to them
+- Option 2 :
+    `vars::NamedTuple/Dict`
+Will plot each pair of symbol and series of plots.
+Note that for multidimensional variable you should pass a Symbol as `Symbol("m[1]")` for example.
+See the docs for some examples.
 ## Keyword arguments
 - `window=1000` : Use a window for plotting the trace
+- `refresh=false` : Restart the plots from scratch everytime `sample` is called again (still WIP)
 """
 TurkieCallback
 
-struct TurkieCallback
+struct TurkieCallback{TN<:NamedTuple,TD<:AbstractDict}
     scene::Scene
     data::Dict{Symbol, MovingWindow}
     axis_dict::Dict
-    vars::Dict{Symbol, Any}
-    params::Dict{Any, Any}
+    vars::TN
+    params::TD
     iter::Observable{Int}
 end
 
@@ -44,38 +56,39 @@ function TurkieCallback(model::Model, plots::Series; kwargs...)
 end
 
 function TurkieCallback(model::Model, plots::AbstractVector = [:histkde, Mean(Float32), Variance(Float32), AutoCov(20, Float32)]; kwargs...)
-    variables = VarInfo(model).metadata
+    vars, vals = _params_to_array([VarInfo(model)])
     return TurkieCallback(
-        Dict(Pair.(keys(variables), Ref(plots)));
+        (;Pair.(vars, Ref(plots))...); # Return a named Tuple
         kwargs...
-        )
+    )
 end
 
-function TurkieCallback(varsdict::Dict; kwargs...)
-    return TurkieCallback(varsdict, Dict{Symbol,Any}(kwargs...))
+function TurkieCallback(vars::Union{Dict, NamedTuple}; kwargs...)
+    return TurkieCallback((;vars...), Dict{Symbol,Any}(kwargs...))
 end
 
-function TurkieCallback(vars::Dict, params::Dict)
+function TurkieCallback(vars::NamedTuple, params::Dict)
 # Create a scene and a layout
     outer_padding = 5
     scene, layout = layoutscene(outer_padding, resolution = (1200, 700))
     window = get!(params, :window, 1000)
     refresh = get!(params, :refresh, false)
     params[:t0] = 0
-    iter = Node(0)
+    iter = Observable(0)
     data = Dict{Symbol, MovingWindow}(:iter => MovingWindow(window, Int))
     obs = Dict{Symbol, Any}()
     axis_dict = Dict()
-    for (i, (variable, plots)) in enumerate(vars)
+    for (i, variable) in enumerate(keys(vars))
+        plots = vars[variable]
         data[variable] = MovingWindow(window, Float32)
         axis_dict[(variable, :varname)] = layout[i, 1, Left()] = Label(scene, string(variable), textsize = 30)
-        axis_dict[(variable, :varname)].padding = (0, 50, 0, 0)
+        axis_dict[(variable, :varname)].padding = (0, 60, 0, 0)   
         onlineplot!(scene, layout, axis_dict, plots, iter, data, variable, i)
     end
     on(iter) do i
         if i > 1 # To deal with autolimits a certain number of samples are needed
-            for (variable, plots) in vars
-                for p in plots
+            for variable in keys(vars)
+                for p in vars[variable]
                     autolimits!(axis_dict[(variable, p)])
                 end
             end
@@ -97,12 +110,10 @@ function (cb::TurkieCallback)(rng, model, sampler, transition, iteration)
         end
         cb.params[:t0] = cb.iter[] 
     end
-    fit!(cb.data[:iter], iteration + cb.params[:t0])
-    for (vals, ks) in values(transition.θ)
-        for (k, val) in zip(ks, vals)
-            if haskey(cb.data, Symbol(k))
-                fit!(cb.data[Symbol(k)], Float32(val))
-            end
+    fit!(cb.data[:iter], iteration + cb.params[:t0]) # Update the iteration value
+    for (variable, val) in zip(_params_to_array([transition])...)
+        if haskey(cb.data, variable) # Check if symbol should be plotted
+            fit!(cb.data[variable], Float32(val)) # Update its value
         end
     end
     cb.iter[] += 1

src/online_stats_plots.jl

@@ -24,11 +24,11 @@ onlineplot!(axis, ::Val{:hist}, args...) = onlineplot!(axis, KHist(50, Float32),
 function onlineplot!(axis, stat::T, iter, data, iterations, i, j) where {T<:OnlineStat}
     window = data.b
     @eval TStat = $(nameof(T))
-    stat = Node(TStat(Float32))
+    stat = Observable(TStat(Float32))
     on(iter) do i
         stat[] = fit!(stat[], last(value(data)))
     end
-    statvals = Node(MovingWindow(window, Float32))
+    statvals = Observable(MovingWindow(window, Float32))
     on(stat) do s
         statvals[] = fit!(statvals[], Float32(value(s)))
     end
@@ -39,15 +39,15 @@ function onlineplot!(axis, stat::T, iter, data, iterations, i, j) where {T<:Onli
 end
 
 function onlineplot!(axis, ::Val{:trace}, iter, data, iterations, i, j)
-    trace = lift(iter; init = [Point2f0(0, 0f0)]) do i
+    trace = lift(iter; init = [Point2f0(0f0, 0f0)]) do i
         Point2f0.(value(iterations), value(data))
     end
     lines!(axis, trace, color = std_colors[i]; linewidth = 3.0)
 end
 
 function onlineplot!(axis, stat::KHist, iter, data, iterations, i, j)
     nbins = stat.k
-    stat = Node(KHist(nbins, Float32))
+    stat = Observable(KHist(nbins, Float32))
     on(iter) do i
         stat[] = fit!(stat[], last(value(data)))
     end
@@ -68,7 +68,7 @@ function expand_extrema(xs)
 end
 
 function onlineplot!(axis, ::Val{:kde}, iter, data, iterations, i, j)
-    interpkde = Node(InterpKDE(kde([1f0])))
+    interpkde = Observable(InterpKDE(kde([1f0])))
     on(iter) do i
         interpkde[] = InterpKDE(kde(value(data)))
     end
@@ -90,7 +90,7 @@ end
 
 function onlineplot!(axis, stat::AutoCov, iter, data, iterations, i, j)
     b = length(stat.cross)
-    stat = Node(AutoCov(b, Float32))
+    stat = Observable(AutoCov(b, Float32))
     on(iter) do i
         stat[] = fit!(stat[], last(value(data)))
     end

test/dev_test.jl

@@ -1,7 +1,9 @@
 using Turing
 using Turkie
 using GLMakie # You could also use CairoMakie or another backend
+using CairoMakie
 GLMakie.activate!()
+CairoMakie.activate!()
 Turing.@model function demo(x) # Some random Turing model
     m0 ~ Normal(0, 2)
     s ~ InverseGamma(2, 3)
@@ -38,4 +40,27 @@ end
 cb = TurkieCallback(Dict(:m0 => [:trace, :mean],
                         :s => [:autocov, :var]))
 
-advancedHMC(sossdemo(), (x=xs,), 100; callback = cb)
+advancedHMC(sossdemo(), (x=xs,), 100; callback = cb)
+
+## Test for array of parameters
+using LinearAlgebra
+D = 1
+N = 20
+Turing.@model function vectordemo(x, y, σ)
+    m ~ Normal(0, 10)
+    β ~ MvNormal(m * ones(D + 1), ones(D + 1))
+    for i in eachindex(y)
+        y[i] ~ Normal(dot(β, vcat(1, x[i])), σ)
+    end
+end
+
+x = [rand(D) for _ in 1:N]
+β = randn(D + 1) * 2
+σ = 0.1
+y = dot.(Ref(β), vcat.(1, x)) .+ σ * randn(N)
+
+m = vectordemo(x, y, σ)
+
+cb = TurkieCallback(m) # Create a callback function to be given to the sample function
+
+chain = sample(m, NUTS(0.65), 200; callback = cb)