Shorter printing of Soss models in stack traces

[DilumAluthge]

Whenever the Soss or SossMLJ model types show up in stack traces, it is quite verbose, and makes the stack trace hard to read.

Here's an example. Look e.g. at [20], [22], [23], etc.

julia> particles = predict_particles(mach, Xnew; response = :y)
ERROR: MethodError: no method matching &(::Particles{Bool,1000}, ::Particles{Bool,1000})
Closest candidates are:
  &(::Any, ::Any, ::Any, ::Any...) at operators.jl:538
  &(::VectorizationBase.Static{M}, ::Number) where M at /Users/dilum/.julia/packages/VectorizationBase/kIoqa/src/static.jl:123
  &(::Number, ::VectorizationBase.Static{N}) where N at /Users/dilum/.julia/packages/VectorizationBase/kIoqa/src/static.jl:124
  ...
Stacktrace:
  [1] _broadcast_getindex_evalf
    @ ./broadcast.jl:648 [inlined]
  [2] _broadcast_getindex
    @ ./broadcast.jl:621 [inlined]
  [3] getindex
    @ ./broadcast.jl:575 [inlined]
  [4] copy
    @ ./broadcast.jl:876 [inlined]
  [5] materialize
    @ ./broadcast.jl:837 [inlined]
  [6] _check_probs_01(probs::Vector{Particles{Float64,1000}})
    @ MLJBase ~/.julia/packages/MLJBase/02f6S/src/univariate_finite/types.jl:66
  [7] _broadcast_getindex_evalf
    @ ./broadcast.jl:648 [inlined]
  [8] _broadcast_getindex
    @ ./broadcast.jl:621 [inlined]
  [9] getindex
    @ ./broadcast.jl:575 [inlined]
 [10] copy
    @ ./broadcast.jl:876 [inlined]
 [11] materialize
    @ ./broadcast.jl:837 [inlined]
 [12] UnivariateFinite(::MLJModelInterface.FullInterface, prob_given_class::OrderedCollections.LittleDict{CategoricalValue{String,UInt8},AbstractVector{Particles{Float64,1000}},Vector{CategoricalValue{String,UInt8}},Vector{AbstractVector{Particles{Float64,1000}}}}; kwargs::Base.Iterators.Pairs{Symbol,Any,Tuple{Symbol,Symbol},NamedTuple{(:pool, :ordered),Tuple{CategoricalPool{String,UInt8,CategoricalValue{String,UInt8}},Bool}}})
    @ MLJBase ~/.julia/packages/MLJBase/02f6S/src/univariate_finite/types.jl:127
 [13] _UnivariateFinite(support::CategoricalVector{String,UInt8,String,CategoricalValue{String,UInt8},Union{}}, probs::Matrix{Particles{Float64,1000}}, N::Int64; augment::Bool, kwargs::Base.Iterators.Pairs{Symbol,Any,Tuple{Symbol,Symbol},NamedTuple{(:pool, :ordered),Tuple{CategoricalPool{String,UInt8,CategoricalValue{String,UInt8}},Bool}}})
    @ MLJBase ~/.julia/packages/MLJBase/02f6S/src/univariate_finite/types.jl:244
 [14] #_UnivariateFinite#36
    @ ~/.julia/packages/MLJBase/02f6S/src/univariate_finite/types.jl:286 [inlined]
 [15] #_UnivariateFinite#39
    @ ~/.julia/packages/MLJBase/02f6S/src/univariate_finite/types.jl:295 [inlined]
 [16] UnivariateFinite(::MLJModelInterface.FullInterface, support::Vector{String}, probs::Matrix{Particles{Float64,1000}}; kwargs::Base.Iterators.Pairs{Symbol,CategoricalPool{String,UInt8,CategoricalValue{String,UInt8}},Tuple{Symbol},NamedTuple{(:pool,),Tuple{CategoricalPool{String,UInt8,CategoricalValue{String,UInt8}}}}})
    @ MLJBase ~/.julia/packages/MLJBase/02f6S/src/univariate_finite/types.jl:211
 [17] UnivariateFinite(support::Vector{String}, probs::Matrix{Particles{Float64,1000}}; kwargs::Base.Iterators.Pairs{Symbol,CategoricalPool{String,UInt8,CategoricalValue{String,UInt8}},Tuple{Symbol},NamedTuple{(:pool,),Tuple{CategoricalPool{String,UInt8,CategoricalValue{String,UInt8}}}}})
    @ MLJModelInterface ~/.julia/packages/MLJModelInterface/Zu24c/src/data_utils.jl:431
 [18] #UnivariateFinite#23
    @ ~/.julia/packages/MLJBase/02f6S/src/univariate_finite/types.jl:33 [inlined]
 [19] macro expansion
    @ ~/.julia/packages/GeneralizedGenerated/wp5nX/src/closure_conv.jl:121 [inlined]
 [20] _particles(#unused#::Type{TypeEncoding(Main)}, _m::Soss.Model{NamedTuple{(:X, :pool, :β),T} where T<:Tuple,TypeEncoding(begin
    η = X * β
    μ = NNlib.softmax(η; dims = 2)
    y_dists = UnivariateFinite(pool.levels, μ; pool = pool)
    n = size(X, 1)
    y ~ For((j->begin
                    y_dists[j]
                end), n)
end),TypeEncoding(Main)}, _args::NamedTuple{(:X, :pool, :β),Tuple{Matrix{Float64},CategoricalPool{String,UInt8,CategoricalValue{String,UInt8}},Matrix{Particles{Float64,1000}}}}, _n::Val{1000})
    @ Soss ~/.julia/packages/GeneralizedGenerated/wp5nX/src/closure_conv.jl:121
 [21] particles
    @ ~/.julia/packages/Soss/FgRzp/src/particles.jl:66 [inlined]
 [22] predict_particles(predictor::SossMLJ.SossMLJPredictor{SossMLJModel{UnivariateFinite,Soss.Model{NamedTuple{(:X, :pool),T} where T<:Tuple,TypeEncoding(begin
    k = length(pool.levels)
    p = size(X, 2)
    β ~ Normal(0.0, 1.0) |> iid(p, k)
    η = X * β
    μ = NNlib.softmax(η; dims = 2)
    y_dists = UnivariateFinite(pool.levels, μ; pool = pool)
    n = size(X, 1)
    y ~ For((j->begin
                    y_dists[j]
                end), n)
end),TypeEncoding(Main)},NamedTuple{(:pool,),Tuple{CategoricalPool{String,UInt8,CategoricalValue{String,UInt8}}}},typeof(dynamicHMC),Symbol,typeof(SossMLJ.default_transform)},Vector{NamedTuple{(:β,),Tuple{Matrix{Float64}}}},Soss.Model{NamedTuple{(:X, :pool, :β),T} where T<:Tuple,TypeEncoding(begin
    η = X * β
    μ = NNlib.softmax(η; dims = 2)
    y_dists = UnivariateFinite(pool.levels, μ; pool = pool)
    n = size(X, 1)
    y ~ For((j->begin
                    y_dists[j]
                end), n)
end),TypeEncoding(Main)},NamedTuple{(:X, :pool),Tuple{Matrix{Float64},CategoricalPool{String,UInt8,CategoricalValue{String,UInt8}}}}}, Xnew::DataFrame)
    @ SossMLJ ~/Downloads/SossMLJ.jl/src/particles.jl:12
 [23] predict_particles(sm::SossMLJModel{UnivariateFinite,Soss.Model{NamedTuple{(:X, :pool),T} where T<:Tuple,TypeEncoding(begin
    k = length(pool.levels)
    p = size(X, 2)
    β ~ Normal(0.0, 1.0) |> iid(p, k)
    η = X * β
    μ = NNlib.softmax(η; dims = 2)
    y_dists = UnivariateFinite(pool.levels, μ; pool = pool)
    n = size(X, 1)
    y ~ For((j->begin
                    y_dists[j]
                end), n)
end),TypeEncoding(Main)},NamedTuple{(:pool,),Tuple{CategoricalPool{String,UInt8,CategoricalValue{String,UInt8}}}},typeof(dynamicHMC),Symbol,typeof(SossMLJ.default_transform)}, fitresult::NamedTuple{(:model, :post),Tuple{SossMLJModel{UnivariateFinite,Soss.Model{NamedTuple{(:X, :pool),T} where T<:Tuple,TypeEncoding(begin
    k = length(pool.levels)
    p = size(X, 2)
    β ~ Normal(0.0, 1.0) |> iid(p, k)
    η = X * β
    μ = NNlib.softmax(η; dims = 2)
    y_dists = UnivariateFinite(pool.levels, μ; pool = pool)
    n = size(X, 1)
    y ~ For((j->begin
                    y_dists[j]
                end), n)
end),TypeEncoding(Main)},NamedTuple{(:pool,),Tuple{CategoricalPool{String,UInt8,CategoricalValue{String,UInt8}}}},typeof(dynamicHMC),Symbol,typeof(SossMLJ.default_transform)},Vector{NamedTuple{(:β,),Tuple{Matrix{Float64}}}}}}, Xnew::DataFrame; response::Symbol)
    @ SossMLJ ~/Downloads/SossMLJ.jl/src/particles.jl:20
 [24] predict_particles(mach::Machine{SossMLJModel{UnivariateFinite,Soss.Model{NamedTuple{(:X, :pool),T} where T<:Tuple,TypeEncoding(begin
    k = length(pool.levels)
    p = size(X, 2)
    β ~ Normal(0.0, 1.0) |> iid(p, k)
    η = X * β
    μ = NNlib.softmax(η; dims = 2)
    y_dists = UnivariateFinite(pool.levels, μ; pool = pool)
    n = size(X, 1)
    y ~ For((j->begin
                    y_dists[j]
                end), n)
end),TypeEncoding(Main)},NamedTuple{(:pool,),Tuple{CategoricalPool{String,UInt8,CategoricalValue{String,UInt8}}}},typeof(dynamicHMC),Symbol,typeof(SossMLJ.default_transform)}}, Xraw::DataFrame; kwargs::Base.Iterators.Pairs{Symbol,Symbol,Tuple{Symbol},NamedTuple{(:response,),Tuple{Symbol}}})
    @ SossMLJ ~/Downloads/SossMLJ.jl/src/machine-operations.jl:10
 [25] top-level scope
    @ REPL[25]:1

