Rename transformed (Update for Measure terminology)

.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+    "julia.environmentPath": "C:\\Users\\Cornelius\\.julia\\environments\\v1.9"
+}

Project.toml

@@ -48,6 +48,7 @@ ParallelProcessingTools = "8e8a01fc-6193-5ca1-a2f1-20776dae4199"
 Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"
 PositiveFactorizations = "85a6dd25-e78a-55b7-8502-1745935b8125"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Random123 = "74087812-796a-5b5d-8853-05524746bad3"
 RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
@@ -127,7 +128,7 @@ LaTeXStrings = "1"
 LinearAlgebra = "1"
 MacroTools = "0.5"
 Markdown = "1"
-MeasureBase = "0.12, 0.13, 0.14"
+MeasureBase = "0.14"
 Measurements = "2"
 NamedArrays = "0.9, 0.10"
 NestedSamplers = "0.8"
@@ -138,6 +139,7 @@ Parameters = "0.12, 0.13"
 Plots = "1"
 PositiveFactorizations = "0.2"
 Printf = "1"
+ProgressMeter = "1"
 Random = "1"
 Random123 = "1.2"
 RecipesBase = "0.7, 0.8, 1.0"

examples/dev-internal/transformed_example.jl

@@ -0,0 +1,67 @@
+using BAT
+using BAT.MeasureBase
+using AffineMaps
+using ChangesOfVariables
+using BAT.LinearAlgebra
+using BAT.Distributions
+using BAT.InverseFunctions
+import BAT: TransformedMCMCIterator, TransformedAdaptiveMHTuning, TransformedRAMTuner, TransformedMHProposal, TransformedNoTransformedMCMCTempering, transformed_mcmc_step!!, MCMCSampleID
+using Random123, PositiveFactorizations
+using AutoDiffOperators
+import AdvancedHMC
+
+import BAT: mcmc_iterate!, transformed_mcmc_iterate!, TransformedMCMCSampling
+
+#ENV["JULIA_DEBUG"] = "BAT"
+
+context = BATContext(ad = ADModule(:ForwardDiff))
+
+posterior = BAT.example_posterior()
+
+my_result = @time BAT.bat_sample_impl(posterior, TransformedMCMCSampling(pre_transform=PriorToGaussian(), nchains=4, nsteps=4*100000), context)
+
+
+density_notrafo = convert(BATMeasure, posterior)
+density, trafo = BAT.transform_and_unshape(PriorToGaussian(), density_notrafo, context)
+
+s = cholesky(Positive, BAT._approx_cov(density)).L
+f = BAT.CustomTransform(Mul(s))
+
+my_result = @time BAT.bat_sample_impl(posterior, TransformedMCMCSampling(pre_transform=PriorToGaussian(), tuning_alg=TransformedAdaptiveMHTuning(), nchains=4, nsteps=4*100000, adaptive_transform=f), context)
+
+my_samples = my_result.result
+
+
+
+using Plots
+plot(my_samples)
+
+r_mh = @time BAT.bat_sample_impl(posterior, MCMCSampling( nchains=4, nsteps=4*100000, store_burnin=true), context)
+
+r_hmc = @time BAT.bat_sample_impl(posterior, MCMCSampling(mcalg=HamiltonianMC(), nchains=4, nsteps=4*20000), context)
+
+plot(bat_sample(posterior).result)
+
+using BAT.Distributions
+using BAT.ValueShapes
+prior2 = NamedTupleDist(ShapedAsNT,
+    b = [4.2, 3.3],
+    a = Exponential(1.0),
+    c = Normal(1.0,3.0),
+    d = product_distribution(Weibull.(ones(2),1)),
+    e = Beta(1.0, 1.0),
+    f = MvNormal([0.3,-2.9],Matrix([1.7 0.5;0.5 2.3]))
+    )
+
+posterior.likelihood.density._log_f(rand(posterior.prior))
+
+posterior.likelihood.density._log_f(rand(prior2))
+
+posterior2 = PosteriorDensity(BAT.logfuncdensity(posterior.likelihood.density._log_f), prior2)
+
+
+@profview r_ram2 = @time BAT.bat_sample_impl(posterior2, TransformedMCMCSampling(pre_transform=PriorToGaussian(), nchains=4, nsteps=4*100000), context)
+
+@profview r_mh2 = @time BAT.bat_sample_impl(posterior2, MCMCSampling( nchains=4, nsteps=4*100000, store_burnin=true), context)
+
+r_hmc2 = @time BAT.bat_sample_impl(posterior2, MCMCSampling(mcalg=HamiltonianMC(), nchains=4, nsteps=4*20000), context)

