From 330972fa762192e26a3cb7a04c4306af588c14d6 Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Sun, 22 Mar 2026 02:56:52 +0000 Subject: [PATCH 1/8] Implement factorised pointwise probabilities --- Project.toml | 6 +- src/accumulators/pointwise_logdensities.jl | 80 ++++++++++++++++------ 2 files changed, 64 insertions(+), 22 deletions(-) diff --git a/Project.toml b/Project.toml index 36a590677..efb1c474e 100644 --- a/Project.toml +++ b/Project.toml @@ -21,6 +21,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09" OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d" +PartitionedDistributions = "569bd051-8d7b-4221-bcb8-d78512b5866a" PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" Preferences = "21216c6a-2e73-6563-6e65-726566657250" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" @@ -40,8 +41,8 @@ ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" [extensions] DynamicPPLEnzymeCoreExt = ["EnzymeCore"] DynamicPPLForwardDiffExt = ["ForwardDiff"] -DynamicPPLMarginalLogDensitiesExt = ["MarginalLogDensities"] DynamicPPLMCMCChainsExt = ["MCMCChains"] +DynamicPPLMarginalLogDensitiesExt = ["MarginalLogDensities"] DynamicPPLMooncakeExt = ["Mooncake", "DifferentiationInterface"] DynamicPPLReverseDiffExt = ["ReverseDiff"] @@ -65,11 +66,12 @@ InteractiveUtils = "1" KernelAbstractions = "0.9.33" LinearAlgebra = "1.6" LogDensityProblems = "2" -MarginalLogDensities = "0.4.3" MCMCChains = "6, 7" MacroTools = "0.5.6" +MarginalLogDensities = "0.4.3" Mooncake = "0.4.147, 0.5" OrderedCollections = "1" +PartitionedDistributions = "0.0.1" PrecompileTools = "1.2.1" Preferences = "1.5.2" Printf = "1.10" diff --git a/src/accumulators/pointwise_logdensities.jl b/src/accumulators/pointwise_logdensities.jl index 9e8a9765c..2050c51cb 100644 --- a/src/accumulators/pointwise_logdensities.jl +++ b/src/accumulators/pointwise_logdensities.jl @@ -1,3 +1,20 @@ +import PartitionedDistributions + +# Force specialisation on D and V. +function _maybe_pointwise_logpdf(dist::D, value::V, ::Val{true}) where {D<:Distribution,V} + return if hasmethod( + PartitionedDistributions.pointwise_conditional_logpdfs, + Tuple{typeof(value),typeof(dist)}, + ) + PartitionedDistributions.pointwise_conditional_logpdfs(value, dist) + else + logpdf(dist, value) + end +end +function _maybe_pointwise_logpdf(dist::Distribution, value, ::Val{false}) + return logpdf(dist, value) +end + """ PointwiseLogProb{Prior,Likelihood} @@ -12,13 +29,13 @@ This struct is used in conjunction with `VNTAccumulator`, via where `Prior` and `Likelihood` are the boolean type parameters. This accumulator will then store the log-probabilities for all tilde-statements in the model. """ -struct PointwiseLogProb{Prior,Likelihood} end +struct PointwiseLogProb{Prior,Likelihood,Factorised} end Base.copy(plp::PointwiseLogProb) = plp -function (plp::PointwiseLogProb{Prior,Likelihood})( +function (plp::PointwiseLogProb{Prior,Likelihood,Factorised})( val, tval, logjac, vn, dist -) where {Prior,Likelihood} - if Prior - return logpdf(dist, val) +) where {Prior,Likelihood,Factorised} + return if Prior + _maybe_pointwise_logpdf(dist, val, Val{Factorised}()) else return DoNotAccumulate() end @@ -33,15 +50,15 @@ end # Have to overload accumulate_assume!! since VNTAccumulator by default does not track # observe statements. function accumulate_observe!!( - acc::VNTAccumulator{POINTWISE_ACCNAME,PointwiseLogProb{Prior,Likelihood}}, + acc::VNTAccumulator{POINTWISE_ACCNAME,PointwiseLogProb{Prior,Likelihood,Factorised}}, right, left, vn, template, -) where {Prior,Likelihood} +) where {Prior,Likelihood,Factorised} # vn could be `nothing`, in which case we can't store it in a VNT. return if Likelihood && vn isa VarName - logp = logpdf(right, left) + logp = _maybe_pointwise_logpdf(right, left, Val{Factorised}()) new_values = DynamicPPL.templated_setindex!!(acc.values, logp, vn, template) return VNTAccumulator{POINTWISE_ACCNAME}(acc.f, new_values) else @@ -55,7 +72,8 @@ end model::Model, init_strat::AbstractInitStrategy, ::Val{Prior}=Val(true), - ::Val{Likelihood}=Val(true), + ::Val{Likelihood}=Val(true); + factorize=false ) where {Prior,Likelihood} Shared internal function that computes pointwise log-densities (either priors, likelihoods, @@ -65,9 +83,10 @@ function _pointwise_logdensities( model::Model, init_strat::AbstractInitStrategy, ::Val{Prior}=Val(true), - ::Val{Likelihood}=Val(true), + ::Val{Likelihood}=Val(true); + factorize=false, ) where {Prior,Likelihood} - acc = VNTAccumulator{POINTWISE_ACCNAME}(PointwiseLogProb{Prior,Likelihood}()) + acc = VNTAccumulator{POINTWISE_ACCNAME}(PointwiseLogProb{Prior,Likelihood,factorize}()) oavi = OnlyAccsVarInfo(acc) oavi = last(init!!(model, oavi, init_strat, UnlinkAll())) return get_pointwise_logprobs(oavi) @@ -76,38 +95,59 @@ end """ DynamicPPL.pointwise_logdensities( model::Model, - init_strat::AbstractInitStrategy + init_strat::AbstractInitStrategy; + factorize=false ) Calculate the pointwise log-densities for the parameters obtained by evaluating the model with the given initialisation strategy. The resulting VarNamedTuple will contain both log-prior probabilities (for random variables) and log-likelihoods (for observed variables). + +If `factorize=true`, additionally attempt to provide factorised log-densities for +distributions that can be partitioned into blocks, using PartitionedDistributions.jl. For +example, if `factorize=true`, then `y ~ MvNormal(...)` will return a vector of +log-densities, one for each element of `y`. If `factorize=false`, then the log-density for +`y ~ MvNormal(...)` will be a single scalar. """ -function pointwise_logdensities(model::Model, init_strat::AbstractInitStrategy) - return _pointwise_logdensities(model, init_strat, Val(true), Val(true)) +function pointwise_logdensities( + model::Model, init_strat::AbstractInitStrategy; factorize=false +) + return _pointwise_logdensities( + model, init_strat, Val(true), Val(true); factorize=factorize + ) end """ DynamicPPL.pointwise_loglikelihoods( model::Model, - init_strat::AbstractInitStrategy + init_strat::AbstractInitStrategy; + factorize=false ) Same as `pointwise_logdensities`, but only returns the log-likelihoods for observed variables. """ -function pointwise_loglikelihoods(model::Model, init_strat::AbstractInitStrategy) - return _pointwise_logdensities(model, init_strat, Val(false), Val(true)) +function pointwise_loglikelihoods( + model::Model, init_strat::AbstractInitStrategy; factorize=false +) + return _pointwise_logdensities( + model, init_strat, Val(false), Val(true); factorize=factorize + ) end """ DynamicPPL.pointwise_prior_logdensities( model::Model, - init_strat::AbstractInitStrategy + init_strat::AbstractInitStrategy; + factorize=false ) Same as `pointwise_logdensities`, but only returns the log-densities for random variables (i.e. the priors). """ -function pointwise_prior_logdensities(model::Model, init_strat::AbstractInitStrategy) - return _pointwise_logdensities(model, init_strat, Val(true), Val(false)) +function pointwise_prior_logdensities( + model::Model, init_strat::AbstractInitStrategy; factorize=false +) + return _pointwise_logdensities( + model, init_strat, Val(true), Val(false); factorize=factorize + ) end From 26032e7186c8d03508271c163faa1ef042d99355 Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Wed, 22 Apr 2026 14:51:14 +0100 Subject: [PATCH 2/8] bump patch & add to MCMCChainsExt as well --- HISTORY.md | 5 +++++ Project.toml | 2 +- ext/DynamicPPLMCMCChainsExt.jl | 38 ++++++++++++++++++++++++---------- 3 files changed, 33 insertions(+), 12 deletions(-) diff --git a/HISTORY.md b/HISTORY.md index 55bdd406d..1cff2d364 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -1,3 +1,8 @@ +# 0.41.4 + +Add a `factorize::Bool` keyword argument for `pointwise_logdensities(model, values)`, which controls whether pointwise logdensities for factorisable distributions (e.g. `MvNormal`, `product_distribution`, etc.) are returned as a single log-density for the whole distribution, or as an array of log-densities for each factor. +The same argument is also added to `pointwise_loglikelihoods` and `pointwise_prior_logdensities`. + # 0.41.3 Add a lower-level constructor for `LogDensityFunction` which directly takes a VNT of `RangeAndTransform`s plus a sample vectorised input. diff --git a/Project.toml b/Project.toml index efb1c474e..f2708cb99 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "DynamicPPL" uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8" -version = "0.41.3" +version = "0.41.4" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" diff --git a/ext/DynamicPPLMCMCChainsExt.jl b/ext/DynamicPPLMCMCChainsExt.jl index d7987ffa5..1b5f7b3eb 100644 --- a/ext/DynamicPPLMCMCChainsExt.jl +++ b/ext/DynamicPPLMCMCChainsExt.jl @@ -422,10 +422,11 @@ function _pointwise_logdensities_chain( model::DynamicPPL.Model, chain::MCMCChains.Chains, ::Val{Prior}=Val(true), - ::Val{Likelihood}=Val(true), + ::Val{Likelihood}=Val(true); + factorize=false, ) where {Prior,Likelihood} acc = DynamicPPL.VNTAccumulator{DynamicPPL.POINTWISE_ACCNAME}( - DynamicPPL.PointwiseLogProb{Prior,Likelihood}() + DynamicPPL.PointwiseLogProb{Prior,Likelihood,factorize}() ) parameter_only_chain = MCMCChains.get_sections(chain, :parameters) # Reevaluating this gives us a VNT of log probs. We can densify and then wrap in @@ -443,12 +444,19 @@ end """ DynamicPPL.pointwise_logdensities( model::DynamicPPL.Model, - chain::MCMCChains.Chains, + chain::MCMCChains.Chains; + factorize=false, ) Runs `model` on each sample in `chain`, returning a new `MCMCChains.Chains` object where the log-density of each variable at each sample is stored (rather than its value). +If `factorize=true`, additionally attempt to provide factorised log-densities for +distributions that can be partitioned into blocks, using PartitionedDistributions.jl. For +example, if `factorize=true`, then `y ~ MvNormal(...)` will return a vector of +log-densities, one for each element of `y`. If `factorize=false`, then the log-density for +`y ~ MvNormal(...)` will be a single scalar. + See also: [`DynamicPPL.pointwise_loglikelihoods`](@ref), [`DynamicPPL.pointwise_prior_logdensities`](@ref). @@ -507,15 +515,18 @@ julia> # The above is the same as: ``` """ function DynamicPPL.pointwise_logdensities( - model::DynamicPPL.Model, chain::MCMCChains.Chains + model::DynamicPPL.Model, chain::MCMCChains.Chains; factorize=false ) - return _pointwise_logdensities_chain(model, chain, Val(true), Val(true)) + return _pointwise_logdensities_chain( + model, chain, Val(true), Val(true); factorize=factorize + ) end """ DynamicPPL.pointwise_loglikelihoods( model::DynamicPPL.Model, - chain::MCMCChains.Chains, + chain::MCMCChains.Chains; + factorize=false, ) Compute the pointwise log-likelihoods of the model given the chain. This is the same as @@ -524,15 +535,18 @@ Compute the pointwise log-likelihoods of the model given the chain. This is the See also: [`DynamicPPL.pointwise_logdensities`](@ref), [`DynamicPPL.pointwise_prior_logdensities`](@ref). """ function DynamicPPL.pointwise_loglikelihoods( - model::DynamicPPL.Model, chain::MCMCChains.Chains + model::DynamicPPL.Model, chain::MCMCChains.Chains; factorize=false ) - return _pointwise_logdensities_chain(model, chain, Val(false), Val(true)) + return _pointwise_logdensities_chain( + model, chain, Val(false), Val(true); factorize=factorize + ) end """ DynamicPPL.pointwise_prior_logdensities( model::DynamicPPL.Model, - chain::MCMCChains.Chains + chain::MCMCChains.Chain; + factorize=false, ) Compute the pointwise log-prior-densities of the model given the chain. This is the same as @@ -541,9 +555,11 @@ Compute the pointwise log-prior-densities of the model given the chain. This is See also: [`DynamicPPL.pointwise_logdensities`](@ref), [`DynamicPPL.pointwise_loglikelihoods`](@ref). """ function DynamicPPL.pointwise_prior_logdensities( - model::DynamicPPL.Model, chain::MCMCChains.Chains + model::DynamicPPL.Model, chain::MCMCChains.Chains; factorize=false ) - return _pointwise_logdensities_chain(model, chain, Val(true), Val(false)) + return _pointwise_logdensities_chain( + model, chain, Val(true), Val(false); factorize=factorize + ) end """ From b48dd3f985583db337dbb969e1b1008f15763b40 Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Wed, 22 Apr 2026 14:58:47 +0100 Subject: [PATCH 3/8] Fix method signature --- src/accumulators/pointwise_logdensities.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/accumulators/pointwise_logdensities.jl b/src/accumulators/pointwise_logdensities.jl index 2050c51cb..2670d5aab 100644 --- a/src/accumulators/pointwise_logdensities.jl +++ b/src/accumulators/pointwise_logdensities.jl @@ -4,9 +4,9 @@ import PartitionedDistributions function _maybe_pointwise_logpdf(dist::D, value::V, ::Val{true}) where {D<:Distribution,V} return if hasmethod( PartitionedDistributions.pointwise_conditional_logpdfs, - Tuple{typeof(value),typeof(dist)}, + Tuple{typeof(dist),typeof(value)}, ) - PartitionedDistributions.pointwise_conditional_logpdfs(value, dist) + PartitionedDistributions.pointwise_conditional_logpdfs(dist, value) else logpdf(dist, value) end From 52f9c2e6a65e878f5983cfa37150c1d08b2e236d Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Wed, 22 Apr 2026 15:08:31 +0100 Subject: [PATCH 4/8] Fix test --- test/accumulators.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/accumulators.jl b/test/accumulators.jl index 1e209e4c5..c7a1d8a50 100644 --- a/test/accumulators.jl +++ b/test/accumulators.jl @@ -33,7 +33,7 @@ TEST_ACCUMULATORS = ( DynamicPPL.DebugRawValueAccumulator(), DynamicPPL.FixedTransformAccumulator(), DynamicPPL.VNTAccumulator{DynamicPPL.POINTWISE_ACCNAME}( - DynamicPPL.PointwiseLogProb{true,true}() + DynamicPPL.PointwiseLogProb{true,true,false}() ), PriorDistributionAccumulator(), DynamicPPL.VectorValueAccumulator(), From 43e9e1a3a515da73c64a83fb2218905aa782e6fd Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Wed, 22 Apr 2026 17:12:09 +0100 Subject: [PATCH 5/8] Add tests and improve docs --- src/accumulators/pointwise_logdensities.jl | 38 ++++++++--- test/pointwise_logdensities.jl | 79 +++++++++++++++++++++- 2 files changed, 103 insertions(+), 14 deletions(-) diff --git a/src/accumulators/pointwise_logdensities.jl b/src/accumulators/pointwise_logdensities.jl index 2670d5aab..73cb3c02a 100644 --- a/src/accumulators/pointwise_logdensities.jl +++ b/src/accumulators/pointwise_logdensities.jl @@ -15,6 +15,21 @@ function _maybe_pointwise_logpdf(dist::Distribution, value, ::Val{false}) return logpdf(dist, value) end +const _FACTORIZE_KWARG_DOC = """ +If `factorize=true`, additionally attempt to provide factorised log-densities for +distributions that can be partitioned into blocks, using PartitionedDistributions.jl. For +example, if `factorize=true`, then `y ~ MvNormal(...)` will return a vector of +log-densities, one for each element of `y`. If `factorize=false`, then the log-density for +`y ~ MvNormal(...)` will be a single scalar. + +Note that the sum of the factorised log-densities may not, in general, be equal to the +log-density of the full distribution: they will only be equal if the original distribution +can be completely factorised into independent components. For example, if `y ~ MvNormal(μ, +Σ)` where `Σ` is diagonal, then each element of `y` is independent and the sum of the +factorised log-densities will be equal to the log-density of the full distribution. In +contrast, if `Σ` has off-diagonal entries, then the elements of `y` are not independent. +""" + """ PointwiseLogProb{Prior,Likelihood} @@ -97,17 +112,13 @@ end model::Model, init_strat::AbstractInitStrategy; factorize=false - ) + )::VarNamedTuple Calculate the pointwise log-densities for the parameters obtained by evaluating the model with the given initialisation strategy. The resulting VarNamedTuple will contain both log-prior probabilities (for random variables) and log-likelihoods (for observed variables). -If `factorize=true`, additionally attempt to provide factorised log-densities for -distributions that can be partitioned into blocks, using PartitionedDistributions.jl. For -example, if `factorize=true`, then `y ~ MvNormal(...)` will return a vector of -log-densities, one for each element of `y`. If `factorize=false`, then the log-density for -`y ~ MvNormal(...)` will be a single scalar. +$(_FACTORIZE_KWARG_DOC) """ function pointwise_logdensities( model::Model, init_strat::AbstractInitStrategy; factorize=false @@ -122,9 +133,12 @@ end model::Model, init_strat::AbstractInitStrategy; factorize=false - ) + )::VarNamedTuple + +Calculate the pointwise log-likelihoods for observed variables, using parameters obtained +from the given initialisation strategy. -Same as `pointwise_logdensities`, but only returns the log-likelihoods for observed variables. +$(_FACTORIZE_KWARG_DOC) """ function pointwise_loglikelihoods( model::Model, init_strat::AbstractInitStrategy; factorize=false @@ -139,10 +153,12 @@ end model::Model, init_strat::AbstractInitStrategy; factorize=false - ) + )::VarNamedTuple + +Calculate the pointwise log-prior probabilities for random variables, using parameters +obtained from the given initialisation strategy. -Same as `pointwise_logdensities`, but only returns the log-densities for random variables -(i.e. the priors). +$(_FACTORIZE_KWARG_DOC) """ function pointwise_prior_logdensities( model::Model, init_strat::AbstractInitStrategy; factorize=false diff --git a/test/pointwise_logdensities.jl b/test/pointwise_logdensities.jl index fe79a622e..91d01405a 100644 --- a/test/pointwise_logdensities.jl +++ b/test/pointwise_logdensities.jl @@ -6,8 +6,9 @@ __now__ = now() using AbstractMCMC: AbstractMCMC using AbstractPPL: AbstractPPL -using Distributions: Normal, Exponential +using Distributions using DynamicPPL +using LinearAlgebra using MCMCChains: MCMCChains using Random: Random using Test @@ -27,7 +28,7 @@ function make_chain_from_prior(model::Model, n_iters::Int) return make_chain_from_prior(Random.default_rng(), model, n_iters) end -@testset "pointwise_logdensities.jl" begin +@testset "pointwise_logdensities with values" begin @testset "$(model.f)" for model in DynamicPPL.TestUtils.DEMO_MODELS example_values = DynamicPPL.TestUtils.rand_prior_true(model) @@ -62,9 +63,39 @@ end logp = sum(sum, values(lps)) @test logp ≈ (logprior_true + loglikelihood_true) end + + @testset "with factorize=true" begin + @model function f(z) + x ~ MvNormal(zeros(2), I) + y ~ product_distribution((; a=Normal(), b=Normal())) + return z ~ MvNormal(x, I) + end + model = f(randn(2)) + + logdensities = pointwise_logdensities(model, InitFromPrior(); factorize=true) + @test logdensities isa DynamicPPL.VarNamedTuple + @test logdensities[@varname(x)] isa Vector{Float64} + @test length(logdensities[@varname(x)]) == 2 + @test logdensities[@varname(z)] isa Vector{Float64} + @test length(logdensities[@varname(z)]) == 2 + + log_prior_densities = pointwise_prior_logdensities( + model, InitFromPrior(); factorize=true + ) + @test logdensities isa DynamicPPL.VarNamedTuple + @test log_prior_densities[@varname(x)] isa Vector{Float64} + @test length(log_prior_densities[@varname(x)]) == 2 + @test !haskey(log_prior_densities, @varname(z)) + + log_likelihoods = pointwise_loglikelihoods(model, InitFromPrior(); factorize=true) + @test log_likelihoods isa DynamicPPL.VarNamedTuple + @test log_likelihoods[@varname(z)] isa Vector{Float64} + @test length(log_likelihoods[@varname(z)]) == 2 + @test !haskey(log_likelihoods, @varname(x)) + end end -@testset "pointwise_logdensities chain" begin +@testset "pointwise_logdensities with chain" begin @testset "correctness" begin model = DynamicPPL.TestUtils.demo_assume_index_observe() vns = DynamicPPL.TestUtils.varnames(model) @@ -135,6 +166,48 @@ end model, chain_m_only ) end + + @testset "with factorize=true" begin + @model function f(z) + x ~ MvNormal(zeros(2), I) + # NOTE: This product_distribution doesn't work with MCMCChains because in the + # chain the elements of `y` are stored as `y.a` and `y.b`. When we call + # pointwise_logdensities(model, chn), DynamicPPL has to reconstruct the values + # and it doesn't realise that a and b should be part of the same distribution. + # Thus it recreates a VNT that looks like + # + # VarNamedTuple + # ├─ x => [-0.37054531061901747, -0.3947539032928444] + # └─ y => VarNamedTuple + # ├─ a => 0.015729054630598562 + # └─ b => -2.152129714223151 + # + # rather than + # + # VarNamedTuple + # ├─ x => [-0.14269319184188936, -1.6481452899699949] + # └─ y => (a = -0.5455556873253069, b = -1.0051579479025197) + # + # This is an inherent failure of MCMCChains and not something that we can fix + # here. + # y ~ product_distribution((; a=Normal(), b=Normal())) + return z ~ MvNormal(x, I) + end + model = f(randn(2)) + chn = make_chain_from_prior(model, 10) + logdensities = pointwise_logdensities(model, chn; factorize=true) + for k in [Symbol("x[1]"), Symbol("x[2]"), Symbol("z[1]"), Symbol("z[2]")] + @test k in keys(logdensities) + end + log_prior_densities = pointwise_prior_logdensities(model, chn; factorize=true) + for k in [Symbol("x[1]"), Symbol("x[2]")] + @test k in keys(logdensities) + end + log_likelihoods = pointwise_loglikelihoods(model, chn; factorize=true) + for k in [Symbol("z[1]"), Symbol("z[2]")] + @test k in keys(logdensities) + end + end end @info "Completed $(@__FILE__) in $(now() - __now__)." From 3035080553cfd8d612cfd75179889871b371d38e Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Wed, 22 Apr 2026 17:13:24 +0100 Subject: [PATCH 6/8] Add docstring to MCMCChainsExt methods as well --- ext/DynamicPPLMCMCChainsExt.jl | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/ext/DynamicPPLMCMCChainsExt.jl b/ext/DynamicPPLMCMCChainsExt.jl index 1b5f7b3eb..0535ce816 100644 --- a/ext/DynamicPPLMCMCChainsExt.jl +++ b/ext/DynamicPPLMCMCChainsExt.jl @@ -451,11 +451,7 @@ end Runs `model` on each sample in `chain`, returning a new `MCMCChains.Chains` object where the log-density of each variable at each sample is stored (rather than its value). -If `factorize=true`, additionally attempt to provide factorised log-densities for -distributions that can be partitioned into blocks, using PartitionedDistributions.jl. For -example, if `factorize=true`, then `y ~ MvNormal(...)` will return a vector of -log-densities, one for each element of `y`. If `factorize=false`, then the log-density for -`y ~ MvNormal(...)` will be a single scalar. +$(DynamicPPL._FACTORIZE_KWARG_DOC) See also: [`DynamicPPL.pointwise_loglikelihoods`](@ref), [`DynamicPPL.pointwise_prior_logdensities`](@ref). @@ -532,6 +528,8 @@ end Compute the pointwise log-likelihoods of the model given the chain. This is the same as `pointwise_logdensities(model, chain)`, but only including the likelihood terms. +$(DynamicPPL._FACTORIZE_KWARG_DOC) + See also: [`DynamicPPL.pointwise_logdensities`](@ref), [`DynamicPPL.pointwise_prior_logdensities`](@ref). """ function DynamicPPL.pointwise_loglikelihoods( @@ -552,6 +550,8 @@ end Compute the pointwise log-prior-densities of the model given the chain. This is the same as `pointwise_logdensities(model, chain)`, but only including the prior terms. +$(DynamicPPL._FACTORIZE_KWARG_DOC) + See also: [`DynamicPPL.pointwise_logdensities`](@ref), [`DynamicPPL.pointwise_loglikelihoods`](@ref). """ function DynamicPPL.pointwise_prior_logdensities( From 5943c1d28f65967a31d2ab43296f834a2c28b6b7 Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Thu, 23 Apr 2026 23:20:07 +0100 Subject: [PATCH 7/8] Explain more in the docstring --- src/accumulators/pointwise_logdensities.jl | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/src/accumulators/pointwise_logdensities.jl b/src/accumulators/pointwise_logdensities.jl index 73cb3c02a..3cebe904c 100644 --- a/src/accumulators/pointwise_logdensities.jl +++ b/src/accumulators/pointwise_logdensities.jl @@ -17,10 +17,16 @@ end const _FACTORIZE_KWARG_DOC = """ If `factorize=true`, additionally attempt to provide factorised log-densities for -distributions that can be partitioned into blocks, using PartitionedDistributions.jl. For -example, if `factorize=true`, then `y ~ MvNormal(...)` will return a vector of -log-densities, one for each element of `y`. If `factorize=false`, then the log-density for -`y ~ MvNormal(...)` will be a single scalar. +distributions that can be partitioned into blocks, using PartitionedDistributions.jl. + +For example, if `factorize=true`, then `y ~ MvNormal(...)` will return a vector of +log-densities, one for each element of `y`. The `i`-th element of this vector will be the +conditional log-probability of `y[i]` given all the other elements of `y` (often denoted +`log p(y_{i} | y_{-i})`): in particular this is exactly the log-density required for +leave-one-out cross-validation. + +In contrast, if `factorize=false`, then the log-density for `y ~ MvNormal(...)` will be a +single scalar corresponding to `logpdf(MvNormal(...), y)`. Note that the sum of the factorised log-densities may not, in general, be equal to the log-density of the full distribution: they will only be equal if the original distribution From 4ee4da3627ff7053d0535e0fb9b8d39bcf7bec5e Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Thu, 23 Apr 2026 23:33:17 +0100 Subject: [PATCH 8/8] Fix typos and other stuff --- ext/DynamicPPLMCMCChainsExt.jl | 2 +- src/accumulators/pointwise_logdensities.jl | 15 ++++++++------- test/pointwise_logdensities.jl | 10 ++++++---- 3 files changed, 15 insertions(+), 12 deletions(-) diff --git a/ext/DynamicPPLMCMCChainsExt.jl b/ext/DynamicPPLMCMCChainsExt.jl index 0535ce816..dbce4edd5 100644 --- a/ext/DynamicPPLMCMCChainsExt.jl +++ b/ext/DynamicPPLMCMCChainsExt.jl @@ -543,7 +543,7 @@ end """ DynamicPPL.pointwise_prior_logdensities( model::DynamicPPL.Model, - chain::MCMCChains.Chain; + chain::MCMCChains.Chains; factorize=false, ) diff --git a/src/accumulators/pointwise_logdensities.jl b/src/accumulators/pointwise_logdensities.jl index 3cebe904c..38f14254e 100644 --- a/src/accumulators/pointwise_logdensities.jl +++ b/src/accumulators/pointwise_logdensities.jl @@ -37,18 +37,19 @@ contrast, if `Σ` has off-diagonal entries, then the elements of `y` are not ind """ """ - PointwiseLogProb{Prior,Likelihood} + PointwiseLogProb{Prior,Likelihood,Factorised} A callable struct that computes the log probability of a given value under a distribution. The `Prior` and `Likelihood` type parameters are used to control whether the log probability -is computed for prior or likelihood terms, respectively. +is computed for prior or likelihood terms, respectively. The `Factorised` type parameter +controls whether to attempt to factorise the log-densities. This struct is used in conjunction with `VNTAccumulator`, via - acc = VNTAccumulator{POINTWISE_ACCNAME}(PointwiseLogProb{Prior,Likelihood}()) + acc = VNTAccumulator{POINTWISE_ACCNAME}(PointwiseLogProb{Prior,Likelihood,Factorised}()) -where `Prior` and `Likelihood` are the boolean type parameters. This accumulator will then -store the log-probabilities for all tilde-statements in the model. +where `Prior`, `Likelihood`, and `Factorised` are the boolean type parameters. This +accumulator will then store the log-probabilities for all tilde-statements in the model. """ struct PointwiseLogProb{Prior,Likelihood,Factorised} end Base.copy(plp::PointwiseLogProb) = plp @@ -58,7 +59,7 @@ function (plp::PointwiseLogProb{Prior,Likelihood,Factorised})( return if Prior _maybe_pointwise_logpdf(dist, val, Val{Factorised}()) else - return DoNotAccumulate() + DoNotAccumulate() end end const POINTWISE_ACCNAME = :PointwiseLogProb @@ -81,7 +82,7 @@ function accumulate_observe!!( return if Likelihood && vn isa VarName logp = _maybe_pointwise_logpdf(right, left, Val{Factorised}()) new_values = DynamicPPL.templated_setindex!!(acc.values, logp, vn, template) - return VNTAccumulator{POINTWISE_ACCNAME}(acc.f, new_values) + VNTAccumulator{POINTWISE_ACCNAME}(acc.f, new_values) else # No need to accumulate likelihoods. acc diff --git a/test/pointwise_logdensities.jl b/test/pointwise_logdensities.jl index 91d01405a..2d9f2c13c 100644 --- a/test/pointwise_logdensities.jl +++ b/test/pointwise_logdensities.jl @@ -82,7 +82,7 @@ end log_prior_densities = pointwise_prior_logdensities( model, InitFromPrior(); factorize=true ) - @test logdensities isa DynamicPPL.VarNamedTuple + @test log_prior_densities isa DynamicPPL.VarNamedTuple @test log_prior_densities[@varname(x)] isa Vector{Float64} @test length(log_prior_densities[@varname(x)]) == 2 @test !haskey(log_prior_densities, @varname(z)) @@ -200,12 +200,14 @@ end @test k in keys(logdensities) end log_prior_densities = pointwise_prior_logdensities(model, chn; factorize=true) + log_likelihoods = pointwise_loglikelihoods(model, chn; factorize=true) for k in [Symbol("x[1]"), Symbol("x[2]")] - @test k in keys(logdensities) + @test k in keys(log_prior_densities) + @test !(k in keys(log_likelihoods)) end - log_likelihoods = pointwise_loglikelihoods(model, chn; factorize=true) for k in [Symbol("z[1]"), Symbol("z[2]")] - @test k in keys(logdensities) + @test k in keys(log_likelihoods) + @test !(k in keys(log_prior_densities)) end end end