From f453abd3be8e13c1841405dfbda3d87a41c40d36 Mon Sep 17 00:00:00 2001 From: Miguel Biron-Lattes Date: Mon, 2 Dec 2024 10:00:39 -0800 Subject: [PATCH 01/32] automatically build prior --- Project.toml | 10 ++++++ .../PigeonsJuliaBUGSExt.jl | 27 ++++++++++++++++ ext/PigeonsJuliaBUGSExt/interface.jl | 32 +++++++++++++++++++ src/Pigeons.jl | 5 ++- src/includes.jl | 1 + src/targets/JuliaBUGSLogPotential.jl | 14 ++++++++ test/Project.toml | 1 + 7 files changed, 87 insertions(+), 3 deletions(-) create mode 100644 ext/PigeonsJuliaBUGSExt/PigeonsJuliaBUGSExt.jl create mode 100644 ext/PigeonsJuliaBUGSExt/interface.jl create mode 100644 src/targets/JuliaBUGSLogPotential.jl diff --git a/Project.toml b/Project.toml index 106bef954..58cfdd782 100644 --- a/Project.toml +++ b/Project.toml @@ -37,12 +37,15 @@ StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" ZipFile = "a5390f91-8eb1-5f08-bee0-b1d1ffed6cea" [weakdeps] +AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" BridgeStan = "c88b6f0a-829e-4b0b-94b7-f06ab5908f5a" DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" HypothesisTests = "09f84164-cd44-5f33-b23f-e6b0d136a0d5" +JuliaBUGS = "ba9fb4c0-828e-4473-b6a1-cd2560fee5bf" MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d" +MetaGraphsNext = "fa8bd995-216d-47f1-8a91-f3b68fbeb377" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" [extensions] @@ -51,10 +54,12 @@ PigeonsDynamicPPLExt = "DynamicPPL" PigeonsEnzymeExt = "Enzyme" PigeonsForwardDiffExt = "ForwardDiff" PigeonsHypothesisTestsExt = "HypothesisTests" +PigeonsJuliaBUGSExt = ["AbstractPPL", "MetaGraphsNext", "JuliaBUGS"] PigeonsMCMCChainsExt = "MCMCChains" PigeonsReverseDiffExt = "ReverseDiff" [compat] +AbstractPPL = "0.9" BridgeStan = "2" DataFrames = "1" Distributions = "0.25" @@ -68,10 +73,12 @@ Graphs = "1" HypothesisTests = "0.11" Interpolations = "0.14, 0.15" JSON = "0.21" +JuliaBUGS = "0.7" LogDensityProblems = "2" LogDensityProblemsAD = "1" LogExpFunctions = "0.3" MCMCChains = "6" +MetaGraphsNext = "0.7" MPI = "0.20" MPIPreferences = "0.1" MacroTools = "0.5" @@ -90,10 +97,13 @@ ZipFile = "0.10" julia = "1.8" [extras] +AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" BridgeStan = "c88b6f0a-829e-4b0b-94b7-f06ab5908f5a" DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" HypothesisTests = "09f84164-cd44-5f33-b23f-e6b0d136a0d5" +JuliaBUGS = "ba9fb4c0-828e-4473-b6a1-cd2560fee5bf" MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d" +MetaGraphsNext = "fa8bd995-216d-47f1-8a91-f3b68fbeb377" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" diff --git a/ext/PigeonsJuliaBUGSExt/PigeonsJuliaBUGSExt.jl b/ext/PigeonsJuliaBUGSExt/PigeonsJuliaBUGSExt.jl new file mode 100644 index 000000000..a486f64c0 --- /dev/null +++ b/ext/PigeonsJuliaBUGSExt/PigeonsJuliaBUGSExt.jl @@ -0,0 +1,27 @@ +module PigeonsJuliaBUGSExt + +using Pigeons +if isdefined(Base, :get_extension) + import JuliaBUGS + using AbstractPPL: getsym + using Graphs + using MetaGraphsNext + using LogDensityProblems + using DocStringExtensions + using SplittableRandoms + using Random +else + import ..JuliaBUGS + using ..AbstractPPL: getsym + using ..Graphs + using ..MetaGraphsNext + using ..LogDensityProblems + using ..DocStringExtensions + using ..SplittableRandoms: SplittableRandom, split + using ..Random +end + +include(joinpath(@__DIR__, "interface.jl")) + + +end diff --git a/ext/PigeonsJuliaBUGSExt/interface.jl b/ext/PigeonsJuliaBUGSExt/interface.jl new file mode 100644 index 000000000..a25843c86 --- /dev/null +++ b/ext/PigeonsJuliaBUGSExt/interface.jl @@ -0,0 +1,32 @@ +Pigeons.default_reference(target::JuliaBUGSLogPotential) = + JuliaBUGSLogPotential(make_prior_model(target.model)) +function make_prior_model(target_model::JuliaBUGS.BUGSModel) + # copy underlying graph and remove all vertices that + # - are stochastic and observed: likelihood statements + # - aren't stochastic nor observed and whose only children are observations + prior_graph = deepcopy(target_model.g) + for (label, (code, node_meta)) in prior_graph.vertex_properties + if node_meta.is_stochastic && node_meta.is_observed + for parent_code in inneighbors(prior_graph, code) + # TODO: this should keep recursing up the parent, removing all unnecessary + # deterministic functions, not just one level + parent_meta = prior_graph.vertex_properties[label_for(prior_graph, parent_code)][2] + if !parent_meta.is_stochastic && !parent_meta.is_observed + rem_vertex!(prior_graph, parent_code) + end + end + rem_vertex!(prior_graph, code) + end + end + + # make the corresponding evaluation environment + prior_parameters = Set(getsym(k) for k in keys(prior_graph.vertex_properties)) + eval_env = target_model.evaluation_env + prior_eval_env = NamedTuple(k => v for (k,v) in zip(keys(eval_env),eval_env) if k in prior_parameters) + + # create prior model and check consistency + prior_model = JuliaBUGS.BUGSModel(prior_graph, prior_eval_env) + @assert Set(prior_model.parameters) == Set(target_model.parameters) + + return prior_model +end \ No newline at end of file diff --git a/src/Pigeons.jl b/src/Pigeons.jl index fb748d912..1b8d01909 100644 --- a/src/Pigeons.jl +++ b/src/Pigeons.jl @@ -52,10 +52,8 @@ include("includes.jl") export pigeons, Inputs, PT, # for running jobs: ChildProcess, MPIProcesses, - # references: - DistributionLogPotential, # targets: - TuringLogPotential, StanLogPotential, + TuringLogPotential, StanLogPotential, DistributionLogPotential, JuliaBUGSLogPotential, # some examples toy_mvn_target, toy_stan_target, # post-processing helpers @@ -89,6 +87,7 @@ end @require Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" include(joinpath(@__DIR__, "../ext/PigeonsEnzymeExt/PigeonsEnzymeExt.jl")) @require ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" include(joinpath(@__DIR__, "../ext/PigeonsForwardDiffExt/PigeonsForwardDiffExt.jl")) @require HypothesisTests = "09f84164-cd44-5f33-b23f-e6b0d136a0d5" include(joinpath(@__DIR__, "../ext/PigeonsHypothesisTestsExt/PigeonsHypothesisTestsExt.jl")) + @require JuliaBUGS = "ba9fb4c0-828e-4473-b6a1-cd2560fee5bf" include(joinpath(@__DIR__, "../ext/PigeonsJuliaBUGSExt/PigeonsJuliaBUGSExt.jl")) @require MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d" include(joinpath(@__DIR__, "../ext/PigeonsMCMCChainsExt/PigeonsMCMCChainsExt.jl")) @require ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" include(joinpath(@__DIR__, "../ext/PigeonsReverseDiffExt/PigeonsReverseDiffExt.jl")) end diff --git a/src/includes.jl b/src/includes.jl index 225a98af9..df0fdd825 100644 --- a/src/includes.jl +++ b/src/includes.jl @@ -87,3 +87,4 @@ include("explorers/AAPS.jl") include("explorers/GradientBasedSampler.jl") include("evidence/stepping_stone.jl") include("api.jl") +include("targets/JuliaBUGSLogPotential.jl") diff --git a/src/targets/JuliaBUGSLogPotential.jl b/src/targets/JuliaBUGSLogPotential.jl new file mode 100644 index 000000000..7579750a2 --- /dev/null +++ b/src/targets/JuliaBUGSLogPotential.jl @@ -0,0 +1,14 @@ +""" +$SIGNATURES + +A thin wrapper around a `JuliaBUGS.BUGSModel`. +To work with Pigeons, `JuliaBUGS` needs to be imported into the current session. + +$FIELDS +""" +@auto struct JuliaBUGSLogPotential + """ + A `JuliaBUGS.BUGSModel`. + """ + model +end diff --git a/test/Project.toml b/test/Project.toml index c38b35c9c..a09c32e67 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -15,6 +15,7 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" HypothesisTests = "09f84164-cd44-5f33-b23f-e6b0d136a0d5" InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" +JuliaBUGS = "ba9fb4c0-828e-4473-b6a1-cd2560fee5bf" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LinearRegression = "92481ed7-9fb7-40fd-80f2-46fd0f076581" LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" From 51f70cb0cf06aba7894e552379c525a68180bc02 Mon Sep 17 00:00:00 2001 From: Miguel Biron-Lattes Date: Mon, 2 Dec 2024 16:40:05 -0800 Subject: [PATCH 02/32] fix prior extraction + test --- ext/PigeonsJuliaBUGSExt/interface.jl | 26 +++++++++----------------- test/test_JuliaBUGS.jl | 19 +++++++++++++++++++ 2 files changed, 28 insertions(+), 17 deletions(-) create mode 100644 test/test_JuliaBUGS.jl diff --git a/ext/PigeonsJuliaBUGSExt/interface.jl b/ext/PigeonsJuliaBUGSExt/interface.jl index a25843c86..f7adeff1a 100644 --- a/ext/PigeonsJuliaBUGSExt/interface.jl +++ b/ext/PigeonsJuliaBUGSExt/interface.jl @@ -1,32 +1,24 @@ Pigeons.default_reference(target::JuliaBUGSLogPotential) = JuliaBUGSLogPotential(make_prior_model(target.model)) + function make_prior_model(target_model::JuliaBUGS.BUGSModel) - # copy underlying graph and remove all vertices that - # - are stochastic and observed: likelihood statements - # - aren't stochastic nor observed and whose only children are observations + # copy the target model graph, then drop any nodes that are not parameters prior_graph = deepcopy(target_model.g) - for (label, (code, node_meta)) in prior_graph.vertex_properties - if node_meta.is_stochastic && node_meta.is_observed - for parent_code in inneighbors(prior_graph, code) - # TODO: this should keep recursing up the parent, removing all unnecessary - # deterministic functions, not just one level - parent_meta = prior_graph.vertex_properties[label_for(prior_graph, parent_code)][2] - if !parent_meta.is_stochastic && !parent_meta.is_observed - rem_vertex!(prior_graph, parent_code) - end - end + model_params_syms = Set(getsym(vn) for vn in target_model.parameters) + for (label, (code, _)) in prior_graph.vertex_properties + if getsym(label) ∉ model_params_syms rem_vertex!(prior_graph, code) end end # make the corresponding evaluation environment - prior_parameters = Set(getsym(k) for k in keys(prior_graph.vertex_properties)) eval_env = target_model.evaluation_env - prior_eval_env = NamedTuple(k => v for (k,v) in zip(keys(eval_env),eval_env) if k in prior_parameters) + prior_eval_env = NamedTuple(k => v for (k,v) in zip(keys(eval_env),eval_env) if k in model_params_syms) # create prior model and check consistency - prior_model = JuliaBUGS.BUGSModel(prior_graph, prior_eval_env) + prior_model = JuliaBUGS.BUGSModel( + prior_graph, prior_eval_env; is_transformed = target_model.transformed) @assert Set(prior_model.parameters) == Set(target_model.parameters) return prior_model -end \ No newline at end of file +end diff --git a/test/test_JuliaBUGS.jl b/test/test_JuliaBUGS.jl new file mode 100644 index 000000000..16d06cd71 --- /dev/null +++ b/test/test_JuliaBUGS.jl @@ -0,0 +1,19 @@ +using JuliaBUGS +using AbstractPPL: getsym + +# good ol' toy unidentifiable model for testing purposes +unid_model_def = @bugs begin + for i in 1:2 + p[i] ~ dunif(0,1) + end + p_prod = p[1]*p[2] + n_heads ~ dbin(p_prod, n_flips) +end +unid_target_model = compile(unid_model_def, (; n_heads=50000, n_flips=100000)) +unid_target = JuliaBUGSLogPotential(unid_target_model) + +@testset "JuliaBUGS: extracting prior model" begin + unid_ref = Pigeons.default_reference(unid_target) + unid_prior_model = unid_ref.model + @test Set(getsym(vn) for vn in unid_prior_model.parameters) == Set((:p,)) +end From 81d6612d19693e8731116fb10801ef97b1241900 Mon Sep 17 00:00:00 2001 From: Miguel Biron-Lattes Date: Mon, 2 Dec 2024 16:41:53 -0800 Subject: [PATCH 03/32] fix missing dep --- test/Project.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/test/Project.toml b/test/Project.toml index a09c32e67..257f24e21 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,4 +1,5 @@ [deps] +AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d" ArgMacros = "dbc42088-9de8-42a0-8ec8-2cd114e1ea3e" Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" From f508bf42755177c0e5a10c1c85a19bae6a4e677a Mon Sep 17 00:00:00 2001 From: Miguel Biron-Lattes Date: Mon, 2 Dec 2024 18:38:01 -0800 Subject: [PATCH 04/32] init + sample_iid + logpotential eval --- .../PigeonsJuliaBUGSExt.jl | 4 +-- ext/PigeonsJuliaBUGSExt/interface.jl | 36 ++++++++++++++++--- test/test_JuliaBUGS.jl | 18 ++++++++-- 3 files changed, 49 insertions(+), 9 deletions(-) diff --git a/ext/PigeonsJuliaBUGSExt/PigeonsJuliaBUGSExt.jl b/ext/PigeonsJuliaBUGSExt/PigeonsJuliaBUGSExt.jl index a486f64c0..06a050172 100644 --- a/ext/PigeonsJuliaBUGSExt/PigeonsJuliaBUGSExt.jl +++ b/ext/PigeonsJuliaBUGSExt/PigeonsJuliaBUGSExt.jl @@ -3,7 +3,7 @@ module PigeonsJuliaBUGSExt using Pigeons if isdefined(Base, :get_extension) import JuliaBUGS - using AbstractPPL: getsym + using AbstractPPL: AbstractPPL using Graphs using MetaGraphsNext using LogDensityProblems @@ -12,7 +12,7 @@ if isdefined(Base, :get_extension) using Random else import ..JuliaBUGS - using ..AbstractPPL: getsym + using ..AbstractPPL: AbstractPPL using ..Graphs using ..MetaGraphsNext using ..LogDensityProblems diff --git a/ext/PigeonsJuliaBUGSExt/interface.jl b/ext/PigeonsJuliaBUGSExt/interface.jl index f7adeff1a..952126617 100644 --- a/ext/PigeonsJuliaBUGSExt/interface.jl +++ b/ext/PigeonsJuliaBUGSExt/interface.jl @@ -1,19 +1,46 @@ +#= +JuliaBUGS state is an "evaluation environment" (NamedTuple), containing only +the parameters of the model (i.e., with observations and other variables deleted). +This is so that the state corresponds to the evaluation_env of the prior +=# +function Pigeons.initialization(target::JuliaBUGSLogPotential, _::AbstractRNG, _::Int64) + ev_env = target.model.evaluation_env + model_params_syms = Set(AbstractPPL.getsym(vn) for vn in target.model.parameters) + NamedTuple(k => v for (k,v) in zip(keys(ev_env), ev_env) if k in model_params_syms) +end + +# iid sampling from the prior is straightforward because state coincides with evaluation_env +function Pigeons.sample_iid!(ref_lp::JuliaBUGSLogPotential, replica, shared) + replica.state = first(JuliaBUGS.evaluate!!(replica.rng, ref_lp.model)) +end + +# log_potential evaluation +# note: `initialize!` merges, so it works even when eval_env has more fields than +# log_potential.model.evaluation_env (which happens when log_potential is the target) +function (log_potential::JuliaBUGSLogPotential)(eval_env) + model = JuliaBUGS.initialize!(log_potential.model, eval_env) + return last(JuliaBUGS.evaluate!!(model)) +end + +# Set the default reference to a JuliaBUGS model for the prior Pigeons.default_reference(target::JuliaBUGSLogPotential) = JuliaBUGSLogPotential(make_prior_model(target.model)) +# Obtain the JuliaBUGS model for the prior by pruning the underlying DAG function make_prior_model(target_model::JuliaBUGS.BUGSModel) # copy the target model graph, then drop any nodes that are not parameters prior_graph = deepcopy(target_model.g) - model_params_syms = Set(getsym(vn) for vn in target_model.parameters) - for (label, (code, _)) in prior_graph.vertex_properties - if getsym(label) ∉ model_params_syms + model_params_syms = Set(AbstractPPL.getsym(vn) for vn in target_model.parameters) + for (vn, (code, _)) in prior_graph.vertex_properties + if AbstractPPL.getsym(vn) ∉ model_params_syms rem_vertex!(prior_graph, code) end end # make the corresponding evaluation environment eval_env = target_model.evaluation_env - prior_eval_env = NamedTuple(k => v for (k,v) in zip(keys(eval_env),eval_env) if k in model_params_syms) + prior_eval_env = NamedTuple( + k => copy(v) for (k,v) in zip(keys(eval_env),eval_env) if k in model_params_syms) # create prior model and check consistency prior_model = JuliaBUGS.BUGSModel( @@ -22,3 +49,4 @@ function make_prior_model(target_model::JuliaBUGS.BUGSModel) return prior_model end + diff --git a/test/test_JuliaBUGS.jl b/test/test_JuliaBUGS.jl index 16d06cd71..20293340f 100644 --- a/test/test_JuliaBUGS.jl +++ b/test/test_JuliaBUGS.jl @@ -1,5 +1,5 @@ using JuliaBUGS -using AbstractPPL: getsym +using AbstractPPL # good ol' toy unidentifiable model for testing purposes unid_model_def = @bugs begin @@ -12,8 +12,20 @@ end unid_target_model = compile(unid_model_def, (; n_heads=50000, n_flips=100000)) unid_target = JuliaBUGSLogPotential(unid_target_model) -@testset "JuliaBUGS: extracting prior model" begin +@testset "Extracting prior model" begin unid_ref = Pigeons.default_reference(unid_target) unid_prior_model = unid_ref.model - @test Set(getsym(vn) for vn in unid_prior_model.parameters) == Set((:p,)) + @test Set(AbstractPPL.getsym(vn) for vn in unid_prior_model.parameters) == Set((:p,)) +end + +@testset "Initialization" begin + true_init_pars = (;p = unid_target_model.evaluation_env.p) + @test Pigeons.initialization(unid_target, SplittableRandom(1), 1) == true_init_pars +end + +@testset "sample_iid!" begin + pt = pigeons(target = unid_target, n_rounds = 0, n_chains = 1) + ref = Pigeons.default_reference(unid_target) + Pigeons.sample_iid!(ref, pt.replicas[1], pt.shared) + @test true end From 308554c4b145a874db91869191896b9be004e55a Mon Sep 17 00:00:00 2001 From: Miguel Biron-Lattes Date: Mon, 2 Dec 2024 18:52:04 -0800 Subject: [PATCH 05/32] logpotential eval test --- test/test_JuliaBUGS.jl | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/test/test_JuliaBUGS.jl b/test/test_JuliaBUGS.jl index 20293340f..6ad192e30 100644 --- a/test/test_JuliaBUGS.jl +++ b/test/test_JuliaBUGS.jl @@ -26,6 +26,18 @@ end @testset "sample_iid!" begin pt = pigeons(target = unid_target, n_rounds = 0, n_chains = 1) ref = Pigeons.default_reference(unid_target) - Pigeons.sample_iid!(ref, pt.replicas[1], pt.shared) - @test true + new_state = Pigeons.sample_iid!(ref, pt.replicas[1], pt.shared) + @test pt.replicas[1].state === new_state +end + +@testset "log_potential eval" begin + # check log_potential evaluation with constrained version (easier, no Jacobian) + unid_target_const = JuliaBUGSLogPotential(JuliaBUGS.settrans(unid_target_model)) + unid_ref_const = Pigeons.default_reference(unid_target_const) + state = (; p = unid_target_model.evaluation_env.p) + @test unid_target_const(state) == + logpdf( + Binomial(unid_target_model.evaluation_env.n_flips,prod(state.p)), + unid_target_model.evaluation_env.n_heads) + @test unid_ref_const(state) == 0 end From c6a9b414d6da131a577dbcb6c81180a04436c86b Mon Sep 17 00:00:00 2001 From: Miguel Biron-Lattes Date: Sat, 7 Dec 2024 07:44:47 -0800 Subject: [PATCH 06/32] rewrite JuliaBUGSExt using tempered_evaluate --- Project.toml | 8 +- .../PigeonsJuliaBUGSExt.jl | 8 +- ext/PigeonsJuliaBUGSExt/deprecated.jl | 37 ++++++ ext/PigeonsJuliaBUGSExt/interface.jl | 106 ++++++++++-------- src/Pigeons.jl | 2 +- src/includes.jl | 2 +- .../JuliaBUGSPath.jl} | 4 +- test/Project.toml | 1 - test/test_JuliaBUGS.jl | 46 ++++---- 9 files changed, 123 insertions(+), 91 deletions(-) create mode 100644 ext/PigeonsJuliaBUGSExt/deprecated.jl rename src/{targets/JuliaBUGSLogPotential.jl => paths/JuliaBUGSPath.jl} (61%) diff --git a/Project.toml b/Project.toml index 58cfdd782..5f360c74d 100644 --- a/Project.toml +++ b/Project.toml @@ -37,7 +37,6 @@ StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" ZipFile = "a5390f91-8eb1-5f08-bee0-b1d1ffed6cea" [weakdeps] -AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" BridgeStan = "c88b6f0a-829e-4b0b-94b7-f06ab5908f5a" DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" @@ -45,7 +44,6 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" HypothesisTests = "09f84164-cd44-5f33-b23f-e6b0d136a0d5" JuliaBUGS = "ba9fb4c0-828e-4473-b6a1-cd2560fee5bf" MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d" -MetaGraphsNext = "fa8bd995-216d-47f1-8a91-f3b68fbeb377" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" [extensions] @@ -54,12 +52,11 @@ PigeonsDynamicPPLExt = "DynamicPPL" PigeonsEnzymeExt = "Enzyme" PigeonsForwardDiffExt = "ForwardDiff" PigeonsHypothesisTestsExt = "HypothesisTests" -PigeonsJuliaBUGSExt = ["AbstractPPL", "MetaGraphsNext", "JuliaBUGS"] +PigeonsJuliaBUGSExt = "JuliaBUGS" PigeonsMCMCChainsExt = "MCMCChains" PigeonsReverseDiffExt = "ReverseDiff" [compat] -AbstractPPL = "0.9" BridgeStan = "2" DataFrames = "1" Distributions = "0.25" @@ -78,7 +75,6 @@ LogDensityProblems = "2" LogDensityProblemsAD = "1" LogExpFunctions = "0.3" MCMCChains = "6" -MetaGraphsNext = "0.7" MPI = "0.20" MPIPreferences = "0.1" MacroTools = "0.5" @@ -97,7 +93,6 @@ ZipFile = "0.10" julia = "1.8" [extras] -AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" BridgeStan = "c88b6f0a-829e-4b0b-94b7-f06ab5908f5a" DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" @@ -105,5 +100,4 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" HypothesisTests = "09f84164-cd44-5f33-b23f-e6b0d136a0d5" JuliaBUGS = "ba9fb4c0-828e-4473-b6a1-cd2560fee5bf" MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d" -MetaGraphsNext = "fa8bd995-216d-47f1-8a91-f3b68fbeb377" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" diff --git a/ext/PigeonsJuliaBUGSExt/PigeonsJuliaBUGSExt.jl b/ext/PigeonsJuliaBUGSExt/PigeonsJuliaBUGSExt.jl index 06a050172..133bbf8e5 100644 --- a/ext/PigeonsJuliaBUGSExt/PigeonsJuliaBUGSExt.jl +++ b/ext/PigeonsJuliaBUGSExt/PigeonsJuliaBUGSExt.jl @@ -3,18 +3,12 @@ module PigeonsJuliaBUGSExt using Pigeons if isdefined(Base, :get_extension) import JuliaBUGS - using AbstractPPL: AbstractPPL - using Graphs - using MetaGraphsNext using LogDensityProblems using DocStringExtensions - using SplittableRandoms + using SplittableRandoms: SplittableRandom, split using Random else import ..JuliaBUGS - using ..AbstractPPL: AbstractPPL - using ..Graphs - using ..MetaGraphsNext using ..LogDensityProblems using ..DocStringExtensions using ..SplittableRandoms: SplittableRandom, split diff --git a/ext/PigeonsJuliaBUGSExt/deprecated.jl b/ext/PigeonsJuliaBUGSExt/deprecated.jl new file mode 100644 index 000000000..bc4d6bbfe --- /dev/null +++ b/ext/PigeonsJuliaBUGSExt/deprecated.jl @@ -0,0 +1,37 @@ +############################################################################### +# InterpolatingPath approach, where we built a separate prior model from the +# target one. +# Not used right now, but keeping it here in case we want to revert back to +# this approach +############################################################################### + +get_unique_params(model::JuliaBUGS.BUGSModel) = Set(AbstractPPL.getsym(vn) for vn in model.parameters) + +# Set the default reference to a JuliaBUGS model for the prior +Pigeons.default_reference(target::JuliaBUGSLogPotential) = + JuliaBUGSLogPotential(make_prior_model(target.model)) + +# Obtain the JuliaBUGS model for the prior by pruning the underlying DAG +function make_prior_model(target_model::JuliaBUGS.BUGSModel) + # copy the target model graph, then drop any nodes that are not parameters + prior_graph = deepcopy(target_model.g) + model_params_syms = unique_params(target_model.parameters) + for (vn, (code, _)) in prior_graph.vertex_properties + if AbstractPPL.getsym(vn) ∉ model_params_syms + rem_vertex!(prior_graph, code) + end + end + + # make the corresponding evaluation environment + eval_env = target_model.evaluation_env + prior_eval_env = NamedTuple( + k => copy(v) for (k,v) in zip(keys(eval_env),eval_env) if k in model_params_syms) + + # create prior model and check consistency + prior_model = JuliaBUGS.BUGSModel( + prior_graph, prior_eval_env; is_transformed = target_model.transformed) + @assert Set(prior_model.parameters) == Set(target_model.parameters) + + return prior_model +end + diff --git a/ext/PigeonsJuliaBUGSExt/interface.jl b/ext/PigeonsJuliaBUGSExt/interface.jl index 952126617..5d5863535 100644 --- a/ext/PigeonsJuliaBUGSExt/interface.jl +++ b/ext/PigeonsJuliaBUGSExt/interface.jl @@ -1,52 +1,68 @@ -#= -JuliaBUGS state is an "evaluation environment" (NamedTuple), containing only -the parameters of the model (i.e., with observations and other variables deleted). -This is so that the state corresponds to the evaluation_env of the prior -=# -function Pigeons.initialization(target::JuliaBUGSLogPotential, _::AbstractRNG, _::Int64) - ev_env = target.model.evaluation_env - model_params_syms = Set(AbstractPPL.getsym(vn) for vn in target.model.parameters) - NamedTuple(k => v for (k,v) in zip(keys(ev_env), ev_env) if k in model_params_syms) +####################################### +# Path interface +####################################### + +# State initialization: state is a flattened vector of the parameters +# Note: JuliaBUGS.getparams creates a new vector on each call, so it is safe +# to call these for different Replicas +Pigeons.initialization(target::JuliaBUGSPath, _::AbstractRNG, _::Int64) = + JuliaBUGS.getparams(target.model) + +# target is already a Path +Pigeons.create_path(target::JuliaBUGSPath, ::Inputs) = target + +####################################### +# Log-potential interface +####################################### + +""" +$SIGNATURES + +A log-potential built from a [`JuliaBUGSPath`](@ref) for a specific inverse +temperature parameter. + +$FIELDS +""" +struct JuliaBUGSLogPotential{TMod<:JuliaBUGS.BUGSModel, TF<:AbstractFloat} + """ + A deep-enough copy of the original model that allows evaluation while + avoiding race conditions between different Replicas. + """ + private_model::TMod + + """ + Tempering parameter. + """ + beta::TF end -# iid sampling from the prior is straightforward because state coincides with evaluation_env -function Pigeons.sample_iid!(ref_lp::JuliaBUGSLogPotential, replica, shared) - replica.state = first(JuliaBUGS.evaluate!!(replica.rng, ref_lp.model)) +# make a log-potential by creating a new model with independent graph and +# evaluation environment. Both of these could be modified during density +# evaluations and/or during Gibbs sampling +function Pigeons.interpolate(path::JuliaBUGSPath, beta) + model = path.model + private_model = JuliaBUGS.BUGSModel( + model, + deepcopy(model.g), + model.parameters, model.flattened_graph_node_data.sorted_nodes, + deepcopy(model.evaluation_env) + ) + JuliaBUGSLogPotential(private_model, beta) end # log_potential evaluation -# note: `initialize!` merges, so it works even when eval_env has more fields than -# log_potential.model.evaluation_env (which happens when log_potential is the target) -function (log_potential::JuliaBUGSLogPotential)(eval_env) - model = JuliaBUGS.initialize!(log_potential.model, eval_env) - return last(JuliaBUGS.evaluate!!(model)) -end +(log_potential::JuliaBUGSLogPotential)(flattened_values) = + last(last(JuliaBUGS._tempered_evaluate!!( + log_potential.private_model, + flattened_values; + temperature=log_potential.beta + ))) -# Set the default reference to a JuliaBUGS model for the prior -Pigeons.default_reference(target::JuliaBUGSLogPotential) = - JuliaBUGSLogPotential(make_prior_model(target.model)) - -# Obtain the JuliaBUGS model for the prior by pruning the underlying DAG -function make_prior_model(target_model::JuliaBUGS.BUGSModel) - # copy the target model graph, then drop any nodes that are not parameters - prior_graph = deepcopy(target_model.g) - model_params_syms = Set(AbstractPPL.getsym(vn) for vn in target_model.parameters) - for (vn, (code, _)) in prior_graph.vertex_properties - if AbstractPPL.getsym(vn) ∉ model_params_syms - rem_vertex!(prior_graph, code) - end - end - - # make the corresponding evaluation environment - eval_env = target_model.evaluation_env - prior_eval_env = NamedTuple( - k => copy(v) for (k,v) in zip(keys(eval_env),eval_env) if k in model_params_syms) - - # create prior model and check consistency - prior_model = JuliaBUGS.BUGSModel( - prior_graph, prior_eval_env; is_transformed = target_model.transformed) - @assert Set(prior_model.parameters) == Set(target_model.parameters) - - return prior_model +# iid sampling +# Note: JuliaBUGS.getparams always allocates a new vector so there is no point +# of copying the result into the Replica's state; just replace it. +function Pigeons.sample_iid!(log_potential::JuliaBUGSLogPotential, replica, shared) + new_env = first(JuliaBUGS.evaluate!!(replica.rng, log_potential.private_model)) # sample a new evaluation environment + JuliaBUGS.initialize!(log_potential.private_model, new_env) # set the private_model's environment to the newly created one + replica.state = JuliaBUGS.getparams(log_potential.private_model) # finally, flatten the eval environment in the model and set that as the replica state end - diff --git a/src/Pigeons.jl b/src/Pigeons.jl index 1b8d01909..8cf17b753 100644 --- a/src/Pigeons.jl +++ b/src/Pigeons.jl @@ -53,7 +53,7 @@ export pigeons, Inputs, PT, # for running jobs: ChildProcess, MPIProcesses, # targets: - TuringLogPotential, StanLogPotential, DistributionLogPotential, JuliaBUGSLogPotential, + TuringLogPotential, StanLogPotential, DistributionLogPotential, JuliaBUGSPath, # some examples toy_mvn_target, toy_stan_target, # post-processing helpers diff --git a/src/includes.jl b/src/includes.jl index df0fdd825..5600eb669 100644 --- a/src/includes.jl +++ b/src/includes.jl @@ -87,4 +87,4 @@ include("explorers/AAPS.jl") include("explorers/GradientBasedSampler.jl") include("evidence/stepping_stone.jl") include("api.jl") -include("targets/JuliaBUGSLogPotential.jl") +include("paths/JuliaBUGSPath.jl") diff --git a/src/targets/JuliaBUGSLogPotential.jl b/src/paths/JuliaBUGSPath.jl similarity index 61% rename from src/targets/JuliaBUGSLogPotential.jl rename to src/paths/JuliaBUGSPath.jl index 7579750a2..1c67ae68d 100644 --- a/src/targets/JuliaBUGSLogPotential.jl +++ b/src/paths/JuliaBUGSPath.jl @@ -1,12 +1,12 @@ """ $SIGNATURES -A thin wrapper around a `JuliaBUGS.BUGSModel`. +A thin wrapper around a `JuliaBUGS.BUGSModel` to provide a prior-posterior path. To work with Pigeons, `JuliaBUGS` needs to be imported into the current session. $FIELDS """ -@auto struct JuliaBUGSLogPotential +@auto struct JuliaBUGSPath """ A `JuliaBUGS.BUGSModel`. """ diff --git a/test/Project.toml b/test/Project.toml index 257f24e21..a09c32e67 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,5 +1,4 @@ [deps] -AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d" ArgMacros = "dbc42088-9de8-42a0-8ec8-2cd114e1ea3e" Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" diff --git a/test/test_JuliaBUGS.jl b/test/test_JuliaBUGS.jl index 6ad192e30..fd8646047 100644 --- a/test/test_JuliaBUGS.jl +++ b/test/test_JuliaBUGS.jl @@ -1,5 +1,4 @@ using JuliaBUGS -using AbstractPPL # good ol' toy unidentifiable model for testing purposes unid_model_def = @bugs begin @@ -10,34 +9,27 @@ unid_model_def = @bugs begin n_heads ~ dbin(p_prod, n_flips) end unid_target_model = compile(unid_model_def, (; n_heads=50000, n_flips=100000)) -unid_target = JuliaBUGSLogPotential(unid_target_model) +unid_target_constrained = JuliaBUGSPath(JuliaBUGS.settrans(unid_target_model)) +struct IdentityExplorer end +function Pigeons.step!(::IdentityExplorer, replica, shared) end -@testset "Extracting prior model" begin - unid_ref = Pigeons.default_reference(unid_target) - unid_prior_model = unid_ref.model - @test Set(AbstractPPL.getsym(vn) for vn in unid_prior_model.parameters) == Set((:p,)) -end - -@testset "Initialization" begin - true_init_pars = (;p = unid_target_model.evaluation_env.p) - @test Pigeons.initialization(unid_target, SplittableRandom(1), 1) == true_init_pars -end +@testset "Basic sampling via independent MH from the prior" begin + pt = pigeons( + target = unid_target_constrained, + n_chains = 2, + explorer = IdentityExplorer(), + record = [traces], + extended_traces = true + ) + # check sample_iid! + @test isapprox(0.5, mean(v[1] for (k,v) in pt.reduced_recorders.traces if first(k)==1), rtol=0.05) + @test isapprox(0.5, mean(v[2] for (k,v) in pt.reduced_recorders.traces if first(k)==1), rtol=0.05) -@testset "sample_iid!" begin - pt = pigeons(target = unid_target, n_rounds = 0, n_chains = 1) - ref = Pigeons.default_reference(unid_target) - new_state = Pigeons.sample_iid!(ref, pt.replicas[1], pt.shared) - @test pt.replicas[1].state === new_state -end - -@testset "log_potential eval" begin # check log_potential evaluation with constrained version (easier, no Jacobian) - unid_target_const = JuliaBUGSLogPotential(JuliaBUGS.settrans(unid_target_model)) - unid_ref_const = Pigeons.default_reference(unid_target_const) - state = (; p = unid_target_model.evaluation_env.p) - @test unid_target_const(state) == + @test all(v for (k,v) in pt.reduced_recorders.traces if first(k)==2) do v logpdf( - Binomial(unid_target_model.evaluation_env.n_flips,prod(state.p)), - unid_target_model.evaluation_env.n_heads) - @test unid_ref_const(state) == 0 + Binomial(unid_target_model.evaluation_env.n_flips, v[1]*v[2]), + unid_target_model.evaluation_env.n_heads + ) == last(v) + end end From 2882141e472407cdc7a110eb87ac0ad5787e19de Mon Sep 17 00:00:00 2001 From: Seren Lee Date: Tue, 10 Dec 2024 09:33:25 -0800 Subject: [PATCH 07/32] Domain error try-catch and sample_iid update --- ext/PigeonsJuliaBUGSExt/interface.jl | 20 +++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/ext/PigeonsJuliaBUGSExt/interface.jl b/ext/PigeonsJuliaBUGSExt/interface.jl index 5d5863535..6667956a8 100644 --- a/ext/PigeonsJuliaBUGSExt/interface.jl +++ b/ext/PigeonsJuliaBUGSExt/interface.jl @@ -51,18 +51,24 @@ function Pigeons.interpolate(path::JuliaBUGSPath, beta) end # log_potential evaluation -(log_potential::JuliaBUGSLogPotential)(flattened_values) = - last(last(JuliaBUGS._tempered_evaluate!!( +function (log_potential::JuliaBUGSLogPotential)(flattened_values) + return try last(last(JuliaBUGS._tempered_evaluate!!( log_potential.private_model, flattened_values; - temperature=log_potential.beta - ))) + temperature=log_potential.beta))) + catch e + return -Inf64 + end +end # iid sampling # Note: JuliaBUGS.getparams always allocates a new vector so there is no point # of copying the result into the Replica's state; just replace it. function Pigeons.sample_iid!(log_potential::JuliaBUGSLogPotential, replica, shared) - new_env = first(JuliaBUGS.evaluate!!(replica.rng, log_potential.private_model)) # sample a new evaluation environment - JuliaBUGS.initialize!(log_potential.private_model, new_env) # set the private_model's environment to the newly created one - replica.state = JuliaBUGS.getparams(log_potential.private_model) # finally, flatten the eval environment in the model and set that as the replica state + values, logp = JuliaBUGS.evaluate!!(replica.rng, log_potential.private_model); + JuliaBUGS.initialize!(log_potential.private_model, values); + #new_env = first(JuliaBUGS.evaluate!!(replica.rng, log_potential.private_model)) # sample a new evaluation environment + #JuliaBUGS.initialize!(log_potential.private_model, new_env) # set the private_model's environment to the newly created one + #replica.state = JuliaBUGS.getparams(log_potential.private_model) # finally, flatten the eval environment in the model and set that as the replica state + replica.state = JuliaBUGS.getparams(log_potential.private_model) end From f25d5ffb7be5aa08d6107fc331cfe9fdc40e3b8c Mon Sep 17 00:00:00 2001 From: Miguel Biron-Lattes Date: Tue, 10 Dec 2024 09:48:30 -0800 Subject: [PATCH 08/32] finer catch + more tests --- ext/PigeonsJuliaBUGSExt/interface.jl | 25 ++++++++++++------------- test/test_JuliaBUGS.jl | 27 +++++++++++++++++++++++++-- 2 files changed, 37 insertions(+), 15 deletions(-) diff --git a/ext/PigeonsJuliaBUGSExt/interface.jl b/ext/PigeonsJuliaBUGSExt/interface.jl index 6667956a8..c2f5cfa34 100644 --- a/ext/PigeonsJuliaBUGSExt/interface.jl +++ b/ext/PigeonsJuliaBUGSExt/interface.jl @@ -51,24 +51,23 @@ function Pigeons.interpolate(path::JuliaBUGSPath, beta) end # log_potential evaluation -function (log_potential::JuliaBUGSLogPotential)(flattened_values) - return try last(last(JuliaBUGS._tempered_evaluate!!( - log_potential.private_model, - flattened_values; - temperature=log_potential.beta))) +(log_potential::JuliaBUGSLogPotential)(flattened_values) = + try + last(last(JuliaBUGS._tempered_evaluate!!( + log_potential.private_model, + flattened_values; + temperature=log_potential.beta + ))) catch e - return -Inf64 + (isa(e, DomainError) || isa(e, BoundsError)) && return -Inf + rethrow(e) end -end # iid sampling # Note: JuliaBUGS.getparams always allocates a new vector so there is no point # of copying the result into the Replica's state; just replace it. function Pigeons.sample_iid!(log_potential::JuliaBUGSLogPotential, replica, shared) - values, logp = JuliaBUGS.evaluate!!(replica.rng, log_potential.private_model); - JuliaBUGS.initialize!(log_potential.private_model, values); - #new_env = first(JuliaBUGS.evaluate!!(replica.rng, log_potential.private_model)) # sample a new evaluation environment - #JuliaBUGS.initialize!(log_potential.private_model, new_env) # set the private_model's environment to the newly created one - #replica.state = JuliaBUGS.getparams(log_potential.private_model) # finally, flatten the eval environment in the model and set that as the replica state - replica.state = JuliaBUGS.getparams(log_potential.private_model) + new_env = first(JuliaBUGS.evaluate!!(replica.rng, log_potential.private_model)) # sample a new evaluation environment + JuliaBUGS.initialize!(log_potential.private_model, new_env) # set the private_model's environment to the newly created one + replica.state = JuliaBUGS.getparams(log_potential.private_model) # finally, flatten the eval environment in the model and set that as the replica state end diff --git a/test/test_JuliaBUGS.jl b/test/test_JuliaBUGS.jl index fd8646047..9301193ca 100644 --- a/test/test_JuliaBUGS.jl +++ b/test/test_JuliaBUGS.jl @@ -1,5 +1,7 @@ using JuliaBUGS +include("supporting/analytic_solutions.jl") + # good ol' toy unidentifiable model for testing purposes unid_model_def = @bugs begin for i in 1:2 @@ -9,6 +11,7 @@ unid_model_def = @bugs begin n_heads ~ dbin(p_prod, n_flips) end unid_target_model = compile(unid_model_def, (; n_heads=50000, n_flips=100000)) +unid_target = JuliaBUGSPath(unid_target_model) unid_target_constrained = JuliaBUGSPath(JuliaBUGS.settrans(unid_target_model)) struct IdentityExplorer end function Pigeons.step!(::IdentityExplorer, replica, shared) end @@ -27,9 +30,29 @@ function Pigeons.step!(::IdentityExplorer, replica, shared) end # check log_potential evaluation with constrained version (easier, no Jacobian) @test all(v for (k,v) in pt.reduced_recorders.traces if first(k)==2) do v - logpdf( + last(v) == logpdf( Binomial(unid_target_model.evaluation_env.n_flips, v[1]*v[2]), unid_target_model.evaluation_env.n_heads - ) == last(v) + ) + end +end + +@testset "SliceSampler on constrained and unconstrained versions" begin + for target in (unid_target, unid_target_constrained) + @show target.model + pt = pigeons(; + target, + explorer = SliceSampler(), + n_chains=7, + n_rounds=7 + ) + @test isapprox( + Pigeons.stepping_stone(pt), + unid_target_exact_logZ( + unid_target_model.evaluation_env.n_flips, + unid_target_model.evaluation_env.n_heads + ), + rtol=0.1 + ) end end From a6d03cf8701218d9169ad460348518dbb53a8269 Mon Sep 17 00:00:00 2001 From: Miguel Biron-Lattes Date: Tue, 10 Dec 2024 11:05:34 -0800 Subject: [PATCH 09/32] reproducible initialization --- ext/PigeonsJuliaBUGSExt/deprecated.jl | 37 --------------------------- ext/PigeonsJuliaBUGSExt/interface.jl | 22 +++++++++------- 2 files changed, 13 insertions(+), 46 deletions(-) delete mode 100644 ext/PigeonsJuliaBUGSExt/deprecated.jl diff --git a/ext/PigeonsJuliaBUGSExt/deprecated.jl b/ext/PigeonsJuliaBUGSExt/deprecated.jl deleted file mode 100644 index bc4d6bbfe..000000000 --- a/ext/PigeonsJuliaBUGSExt/deprecated.jl +++ /dev/null @@ -1,37 +0,0 @@ -############################################################################### -# InterpolatingPath approach, where we built a separate prior model from the -# target one. -# Not used right now, but keeping it here in case we want to revert back to -# this approach -############################################################################### - -get_unique_params(model::JuliaBUGS.BUGSModel) = Set(AbstractPPL.getsym(vn) for vn in model.parameters) - -# Set the default reference to a JuliaBUGS model for the prior -Pigeons.default_reference(target::JuliaBUGSLogPotential) = - JuliaBUGSLogPotential(make_prior_model(target.model)) - -# Obtain the JuliaBUGS model for the prior by pruning the underlying DAG -function make_prior_model(target_model::JuliaBUGS.BUGSModel) - # copy the target model graph, then drop any nodes that are not parameters - prior_graph = deepcopy(target_model.g) - model_params_syms = unique_params(target_model.parameters) - for (vn, (code, _)) in prior_graph.vertex_properties - if AbstractPPL.getsym(vn) ∉ model_params_syms - rem_vertex!(prior_graph, code) - end - end - - # make the corresponding evaluation environment - eval_env = target_model.evaluation_env - prior_eval_env = NamedTuple( - k => copy(v) for (k,v) in zip(keys(eval_env),eval_env) if k in model_params_syms) - - # create prior model and check consistency - prior_model = JuliaBUGS.BUGSModel( - prior_graph, prior_eval_env; is_transformed = target_model.transformed) - @assert Set(prior_model.parameters) == Set(target_model.parameters) - - return prior_model -end - diff --git a/ext/PigeonsJuliaBUGSExt/interface.jl b/ext/PigeonsJuliaBUGSExt/interface.jl index c2f5cfa34..de4e1ce0b 100644 --- a/ext/PigeonsJuliaBUGSExt/interface.jl +++ b/ext/PigeonsJuliaBUGSExt/interface.jl @@ -2,11 +2,19 @@ # Path interface ####################################### -# State initialization: state is a flattened vector of the parameters +# Initialization and iid sampling +# state is a flattened vector of the parameters # Note: JuliaBUGS.getparams creates a new vector on each call, so it is safe -# to call these for different Replicas -Pigeons.initialization(target::JuliaBUGSPath, _::AbstractRNG, _::Int64) = - JuliaBUGS.getparams(target.model) +# to call _sample_iid during initialization (**sequentially**, as done as of time +# of writing) for different Replicas (i.e., they won't share the same state). +function _sample_iid(model::JuliaBUGS.BUGSModel, rng::AbstractRNG) + new_env = first(JuliaBUGS.evaluate!!(rng, model)) # sample a new evaluation environment + JuliaBUGS.initialize!(model, new_env) # set the private_model's environment to the newly created one + return JuliaBUGS.getparams(model) # finally, flatten the unobserved parameters in the model's eval environment and return +end + +Pigeons.initialization(target::JuliaBUGSPath, rng::AbstractRNG, _::Int64) = + _sample_iid(target.model, rng) # target is already a Path Pigeons.create_path(target::JuliaBUGSPath, ::Inputs) = target @@ -64,10 +72,6 @@ end end # iid sampling -# Note: JuliaBUGS.getparams always allocates a new vector so there is no point -# of copying the result into the Replica's state; just replace it. function Pigeons.sample_iid!(log_potential::JuliaBUGSLogPotential, replica, shared) - new_env = first(JuliaBUGS.evaluate!!(replica.rng, log_potential.private_model)) # sample a new evaluation environment - JuliaBUGS.initialize!(log_potential.private_model, new_env) # set the private_model's environment to the newly created one - replica.state = JuliaBUGS.getparams(log_potential.private_model) # finally, flatten the eval environment in the model and set that as the replica state + replica.state = _sample_iid(log_potential.private_model, replica.rng) end From 2908feac61c4491fd87991dd362a19756a41b41c Mon Sep 17 00:00:00 2001 From: Miguel Biron-Lattes Date: Wed, 11 Dec 2024 08:50:51 -0800 Subject: [PATCH 10/32] invariance test --- ext/PigeonsDynamicPPLExt/invariance_test.jl | 5 ++- .../PigeonsHypothesisTestsExt.jl | 4 +-- .../PigeonsJuliaBUGSExt.jl | 2 +- ext/PigeonsJuliaBUGSExt/interface.jl | 19 ++++++---- ext/PigeonsJuliaBUGSExt/invariance_test.jl | 35 +++++++++++++++++++ src/explorers/invariance_test.jl | 5 ++- test/test_JuliaBUGS.jl | 21 ++++++----- 7 files changed, 66 insertions(+), 25 deletions(-) create mode 100644 ext/PigeonsJuliaBUGSExt/invariance_test.jl diff --git a/ext/PigeonsDynamicPPLExt/invariance_test.jl b/ext/PigeonsDynamicPPLExt/invariance_test.jl index 6bb7d9a48..0168bf5a6 100644 --- a/ext/PigeonsDynamicPPLExt/invariance_test.jl +++ b/ext/PigeonsDynamicPPLExt/invariance_test.jl @@ -10,9 +10,8 @@ initial and final states. """ function Pigeons.forward_sample_condition_and_explore( model::DynamicPPL.Model, - explorer, rng::SplittableRandom; - run_explorer::Bool = true, + explorer = nothing, condition_on::NTuple{N,Symbol} ) where {N} # forward simulation @@ -37,7 +36,7 @@ function Pigeons.forward_sample_condition_and_explore( DynamicPPL.link!!(state, DynamicPPL.SampleFromPrior(), conditioned_model) # maybe take a step with explorer - if run_explorer + if !isnothing(explorer) state = Pigeons.explorer_step(rng, TuringLogPotential(conditioned_model), explorer, state) end diff --git a/ext/PigeonsHypothesisTestsExt/PigeonsHypothesisTestsExt.jl b/ext/PigeonsHypothesisTestsExt/PigeonsHypothesisTestsExt.jl index 18bf303e2..6444a04bf 100644 --- a/ext/PigeonsHypothesisTestsExt/PigeonsHypothesisTestsExt.jl +++ b/ext/PigeonsHypothesisTestsExt/PigeonsHypothesisTestsExt.jl @@ -34,9 +34,9 @@ function Pigeons.invariance_test( # iterate iid samples for n in eachindex(initial_values) initial_values[n] = Pigeons.forward_sample_condition_and_explore( - target, explorer, rng; run_explorer=false, simulator_kwargs...) + target, rng; simulator_kwargs...) final_values[n] = Pigeons.forward_sample_condition_and_explore( - target, explorer, rng; simulator_kwargs...) + target, rng; explorer, simulator_kwargs...) end # transform vector of vectors to matrices so that iterating dimensions == iterating columns => faster diff --git a/ext/PigeonsJuliaBUGSExt/PigeonsJuliaBUGSExt.jl b/ext/PigeonsJuliaBUGSExt/PigeonsJuliaBUGSExt.jl index 133bbf8e5..dd6dc4f55 100644 --- a/ext/PigeonsJuliaBUGSExt/PigeonsJuliaBUGSExt.jl +++ b/ext/PigeonsJuliaBUGSExt/PigeonsJuliaBUGSExt.jl @@ -16,6 +16,6 @@ else end include(joinpath(@__DIR__, "interface.jl")) - +include(joinpath(@__DIR__, "invariance_test.jl")) end diff --git a/ext/PigeonsJuliaBUGSExt/interface.jl b/ext/PigeonsJuliaBUGSExt/interface.jl index de4e1ce0b..4cc3e961d 100644 --- a/ext/PigeonsJuliaBUGSExt/interface.jl +++ b/ext/PigeonsJuliaBUGSExt/interface.jl @@ -3,16 +3,21 @@ ####################################### # Initialization and iid sampling -# state is a flattened vector of the parameters +function evaluate_and_initialize(model::JuliaBUGS.BUGSModel, rng::AbstractRNG) + new_env = first(JuliaBUGS.evaluate!!(rng, model)) # sample a new evaluation environment + return JuliaBUGS.initialize!(model, new_env) # set the private_model's environment to the newly created one +end + +# used for both initializing and iid sampling +# Note: state is a flattened vector of the parameters +# Also, the vector is **concretely typed**. This means that if the evaluation +# environment contains floats and integers, the latter will be cast to float. +_sample_iid(model::JuliaBUGS.BUGSModel, rng::AbstractRNG) = + JuliaBUGS.getparams(evaluate_and_initialize(model, rng)) # flatten the unobserved parameters in the model's eval environment and return + # Note: JuliaBUGS.getparams creates a new vector on each call, so it is safe # to call _sample_iid during initialization (**sequentially**, as done as of time # of writing) for different Replicas (i.e., they won't share the same state). -function _sample_iid(model::JuliaBUGS.BUGSModel, rng::AbstractRNG) - new_env = first(JuliaBUGS.evaluate!!(rng, model)) # sample a new evaluation environment - JuliaBUGS.initialize!(model, new_env) # set the private_model's environment to the newly created one - return JuliaBUGS.getparams(model) # finally, flatten the unobserved parameters in the model's eval environment and return -end - Pigeons.initialization(target::JuliaBUGSPath, rng::AbstractRNG, _::Int64) = _sample_iid(target.model, rng) diff --git a/ext/PigeonsJuliaBUGSExt/invariance_test.jl b/ext/PigeonsJuliaBUGSExt/invariance_test.jl new file mode 100644 index 000000000..2150693b8 --- /dev/null +++ b/ext/PigeonsJuliaBUGSExt/invariance_test.jl @@ -0,0 +1,35 @@ +""" +$SIGNATURES + +Implements `Pigeons.forward_sample_condition_and_explore` for running invariance +tests using a [`JuliaBUGSPath`](@ref) as target. +""" +function Pigeons.forward_sample_condition_and_explore( + model::JuliaBUGS.BUGSModel, + rng::SplittableRandom; + explorer = nothing, + condition_on = () + ) + # forward simulation (new values stored in model.evaluation_env) + model = evaluate_and_initialize(model, rng) + + # maybe condition the model using the sampled observations + conditioned_model = if length(condition_on) > 0 + var_group = [JuliaBUGS.VarName{sym}() for sym in condition_on] # transform Symbols into VarNames + JuliaBUGS.condition(model, var_group) + else + model + end + + # maybe take a step with explorer + state = JuliaBUGS.getparams(conditioned_model) + return if !isnothing(explorer) + Pigeons.explorer_step(rng, JuliaBUGSPath(conditioned_model), explorer, state) + else + state + end +end + +Pigeons.forward_sample_condition_and_explore(target::JuliaBUGSPath, args...; kwargs...) = + Pigeons.forward_sample_condition_and_explore(target.model, args...; kwargs...) + \ No newline at end of file diff --git a/src/explorers/invariance_test.jl b/src/explorers/invariance_test.jl index 1465bc010..5b75a2061 100644 --- a/src/explorers/invariance_test.jl +++ b/src/explorers/invariance_test.jl @@ -65,12 +65,11 @@ allows direct iid sampling from the target, conditioning is not necessary. """ function forward_sample_condition_and_explore( target::ScaledPrecisionNormalPath, - explorer, rng::SplittableRandom; - run_explorer::Bool = true + explorer = nothing ) state = initialization(target, rng, 1) # forward simulation - if run_explorer + if !isnothing(explorer) state = explorer_step(rng, target, explorer, state) end return state diff --git a/test/test_JuliaBUGS.jl b/test/test_JuliaBUGS.jl index 9301193ca..c7df6fed4 100644 --- a/test/test_JuliaBUGS.jl +++ b/test/test_JuliaBUGS.jl @@ -38,21 +38,24 @@ function Pigeons.step!(::IdentityExplorer, replica, shared) end end @testset "SliceSampler on constrained and unconstrained versions" begin + exact_logZ = unid_target_exact_logZ( + unid_target_model.evaluation_env.n_flips, + unid_target_model.evaluation_env.n_heads + ) for target in (unid_target, unid_target_constrained) @show target.model pt = pigeons(; target, explorer = SliceSampler(), n_chains=7, - n_rounds=7 - ) - @test isapprox( - Pigeons.stepping_stone(pt), - unid_target_exact_logZ( - unid_target_model.evaluation_env.n_flips, - unid_target_model.evaluation_env.n_heads - ), - rtol=0.1 + n_rounds=5 ) + @test isapprox(Pigeons.stepping_stone(pt), exact_logZ, rtol=0.1) end end + +@testset "Invariance test" begin + res = Pigeons.invariance_test(target, SliceSampler(), rng; condition_on=(:n_heads,)) + @show res.pvalues + @test res.passed +end From 536930d0584a91d62a603a8a2930d06b258039be Mon Sep 17 00:00:00 2001 From: Miguel Biron-Lattes Date: Wed, 11 Dec 2024 09:53:11 -0800 Subject: [PATCH 11/32] fix bug --- test/test_JuliaBUGS.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_JuliaBUGS.jl b/test/test_JuliaBUGS.jl index c7df6fed4..bc62433ee 100644 --- a/test/test_JuliaBUGS.jl +++ b/test/test_JuliaBUGS.jl @@ -55,7 +55,7 @@ end end @testset "Invariance test" begin - res = Pigeons.invariance_test(target, SliceSampler(), rng; condition_on=(:n_heads,)) + res = Pigeons.invariance_test(unid_target, SliceSampler(), rng; condition_on=(:n_heads,)) @show res.pvalues @test res.passed end From d1f29ce768beaa0039445a91de5b8fb04281f3c2 Mon Sep 17 00:00:00 2001 From: Miguel Biron-Lattes Date: Wed, 11 Dec 2024 10:28:57 -0800 Subject: [PATCH 12/32] fix --- test/test_JuliaBUGS.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test/test_JuliaBUGS.jl b/test/test_JuliaBUGS.jl index bc62433ee..8a082d27c 100644 --- a/test/test_JuliaBUGS.jl +++ b/test/test_JuliaBUGS.jl @@ -55,7 +55,8 @@ end end @testset "Invariance test" begin - res = Pigeons.invariance_test(unid_target, SliceSampler(), rng; condition_on=(:n_heads,)) + uncond_target = JuliaBUGSPath(compile(unid_model_def, (;n_flips=100000))) + res = Pigeons.invariance_test(uncond_target, SliceSampler(); condition_on=(:n_heads,)) @show res.pvalues @test res.passed end From b427da788d6664380eb7b08f9798ab5e0a3ed3a2 Mon Sep 17 00:00:00 2001 From: Miguel Biron-Lattes Date: Wed, 11 Dec 2024 20:38:59 -0800 Subject: [PATCH 13/32] back to state as namedtuple --- ext/PigeonsJuliaBUGSExt/interface.jl | 50 ++++++++++----- ext/PigeonsJuliaBUGSExt/invariance_test.jl | 5 +- src/paths/JuliaBUGSPath.jl | 5 ++ test/test_JuliaBUGS.jl | 73 ++++++++++++---------- 4 files changed, 82 insertions(+), 51 deletions(-) diff --git a/ext/PigeonsJuliaBUGSExt/interface.jl b/ext/PigeonsJuliaBUGSExt/interface.jl index 4cc3e961d..8a557c676 100644 --- a/ext/PigeonsJuliaBUGSExt/interface.jl +++ b/ext/PigeonsJuliaBUGSExt/interface.jl @@ -2,22 +2,22 @@ # Path interface ####################################### -# Initialization and iid sampling -function evaluate_and_initialize(model::JuliaBUGS.BUGSModel, rng::AbstractRNG) - new_env = first(JuliaBUGS.evaluate!!(rng, model)) # sample a new evaluation environment - return JuliaBUGS.initialize!(model, new_env) # set the private_model's environment to the newly created one +get_symbol(::JuliaBUGS.VarName{sym}) where sym = sym + +function Pigeons.JuliaBUGSPath(model::JuliaBUGS.BUGSModel) + Pigeons.JuliaBUGSPath( + model, + Set(get_symbol(vn) for vn in model.parameters) + ) end # used for both initializing and iid sampling -# Note: state is a flattened vector of the parameters -# Also, the vector is **concretely typed**. This means that if the evaluation -# environment contains floats and integers, the latter will be cast to float. +# sample a new evaluation environment (without resampling observed data) +# Note: JuliaBUGS.evaluate!! deepcopies model.evaluation_env, so the +# new environment is completely independent of model.evaluation_env _sample_iid(model::JuliaBUGS.BUGSModel, rng::AbstractRNG) = - JuliaBUGS.getparams(evaluate_and_initialize(model, rng)) # flatten the unobserved parameters in the model's eval environment and return + first(JuliaBUGS.evaluate!!(rng, model; sample_all=false)) -# Note: JuliaBUGS.getparams creates a new vector on each call, so it is safe -# to call _sample_iid during initialization (**sequentially**, as done as of time -# of writing) for different Replicas (i.e., they won't share the same state). Pigeons.initialization(target::JuliaBUGSPath, rng::AbstractRNG, _::Int64) = _sample_iid(target.model, rng) @@ -36,7 +36,7 @@ temperature parameter. $FIELDS """ -struct JuliaBUGSLogPotential{TMod<:JuliaBUGS.BUGSModel, TF<:AbstractFloat} +mutable struct JuliaBUGSLogPotential{TMod<:JuliaBUGS.BUGSModel, TF<:AbstractFloat, TPars<:Set} """ A deep-enough copy of the original model that allows evaluation while avoiding race conditions between different Replicas. @@ -47,6 +47,11 @@ struct JuliaBUGSLogPotential{TMod<:JuliaBUGS.BUGSModel, TF<:AbstractFloat} Tempering parameter. """ beta::TF + + """ + See [`JuliaBUGSPath`](@ref). + """ + parameter_names::TPars end # make a log-potential by creating a new model with independent graph and @@ -60,12 +65,20 @@ function Pigeons.interpolate(path::JuliaBUGSPath, beta) model.parameters, model.flattened_graph_node_data.sorted_nodes, deepcopy(model.evaluation_env) ) - JuliaBUGSLogPotential(private_model, beta) + JuliaBUGSLogPotential(private_model, beta, path.parameter_names) end # log_potential evaluation -(log_potential::JuliaBUGSLogPotential)(flattened_values) = - try +(log_potential::JuliaBUGSLogPotential)(new_env) = + try + # update model evaluation_env with the one passed + log_potential.private_model = JuliaBUGS.initialize!(log_potential.private_model, new_env) + + # FIXME (temporary hack): ideally we would just call + # `_tempered_evaluate!!(log_potential.private_model; temperature=log_potential.beta) + # but the method does not exist yet, and so we must flatten first + # see https://github.com/TuringLang/JuliaBUGS.jl/issues/260 + flattened_values = JuliaBUGS.getparams(log_potential.private_model) last(last(JuliaBUGS._tempered_evaluate!!( log_potential.private_model, flattened_values; @@ -80,3 +93,10 @@ end function Pigeons.sample_iid!(log_potential::JuliaBUGSLogPotential, replica, shared) replica.state = _sample_iid(log_potential.private_model, replica.rng) end + +# custom sample extraction, needed because +# 1) evaluation_env contains both observations and parameters (only want the latter) +# 2) there is no copy method for NamedTuples +Pigeons.extract_sample(state::NamedTuple, log_potential::JuliaBUGSLogPotential) = + NamedTuple(k => copy(v) for (k,v) in zip(keys(state), state) if k in log_potential.parameter_names) + diff --git a/ext/PigeonsJuliaBUGSExt/invariance_test.jl b/ext/PigeonsJuliaBUGSExt/invariance_test.jl index 2150693b8..1eafa75fa 100644 --- a/ext/PigeonsJuliaBUGSExt/invariance_test.jl +++ b/ext/PigeonsJuliaBUGSExt/invariance_test.jl @@ -10,8 +10,9 @@ function Pigeons.forward_sample_condition_and_explore( explorer = nothing, condition_on = () ) - # forward simulation (new values stored in model.evaluation_env) - model = evaluate_and_initialize(model, rng) + # sample and then store results in model evaluation_env + new_env = _sample_iid(target.model, rng) + model = JuliaBUGS.initialize!(model, new_env) # maybe condition the model using the sampled observations conditioned_model = if length(condition_on) > 0 diff --git a/src/paths/JuliaBUGSPath.jl b/src/paths/JuliaBUGSPath.jl index 1c67ae68d..3ca1c0282 100644 --- a/src/paths/JuliaBUGSPath.jl +++ b/src/paths/JuliaBUGSPath.jl @@ -11,4 +11,9 @@ $FIELDS A `JuliaBUGS.BUGSModel`. """ model + + """ + Set of names of unobserved parameters in the model. + """ + parameter_names end diff --git a/test/test_JuliaBUGS.jl b/test/test_JuliaBUGS.jl index 8a082d27c..611adc63a 100644 --- a/test/test_JuliaBUGS.jl +++ b/test/test_JuliaBUGS.jl @@ -15,48 +15,53 @@ unid_target = JuliaBUGSPath(unid_target_model) unid_target_constrained = JuliaBUGSPath(JuliaBUGS.settrans(unid_target_model)) struct IdentityExplorer end function Pigeons.step!(::IdentityExplorer, replica, shared) end +exact_logZ = unid_target_exact_logZ( + unid_target_model.evaluation_env.n_flips, + unid_target_model.evaluation_env.n_heads +) @testset "Basic sampling via independent MH from the prior" begin pt = pigeons( target = unid_target_constrained, - n_chains = 2, + n_chains = 8, explorer = IdentityExplorer(), record = [traces], extended_traces = true ) - # check sample_iid! - @test isapprox(0.5, mean(v[1] for (k,v) in pt.reduced_recorders.traces if first(k)==1), rtol=0.05) - @test isapprox(0.5, mean(v[2] for (k,v) in pt.reduced_recorders.traces if first(k)==1), rtol=0.05) - - # check log_potential evaluation with constrained version (easier, no Jacobian) - @test all(v for (k,v) in pt.reduced_recorders.traces if first(k)==2) do v - last(v) == logpdf( - Binomial(unid_target_model.evaluation_env.n_flips, v[1]*v[2]), - unid_target_model.evaluation_env.n_heads - ) - end -end - -@testset "SliceSampler on constrained and unconstrained versions" begin - exact_logZ = unid_target_exact_logZ( - unid_target_model.evaluation_env.n_flips, - unid_target_model.evaluation_env.n_heads + # check normalizing constant + @test isapprox(Pigeons.stepping_stone(pt), exact_logZ, rtol=0.1) + + # check state extraction + # check sample_iid! produces p1,p2 ~iid U(0,1) => E[p1p2]=0.25 + @test isapprox( + mean(prod(v.p) for (k,v) in pt.reduced_recorders.traces if first(k)==1), + 0.25, + rtol=0.05 + ) + # check posterior E[p1*p2] ~ 0.5 = n_heads/n_flips + @test isapprox( + mean(prod(v.p) for (k,v) in pt.reduced_recorders.traces if first(k)==2), + 0.5, + rtol=0.05 ) - for target in (unid_target, unid_target_constrained) - @show target.model - pt = pigeons(; - target, - explorer = SliceSampler(), - n_chains=7, - n_rounds=5 - ) - @test isapprox(Pigeons.stepping_stone(pt), exact_logZ, rtol=0.1) - end end -@testset "Invariance test" begin - uncond_target = JuliaBUGSPath(compile(unid_model_def, (;n_flips=100000))) - res = Pigeons.invariance_test(uncond_target, SliceSampler(); condition_on=(:n_heads,)) - @show res.pvalues - @test res.passed -end +# @testset "SliceSampler on constrained and unconstrained versions" begin +# for target in (unid_target, unid_target_constrained) +# @show target.model +# pt = pigeons(; +# target, +# explorer = SliceSampler(), +# n_chains=7, +# n_rounds=5 +# ) +# @test isapprox(Pigeons.stepping_stone(pt), exact_logZ, rtol=0.1) +# end +# end + +# @testset "Invariance test" begin +# uncond_target = JuliaBUGSPath(compile(unid_model_def, (;n_flips=100000))) +# res = Pigeons.invariance_test(uncond_target, SliceSampler(); condition_on=(:n_heads,)) +# @show res.pvalues +# @test res.passed +# end From cfe5491f27bda2c1bbec8d365f4764f39c6d90bf Mon Sep 17 00:00:00 2001 From: Miguel Biron-Lattes Date: Tue, 10 Dec 2024 17:08:08 -0800 Subject: [PATCH 14/32] bump version --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 5f360c74d..4960f19ea 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "Pigeons" uuid = "0eb8d820-af6a-4919-95ae-11206f830c31" authors = ["Alexandre Bouchard-Côté , Nikola Surjanovic , Paul Tiede , Trevor Campbell, Miguel Biron-Lattes, Saifuddin Syed"] -version = "0.4.7" +version = "0.4.8" [deps] DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" From 7e4aade90e52d330dcdad9fd71a23ee0f93653fa Mon Sep 17 00:00:00 2001 From: Miguel Biron-Lattes Date: Thu, 12 Dec 2024 08:48:46 -0800 Subject: [PATCH 15/32] Revert "fix" This reverts commit d1f29ce768beaa0039445a91de5b8fb04281f3c2. --- test/test_JuliaBUGS.jl | 40 ++++++++++++++++++++++------------------ 1 file changed, 22 insertions(+), 18 deletions(-) diff --git a/test/test_JuliaBUGS.jl b/test/test_JuliaBUGS.jl index 611adc63a..09cc6b829 100644 --- a/test/test_JuliaBUGS.jl +++ b/test/test_JuliaBUGS.jl @@ -46,22 +46,26 @@ exact_logZ = unid_target_exact_logZ( ) end -# @testset "SliceSampler on constrained and unconstrained versions" begin -# for target in (unid_target, unid_target_constrained) -# @show target.model -# pt = pigeons(; -# target, -# explorer = SliceSampler(), -# n_chains=7, -# n_rounds=5 -# ) -# @test isapprox(Pigeons.stepping_stone(pt), exact_logZ, rtol=0.1) -# end -# end +@testset "SliceSampler on constrained and unconstrained versions" begin + exact_logZ = unid_target_exact_logZ( + unid_target_model.evaluation_env.n_flips, + unid_target_model.evaluation_env.n_heads + ) + for target in (unid_target, unid_target_constrained) + @show target.model + pt = pigeons(; + target, + explorer = SliceSampler(), + n_chains=7, + n_rounds=5 + ) + @test isapprox(Pigeons.stepping_stone(pt), exact_logZ, rtol=0.1) + end +end -# @testset "Invariance test" begin -# uncond_target = JuliaBUGSPath(compile(unid_model_def, (;n_flips=100000))) -# res = Pigeons.invariance_test(uncond_target, SliceSampler(); condition_on=(:n_heads,)) -# @show res.pvalues -# @test res.passed -# end +@testset "Invariance test" begin + uncond_target = JuliaBUGSPath(compile(unid_model_def, (;n_flips=100000))) + res = Pigeons.invariance_test(uncond_target, SliceSampler(); condition_on=(:n_heads,)) + @show res.pvalues + @test res.passed +end From ece62905758b1ab8e3e29ecb400614176e167a47 Mon Sep 17 00:00:00 2001 From: Miguel Biron-Lattes Date: Thu, 12 Dec 2024 08:59:10 -0800 Subject: [PATCH 16/32] Reapply "fix" This reverts commit 7e4aade90e52d330dcdad9fd71a23ee0f93653fa. --- test/test_JuliaBUGS.jl | 40 ++++++++++++++++++---------------------- 1 file changed, 18 insertions(+), 22 deletions(-) diff --git a/test/test_JuliaBUGS.jl b/test/test_JuliaBUGS.jl index 09cc6b829..611adc63a 100644 --- a/test/test_JuliaBUGS.jl +++ b/test/test_JuliaBUGS.jl @@ -46,26 +46,22 @@ exact_logZ = unid_target_exact_logZ( ) end -@testset "SliceSampler on constrained and unconstrained versions" begin - exact_logZ = unid_target_exact_logZ( - unid_target_model.evaluation_env.n_flips, - unid_target_model.evaluation_env.n_heads - ) - for target in (unid_target, unid_target_constrained) - @show target.model - pt = pigeons(; - target, - explorer = SliceSampler(), - n_chains=7, - n_rounds=5 - ) - @test isapprox(Pigeons.stepping_stone(pt), exact_logZ, rtol=0.1) - end -end +# @testset "SliceSampler on constrained and unconstrained versions" begin +# for target in (unid_target, unid_target_constrained) +# @show target.model +# pt = pigeons(; +# target, +# explorer = SliceSampler(), +# n_chains=7, +# n_rounds=5 +# ) +# @test isapprox(Pigeons.stepping_stone(pt), exact_logZ, rtol=0.1) +# end +# end -@testset "Invariance test" begin - uncond_target = JuliaBUGSPath(compile(unid_model_def, (;n_flips=100000))) - res = Pigeons.invariance_test(uncond_target, SliceSampler(); condition_on=(:n_heads,)) - @show res.pvalues - @test res.passed -end +# @testset "Invariance test" begin +# uncond_target = JuliaBUGSPath(compile(unid_model_def, (;n_flips=100000))) +# res = Pigeons.invariance_test(uncond_target, SliceSampler(); condition_on=(:n_heads,)) +# @show res.pvalues +# @test res.passed +# end From ae3c821b036f4359d1c5f79611b7c8a7e529f806 Mon Sep 17 00:00:00 2001 From: Miguel Biron-Lattes Date: Thu, 12 Dec 2024 08:59:28 -0800 Subject: [PATCH 17/32] Revert "bump version" This reverts commit cfe5491f27bda2c1bbec8d365f4764f39c6d90bf. --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 4960f19ea..5f360c74d 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "Pigeons" uuid = "0eb8d820-af6a-4919-95ae-11206f830c31" authors = ["Alexandre Bouchard-Côté , Nikola Surjanovic , Paul Tiede , Trevor Campbell, Miguel Biron-Lattes, Saifuddin Syed"] -version = "0.4.8" +version = "0.4.7" [deps] DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" From 0973ba411f682d6231754fbbfb6bb994704e21ff Mon Sep 17 00:00:00 2001 From: Miguel Biron-Lattes Date: Thu, 12 Dec 2024 08:59:30 -0800 Subject: [PATCH 18/32] Revert "back to state as namedtuple" This reverts commit b427da788d6664380eb7b08f9798ab5e0a3ed3a2. --- ext/PigeonsJuliaBUGSExt/interface.jl | 50 +++++---------- ext/PigeonsJuliaBUGSExt/invariance_test.jl | 5 +- src/paths/JuliaBUGSPath.jl | 5 -- test/test_JuliaBUGS.jl | 73 ++++++++++------------ 4 files changed, 51 insertions(+), 82 deletions(-) diff --git a/ext/PigeonsJuliaBUGSExt/interface.jl b/ext/PigeonsJuliaBUGSExt/interface.jl index 8a557c676..4cc3e961d 100644 --- a/ext/PigeonsJuliaBUGSExt/interface.jl +++ b/ext/PigeonsJuliaBUGSExt/interface.jl @@ -2,22 +2,22 @@ # Path interface ####################################### -get_symbol(::JuliaBUGS.VarName{sym}) where sym = sym - -function Pigeons.JuliaBUGSPath(model::JuliaBUGS.BUGSModel) - Pigeons.JuliaBUGSPath( - model, - Set(get_symbol(vn) for vn in model.parameters) - ) +# Initialization and iid sampling +function evaluate_and_initialize(model::JuliaBUGS.BUGSModel, rng::AbstractRNG) + new_env = first(JuliaBUGS.evaluate!!(rng, model)) # sample a new evaluation environment + return JuliaBUGS.initialize!(model, new_env) # set the private_model's environment to the newly created one end # used for both initializing and iid sampling -# sample a new evaluation environment (without resampling observed data) -# Note: JuliaBUGS.evaluate!! deepcopies model.evaluation_env, so the -# new environment is completely independent of model.evaluation_env +# Note: state is a flattened vector of the parameters +# Also, the vector is **concretely typed**. This means that if the evaluation +# environment contains floats and integers, the latter will be cast to float. _sample_iid(model::JuliaBUGS.BUGSModel, rng::AbstractRNG) = - first(JuliaBUGS.evaluate!!(rng, model; sample_all=false)) + JuliaBUGS.getparams(evaluate_and_initialize(model, rng)) # flatten the unobserved parameters in the model's eval environment and return +# Note: JuliaBUGS.getparams creates a new vector on each call, so it is safe +# to call _sample_iid during initialization (**sequentially**, as done as of time +# of writing) for different Replicas (i.e., they won't share the same state). Pigeons.initialization(target::JuliaBUGSPath, rng::AbstractRNG, _::Int64) = _sample_iid(target.model, rng) @@ -36,7 +36,7 @@ temperature parameter. $FIELDS """ -mutable struct JuliaBUGSLogPotential{TMod<:JuliaBUGS.BUGSModel, TF<:AbstractFloat, TPars<:Set} +struct JuliaBUGSLogPotential{TMod<:JuliaBUGS.BUGSModel, TF<:AbstractFloat} """ A deep-enough copy of the original model that allows evaluation while avoiding race conditions between different Replicas. @@ -47,11 +47,6 @@ mutable struct JuliaBUGSLogPotential{TMod<:JuliaBUGS.BUGSModel, TF<:AbstractFloa Tempering parameter. """ beta::TF - - """ - See [`JuliaBUGSPath`](@ref). - """ - parameter_names::TPars end # make a log-potential by creating a new model with independent graph and @@ -65,20 +60,12 @@ function Pigeons.interpolate(path::JuliaBUGSPath, beta) model.parameters, model.flattened_graph_node_data.sorted_nodes, deepcopy(model.evaluation_env) ) - JuliaBUGSLogPotential(private_model, beta, path.parameter_names) + JuliaBUGSLogPotential(private_model, beta) end # log_potential evaluation -(log_potential::JuliaBUGSLogPotential)(new_env) = - try - # update model evaluation_env with the one passed - log_potential.private_model = JuliaBUGS.initialize!(log_potential.private_model, new_env) - - # FIXME (temporary hack): ideally we would just call - # `_tempered_evaluate!!(log_potential.private_model; temperature=log_potential.beta) - # but the method does not exist yet, and so we must flatten first - # see https://github.com/TuringLang/JuliaBUGS.jl/issues/260 - flattened_values = JuliaBUGS.getparams(log_potential.private_model) +(log_potential::JuliaBUGSLogPotential)(flattened_values) = + try last(last(JuliaBUGS._tempered_evaluate!!( log_potential.private_model, flattened_values; @@ -93,10 +80,3 @@ end function Pigeons.sample_iid!(log_potential::JuliaBUGSLogPotential, replica, shared) replica.state = _sample_iid(log_potential.private_model, replica.rng) end - -# custom sample extraction, needed because -# 1) evaluation_env contains both observations and parameters (only want the latter) -# 2) there is no copy method for NamedTuples -Pigeons.extract_sample(state::NamedTuple, log_potential::JuliaBUGSLogPotential) = - NamedTuple(k => copy(v) for (k,v) in zip(keys(state), state) if k in log_potential.parameter_names) - diff --git a/ext/PigeonsJuliaBUGSExt/invariance_test.jl b/ext/PigeonsJuliaBUGSExt/invariance_test.jl index 1eafa75fa..2150693b8 100644 --- a/ext/PigeonsJuliaBUGSExt/invariance_test.jl +++ b/ext/PigeonsJuliaBUGSExt/invariance_test.jl @@ -10,9 +10,8 @@ function Pigeons.forward_sample_condition_and_explore( explorer = nothing, condition_on = () ) - # sample and then store results in model evaluation_env - new_env = _sample_iid(target.model, rng) - model = JuliaBUGS.initialize!(model, new_env) + # forward simulation (new values stored in model.evaluation_env) + model = evaluate_and_initialize(model, rng) # maybe condition the model using the sampled observations conditioned_model = if length(condition_on) > 0 diff --git a/src/paths/JuliaBUGSPath.jl b/src/paths/JuliaBUGSPath.jl index 3ca1c0282..1c67ae68d 100644 --- a/src/paths/JuliaBUGSPath.jl +++ b/src/paths/JuliaBUGSPath.jl @@ -11,9 +11,4 @@ $FIELDS A `JuliaBUGS.BUGSModel`. """ model - - """ - Set of names of unobserved parameters in the model. - """ - parameter_names end diff --git a/test/test_JuliaBUGS.jl b/test/test_JuliaBUGS.jl index 611adc63a..8a082d27c 100644 --- a/test/test_JuliaBUGS.jl +++ b/test/test_JuliaBUGS.jl @@ -15,53 +15,48 @@ unid_target = JuliaBUGSPath(unid_target_model) unid_target_constrained = JuliaBUGSPath(JuliaBUGS.settrans(unid_target_model)) struct IdentityExplorer end function Pigeons.step!(::IdentityExplorer, replica, shared) end -exact_logZ = unid_target_exact_logZ( - unid_target_model.evaluation_env.n_flips, - unid_target_model.evaluation_env.n_heads -) @testset "Basic sampling via independent MH from the prior" begin pt = pigeons( target = unid_target_constrained, - n_chains = 8, + n_chains = 2, explorer = IdentityExplorer(), record = [traces], extended_traces = true ) - # check normalizing constant - @test isapprox(Pigeons.stepping_stone(pt), exact_logZ, rtol=0.1) - - # check state extraction - # check sample_iid! produces p1,p2 ~iid U(0,1) => E[p1p2]=0.25 - @test isapprox( - mean(prod(v.p) for (k,v) in pt.reduced_recorders.traces if first(k)==1), - 0.25, - rtol=0.05 - ) - # check posterior E[p1*p2] ~ 0.5 = n_heads/n_flips - @test isapprox( - mean(prod(v.p) for (k,v) in pt.reduced_recorders.traces if first(k)==2), - 0.5, - rtol=0.05 - ) + # check sample_iid! + @test isapprox(0.5, mean(v[1] for (k,v) in pt.reduced_recorders.traces if first(k)==1), rtol=0.05) + @test isapprox(0.5, mean(v[2] for (k,v) in pt.reduced_recorders.traces if first(k)==1), rtol=0.05) + + # check log_potential evaluation with constrained version (easier, no Jacobian) + @test all(v for (k,v) in pt.reduced_recorders.traces if first(k)==2) do v + last(v) == logpdf( + Binomial(unid_target_model.evaluation_env.n_flips, v[1]*v[2]), + unid_target_model.evaluation_env.n_heads + ) + end end -# @testset "SliceSampler on constrained and unconstrained versions" begin -# for target in (unid_target, unid_target_constrained) -# @show target.model -# pt = pigeons(; -# target, -# explorer = SliceSampler(), -# n_chains=7, -# n_rounds=5 -# ) -# @test isapprox(Pigeons.stepping_stone(pt), exact_logZ, rtol=0.1) -# end -# end +@testset "SliceSampler on constrained and unconstrained versions" begin + exact_logZ = unid_target_exact_logZ( + unid_target_model.evaluation_env.n_flips, + unid_target_model.evaluation_env.n_heads + ) + for target in (unid_target, unid_target_constrained) + @show target.model + pt = pigeons(; + target, + explorer = SliceSampler(), + n_chains=7, + n_rounds=5 + ) + @test isapprox(Pigeons.stepping_stone(pt), exact_logZ, rtol=0.1) + end +end -# @testset "Invariance test" begin -# uncond_target = JuliaBUGSPath(compile(unid_model_def, (;n_flips=100000))) -# res = Pigeons.invariance_test(uncond_target, SliceSampler(); condition_on=(:n_heads,)) -# @show res.pvalues -# @test res.passed -# end +@testset "Invariance test" begin + uncond_target = JuliaBUGSPath(compile(unid_model_def, (;n_flips=100000))) + res = Pigeons.invariance_test(uncond_target, SliceSampler(); condition_on=(:n_heads,)) + @show res.pvalues + @test res.passed +end From 7db17095ba4494e93fc1a543464a5d67a5ef2d14 Mon Sep 17 00:00:00 2001 From: Miguel Biron-Lattes Date: Thu, 12 Dec 2024 09:22:48 -0800 Subject: [PATCH 19/32] mixed eltype vector --- Project.toml | 8 +++- .../PigeonsJuliaBUGSExt.jl | 7 +++- ext/PigeonsJuliaBUGSExt/interface.jl | 2 +- ext/PigeonsJuliaBUGSExt/invariance_test.jl | 2 +- ext/PigeonsJuliaBUGSExt/utils.jl | 38 +++++++++++++++++++ 5 files changed, 52 insertions(+), 5 deletions(-) create mode 100644 ext/PigeonsJuliaBUGSExt/utils.jl diff --git a/Project.toml b/Project.toml index 5f360c74d..f609cd50d 100644 --- a/Project.toml +++ b/Project.toml @@ -37,6 +37,8 @@ StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" ZipFile = "a5390f91-8eb1-5f08-bee0-b1d1ffed6cea" [weakdeps] +AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" +Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" BridgeStan = "c88b6f0a-829e-4b0b-94b7-f06ab5908f5a" DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" @@ -52,11 +54,13 @@ PigeonsDynamicPPLExt = "DynamicPPL" PigeonsEnzymeExt = "Enzyme" PigeonsForwardDiffExt = "ForwardDiff" PigeonsHypothesisTestsExt = "HypothesisTests" -PigeonsJuliaBUGSExt = "JuliaBUGS" +PigeonsJuliaBUGSExt = ["JuliaBUGS", "AbstractPPL", "Bijectors"] PigeonsMCMCChainsExt = "MCMCChains" PigeonsReverseDiffExt = "ReverseDiff" [compat] +AbstractPPL = "0.9" +Bijectors = "0.14" BridgeStan = "2" DataFrames = "1" Distributions = "0.25" @@ -93,6 +97,8 @@ ZipFile = "0.10" julia = "1.8" [extras] +AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" +Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" BridgeStan = "c88b6f0a-829e-4b0b-94b7-f06ab5908f5a" DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" diff --git a/ext/PigeonsJuliaBUGSExt/PigeonsJuliaBUGSExt.jl b/ext/PigeonsJuliaBUGSExt/PigeonsJuliaBUGSExt.jl index dd6dc4f55..d797e36e5 100644 --- a/ext/PigeonsJuliaBUGSExt/PigeonsJuliaBUGSExt.jl +++ b/ext/PigeonsJuliaBUGSExt/PigeonsJuliaBUGSExt.jl @@ -3,18 +3,21 @@ module PigeonsJuliaBUGSExt using Pigeons if isdefined(Base, :get_extension) import JuliaBUGS - using LogDensityProblems + using AbstractPPL # only need because we rewrite JuliaBUGS.getparams + using Bijectors # only need because we rewrite JuliaBUGS.getparams using DocStringExtensions using SplittableRandoms: SplittableRandom, split using Random else import ..JuliaBUGS - using ..LogDensityProblems + using ..AbstractPPL # only need because we rewrite JuliaBUGS.getparams + using ..Bijectors # only need because we rewrite JuliaBUGS.getparams using ..DocStringExtensions using ..SplittableRandoms: SplittableRandom, split using ..Random end +include(joinpath(@__DIR__, "utils.jl")) include(joinpath(@__DIR__, "interface.jl")) include(joinpath(@__DIR__, "invariance_test.jl")) diff --git a/ext/PigeonsJuliaBUGSExt/interface.jl b/ext/PigeonsJuliaBUGSExt/interface.jl index 4cc3e961d..4c3c0d061 100644 --- a/ext/PigeonsJuliaBUGSExt/interface.jl +++ b/ext/PigeonsJuliaBUGSExt/interface.jl @@ -13,7 +13,7 @@ end # Also, the vector is **concretely typed**. This means that if the evaluation # environment contains floats and integers, the latter will be cast to float. _sample_iid(model::JuliaBUGS.BUGSModel, rng::AbstractRNG) = - JuliaBUGS.getparams(evaluate_and_initialize(model, rng)) # flatten the unobserved parameters in the model's eval environment and return + getparams(evaluate_and_initialize(model, rng)) # flatten the unobserved parameters in the model's eval environment and return # Note: JuliaBUGS.getparams creates a new vector on each call, so it is safe # to call _sample_iid during initialization (**sequentially**, as done as of time diff --git a/ext/PigeonsJuliaBUGSExt/invariance_test.jl b/ext/PigeonsJuliaBUGSExt/invariance_test.jl index 2150693b8..fa44a0731 100644 --- a/ext/PigeonsJuliaBUGSExt/invariance_test.jl +++ b/ext/PigeonsJuliaBUGSExt/invariance_test.jl @@ -22,7 +22,7 @@ function Pigeons.forward_sample_condition_and_explore( end # maybe take a step with explorer - state = JuliaBUGS.getparams(conditioned_model) + state = getparams(conditioned_model) return if !isnothing(explorer) Pigeons.explorer_step(rng, JuliaBUGSPath(conditioned_model), explorer, state) else diff --git a/ext/PigeonsJuliaBUGSExt/utils.jl b/ext/PigeonsJuliaBUGSExt/utils.jl new file mode 100644 index 000000000..49d075e3c --- /dev/null +++ b/ext/PigeonsJuliaBUGSExt/utils.jl @@ -0,0 +1,38 @@ +#= +Tweak of JuliaBUGS.getparams to allow for flattened vectors of mixed type +=# +function getparams(model::JuliaBUGS.BUGSModel) + param_length = if model.transformed + model.transformed_param_length + else + model.untransformed_param_length + end + + param_vals = Vector{Real}(undef, param_length) # NB: use mixed type vector for correct dispatch in SliceSampler + pos = 1 + for v in model.parameters + if !model.transformed + val = AbstractPPL.get(model.evaluation_env, v) + len = model.untransformed_var_lengths[v] + if val isa AbstractArray + param_vals[pos:(pos + len - 1)] .= vec(val) + else + param_vals[pos] = val + end + else + (; node_function, loop_vars) = model.g[v] + dist = node_function(model.evaluation_env, loop_vars) + transformed_value = Bijectors.transform( + Bijectors.bijector(dist), AbstractPPL.get(model.evaluation_env, v) + ) + len = model.transformed_var_lengths[v] + if transformed_value isa AbstractArray + param_vals[pos:(pos + len - 1)] .= vec(transformed_value) + else + param_vals[pos] = transformed_value + end + end + pos += len + end + return param_vals +end \ No newline at end of file From 4c6c6e5c431776dc590f0459279194ea69b50fed Mon Sep 17 00:00:00 2001 From: Miguel Biron-Lattes Date: Thu, 12 Dec 2024 09:40:29 -0800 Subject: [PATCH 20/32] test incomplete count data model --- examples/JuliaBUGS.jl | 23 +++++++++++++++++++++++ test/test_JuliaBUGS.jl | 6 ++++++ 2 files changed, 29 insertions(+) create mode 100644 examples/JuliaBUGS.jl diff --git a/examples/JuliaBUGS.jl b/examples/JuliaBUGS.jl new file mode 100644 index 000000000..4e1c7109f --- /dev/null +++ b/examples/JuliaBUGS.jl @@ -0,0 +1,23 @@ +function incomplete_count_data() + model_def = @bugs("model{ + for (i in 1:n) { + r[i] ~ dbern(pr[i]) + pr[i] <- ilogit(y[i] * alpha1 + alpha0) + y[i] ~ dpois(mu) + } + mu ~ dgamma(1,1) + alpha0 ~ dnorm(0, 0.1) + alpha1 ~ dnorm(0, tau) + }",false,false + ) + data = ( + y = [ + 6,missing,missing,missing,missing,missing,missing,5,1,missing,1,missing, + missing,missing,2,missing,missing,0,missing,1,2,1,7,4,6,missing,missing, + missing,5,missing + ], + r = [1,0,0,0,0,0,0,1,1,0,1,0,0,0,1,0,0,1,0,1,1,1,1,1,1,0,0,0,1,0], + n = 30, tau = 4 + ) + return JuliaBUGSPath(compile(model_def, data)) +end diff --git a/test/test_JuliaBUGS.jl b/test/test_JuliaBUGS.jl index 8a082d27c..d92c24ae3 100644 --- a/test/test_JuliaBUGS.jl +++ b/test/test_JuliaBUGS.jl @@ -1,6 +1,7 @@ using JuliaBUGS include("supporting/analytic_solutions.jl") +include("../examples/JuliaBUGS.jl") # good ol' toy unidentifiable model for testing purposes unid_model_def = @bugs begin @@ -60,3 +61,8 @@ end @show res.pvalues @test res.passed end + +@testset "Model with mixed state types" begin + pt = pigeons(;target = incomplete_count_data(), n_chains = 5, n_rounds = 5) + @test true +end \ No newline at end of file From 6b4bf0a37fa6888ddfefa7fa6c065a87fecd9ee5 Mon Sep 17 00:00:00 2001 From: Miguel Biron-Lattes Date: Thu, 12 Dec 2024 10:23:51 -0800 Subject: [PATCH 21/32] bump JuliaBUGS compat --- Project.toml | 2 +- examples/JuliaBUGS.jl | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/Project.toml b/Project.toml index f609cd50d..7e6d5de0c 100644 --- a/Project.toml +++ b/Project.toml @@ -74,7 +74,7 @@ Graphs = "1" HypothesisTests = "0.11" Interpolations = "0.14, 0.15" JSON = "0.21" -JuliaBUGS = "0.7" +JuliaBUGS = "0.8" LogDensityProblems = "2" LogDensityProblemsAD = "1" LogExpFunctions = "0.3" diff --git a/examples/JuliaBUGS.jl b/examples/JuliaBUGS.jl index 4e1c7109f..0390a5461 100644 --- a/examples/JuliaBUGS.jl +++ b/examples/JuliaBUGS.jl @@ -1,4 +1,4 @@ -function incomplete_count_data() +function incomplete_count_data_model() model_def = @bugs("model{ for (i in 1:n) { r[i] ~ dbern(pr[i]) @@ -19,5 +19,6 @@ function incomplete_count_data() r = [1,0,0,0,0,0,0,1,1,0,1,0,0,0,1,0,0,1,0,1,1,1,1,1,1,0,0,0,1,0], n = 30, tau = 4 ) - return JuliaBUGSPath(compile(model_def, data)) + return compile(model_def, data) end +incomplete_count_data() = JuliaBUGSPath(incomplete_count_data_model()) \ No newline at end of file From 4d73dc39b7951f8b6bc0ee90326681d3f9c995c9 Mon Sep 17 00:00:00 2001 From: Miguel Biron-Lattes Date: Thu, 12 Dec 2024 12:37:05 -0800 Subject: [PATCH 22/32] match new version number --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 7e6d5de0c..11eb4819f 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "Pigeons" uuid = "0eb8d820-af6a-4919-95ae-11206f830c31" authors = ["Alexandre Bouchard-Côté , Nikola Surjanovic , Paul Tiede , Trevor Campbell, Miguel Biron-Lattes, Saifuddin Syed"] -version = "0.4.7" +version = "0.4.8" [deps] DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" From 2a190958ea75363c20e94685e6b2be8d9fedac89 Mon Sep 17 00:00:00 2001 From: Seren Lee Date: Sat, 14 Dec 2024 03:38:18 -0800 Subject: [PATCH 23/32] Bijectors 0.13 version compat --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 11eb4819f..8c95fa857 100644 --- a/Project.toml +++ b/Project.toml @@ -60,7 +60,7 @@ PigeonsReverseDiffExt = "ReverseDiff" [compat] AbstractPPL = "0.9" -Bijectors = "0.14" +Bijectors = "0.13, 0.14" BridgeStan = "2" DataFrames = "1" Distributions = "0.25" From e9224810e97b375c595f4699569a17688188c169 Mon Sep 17 00:00:00 2001 From: Seren Lee Date: Sat, 14 Dec 2024 03:42:37 -0800 Subject: [PATCH 24/32] AbstractPPL ver --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 8c95fa857..53749f127 100644 --- a/Project.toml +++ b/Project.toml @@ -59,7 +59,7 @@ PigeonsMCMCChainsExt = "MCMCChains" PigeonsReverseDiffExt = "ReverseDiff" [compat] -AbstractPPL = "0.9" +AbstractPPL = "0.8.4, 0.9" Bijectors = "0.13, 0.14" BridgeStan = "2" DataFrames = "1" From 722dde8f2a6937a5d320582043e421881f24e4ed Mon Sep 17 00:00:00 2001 From: Miguel Biron-Lattes Date: Sun, 15 Dec 2024 21:31:28 -0800 Subject: [PATCH 25/32] finer mix type --- ext/PigeonsJuliaBUGSExt/utils.jl | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/ext/PigeonsJuliaBUGSExt/utils.jl b/ext/PigeonsJuliaBUGSExt/utils.jl index 49d075e3c..4bbda7fc4 100644 --- a/ext/PigeonsJuliaBUGSExt/utils.jl +++ b/ext/PigeonsJuliaBUGSExt/utils.jl @@ -1,6 +1,7 @@ #= Tweak of JuliaBUGS.getparams to allow for flattened vectors of mixed type =# +type_join_eval_env(env) = typejoin(Set(eltype(v) for v in env)...) function getparams(model::JuliaBUGS.BUGSModel) param_length = if model.transformed model.transformed_param_length @@ -8,7 +9,11 @@ function getparams(model::JuliaBUGS.BUGSModel) model.untransformed_param_length end - param_vals = Vector{Real}(undef, param_length) # NB: use mixed type vector for correct dispatch in SliceSampler + # search for an umbrella type for all parameters in the model to avoid + # promotion of e.g. ints to floats. For models with a unique parameter + # type T, it holds that TMix=T. + TMix = type_join_eval_env(model.evaluation_env) + param_vals = Vector{TMix}(undef, param_length) pos = 1 for v in model.parameters if !model.transformed From 0e4bb7203a324904bdecc63fdb45b50b80157764 Mon Sep 17 00:00:00 2001 From: Miguel Biron-Lattes Date: Mon, 16 Dec 2024 20:17:08 -0800 Subject: [PATCH 26/32] fix private model builder --- ext/PigeonsJuliaBUGSExt/interface.jl | 7 +------ ext/PigeonsJuliaBUGSExt/utils.jl | 22 +++++++++++++++++++++- test/test_JuliaBUGS.jl | 22 +++++++++++++++++++--- 3 files changed, 41 insertions(+), 10 deletions(-) diff --git a/ext/PigeonsJuliaBUGSExt/interface.jl b/ext/PigeonsJuliaBUGSExt/interface.jl index 4c3c0d061..fa2b46bef 100644 --- a/ext/PigeonsJuliaBUGSExt/interface.jl +++ b/ext/PigeonsJuliaBUGSExt/interface.jl @@ -54,12 +54,7 @@ end # evaluations and/or during Gibbs sampling function Pigeons.interpolate(path::JuliaBUGSPath, beta) model = path.model - private_model = JuliaBUGS.BUGSModel( - model, - deepcopy(model.g), - model.parameters, model.flattened_graph_node_data.sorted_nodes, - deepcopy(model.evaluation_env) - ) + private_model = make_private_model_copy(model) JuliaBUGSLogPotential(private_model, beta) end diff --git a/ext/PigeonsJuliaBUGSExt/utils.jl b/ext/PigeonsJuliaBUGSExt/utils.jl index 4bbda7fc4..0a1ebedc1 100644 --- a/ext/PigeonsJuliaBUGSExt/utils.jl +++ b/ext/PigeonsJuliaBUGSExt/utils.jl @@ -40,4 +40,24 @@ function getparams(model::JuliaBUGS.BUGSModel) pos += len end return param_vals -end \ No newline at end of file +end + +function make_private_model_copy(model::JuliaBUGS.BUGSModel) + g = deepcopy(model.g) + parameters = model.parameters + sorted_nodes = model.flattened_graph_node_data.sorted_nodes + return JuliaBUGS.BUGSModel( + model.transformed, + sum(model.untransformed_var_lengths[v] for v in parameters), + sum(model.transformed_var_lengths[v] for v in parameters), + model.untransformed_var_lengths, + model.transformed_var_lengths, + deepcopy(model.evaluation_env), + parameters, + JuliaBUGS.FlattenedGraphNodeData(g, sorted_nodes), + g, + nothing, + model.model_def, + model.data + ) +end diff --git a/test/test_JuliaBUGS.jl b/test/test_JuliaBUGS.jl index d92c24ae3..992062742 100644 --- a/test/test_JuliaBUGS.jl +++ b/test/test_JuliaBUGS.jl @@ -1,6 +1,7 @@ using JuliaBUGS include("supporting/analytic_solutions.jl") +include("supporting/mpi_test_utils.jl") include("../examples/JuliaBUGS.jl") # good ol' toy unidentifiable model for testing purposes @@ -62,7 +63,22 @@ end @test res.passed end -@testset "Model with mixed state types" begin - pt = pigeons(;target = incomplete_count_data(), n_chains = 5, n_rounds = 5) +@testset "Model with mixed state types using MPI" begin + target=incomplete_count_data() + r = pigeons(; + target, + n_rounds = 5, + n_chains = 4, + checkpoint = true, + # checked_round = 4, # NB: doesn't work yet, need a fine-tuned equality check for JuliaBUGS.BUGSModel + multithreaded = true, + on = ChildProcess( + n_local_mpi_processes = set_n_mpis_to_one_on_windows(2), + n_threads = 2, + mpiexec_args = extra_mpi_args(), + dependencies = [JuliaBUGS] + ) + ) + pt = Pigeons.load(r) @test true -end \ No newline at end of file +end From 6e0c0cbbe0f65dc6d15290c855a381a00bd0b0d4 Mon Sep 17 00:00:00 2001 From: Seren Lee Date: Wed, 18 Dec 2024 10:56:53 -0800 Subject: [PATCH 27/32] Add test --- examples/JuliaBUGS.jl | 6 +++--- test/test_JuliaBUGS.jl | 11 ++++++++++- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/examples/JuliaBUGS.jl b/examples/JuliaBUGS.jl index 0390a5461..20520c95a 100644 --- a/examples/JuliaBUGS.jl +++ b/examples/JuliaBUGS.jl @@ -1,4 +1,4 @@ -function incomplete_count_data_model() +function incomplete_count_data_model(tau::Float64) model_def = @bugs("model{ for (i in 1:n) { r[i] ~ dbern(pr[i]) @@ -17,8 +17,8 @@ function incomplete_count_data_model() missing,5,missing ], r = [1,0,0,0,0,0,0,1,1,0,1,0,0,0,1,0,0,1,0,1,1,1,1,1,1,0,0,0,1,0], - n = 30, tau = 4 + n = 30, tau = tau ) return compile(model_def, data) end -incomplete_count_data() = JuliaBUGSPath(incomplete_count_data_model()) \ No newline at end of file +incomplete_count_data(tau::Float64) = JuliaBUGSPath(incomplete_count_data_model(tau)) \ No newline at end of file diff --git a/test/test_JuliaBUGS.jl b/test/test_JuliaBUGS.jl index 992062742..09b17d5b6 100644 --- a/test/test_JuliaBUGS.jl +++ b/test/test_JuliaBUGS.jl @@ -64,7 +64,7 @@ end end @testset "Model with mixed state types using MPI" begin - target=incomplete_count_data() + target=incomplete_count_data(4) r = pigeons(; target, n_rounds = 5, @@ -82,3 +82,12 @@ end pt = Pigeons.load(r) @test true end + +@testset "Model with incomplete count data" begin + target=incomplete_count_data(0.01) + pt = pigeons(target = target, + n_rounds = 4, + n_chains = 4 + ) + @test true +end From 4991e63c4a9558a3ae5fb754357d578f6cfbfdaa Mon Sep 17 00:00:00 2001 From: Seren Lee Date: Thu, 19 Dec 2024 11:36:46 -0800 Subject: [PATCH 28/32] Fix Float64 -> Real --- examples/JuliaBUGS.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/JuliaBUGS.jl b/examples/JuliaBUGS.jl index 20520c95a..760277d05 100644 --- a/examples/JuliaBUGS.jl +++ b/examples/JuliaBUGS.jl @@ -1,4 +1,4 @@ -function incomplete_count_data_model(tau::Float64) +function incomplete_count_data_model(tau::Real) model_def = @bugs("model{ for (i in 1:n) { r[i] ~ dbern(pr[i]) @@ -21,4 +21,4 @@ function incomplete_count_data_model(tau::Float64) ) return compile(model_def, data) end -incomplete_count_data(tau::Float64) = JuliaBUGSPath(incomplete_count_data_model(tau)) \ No newline at end of file +incomplete_count_data(tau::Float64) = JuliaBUGSPath(incomplete_count_data_model(tau)) From d4a60711a49618bcd6e21e1ff93ee87495fd7ad3 Mon Sep 17 00:00:00 2001 From: Miguel Biron-Lattes Date: Thu, 19 Dec 2024 15:18:55 -0800 Subject: [PATCH 29/32] fix NaN --- examples/JuliaBUGS.jl | 4 ++-- ext/PigeonsJuliaBUGSExt/interface.jl | 14 +++++++++----- src/log_potentials/log_potentials.jl | 11 +++++++++-- test/test_JuliaBUGS.jl | 11 ++++------- 4 files changed, 24 insertions(+), 16 deletions(-) diff --git a/examples/JuliaBUGS.jl b/examples/JuliaBUGS.jl index 760277d05..65d29a7f1 100644 --- a/examples/JuliaBUGS.jl +++ b/examples/JuliaBUGS.jl @@ -1,4 +1,4 @@ -function incomplete_count_data_model(tau::Real) +function incomplete_count_data_model(;tau::Real=4) model_def = @bugs("model{ for (i in 1:n) { r[i] ~ dbern(pr[i]) @@ -21,4 +21,4 @@ function incomplete_count_data_model(tau::Real) ) return compile(model_def, data) end -incomplete_count_data(tau::Float64) = JuliaBUGSPath(incomplete_count_data_model(tau)) +incomplete_count_data(;kwargs...) = JuliaBUGSPath(incomplete_count_data_model(;kwargs...)) diff --git a/ext/PigeonsJuliaBUGSExt/interface.jl b/ext/PigeonsJuliaBUGSExt/interface.jl index fa2b46bef..aaa5a59c9 100644 --- a/ext/PigeonsJuliaBUGSExt/interface.jl +++ b/ext/PigeonsJuliaBUGSExt/interface.jl @@ -61,11 +61,15 @@ end # log_potential evaluation (log_potential::JuliaBUGSLogPotential)(flattened_values) = try - last(last(JuliaBUGS._tempered_evaluate!!( - log_potential.private_model, - flattened_values; - temperature=log_potential.beta - ))) + log_prior, _, tempered_log_joint = last( + JuliaBUGS._tempered_evaluate!!( + log_potential.private_model, + flattened_values; + temperature=log_potential.beta + ) + ) + # avoid potential 0*Inf (= NaN) + return iszero(log_potential.beta) ? log_prior : tempered_log_joint catch e (isa(e, DomainError) || isa(e, BoundsError)) && return -Inf rethrow(e) diff --git a/src/log_potentials/log_potentials.jl b/src/log_potentials/log_potentials.jl index 28226df05..e5a049010 100644 --- a/src/log_potentials/log_potentials.jl +++ b/src/log_potentials/log_potentials.jl @@ -40,5 +40,12 @@ Assumes the input `log_potentials` is a vector where each element is a [`log_pot This default implementation is sufficient in most cases, but in less standard scenarios, e.g. where the state space is infinite dimensional, this can be overridden. """ -log_unnormalized_ratio(log_potentials::AbstractVector, numerator::Int, denominator::Int, state) = - log_potentials[numerator](state) - log_potentials[denominator](state) +function log_unnormalized_ratio(log_potentials::AbstractVector, numerator::Int, denominator::Int, state) + lp_num = log_potentials[numerator](state) + lp_den = log_potentials[denominator](state) + ans = lp_num-lp_den + if isnan(ans) + error("Got NaN log-unnormalized ratio; Dumping information:\n\tlp_num=$lp_num\n\tlp_den=$lp_den\n\tState=$state") + end + return ans +end diff --git a/test/test_JuliaBUGS.jl b/test/test_JuliaBUGS.jl index 09b17d5b6..3ec4d7c9f 100644 --- a/test/test_JuliaBUGS.jl +++ b/test/test_JuliaBUGS.jl @@ -64,7 +64,7 @@ end end @testset "Model with mixed state types using MPI" begin - target=incomplete_count_data(4) + target=incomplete_count_data() r = pigeons(; target, n_rounds = 5, @@ -83,11 +83,8 @@ end @test true end -@testset "Model with incomplete count data" begin - target=incomplete_count_data(0.01) - pt = pigeons(target = target, - n_rounds = 4, - n_chains = 4 - ) +@testset "Check no NaN log potentials" begin # https://github.com/Julia-Tempering/Pigeons.jl/pull/303#issuecomment-2547306248 + target=incomplete_count_data(tau=0.01) + pt = pigeons(target = target, n_rounds = 4, n_chains = 4) @test true end From 0987925b70702427ef72bd9f29f3e48af62ba6e8 Mon Sep 17 00:00:00 2001 From: Miguel Biron-Lattes Date: Thu, 19 Dec 2024 15:36:25 -0800 Subject: [PATCH 30/32] sample names --- ext/PigeonsJuliaBUGSExt/interface.jl | 3 +++ 1 file changed, 3 insertions(+) diff --git a/ext/PigeonsJuliaBUGSExt/interface.jl b/ext/PigeonsJuliaBUGSExt/interface.jl index aaa5a59c9..d44efe5b5 100644 --- a/ext/PigeonsJuliaBUGSExt/interface.jl +++ b/ext/PigeonsJuliaBUGSExt/interface.jl @@ -79,3 +79,6 @@ end function Pigeons.sample_iid!(log_potential::JuliaBUGSLogPotential, replica, shared) replica.state = _sample_iid(log_potential.private_model, replica.rng) end + +Pigeons.sample_names(::Vector, log_potential::JuliaBUGSLogPotential) = + [(Symbol(string(vn)) for vn in log_potential.private_model.parameters)...,:log_density] From 0a22c4aa1ece0927f243e884b6b5d181269502ff Mon Sep 17 00:00:00 2001 From: Miguel Biron-Lattes Date: Thu, 19 Dec 2024 18:05:09 -0800 Subject: [PATCH 31/32] simple check for names --- test/test_JuliaBUGS.jl | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/test/test_JuliaBUGS.jl b/test/test_JuliaBUGS.jl index 3ec4d7c9f..36263f48f 100644 --- a/test/test_JuliaBUGS.jl +++ b/test/test_JuliaBUGS.jl @@ -85,6 +85,7 @@ end @testset "Check no NaN log potentials" begin # https://github.com/Julia-Tempering/Pigeons.jl/pull/303#issuecomment-2547306248 target=incomplete_count_data(tau=0.01) - pt = pigeons(target = target, n_rounds = 4, n_chains = 4) - @test true + pt = pigeons(target = target, n_rounds = 4, n_chains = 4, record=[traces]) + chns = Chains(pt) + @test first(names(chns)) != Symbol("param_1") # check we're not using the default array-state name builder end From 1ba7742b6160025e21eb0e4b55dd9ac2c6e577b0 Mon Sep 17 00:00:00 2001 From: Miguel Biron-Lattes Date: Fri, 20 Dec 2024 07:44:40 -0800 Subject: [PATCH 32/32] parallelism invariance --- ext/PigeonsJuliaBUGSExt/interface.jl | 13 +++++++++++++ test/test_JuliaBUGS.jl | 6 +++--- 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/ext/PigeonsJuliaBUGSExt/interface.jl b/ext/PigeonsJuliaBUGSExt/interface.jl index d44efe5b5..2b77927d3 100644 --- a/ext/PigeonsJuliaBUGSExt/interface.jl +++ b/ext/PigeonsJuliaBUGSExt/interface.jl @@ -80,5 +80,18 @@ function Pigeons.sample_iid!(log_potential::JuliaBUGSLogPotential, replica, shar replica.state = _sample_iid(log_potential.private_model, replica.rng) end +# parameter names Pigeons.sample_names(::Vector, log_potential::JuliaBUGSLogPotential) = [(Symbol(string(vn)) for vn in log_potential.private_model.parameters)...,:log_density] + +# Parallelism invariance +Pigeons.recursive_equal(a::Union{JuliaBUGSPath,JuliaBUGSLogPotential}, b) = + Pigeons._recursive_equal(a,b) +function Pigeons.recursive_equal(a::T, b) where T <: JuliaBUGS.BUGSModel + included = (:transformed, :model_def, :data) + excluded = Tuple(setdiff(fieldnames(T), included)) + Pigeons._recursive_equal(a,b,excluded) +end +# just check the betas match, the model is already checked within path +Pigeons.recursive_equal(a::AbstractVector{<:JuliaBUGSLogPotential}, b) = + all(lp1.beta == lp2.beta for (lp1,lp2) in zip(a,b)) diff --git a/test/test_JuliaBUGS.jl b/test/test_JuliaBUGS.jl index 36263f48f..b0fad3a6e 100644 --- a/test/test_JuliaBUGS.jl +++ b/test/test_JuliaBUGS.jl @@ -63,14 +63,14 @@ end @test res.passed end -@testset "Model with mixed state types using MPI" begin +@testset "Parallelism invariance using MPI" begin target=incomplete_count_data() r = pigeons(; - target, + target=unid_target, n_rounds = 5, n_chains = 4, checkpoint = true, - # checked_round = 4, # NB: doesn't work yet, need a fine-tuned equality check for JuliaBUGS.BUGSModel + checked_round = 4, multithreaded = true, on = ChildProcess( n_local_mpi_processes = set_n_mpis_to_one_on_windows(2),