I wonder if we could make these type names more concise, at least when they appear in stack traces.

[cscherrer]

That would be great! I guess this involves changing how the type is displayed, right? Any suggestions on a clean representation?

[DilumAluthge]

The big part, I think, will be to not print the model itself (all the ~ statements) when printing the model type.

Also, we could shorten NamedTuple{...} to just NamedTuple.

[millerjoey]

I was thinking this too recently and it's easier than I thought to define custom show for types:

using Soss

m = @model μ begin 
    x ~ For(3) do i
        Normal(μ)
    end
    y = sum(x)
end

import Base.show

function show(io::IO, t::Type{Model{A,B,M}}) where {A,B,M}
    m = Soss.type2model(t)
    args = join((v for v in arguments(m)), ",")
    vars = join((v for v in parameters(m)), ",")
    print(io, "Model: ", args, " -> ", "JointDistribution on $vars")
end

function show(io::IO, t::Type{Soss.JointDistribution{A0,A,B,M}}) where {A0,A,B,M}
    m = type2model(t)
    vars = join((v for v in parameters(m)), ",")
    print(io, "JointDistribution on ", vars)
end

function type2model(::Type{Soss.JointDistribution{A0,A,B,M}}) where {A0,A,B,M}
    args = [Soss.fieldnames(A)...]
    body = Soss.from_type(B)
    Model(Soss.from_type(M), convert(Vector{Symbol},args), body)
