Add logevidence to chain

Author: penelopeysm

Project.toml

@@ -41,7 +41,7 @@ DynamicPPLChainRulesCoreExt = ["ChainRulesCore"]
 DynamicPPLEnzymeCoreExt = ["EnzymeCore"]
 DynamicPPLForwardDiffExt = ["ForwardDiff"]
 DynamicPPLJETExt = ["JET"]
-DynamicPPLMCMCChainsExt = ["MCMCChains"]
+DynamicPPLMCMCChainsExt = ["MCMCChains", "Statistics"]
 DynamicPPLMarginalLogDensitiesExt = ["MarginalLogDensities"]
 DynamicPPLMooncakeExt = ["Mooncake"]
 

ext/DynamicPPLMCMCChainsExt.jl

@@ -2,6 +2,7 @@ module DynamicPPLMCMCChainsExt
 
 using DynamicPPL: DynamicPPL, AbstractPPL, AbstractMCMC
 using MCMCChains: MCMCChains
+using Statistics: mean
 
 _has_varname_to_symbol(info::NamedTuple{names}) where {names} = :varname_to_symbol in names
 
@@ -165,12 +166,20 @@ function AbstractMCMC.bundle_samples(
         info = merge(info, (start_time=stats.start, stop_time=stats.stop))
     end
 
+    # Calculate logevidence, if it's available
+    logevidence = if :logevidence in names(bare_chain)
+        mean(bare_chain[:logevidence])
+    else
+        missing
+    end
+
     # Reconstruct the chain with the extra information
     # Yeah, this is quite ugly. Blame MCMCChains.
     chain = MCMCChains.Chains(
         bare_chain.value.data,
         names(bare_chain),
         bare_chain.name_map;
+        evidence=logevidence,
         info=info,
         start=discard_initial + 1,
         thin=thinning,