From 446f06d07ce7475bd1c413568219bdc56bfd3e47 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Mon, 17 Jul 2023 19:44:01 +0100
Subject: [PATCH] Added missing `posterior_mean` and `rand` for test models
 (#499)

* added missing posterior_mean and rand implementation, in simple checks
to make sure these things do not go unnoticied

* bump patch version

* Update src/test_utils.jl

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

* Update test/model.jl

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

* fixed tests

* fixed MatrixvariateDemoModels

* Apply suggestions from code review

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

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
---
 Project.toml      |  2 +-
 src/test_utils.jl | 30 ++++++++++++++++++++++++++++++
 test/model.jl     | 22 ++++++++++++++++++++++
 3 files changed, 53 insertions(+), 1 deletion(-)

diff --git a/Project.toml b/Project.toml
index 4c1d5c31b..0585d3856 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,6 +1,6 @@
 name = "DynamicPPL"
 uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
-version = "0.23.3"
+version = "0.23.4"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
diff --git a/src/test_utils.jl b/src/test_utils.jl
index 3f6f4274e..cc2d27803 100644
--- a/src/test_utils.jl
+++ b/src/test_utils.jl
@@ -669,6 +669,36 @@ function Random.rand(
     return vals
 end
 
+const MatrixvariateAssumeDemoModels = Union{
+    Model{typeof(demo_assume_matrix_dot_observe_matrix)}
+}
+function posterior_mean(model::MatrixvariateAssumeDemoModels)
+    # Get some containers to fill.
+    vals = Random.rand(model)
+
+    vals.s[1, 1] = 19 / 8
+    vals.m[1] = 3 / 4
+
+    vals.s[1, 2] = 8 / 3
+    vals.m[2] = 1
+
+    return vals
+end
+function Base.rand(
+    rng::Random.AbstractRNG, ::Type{NamedTuple}, model::MatrixvariateAssumeDemoModels
+)
+    # Get template values from `model`.
+    retval = model(rng)
+    vals = (s=retval.s, m=retval.m)
+    # Fill containers with realizations from prior.
+    for i in LinearIndices(vals.s)
+        vals.s[i] = rand(rng, InverseGamma(2, 3))
+        vals.m[i] = rand(rng, Normal(0, sqrt(vals.s[i])))
+    end
+
+    return vals
+end
+
 """
 A collection of models corresponding to the posterior distribution defined by
 the generative process
diff --git a/test/model.jl b/test/model.jl
index 85c44a025..b9ad9fc7a 100644
--- a/test/model.jl
+++ b/test/model.jl
@@ -185,4 +185,26 @@ end
             end
         end
     end
+
+    @testset "TestUtils" begin
+        @testset "$(model.f)" for model in DynamicPPL.TestUtils.DEMO_MODELS
+            x = rand(model)
+            # Ensure log-probability computations are implemented.
+            @test logprior(model, x) ≈ DynamicPPL.TestUtils.logprior_true(model, x...)
+            @test loglikelihood(model, x) ≈
+                DynamicPPL.TestUtils.loglikelihood_true(model, x...)
+            @test logjoint(model, x) ≈ DynamicPPL.TestUtils.logjoint_true(model, x...)
+            @test logjoint(model, x) !=
+                DynamicPPL.TestUtils.logjoint_true_with_logabsdet_jacobian(model, x...)
+            # Ensure `varnames` is implemented.
+            vi = last(
+                DynamicPPL.evaluate!!(
+                    model, SimpleVarInfo(OrderedDict()), SamplingContext()
+                ),
+            )
+            @test all(collect(keys(vi)) .== DynamicPPL.TestUtils.varnames(model))
+            # Ensure `posterior_mean` is implemented.
+            @test DynamicPPL.TestUtils.posterior_mean(model) isa typeof(x)
+        end
+    end
 end