diff --git a/.github/dependabot.yml b/.github/dependabot.yml new file mode 100644 index 000000000..0d08e261a --- /dev/null +++ b/.github/dependabot.yml @@ -0,0 +1,11 @@ +# To get started with Dependabot version updates, you'll need to specify which +# package ecosystems to update and where the package manifests are located. +# Please see the documentation for all configuration options: +# https://docs.github.com/code-security/dependabot/dependabot-version-updates/configuration-options-for-the-dependabot.yml-file + +version: 2 +updates: + - package-ecosystem: "github-actions" # See documentation for possible values + directory: "/" # Location of package manifests + schedule: + interval: "weekly" diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 1fcbee3ac..e1c21f892 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -39,16 +39,18 @@ jobs: with: distribution: 'temurin' java-version: '11' - - uses: julia-actions/setup-julia@latest + - uses: julia-actions/setup-julia@v2 with: version: ${{ matrix.version }} - uses: julia-actions/cache@v2 - - uses: julia-actions/julia-buildpkg@latest - - uses: julia-actions/julia-runtest@latest - - uses: julia-actions/julia-processcoverage@latest + - uses: julia-actions/julia-buildpkg@v1 + - uses: julia-actions/julia-runtest@v1 + - uses: julia-actions/julia-processcoverage@v1 - uses: codecov/codecov-action@v4 with: - files: lcov.info + token: ${{ secrets.CODECOV_TOKEN }} + fail_ci_if_error: false # or true if you want CI to fail when Codecov fails + file: lcov.info # test OpenMPI by requesting it with MPIPreferences @@ -71,7 +73,7 @@ jobs: - name: Checkout uses: actions/checkout@v4 - - uses: julia-actions/setup-julia@latest + - uses: julia-actions/setup-julia@v2 with: version: ${{ matrix.version }} - uses: julia-actions/cache@v2 @@ -85,7 +87,7 @@ jobs: MPIPreferences.use_jll_binary("OpenMPI_jll") rm("test/Manifest.toml") - - uses: julia-actions/julia-runtest@latest + - uses: julia-actions/julia-runtest@v1 # # CI is getting too slow! @@ -98,7 +100,7 @@ jobs: fail-fast: false matrix: version: - - '1' + - '1.10' os: - macos-13 mpi: @@ -119,7 +121,7 @@ jobs: distribution: 'temurin' java-version: '11' - - uses: julia-actions/setup-julia@latest + - uses: julia-actions/setup-julia@v2 with: version: ${{ matrix.version }} @@ -137,7 +139,7 @@ jobs: run(`sed -i.bu 's/unknown/MPICH/' test/LocalPreferences.toml`) # fix wrong abi detection for mpich rm("test/Manifest.toml") - - uses: julia-actions/julia-runtest@latest + - uses: julia-actions/julia-runtest@v1 docs: @@ -151,11 +153,11 @@ jobs: with: distribution: 'temurin' java-version: '11' - - uses: julia-actions/setup-julia@latest + - uses: julia-actions/setup-julia@v2 with: - version: '1.9' - - uses: julia-actions/julia-buildpkg@latest - - uses: julia-actions/julia-docdeploy@latest + version: '1.10' + - uses: julia-actions/cache@v2 + - uses: julia-actions/julia-docdeploy@v1 env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} - run: | diff --git a/.github/workflows/CompatHelper.yml b/.github/workflows/CompatHelper.yml index 757087482..dae5d1fc9 100644 --- a/.github/workflows/CompatHelper.yml +++ b/.github/workflows/CompatHelper.yml @@ -15,7 +15,7 @@ jobs: run: which julia continue-on-error: true - name: Install Julia, but only if it is not already available in the PATH - uses: julia-actions/setup-julia@v1 + uses: julia-actions/setup-julia@v2 with: version: '1' arch: ${{ runner.arch }} diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 000000000..ee530c2ed --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,22 @@ +# Contributing Guide + +Thank you for your interest in contributing to Pigeons.jl! There are multiple ways to contribute, and we appreciate all contributions. + +## Ways to Contribute + +### Bug Reports and Feature Requests + +If you are working with Pigeons, and run into a bug, or require additional features, we definitely want to know about it. Please let us know by creating a [GitHub Issue](https://github.com/PigeonsAD/Pigeons.jl/issues). + +### How to Submit a Pull Request + +If you would like to contribute, please feel free to open a pull request. The pull request should: + +- Include a small unit test +- Not contain any unrelated changes +- Be an isolated change. Independent changes should be submitted as separate pull requests as this makes reviewing easier. + +## Code of Conduct + +As a Julia Project we expect all participants to abide by the [Julia Community Standards](https://julialang.org/community/standards/). + diff --git a/Project.toml b/Project.toml index 004a7e64d..810e4cf16 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.5" +version = "0.4.7" [deps] DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" @@ -59,7 +59,7 @@ BridgeStan = "2" DataFrames = "1" Distributions = "0.25" DocStringExtensions = "0.9" -DynamicPPL = "0.23, 0.24, 0.25, 0.26, 0.27, 0.28" +DynamicPPL = "0.23, 0.24, 0.25, 0.26, 0.27, 0.28, 0.30" Enzyme = "0.12, 0.13" Expect = "0.3" FillArrays = "1" diff --git a/docs/img/alex.png b/docs/img/alex.png new file mode 100644 index 000000000..d324c6b90 Binary files /dev/null and b/docs/img/alex.png differ diff --git a/docs/img/lucas.jpg b/docs/img/lucas.jpg new file mode 100644 index 000000000..3cfbdf820 Binary files /dev/null and b/docs/img/lucas.jpg differ diff --git a/docs/img/miguel.jpg b/docs/img/miguel.jpg new file mode 100644 index 000000000..90f425c73 Binary files /dev/null and b/docs/img/miguel.jpg differ diff --git a/docs/img/nikola.jpg b/docs/img/nikola.jpg new file mode 100644 index 000000000..7bd47e586 Binary files /dev/null and b/docs/img/nikola.jpg differ diff --git a/docs/img/paul.jpg b/docs/img/paul.jpg new file mode 100644 index 000000000..2b37dd008 Binary files /dev/null and b/docs/img/paul.jpg differ diff --git a/docs/img/saif.jpg b/docs/img/saif.jpg new file mode 100644 index 000000000..c98dfe697 Binary files /dev/null and b/docs/img/saif.jpg differ diff --git a/docs/img/son.jpg b/docs/img/son.jpg new file mode 100644 index 000000000..c9b652f06 Binary files /dev/null and b/docs/img/son.jpg differ diff --git a/docs/img/trevor.jpg b/docs/img/trevor.jpg new file mode 100644 index 000000000..05de5af10 Binary files /dev/null and b/docs/img/trevor.jpg differ diff --git a/docs/img/william.png b/docs/img/william.png new file mode 100644 index 000000000..a14e15d21 Binary files /dev/null and b/docs/img/william.png differ diff --git a/docs/img/yujia.jpg b/docs/img/yujia.jpg new file mode 100644 index 000000000..ab39e32a2 Binary files /dev/null and b/docs/img/yujia.jpg differ diff --git a/docs/make.jl b/docs/make.jl index eef78138a..0fc71061a 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -72,6 +72,7 @@ InferenceReport.headless() do "More on distributed PT" => "distributed.md", "Interfaces" => Pigeons.informal_doc(@__DIR__, Pigeons), "Reference" => "reference.md", + "About Us" => "about-us.md", "Google Summer of Code" => "gsoc.md" ], ) diff --git a/docs/src/about-us.md b/docs/src/about-us.md new file mode 100644 index 000000000..f33efd128 --- /dev/null +++ b/docs/src/about-us.md @@ -0,0 +1,103 @@ +```@meta +CurrentModule = Pigeons +``` +# [The Pigeons Team - About Us](@id about) + +Planning for the Pigeons project started taking place in the summer of 2022. +Since then, we have grown as a team and have added various features to the software. +We still have great ambitions and a swath of features to add; if you would like +to become a part of our team feel free to reach out! + +## Team Members (chronological order) + +```@raw html +
+ Alexandre Bouchard-Côté +
+ Alexandre Bouchard-Côté +
+ University of British Columbia +
+
+ +
+ Nikola Surjanovic +
+ Nikola Surjanovic +
+ University of British Columbia +
+
+ +
+ Paul Tiede +
+ Paul Tiede +
+ Harvard University, Event Horizon Telescope +
+
+ +
+ Saifuddin Syed +
+ Saifuddin Syed +
+ University of Oxford +
+
+ +
+ Trevor Campbell +
+ Trevor Campbell +
+ University of British Columbia +
+
+ +
+ Miguel Biron-Lattes +
+ Miguel Biron-Lattes +
+ University of British Columbia +
+
+ +
+ William Thompson +
+ William Thompson +
+ Herzberg Astronomy and Astrophysics Research Centre +
+
+ +
+ Yujia Ma +
+ Yujia Ma +
+ University of British Columbia +
+
+ +
+ Lucas Wong +
+ Lucas Wong +
+ University of British Columbia +
+
+ +
+ Son Luu +
+ Son Luu +
+ University of British Columbia +
+
+``` \ No newline at end of file diff --git a/ext/PigeonsEnzymeExt/PigeonsEnzymeExt.jl b/ext/PigeonsEnzymeExt/PigeonsEnzymeExt.jl index eab75f439..ad23a2862 100644 --- a/ext/PigeonsEnzymeExt/PigeonsEnzymeExt.jl +++ b/ext/PigeonsEnzymeExt/PigeonsEnzymeExt.jl @@ -13,26 +13,32 @@ else using ..LogDensityProblemsAD end -# TODO: currently, the concrete versions of ADGradientWrapper are defined only -# in the extensions of LogDensityProblemsAD. Therefore, it is impossible to -# dispatch on them; see -# https://github.com/tpapp/LogDensityProblemsAD.jl/issues/32 -# This is a HACK to extract that type -const EnzymeGradientLogDensity = if isdefined(Base, :get_extension) - Base.get_extension(LogDensityProblemsAD, :LogDensityProblemsADEnzymeExt).EnzymeGradientLogDensity -else - LogDensityProblemsAD.LogDensityProblemsADEnzymeExt.EnzymeGradientLogDensity +# A simpler version of the wrapper defined in LogDensityProblemsAD's extension +struct EnzymeWrapper{TLP} <: Pigeons.ADWrapper + log_potential::TLP +end + +# special ADgradient constructor for Enzyme +function LogDensityProblemsAD.ADgradient( + kind::Val{:Enzyme}, + log_potential, + buffers::Pigeons.Augmentation + ) + d = LogDensityProblems.dimension(log_potential) + buffer = Pigeons.get_buffer(buffers, :gradient_buffer, d) + enclosed = EnzymeWrapper(log_potential) + Pigeons.BufferedAD(enclosed, buffer, nothing, nothing) end # adapted from LogDensityProblemsAD to use the Replica's buffer function LogDensityProblems.logdensity_and_gradient( - b::Pigeons.BufferedAD{<:EnzymeGradientLogDensity}, + b::Pigeons.BufferedAD{<:EnzymeWrapper}, x::AbstractVector ) ∂ℓ_∂x = fill!(b.buffer, zero(eltype(b.buffer))) # NB: Enzyme gives erroneous answer if buffer is not zeroed first _, y = Enzyme.autodiff( Enzyme.ReverseWithPrimal, LogDensityProblems.logdensity, Enzyme.Active, - Enzyme.Const(b.enclosed.ℓ), Enzyme.Duplicated(x, ∂ℓ_∂x) + Enzyme.Const(b.enclosed.log_potential), Enzyme.Duplicated(x, ∂ℓ_∂x) ) y, ∂ℓ_∂x end diff --git a/ext/PigeonsForwardDiffExt/PigeonsForwardDiffExt.jl b/ext/PigeonsForwardDiffExt/PigeonsForwardDiffExt.jl index 580599a1f..e64aafe6c 100644 --- a/ext/PigeonsForwardDiffExt/PigeonsForwardDiffExt.jl +++ b/ext/PigeonsForwardDiffExt/PigeonsForwardDiffExt.jl @@ -15,15 +15,10 @@ else import ..ForwardDiff: DiffResults end -# TODO: currently, the concrete versions of ADGradientWrapper are defined only -# in the extensions of LogDensityProblemsAD. Therefore, it is impossible to -# dispatch on them; see -# https://github.com/tpapp/LogDensityProblemsAD.jl/issues/32 -# This is a HACK to extract that type -const ForwardDiffLogDensity = if isdefined(Base, :get_extension) - Base.get_extension(LogDensityProblemsAD, :LogDensityProblemsADForwardDiffExt).ForwardDiffLogDensity -else - LogDensityProblemsAD.LogDensityProblemsADForwardDiffExt.ForwardDiffLogDensity +# A simpler version of the wrapper defined in LogDensityProblemsAD's extension +struct ForwardDiffWrapper{TLP, TGC <: ForwardDiff.GradientConfig} <: Pigeons.ADWrapper + log_potential::TLP + gradient_config::TGC end # special ADgradient constructor for ForwardDiff @@ -33,21 +28,22 @@ function LogDensityProblemsAD.ADgradient( buffers::Pigeons.Augmentation ) d = LogDensityProblems.dimension(log_potential) - buffer = Pigeons.get_buffer(buffers, :gradient_buffer, d) - enclosed = ADgradient(kind, log_potential; x = buffer) + buffer = Pigeons.get_buffer(buffers, :gradient_buffer, d) + lp_fix = Base.Fix1(LogDensityProblems.logdensity, log_potential) + gradient_config = ForwardDiff.GradientConfig(lp_fix, buffer, ForwardDiff.Chunk(d)) + enclosed = ForwardDiffWrapper(log_potential, gradient_config) diff_result = DiffResults.MutableDiffResult(zero(eltype(buffer)), (buffer, )) Pigeons.BufferedAD(enclosed, diff_result, nothing, nothing) end # adapted from LogDensityProblemsAD to use the Replica's buffer function LogDensityProblems.logdensity_and_gradient( - b::Pigeons.BufferedAD{<:ForwardDiffLogDensity}, + b::Pigeons.BufferedAD{<:ForwardDiffWrapper}, x::AbstractVector ) diff_result = b.buffer - ℓ_fix = Base.Fix1(LogDensityProblems.logdensity, b.enclosed.ℓ) - ForwardDiff.gradient!(diff_result, ℓ_fix, x, b.enclosed.gradient_config) - + lp_fix = Base.Fix1(LogDensityProblems.logdensity, b.enclosed.log_potential) + ForwardDiff.gradient!(diff_result, lp_fix, x, b.enclosed.gradient_config) return (DiffResults.value(diff_result), DiffResults.gradient(diff_result)) end diff --git a/ext/PigeonsReverseDiffExt/PigeonsReverseDiffExt.jl b/ext/PigeonsReverseDiffExt/PigeonsReverseDiffExt.jl index e2a10b8bb..febb61fc1 100644 --- a/ext/PigeonsReverseDiffExt/PigeonsReverseDiffExt.jl +++ b/ext/PigeonsReverseDiffExt/PigeonsReverseDiffExt.jl @@ -15,16 +15,23 @@ else import ..ReverseDiff: DiffResults end -# TODO: currently, the concrete versions of ADGradientWrapper are defined only -# in the extensions of LogDensityProblemsAD. Therefore, it is impossible to -# dispatch on them; see -# https://github.com/tpapp/LogDensityProblemsAD.jl/issues/32 -# This is a HACK to extract that type -const ReverseDiffLogDensity = if isdefined(Base, :get_extension) - Base.get_extension(LogDensityProblemsAD, :LogDensityProblemsADReverseDiffExt).ReverseDiffLogDensity -else - LogDensityProblemsAD.LogDensityProblemsADReverseDiffExt.ReverseDiffLogDensity +# A simpler version of the wrapper defined in LogDensityProblemsAD's extension +struct ReverseDiffWrapper{TLP, TCT} <: Pigeons.ADWrapper + log_potential::TLP + compiled_tape::TCT +end +function make_compiled_tape(log_potential, x) + lp_fix = Base.Fix1(LogDensityProblems.logdensity, log_potential) + tape = ReverseDiff.GradientTape(lp_fix, x) + return ReverseDiff.compile(tape) end +compute_gradient!(rdw::ReverseDiffWrapper, diff_result, x) = + ReverseDiff.gradient!( + diff_result, + Base.Fix1(LogDensityProblems.logdensity, rdw.log_potential), x + ) +compute_gradient!(rdw::ReverseDiffWrapper{<:Any,<:ReverseDiff.CompiledTape}, diff_result, x) = + ReverseDiff.gradient!(diff_result, rdw.compiled_tape, x) # special ADgradient constructor for ReverseDiff function LogDensityProblemsAD.ADgradient( @@ -34,9 +41,8 @@ function LogDensityProblemsAD.ADgradient( ) d = LogDensityProblems.dimension(log_potential) buffer = Pigeons.get_buffer(buffers, :gradient_buffer, d) - compile_tape = Pigeons.get_tape_compilation_strategy() - - if compile_tape + should_compile = Pigeons.get_tape_compilation_strategy() + if should_compile @info """ Using ReverseDiff with tape compilation, which usually results in huge performance gains. @@ -51,24 +57,19 @@ function LogDensityProblemsAD.ADgradient( by calling `Pigeons.set_tape_compilation_strategy!(true)`. """ maxlog=1 end - - enclosed = ADgradient(kind, log_potential; x = buffer, compile=Val{compile_tape}()) + compiled_tape = should_compile ? make_compiled_tape(log_potential, buffer) : nothing + enclosed = ReverseDiffWrapper(log_potential, compiled_tape) diff_result = DiffResults.MutableDiffResult(zero(eltype(buffer)), (buffer, )) Pigeons.BufferedAD(enclosed, diff_result, nothing, nothing) end # adapted from LogDensityProblemsAD to use the Replica's buffer function LogDensityProblems.logdensity_and_gradient( - b::Pigeons.BufferedAD{<:ReverseDiffLogDensity}, + b::Pigeons.BufferedAD{<:ReverseDiffWrapper}, x::AbstractVector ) diff_result = b.buffer - compiled_tape = b.enclosed.compiledtape - if compiled_tape === nothing - ReverseDiff.gradient!(diff_result, Base.Fix1(LogDensityProblems.logdensity, b.enclosed.ℓ), x) - else - ReverseDiff.gradient!(diff_result, compiled_tape, x) - end + compute_gradient!(b.enclosed, diff_result, x) return (DiffResults.value(diff_result), DiffResults.gradient(diff_result)) end diff --git a/src/explorers/Augmentation.jl b/src/explorers/Augmentation.jl index d40548015..ee9da50af 100644 --- a/src/explorers/Augmentation.jl +++ b/src/explorers/Augmentation.jl @@ -8,7 +8,7 @@ avoiding race-conditions between replicas. For an application, see [`buffers`](@ $FIELDS """ -struct Augmentation{T} +struct Augmentation{T} # For serialization purpose T should support zero-arg constructor T() """ The payload. Can be `nothing` for efficiency purposes. """ @@ -43,7 +43,7 @@ end function Serialization.deserialize(s::AbstractSerializer, ::Type{Augmentation{T}}) where {T} serialize_field = Serialization.deserialize(s) - contents = serialize_field ? Serialization.deserialize(s) : nothing + contents = serialize_field ? Serialization.deserialize(s) : T() return Augmentation{T}(contents, serialize_field) end diff --git a/src/explorers/BufferedAD.jl b/src/explorers/BufferedAD.jl index 04e623588..576a2b8d1 100644 --- a/src/explorers/BufferedAD.jl +++ b/src/explorers/BufferedAD.jl @@ -162,3 +162,10 @@ Currently this is only used for [ReverseDiff](https://github.com/JuliaDiff/Rever function set_tape_compilation_strategy!(compile::Bool) COMPILE_TAPE[] = compile end + +# used in the AD extensions +abstract type ADWrapper end +LogDensityProblems.logdensity(adw::ADWrapper, x::AbstractVector) = + LogDensityProblems.logdensity(adw.log_potential, x) +LogDensityProblems.dimension(adw::ADWrapper) = + LogDensityProblems.dimension(adw.log_potential) diff --git a/src/pt/process_sample.jl b/src/pt/process_sample.jl index 55fba60e3..0a12d3dee 100644 --- a/src/pt/process_sample.jl +++ b/src/pt/process_sample.jl @@ -19,12 +19,12 @@ and pair plots (via [PairPlots](https://sefffal.github.io/PairPlots.jl/dev/chain function sample_array(pt::PT) chains = chains_with_samples(pt) dim, size = sample_dim_size(pt, chains) - result = zeros(size, dim, length(chains)) - for chain_index in eachindex(chains) - t = chains[chain_index] + n_chains_with_samples = count(!isnothing, chains) # iterators have no `length` method + result = zeros(size, dim, n_chains_with_samples) + for (chain_index, t) in enumerate(chains) sample = get_sample(pt, t) - for i in 1:size - vector = sample[i] + for i in 1:size + vector = sample[i] result[i, :, chain_index] .= vector end end @@ -34,13 +34,13 @@ end chains_with_samples(pt) = pt.inputs.extended_traces ? (1:n_chains(pt.inputs)) : target_chains(pt) function sample_dim_size(pt::PT, chains) - sample = get_sample(pt, chains[1]) - return length(sample[1]), length(sample) + sample = get_sample(pt, first(chains)) + return length(first(sample)), length(sample) end function target_chains(pt::PT) n = n_chains(pt.inputs) - return filter(i -> is_target(pt.shared.tempering.swap_graphs, i), 1:n) + return (i for i in 1:n if is_target(pt.shared.tempering.swap_graphs, i)) end """ @@ -84,7 +84,7 @@ The `chain` option can be omitted and by default the first chain targetting the distribution of interest will be used (in many cases, there will be only one, in variational cases, two). """ -get_sample(pt::PT, chain = target_chains(pt)[1]) = SampleArray(pt, chain) +get_sample(pt::PT, chain = first(target_chains(pt))) = SampleArray(pt, chain) function Base.show(io::IO, s::SampleArray{T,PT}) where {T,PT} println(io, "SampleArray{$T}") diff --git a/src/swap/VariationalDEO.jl b/src/swap/VariationalDEO.jl index 07fd31553..1a2b23960 100644 --- a/src/swap/VariationalDEO.jl +++ b/src/swap/VariationalDEO.jl @@ -18,4 +18,4 @@ create_swap_graph(deo::VariationalDEO, shared) = iseven(shared.iterators.scan) ? even(deo.n_chains_fixed, deo.n_chains_var) : odd(deo.n_chains_fixed, deo.n_chains_var) is_reference(deo::VariationalDEO, chain::Int) = (chain == 1) || (chain == n_chains(deo)) -is_target(deo::VariationalDEO, chain::Int) = (chain == deo.n_chains_fixed) || (chain == deo.n_chains_fixed + 1) \ No newline at end of file +is_target(deo::VariationalDEO, chain::Int) = (chain == deo.n_chains_var) || (chain == deo.n_chains_var + 1) \ No newline at end of file diff --git a/test/supporting/HetPrecisionNormalLogPotential.jl b/test/supporting/HetPrecisionNormalLogPotential.jl index 1f61aba74..a50cd2677 100644 --- a/test/supporting/HetPrecisionNormalLogPotential.jl +++ b/test/supporting/HetPrecisionNormalLogPotential.jl @@ -29,5 +29,3 @@ end LogDensityProblems.logdensity(log_potential::HetPrecisionNormalLogPotential, x) = log_potential(x) LogDensityProblems.dimension(log_potential::HetPrecisionNormalLogPotential) = length(log_potential.precisions) -LogDensityProblemsAD.ADgradient(kind::Symbol, log_potential::HetPrecisionNormalLogPotential, buffers::Pigeons.Augmentation) = - LogDensityProblemsAD.ADgradient(kind, log_potential) \ No newline at end of file diff --git a/test/supporting/setup.jl b/test/supporting/setup.jl index c8b3dd27c..dcfee1d41 100644 --- a/test/supporting/setup.jl +++ b/test/supporting/setup.jl @@ -7,6 +7,7 @@ using Pigeons, Distributions, DynamicPPL, Enzyme, + FillArrays, ForwardDiff, HypothesisTests, LinearAlgebra, diff --git a/test/test_checkpoint.jl b/test/test_checkpoint.jl index 42ce2f943..b218d92ee 100644 --- a/test/test_checkpoint.jl +++ b/test/test_checkpoint.jl @@ -15,3 +15,10 @@ end compare_pts(p1, p3) end end + +@testset "Checkpoints-with-AD" begin + pigeons(target = toy_mvn_target(2), n_rounds = 2, checkpoint = true, explorer = AutoMALA()) + pt = PT("results/latest") + pt = increment_n_rounds!(pt, 2) + pigeons(pt) +end \ No newline at end of file diff --git a/test/test_custom_gradient.jl b/test/test_custom_gradient.jl new file mode 100644 index 000000000..aa9ae5486 --- /dev/null +++ b/test/test_custom_gradient.jl @@ -0,0 +1,34 @@ +struct CustomGradientLogPotential + precision::Float64 + dim::Int +end +function (log_potential::CustomGradientLogPotential)(x) + -0.5 * log_potential.precision * sum(abs2, x) +end + +Pigeons.initialization(lp::CustomGradientLogPotential, ::AbstractRNG, ::Int) = zeros(lp.dim) + +LogDensityProblems.dimension(lp::CustomGradientLogPotential) = lp.dim +LogDensityProblems.logdensity(lp::CustomGradientLogPotential, x) = lp(x) + +LogDensityProblemsAD.ADgradient(::Val, log_potential::CustomGradientLogPotential, replica::Pigeons.Replica) = + Pigeons.BufferedAD(log_potential, replica.recorders.buffers) + +const check_custom_grad_called = Ref(false) + +function LogDensityProblems.logdensity_and_gradient(log_potential::Pigeons.BufferedAD{CustomGradientLogPotential}, x) + logdens = log_potential.enclosed(x) + global check_custom_grad_called[] = true + log_potential.buffer .= -log_potential.enclosed.precision .* x + return logdens, log_potential.buffer +end + +@testset "Custom-gradient" begin + pigeons( + target = CustomGradientLogPotential(2.1, 4), + reference = CustomGradientLogPotential(1.1, 4), + n_chains = 1, + explorer = AutoMALA()) + + @assert check_custom_grad_called[] +end \ No newline at end of file diff --git a/test/test_mala.jl b/test/test_mala.jl index d8a880f9c..45e8462a9 100644 --- a/test/test_mala.jl +++ b/test/test_mala.jl @@ -8,6 +8,7 @@ which also cover checks of the Hamiltonian dynamics. target = toy_mvn_target(2), n_chains = 2, explorer = MALA(), + n_chains_variational = 4, record = [Pigeons.online], n_rounds = 10); for var_name in Pigeons.continuous_variables(pt) diff --git a/test/test_two_legs.jl b/test/test_two_legs.jl index 2288ebf25..8c1168d02 100644 --- a/test/test_two_legs.jl +++ b/test/test_two_legs.jl @@ -1,10 +1,11 @@ -@testset "Two legs schedule adaptation" begin +@testset "Test StabilizedPT machinery" begin n_rounds = 10 - n_chains = 10 + n_chains = 8 + n_chains_variational = 7 pt_2_legs = pigeons(; target = Pigeons.toy_turing_unid_target(), variational = GaussianReference(first_tuning_round = n_rounds + 1), # never activate - n_chains_variational = n_chains, + n_chains_variational = n_chains_variational, n_chains, n_rounds) @@ -15,14 +16,26 @@ n_chains, n_rounds) - - @show gcb_1 = Pigeons.global_barrier(pt_1_leg.shared.tempering) - @show gcb_2_1 = Pigeons.global_barrier(pt_2_legs.shared.tempering) - @show gcb_2_2 = Pigeons.global_barrier_variational(pt_2_legs.shared.tempering) + @testset "Two legs schedule adaptation" begin + @show gcb_1 = Pigeons.global_barrier(pt_1_leg.shared.tempering) + @show gcb_2_1 = Pigeons.global_barrier(pt_2_legs.shared.tempering) + @show gcb_2_2 = Pigeons.global_barrier_variational(pt_2_legs.shared.tempering) + truth = 3.5 # based on 15 rounds + for approx in [gcb_1, gcb_2_1, gcb_2_2] + @test isapprox(approx, truth, rtol = 0.1) + end + end - truth = 3.5 # based on 15 rounds - - for approx in [gcb_1, gcb_2_1, gcb_2_2] - @test isapprox(approx, truth, atol = 0.1) + @testset "Issue #290" begin + n = Pigeons.n_chains(pt_2_legs.inputs) + idxs_targets = Pigeons.target_chains(pt_2_legs) + idxs_refs = (i for i in 1:n if Pigeons.is_reference(pt_2_legs.shared.tempering.swap_graphs, i)) + indexer = pt_2_legs.shared.tempering.indexer + @test isempty(intersect(idxs_refs, idxs_targets)) # targets and references should be different + # test: if multiple references/targets exist, they should live on different legs + for idxs in (idxs_targets, idxs_refs) + n_distinct_legs = length(Set(last(indexer.i2t[idx]) for idx in idxs)) + @test length(collect(idxs)) == n_distinct_legs + end end end \ No newline at end of file