end

then

julia> typeof(m)
Model: μ -> JointDistribution on y,x

julia> typeof(m(μ=2))
JointDistribution on y,x

julia> logcdf(m(μ=2))
ERROR: LoadError: MethodError: no method matching logcdf(::JointDistribution on y,x)
Closest candidates are:
  logcdf(::Hypergeometric, ::Int64) at /Users/jx1525/.julia/packages/Distributions/jFoHB/src/univariates.jl:564
  logcdf(::Binomial, ::Int64) at /Users/jx1525/.julia/packages/Distributions/jFoHB/src/univariates.jl:564
  logcdf(::Geometric{T}, ::Int64) where T<:Real at /Users/jx1525/.julia/packages/Distributions/jFoHB/src/univariate/discrete/geometric.jl:115
  ...
Stacktrace:
 [1] top-level scope at /Users/jx1525/Desktop/Soss.jl/scratch.jl:34
 [2] include_string(::Function, ::Module, ::String, ::String) at ./loading.jl:1088
in expression starting at /Users/jx1525/Desktop/Soss.jl/scratch.jl:34

How can we make that better? We may not need it for Models. I think just showing the variables (and maybe dimensions of variables) should be specific enough.

[cscherrer]

How about context dependence, like the way Vector{Bool}s are displayed?

julia> true
true

julia> [true, false, false]
3-element Array{Bool,1}:
 1
 0
 0

julia> [true, false, false][1]
true

[millerjoey]

From https://docs.julialang.org/en/v1/manual/types/#man-custom-pretty-printing It seems like that's supposed to be done with defining the extra MIME argument

function Base.show(io::IO, ::MIME"text/plain", t::Type{Soss.JointDistribution{A0,A,B,M}}) where {A0,A,B,M} 
    print(io, "Verbose goes here")
end

function show(io::IO, t::Type{Model{A,B,M}}) where {A,B,M}
    m = Soss.type2model(t)
    args = join((v for v in arguments(m)), ",")
    vars = join((v for v in parameters(m)), ",")
    print(io, "Model: ", args, " -> ", "JointDistribution on $vars")
end

I think I messed it up though, since it's not working. Also, we'd probably want to use the default type show method for the verbose one. I'll give it some more time later and see if I can figure it out.