ext/BATAdvancedHMCExt.jl

@@ -14,6 +14,7 @@ BAT.pkgext(::Val{:AdvancedHMC}) = BAT.PackageExtension{:AdvancedHMC}()
 using Random
 using DensityInterface
 using HeterogeneousComputing, AutoDiffOperators
+using BAT.ChangesOfVariables 
 
 using BAT: MeasureLike, BATMeasure
 

ext/BATHDF5Ext.jl

@@ -88,6 +88,17 @@ _to_flat_array(A::AbstractArray{<:AbstractArray{<:Real}}) = _to_flat_array(Array
 
 const _AnyRealArrayOrArrays = Union{AbstractArray{<:Real},AbstractArray{<:AbstractArray{<:Real}}}
 
+
+# TODO: MD Discuss, is to handle "nothing" entries in MCMCSampleIDVector objects
+function _h5io_write(datastore::H5DataStore, path::AbstractString, data::Vector{Union{Nothing, Int64}})
+    if any(isnothing.(data))
+        data_tmp = fill(0, length(data))
+    else
+        data_tmp = convert(Vector{Int64}, data)
+    end
+    _h5io_write(datastore, path, data_tmp)
+end
+
 function _h5io_write(datastore::H5DataStore, path::AbstractString, data::_AnyRealArrayOrArrays)
     @nospecialize datastore, path, data
     group = _h5io__get_or_create_group(datastore, dirname(path))

ext/ahmc_impl/ahmc_sampler_impl.jl

@@ -15,6 +15,55 @@ BAT.bat_default(::Type{MCMCSampling}, ::Val{:burnin}, mcalg::HamiltonianMC, traf
 BAT.get_mcmc_tuning(algorithm::HamiltonianMC) = algorithm.tuning
 
 
+
+"""
+    BAT.TransformedHMCProposal
+
+*BAT-internal, not part of stable public API.*
+"""
+mutable struct TransformedHMCProposal{
+    HA<:AdvancedHMC.Hamiltonian,
+    TR<:AdvancedHMC.Transition,
+    KRNL<:AdvancedHMC.HMCKernel
+}<: BAT.TransformedMCMCProposal
+    hamiltonian::HA
+    transition::TR
+    kernel::KRNL
+end
+
+function TransformedHMCProposal(algorithm::HamiltonianMC, target::BATMeasure, context::BATContext, v_init::AbstractVector)
+    adsel = get_adselector(context)
+    rng = get_rng(context)
+    f = checked_logdensityof(target)
+    metric = ahmc_metric(algorithm.metric, v_init)
+    fg = valgrad_func(f, adsel)
+
+    init_hamiltonian = AdvancedHMC.Hamiltonian(metric, f, fg)
+    hamiltonian, init_transition = AdvancedHMC.sample_init(rng, init_hamiltonian, v_init)
+    integrator = _ahmc_set_step_size(algorithm.integrator, hamiltonian, v_init)
+    termination = _ahmc_convert_termination(algorithm.termination, v_init)
+    kernel = HMCKernel(Trajectory{MultinomialTS}(integrator, termination))
+
+    # Perform a dummy step to get type-stable transition value:
+    transition = AdvancedHMC.transition(deepcopy(rng), deepcopy(hamiltonian), deepcopy(kernel), init_transition.z)
+
+    TransformedHMCProposal(hamiltonian, transition, kernel)
+end
+
+BAT._get_proposal(alg::HamiltonianMC, target::BATMeasure, context::BATContext, v_init::AbstractVector) = TransformedHMCProposal(alg, target, context, v_init)
+BAT._get_adaptive_transform(alg::HamiltonianMC) = BAT.default_adaptive_transform(alg)
+
+function MCMCSampleID(iter::TransformedMCMCIterator{<:Any, <:Any, <:Any, <:TransformedHMCProposal})
+    stat = AdvancedHMC.stat(iter.proposal.transition)
+
+    # TODO MD: Handle proposal-dependent tstat (only NUTS has tree_depth):
+    AHMCSampleID(
+        iter.info.id, iter.info.cycle, iter.stepno, CURRENT_SAMPLE,
+        stat.hamiltonian_energy, stat.tree_depth,
+        stat.numerical_error, stat.step_size
+    )
+end
+
 # MCMCIterator subtype for HamiltonianMC
 mutable struct AHMCIterator{
     AL<:HamiltonianMC,
@@ -297,3 +346,30 @@ end
 
 
 BAT.eff_acceptance_ratio(chain::AHMCIterator) = nsamples(chain) / nsteps(chain)
+
+
+function BAT.propose_mcmc(
+    iter::TransformedMCMCIterator{<:Any, <:Any, <:Any, <:TransformedHMCProposal}
+)
+    μ, f_transform, proposal, samples, sample_z, context = iter.μ, iter.f_transform, iter.proposal, iter.samples, iter.sample_z, iter.context
+    rng = get_rng(context)
+    sample_x = last(samples)
+    x, logd_x = sample_x.v, sample_x.logd
+    z, logd_z = sample_z.v, sample_z.logd
+
+    n = size(z, 1)
+
+    proposal.transition = AdvancedHMC.transition(rng, proposal.hamiltonian, proposal.kernel, proposal.transition.z)
+
+    z_proposed = proposal.transition.z.θ
+    x_proposed, ladj = ChangesOfVariables.with_logabsdet_jacobian(f_transform, z_proposed)
+    logd_x_proposed = BAT.checked_logdensityof(μ, x_proposed)
+    logd_z_proposed = logd_x_proposed + ladj
+
+    p_accept = clamp(exp(logd_z_proposed-logd_z), 0, 1)
+
+    sample_z_proposed = BAT._rebuild_density_sample(sample_z, z_proposed, logd_z_proposed)
+    sample_x_proposed = BAT._rebuild_density_sample(sample_x, x_proposed, logd_x_proposed)
+
+    return sample_x_proposed, sample_z_proposed, p_accept
+end

ext/ahmc_impl/ahmc_tuner_impl.jl

@@ -6,26 +6,61 @@ mutable struct AHMCTuner{A<:AdvancedHMC.AbstractAdaptor} <: AbstractMCMCTunerIns
     adaptor::A
 end
 
-function (tuning::HMCTuningAlgorithm)(chain::MCMCIterator)
+
+function BAT.get_tuner(tuning::HMCTuningAlgorithm, chain::TransformedMCMCIterator) 
+    θ = first(chain.samples).v
+    adaptor = ahmc_adaptor(tuning, chain.proposal.hamiltonian.metric, chain.proposal.kernel.τ.integrator, θ)
+    AHMCTuner(tuning.target_acceptance, adaptor)
+end
+
+
+function (tuning::HMCTuningAlgorithm)(chain::TransformedMCMCIterator)
     θ = first(chain.samples).v
-    adaptor = ahmc_adaptor(tuning, chain.hamiltonian.metric, chain.kernel.τ.integrator, θ)
+    adaptor = ahmc_adaptor(tuning, chain.proposal.hamiltonian.metric, chain.proposal.kernel.τ.integrator, θ)
     AHMCTuner(tuning.target_acceptance, adaptor)
 end
 
+# function (tuning::HMCTuningAlgorithm)(chain::MCMCIterator)
+#     θ = first(chain.samples).v
+#     adaptor = ahmc_adaptor(tuning, chain.hamiltonian.metric, chain.kernel.τ.integrator, θ)
+#     AHMCTuner(tuning.target_acceptance, adaptor)
+# end
 
-function BAT.tuning_init!(tuner::AHMCTuner, chain::MCMCIterator, max_nsteps::Integer)
+
+function BAT.tuning_init!(tuner::AHMCTuner, chain::TransformedMCMCIterator, max_nsteps::Integer)
     AdvancedHMC.Adaptation.initialize!(tuner.adaptor, Int(max_nsteps - 1))
     nothing
 end
 
-BAT.tuning_postinit!(tuner::AHMCTuner, chain::MCMCIterator, samples::DensitySampleVector) = nothing
+# function BAT.tuning_init!(tuner::AHMCTuner, chain::MCMCIterator, max_nsteps::Integer)
+#     AdvancedHMC.Adaptation.initialize!(tuner.adaptor, Int(max_nsteps - 1))
+#     nothing
+# end
+
+
+
+BAT.tuning_postinit!(tuner::AHMCTuner, chain::TransformedMCMCIterator, samples::DensitySampleVector) = nothing
+
+# BAT.tuning_postinit!(tuner::AHMCTuner, chain::MCMCIterator, samples::DensitySampleVector) = nothing
 
-function BAT.tuning_reinit!(tuner::AHMCTuner, chain::MCMCIterator, max_nsteps::Integer)
+
+
+function BAT.tuning_reinit!(tuner::AHMCTuner, chain::TransformedMCMCIterator, max_nsteps::Integer)
     AdvancedHMC.Adaptation.initialize!(tuner.adaptor, Int(max_nsteps - 1))
     nothing
 end
 
-function BAT.tuning_update!(tuner::AHMCTuner, chain::MCMCIterator, samples::DensitySampleVector)
+
+# function BAT.tuning_reinit!(tuner::AHMCTuner, chain::MCMCIterator, max_nsteps::Integer)
+#     AdvancedHMC.Adaptation.initialize!(tuner.adaptor, Int(max_nsteps - 1))
+#     nothing
+# end
+
+BAT.default_adaptive_transform(algorithm::HamiltonianMC) = BAT.TriangularAffineTransform()
+BAT.default_adaptive_transform(tuning::HMCTuningAlgorithm) = BAT.TriangularAffineTransform()
+
+
+function BAT.tuning_update!(tuner::AHMCTuner, chain::TransformedMCMCIterator, samples::DensitySampleVector)
     max_log_posterior = maximum(samples.logd)
     accept_ratio = eff_acceptance_ratio(chain)
     if accept_ratio >= 0.9 * tuner.target_acceptance
@@ -38,11 +73,11 @@ function BAT.tuning_update!(tuner::AHMCTuner, chain::MCMCIterator, samples::Dens
     nothing
 end
 
-function BAT.tuning_finalize!(tuner::AHMCTuner, chain::MCMCIterator)
+function BAT.tuning_finalize!(tuner::AHMCTuner, chain::TransformedMCMCIterator)
     adaptor = tuner.adaptor
     AdvancedHMC.finalize!(adaptor)
-    chain.hamiltonian = AdvancedHMC.update(chain.hamiltonian, adaptor)
-    chain.kernel = AdvancedHMC.update(chain.kernel, adaptor)
+    chain.proposal.hamiltonian = AdvancedHMC.update(chain.proposal.hamiltonian, adaptor)
+    chain.proposal.kernel = AdvancedHMC.update(chain.proposal.kernel, adaptor)
     nothing
 end
 
@@ -66,3 +101,16 @@ function (callback::AHMCTunerCallback)(::Val{:mcmc_step}, chain::AHMCIterator)
 
     nothing
 end
+
+function BAT.tune_mcmc_transform!!(
+    tuner::AHMCTuner,
+    transform::Any, #AffineMaps.AbstractAffineMap,#{<:typeof(*), <:LowerTriangular{<:Real}},
+    p_accept::Real,
+    z_proposed::Vector{<:Float64}, #TODO: use DensitySamples instead
+    z_current::Vector{<:Float64},
+    stepno::Int,
+    context::BATContext
+)
+
+    return (tuner, transform, false)
+end

src/BAT.jl

@@ -57,6 +57,7 @@ import EmpiricalDistributions
 import HypothesisTests
 import Measurements
 import NamedArrays
+import ProgressMeter
 import Random123
 import Sobol
 import StableRNGs

src/measures/bat_pushfwd_measure.jl

@@ -61,7 +61,7 @@ MeasureBase.pullback(f, m::BATMeasure) = _bat_pulbck(f, m, KeepRootMeasure())
 MeasureBase.pullback(f, m::BATMeasure, volcorr::KeepRootMeasure) = _bat_pulbck(f, m, volcorr)
 MeasureBase.pullback(f, m::BATMeasure, volcorr::ChangeRootMeasure) = _bat_pulbck(f, m, volcorr)
 
-_bat_pulbck(f, m::BATMeasure, volcorr::PushFwdStyle) = pushfwd(inverse(f), m, volcorr)
+_bat_pulbck(f, m::BATMeasure, volcorr::PushFwdStyle) = MeasureBase.pushfwd(inverse(f), m, volcorr)
 
 
 # ToDo: remove
@@ -84,18 +84,18 @@ function DensityInterface.logdensityof(@nospecialize(m::_NonBijectiveBATPusfwdMe
 end
 
 function DensityInterface.logdensityof(m::BATPushFwdMeasure{F,I,M,ChangeRootMeasure}, v::Any) where {F,I,M}
-    v_orig = inverse(m.trafo)(v)
-    logdensityof(parent(m), v_orig)
+    v_orig = m.finv(v)
+    logdensityof(m.origin, v_orig)
 end
 
 function checked_logdensityof(m::BATPushFwdMeasure{F,I,M,ChangeRootMeasure}, v::Any) where {F,I,M}
-    v_orig = inverse(m.trafo)(v)
-    checked_logdensityof(parent(m), v_orig)
+    v_orig = m.finv(v)
+    checked_logdensityof(m.origin, v_orig)
 end
 
 
 function _v_orig_and_ladj(m::BATPushFwdMeasure, v::Any)
-    with_logabsdet_jacobian(inverse(m.trafo), v)
+    with_logabsdet_jacobian(m.finv, v)
 end
 
 # TODO: Would profit from custom pullback:
@@ -123,13 +123,13 @@ end
 
 function DensityInterface.logdensityof(m::BATPushFwdMeasure{F,I,M,KeepRootMeasure}, v::Any) where {F,I,M}
     v_orig, ladj = _v_orig_and_ladj(m, v)
-    logd_orig = logdensityof(parent(m), v_orig)
+    logd_orig = logdensityof(m.origin, v_orig)
     _combine_logd_with_ladj(logd_orig, ladj)
 end
 
 function checked_logdensityof(m::BATPushFwdMeasure{F,I,M,KeepRootMeasure}, v::Any) where {F,I,M}
     v_orig, ladj = _v_orig_and_ladj(m, v)
-    logd_orig = logdensityof(parent(m), v_orig)
+    logd_orig = logdensityof(m.origin, v_orig)
     isnan(logd_orig) && @throw_logged EvalException(logdensityof, m, v, 0)
     _combine_logd_with_ladj(logd_orig, ladj)
 end