Reversible Jump MCMC

[zenna]

To implement reversible jump mcmc, we need the following things

• Given a probabilistic program, that may or may not be conditioned, we need to be able to compute the target density (which is the posterior) of an ω object.
  • What does this mean exactly?
• An invertible transformation that h : Ω -> Ω such that h(ω) = ω'
• To be able to compute the jacobian correction for h

Problems

• To find consistent ω requires solving, either through sketch, inversion or some other method. What if this solving process is not invertible?

What is ω?

ω is a mapping from random variables to values An ω` is consistent if the values it assigns to each random variably are jointly possible

Reversible Transforms

One of the difficulties is in ensuring that h is invertible. It is hard because in general h has to solve constraints. One idea is as follows:

• Assume h(ω, θ) is some arbitrary function that produces some other valid ω, but is not necessarily invertible
• From h construct a function h' which tracks all the non-invertible function choices so that we can invert them. Can do this by applying the pgf transform h' = pgf(h') such that h'(ω, θ) = (ω', s) where s are the parameters for the parametric inverse.
• h is not then a diffeomorphism, s and θ will generally not have same dimensionality, or even same type.

[zenna]

module ProposalDemo
using Cassette
using Distributions

"std-normal inverse cdf"
q(x) = quantile(Normal(0, 1), x)

"StdUnif(i) is ith element of sequence of iid random variables that are uniformly distributed on [0, 1]"
struct StdUnif
  id::Int
end

"Mapping from Random variables to values"
mutable struct Ω{T}
  data::T
end

(x::StdUnif)(ω) = ω.data[x]

## FIXME - disconcerting that one is mutating and other isnt

function propose!(x::StdUnif, ω)
  if x ∉ keys(ω.data)
    ω.data[x] = rand()
  end
end

function propose(::typeof(x), x_, ω)
  if μ in keys(ω.data) && StdUnif(2) in keys(ω.data)
    return nothing
    # @assert false
  
  ## FIXME< seem to be some numerical issues
  elseif μ in keys(ω.data)
    c(x) = cdf(Normal(0, 1), x)
    Dict{Any, Any}(StdUnif(2) => c(x_ - ω.data[μ]))
  elseif StdUnif(2) in keys(ω.data)
    nothing
  end
    ## Provide a value for μ
    # @assert false
end

propose(T, x_, ω) = nothing

"Make a proposal for random variable `f`"
function propose!(f, ω)
  @show ω.data
  # FIXME: don't want to make proposal more than once do we?
  for (x, x_) in ω.data
    @show x, x_
    proposal_ = propose(x, x_, ω)
    if !isnothing(proposal_)
      ω.data = merge(ω.data, proposal_)
    end
  end
  ω
end

## Where am I making these proposals?

Cassette.@context ProposalCtx
function Cassette.posthook(::ProposalCtx, ret, f, ω::Ω)
  ## Add f to ω, since some proposal might depend on it
  if f in keys(ω.data)
    @assert ω.data[f] == ret
  else
    ω.data[f] = ret
  end
  @show f
  propose!(f, ω)
end

function Cassette.prehook(::ProposalCtx, f::StdUnif, ω::Ω)
  propose!(f, ω)
end

# FIXME: Proposa; probability

"""
Generate proposal for random variables within `f`, returns ω::Ω
# Input
- `f`   Random variable.  Proposal will generate values for depenends of `f`
- `ω` Initial mapping, to condition variables add to here
# Returns
- `ω::Ω` mapping from random variables ot values
"""
propose(f, ω) = (Cassette.overdub(ProposalCtx(), ω -> f(ω), ω); ω)

## Example
function test()
  x_ = 3.5
  ω = Ω(Dict{Any, Any}(x => x_))
  propose(μ, ω)
end

# I;m not sure about calling propose! in pre and host hools

end