Updates to init (#489)

* fixed init for transformable distributions

* imrpove transformable a bit

* fixed tests

* added proper testing for init

* added coverage for matrix distributions

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* forgot to commit the rand implementation

* Update src/test_utils.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* fixed typo

* actually fixed it

* fixed new testutils model

* formatting

* bump Bijectors version so we can use output_size and output_length

* Apply suggestions from code review

Co-authored-by: David Widmann <[email protected]>

* Update test/Project.toml

Co-authored-by: David Widmann <[email protected]>

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Hong Ge <[email protected]>
Co-authored-by: David Widmann <[email protected]>

Author: 4 people

Project.toml

@@ -24,7 +24,7 @@ ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 AbstractMCMC = "2, 3.0, 4"
 AbstractPPL = "0.5.3"
 BangBang = "0.3"
-Bijectors = "0.12.4"
+Bijectors = "0.13"
 ChainRulesCore = "0.9.7, 0.10, 1"
 ConstructionBase = "1"
 Distributions = "0.23.8, 0.24, 0.25"

src/test_utils.jl

@@ -546,6 +546,53 @@ function varnames(model::Model{typeof(demo_dot_assume_matrix_dot_observe_matrix)
     return [@varname(s[:, 1]), @varname(s[:, 2]), @varname(m)]
 end
 
+@model function demo_assume_matrix_dot_observe_matrix(
+    x=transpose([1.5 2.0;]), ::Type{TV}=Array{Float64}
+) where {TV}
+    n = length(x)
+    d = n ÷ 2
+    s ~ reshape(product_distribution(fill(InverseGamma(2, 3), n)), d, 2)
+    s_vec = vec(s)
+    m ~ MvNormal(zeros(n), Diagonal(s_vec))
+
+    # Dotted observe for `Matrix`.
+    x .~ MvNormal(m, Diagonal(s_vec))
+
+    return (; s=s, m=m, x=x, logp=getlogp(__varinfo__))
+end
+function logprior_true(model::Model{typeof(demo_assume_matrix_dot_observe_matrix)}, s, m)
+    n = length(model.args.x)
+    s_vec = vec(s)
+    return loglikelihood(InverseGamma(2, 3), s_vec) +
+           logpdf(MvNormal(zeros(n), Diagonal(s_vec)), m)
+end
+function loglikelihood_true(
+    model::Model{typeof(demo_assume_matrix_dot_observe_matrix)}, s, m
+)
+    return loglikelihood(MvNormal(m, Diagonal(vec(s))), model.args.x)
+end
+function logprior_true_with_logabsdet_jacobian(
+    model::Model{typeof(demo_assume_matrix_dot_observe_matrix)}, s, m
+)
+    return _demo_logprior_true_with_logabsdet_jacobian(model, s, m)
+end
+function varnames(model::Model{typeof(demo_assume_matrix_dot_observe_matrix)})
+    return [@varname(s), @varname(m)]
+end
+
+function Random.rand(
+    rng::Random.AbstractRNG,
+    ::Type{NamedTuple},
+    model::Model{typeof(demo_assume_matrix_dot_observe_matrix)},
+)
+    n = length(model.args.x)
+    s = reshape(rand(rng, InverseGamma(2, 3), n), n ÷ 2, 2)
+    s_vec = vec(s)
+    m = rand(rng, MvNormal(zeros(n), Diagonal(s_vec)))
+
+    return (s=s, m=m)
+end
+
 const DemoModels = Union{
     Model{typeof(demo_dot_assume_dot_observe)},
     Model{typeof(demo_assume_index_observe)},
@@ -559,6 +606,7 @@ const DemoModels = Union{
     Model{typeof(demo_dot_assume_observe_submodel)},
     Model{typeof(demo_dot_assume_dot_observe_matrix)},
     Model{typeof(demo_dot_assume_matrix_dot_observe_matrix)},
+    Model{typeof(demo_assume_matrix_dot_observe_matrix)},
 }
 
 # We require demo models to have explict impleentations of `rand` since we want
@@ -668,6 +716,7 @@ const DEMO_MODELS = (
     demo_dot_assume_observe_submodel(),
     demo_dot_assume_dot_observe_matrix(),
     demo_dot_assume_matrix_dot_observe_matrix(),
+    demo_assume_matrix_dot_observe_matrix(),
 )
 
 # Model to test `StaticTransformation` with.

src/utils.jl

@@ -275,26 +275,24 @@ end
 randrealuni(rng::Random.AbstractRNG) = 4 * rand(rng) - 2
 randrealuni(rng::Random.AbstractRNG, args...) = 4 .* rand(rng, args...) .- 2
 
-const Transformable = Union{
-    PositiveDistribution,
-    UnitDistribution,
-    TransformDistribution,
-    SimplexDistribution,
-    PDMatDistribution,
-}
-istransformable(dist) = false
-istransformable(::Transformable) = true
+istransformable(dist) = link_transform(dist) !== identity
 
 #################################
 # Single-sample initialisations #
 #################################
 
 inittrans(rng, dist::UnivariateDistribution) = Bijectors.invlink(dist, randrealuni(rng))
 function inittrans(rng, dist::MultivariateDistribution)
-    return Bijectors.invlink(dist, randrealuni(rng, size(dist)[1]))
+    # Get the length of the unconstrained vector
+    b = link_transform(dist)
+    d = Bijectors.output_length(b, length(dist))
+    return Bijectors.invlink(dist, randrealuni(rng, d))
 end
 function inittrans(rng, dist::MatrixDistribution)
-    return Bijectors.invlink(dist, randrealuni(rng, size(dist)...))
+    # Get the size of the unconstrained vector
+    b = link_transform(dist)
+    sz = Bijectors.output_size(b, size(dist))
+    return Bijectors.invlink(dist, randrealuni(rng, sz...))
 end
 
 ################################

test/Project.toml

@@ -23,7 +23,7 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 [compat]
 AbstractMCMC = "2.1, 3.0, 4"
 AbstractPPL = "0.5"
-Bijectors = "0.11, 0.12"
+Bijectors = "0.13"
 Distributions = "0.25"
 DistributionsAD = "0.6.3"
 Documenter = "0.26.1, 0.27"

test/sampler.jl

@@ -24,12 +24,37 @@
         @test chains isa Vector{<:VarInfo}
         @test length(chains) == N
 
-        # Expected value of ``X`` where ``X ~ U[-2, 2]`` is ≈ 0.
-        @test mean(vi[@varname(m)] for vi in chains) ≈ 0 atol = 0.1
+        # `m` is Gaussian, i.e. no transformation is used, so it
+        # should have a mean equal to its prior, i.e. 2.
+        @test mean(vi[@varname(m)] for vi in chains) ≈ 2 atol = 0.1
 
         # Expected value of ``exp(X)`` where ``X ~ U[-2, 2]`` is ≈ 1.8.
         @test mean(vi[@varname(s)] for vi in chains) ≈ 1.8 atol = 0.1
     end
+
+    @testset "init" begin
+        @testset "$(model.f)" for model in DynamicPPL.TestUtils.DEMO_MODELS
+            @testset "$(model.f)" for model in DynamicPPL.TestUtils.DEMO_MODELS
+                N = 1000
+                chain_init = sample(model, SampleFromUniform(), N; progress=false)
+
+                for vn in keys(first(chain_init))
+                    if AbstractPPL.subsumes(@varname(s), vn)
+                        # `s ~ InverseGamma(2, 3)` and its unconstrained value will be sampled from Unif[-2,2].
+                        dist = InverseGamma(2, 3)
+                        b = DynamicPPL.link_transform(dist)
+                        @test mean(mean(b(vi[vn])) for vi in chain_init) ≈ 0 atol = 0.11
+                    elseif AbstractPPL.subsumes(@varname(m), vn)
+                        # `m ~ Normal(0, sqrt(s))` and its constrained value is the same.
+                        @test mean(mean(vi[vn]) for vi in chain_init) ≈ 0 atol = 0.11
+                    else
+                        error("Unknown variable name: $vn")
+                    end
+                end
+            end
+        end
+    end
+
     @testset "Initial parameters" begin
         # dummy algorithm that just returns initial value and does not perform any sampling
         abstract type OnlyInitAlg end