tpapp / DynamicHMC.jl

Implementation of robust dynamic Hamiltonian Monte Carlo methods (NUTS) in Julia.

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Struggling to Construct a working sampler

john-heyer opened this issue · comments

Hi Tamas,

I am developing a pretty complex sampler that I am having a hard time getting to work. I have had a bunch of issues, because I am sampling over parameters that could have produced an observed image. Basically, there is a point spread function applied with respect to the spatial locations of sources, that produces a "rate image". What that means is that each pixel in the rate image can be thought of as an independent poisson distribution with rate \lambda = value of the image at that pixel. The observed image is similar, yet has actual integer "event" values that are considered samples from these "rate images". So, the log density function is pretty complex.

Anyway, I have tried to set up the problem by following some of the examples you have provided.

struct NustarModel{T}
    observed_image::T
end

function (model::NustarModel)(θ)
    """
    Model callable to compute the conditional log likelihood of the sampled
    model parameters given the observed image.

    Params:
    - θ (model parameters): Array of 3*n_sources Real numbers, when unpacked
      is Arrays of [X_i, ..., X_n], [Y_i, ..., Y_n], [B_i, ..., B_n]

    Returns: likelihood
    """
    @unpack Xs, Ys, Bs = θ
    model_rate_image = TransformPSF.compose_mean_image(Xs, Ys, Bs)
    model_rate_image .+= 1
    likelihood = sum(
        [poisson_log_prob(model_rate_image[i], model.observed_image[i])
            for i in 1:length(model.observed_image)
        ]
    )
    return likelihood  # currently no priors
end

function parameter_transformation(model::NustarModel)
    as((Xs = as(Array, n_sources), Ys = as(Array, n_sources), Bs = as(Array, asℝ₊, n_sources)))
end

model = NustarModel(observed_image)
θs_truth = (Xs = sources_x, Ys = sources_y, Bs = sources_b)

println("log density of ground truth")
println(model(θs_truth))  # log density of true parameters, seemingly works as a real number is returned

transformation = parameter_transformation(model)
transformed_model = TransformedLogDensity(transformation, model)
∇M = ADgradient(:ForwardDiff, transformed_model)
results = mcmc_with_warmup(Random.GLOBAL_RNG, ∇M, 100)

Running this gives me a long method error stating "no method matching" with an extremely long and nested type. I'm not sure what's going wrong. I have gotten through some issues successfully but am not at a stand.

At first, I was getting a "initial slope must be finite" error, because some of the rates in the rate image of poissons which we were drawing from were zero, and so the log density was -inf. So, I added a constant value to each of the pixels (rates). Then, the sampler would appear to be running for a while, only to eventually return "Reached maximum number of iterations while bisecting interval for ϵ."

I hypothesized that perhaps adding a constant to each of the pixels made the gradient across these locations zero, and so instead I took the approach of adding to each pixel according to a power law 1/(distance to source). This is the point which I am currently at, where I receive a method error. Full error:

"LoadError: MethodError: no method matching Float32(::ForwardDiff.Dual{ForwardDiff.Tag{getfield(LogDensityProblems, Symbol("##34#35")){LogDensityProblems.TransformedLogDensity{TransformVariables.TransformTuple{NamedTuple{(:Xs, :Ys, :Bs),Tuple{TransformVariables.ArrayTransform{TransformVariables.Identity,1},TransformVariables.ArrayTransform{TransformVariables.Identity,1},TransformVariables.ArrayTransform{TransformVariables.ShiftedExp{true,Float64},1}}}},Main.Sampler.NustarModel{Array{Int64,2}}}},Float64},Float64,9})"

I am having trouble debugging such an error. I know that this post is extremely long and complicated, and that you may not be able to answer this without more information or taking a deep dive, but I didn't know where else to turn other than this forum. Also, I am new to Julia so I am struggling in this domain as well. Thanks for your contributions, let me know if you see something immediately wrong with the way I've set up the problem.

  • John

Hi John, I am very happy to help you but I need something self-contained I can run. It should include functions to generate fake data, code to load the libraries you use, and access to the later if they are not publicly available.

If you cannot provide this, at least a full stacktrace would be useful. The code you posted looks OK (as much as I can tell without running it), so the problem may be in TransformPSF.compose_mean_image. In particular, the error message you are posting is typical when trying to assign to a pre-allocated container with a type that is too narrow. Using pre-allocated arrays with ForwardDiff is tricky.

Alternatively, you can try AD via Flux.

Hi Tamas,

Thank you for the quick response!

There is really only one other module involved that contains all of the logic for `compose_mean_image'. It is publicly available here: https://github.com/john-heyer/NustarPointSourceInference. However, the current state is a bit of a mess, and I don't want to send you on a goose chase (some stuff is commented out, particularly translating an image, as I thought rounding might not be differentiable). I run the code from the root of the repo as the working directory.

If this isn't appropriate or immediately telling, I can work on getting a better working example tomorrow (cuz it's 4 AM my time 😅, but I see we are in different time zones). Thanks for your help again Tamas.