Make NestedModel a bit more general

[torfjelde]

This PR makes the following changes:

• NestedModel now only has a single field: prior_transform_and_loglikelihood, which maps u ↦ (v, logl), where u is a uniform random variable, v is u transformed, and logl is the loglikelihood at v.
• Introduces some new methods
  • prior_transform(model, u): replaces the nested_model.prior_transform, allowing us to overload for other model types too, if we so desire.
  • loglikelihood(model, u): replaces the nested_model.loglike, allowing us to overload for other model types too, if we so desire.
  • prior_transform_and_loglikelihood(model, u): combines both prior_transfrom and loglikelihood into a single function call. This is useful because these two methods will often share computation in non-trivial models. All uses of prior_transfrom and loglikelihood in #master follows the pattern v = prior_transform(u); logl = loglikelihood(v), and so combining these into a single function doesn't really change much internally.
• model is passed around instead of functions prior_transform and loglike, and then we instead use the newly introduced methods, e.g. prior_transform_and_loglikelihood(model, u) instead.

As an example/motivational use-cases, this PR makes it quite easy to make a DynamicPPL.Model compatible with NestedSamplers.jl:

struct NestedSampler <: AbstractMCMC.AbstractSampler end

function DynamicPPL.assume(
    rng::Random.AbstractRNG,
    sampler::DynamicPPL.Sampler{<:NestedSampler},
    dist::Distribution,
    vn::DynamicPPL.VarName,
    vi::DynamicPPL.AbstractVarInfo,
)
    if haskey(vi, vn)
        # If it's already present, we assume it's in unit-cube space
        # but we want to return it in the original space.
        # TODO: We should probably specify this somehow, e.g. using
        # a field in `sampler` or something.
        val = invlogcdf(dist, log(vi[vn]))
        vi = DynamicPPL.setindex!!(vi, val, vn)
    else
        val = rand(rng, dist)
        vi = DynamicPPL.push!!(vi, vn, val, dist, sampler)
    end

    return val, zero(eltype(val)), vi
end

function NestedSamplers.NestedModel(model::DynamicPPL.Model)
    return NestedSamplers.NestedModel(Random.GLOBAL_RNG, model)
end

function NestedSamplers.NestedModel(rng::Random.AbstractRNG, model::DynamicPPL.Model)
    sampler = DynamicPPL.Sampler(NestedSampler())
    vi_base = DynamicPPL.VarInfo(rng, model, sampler)
    function prior_transform_and_loglikelihod_model(u)
        # Update in unit-cube space.
        vi_new = DynamicPPL.setindex!!(vi_base, u, sampler)
        # Evaluate model, computing the transformed variables and the loglikelihood.
        _, vi_new = DynamicPPL.evaluate!!(model, vi_new, sampler)
        # Return the new samples and loglikelihood.
        return vi_new[sampler], DynamicPPL.getlogp(vi_new)
    end

    return NestedSamplers.NestedModel(prior_transform_and_loglikelihod_model)
end

then we can just convert a model::DynamicPPL.Model into a NestedModel by calling NestedModel(model).

With a bit more code, we can also easily make a NestedSampler which is compatible with Turing rather than just converting into a NestedModel, e.g. we get parameter names: https://gist.github.com/torfjelde/5dd1ed93a81759c98ff0ef3feeb24237.

[ParadaCarleton]

@torfjelde do you know what's breaking the tests?

[mileslucas]

This change seems good, I'm curious about the choice of creating a custom struct for PriorTransformAndLogLike- can you explain your thinking a bit?

Also, for changing the proposals to pass the whole NestedModel- could you just pass a prior_and_loglike function instead? This might make tests easier since you don't have to construct a NestedModel just to run the MCMC proposal. The current way the proposals are written is pretty agnostic of the NestedModel and Nested structs.

Also, the tests are failing because you need to update these tests to use the new proposal interface:

NestedSamplers.jl/test/proposals/proposals.jl

Lines 16 to 17 in 2b3423c

	logl(X) = -sum(x->x^2, X)
	prior(u) = 2u .- 1 # Uniform -1, to 1

[torfjelde]

This change seems good, I'm curious about the choice of creating a custom struct for PriorTransformAndLogLike- can you explain your thinking a bit?

It's really just to allow the user to specify the prior_transform and the loglikelihood independetly rather than as a single function. Of course this could be done as a anonymous function, but a) it's not very interpretable vs. PriorTransformAndLogLikelihood, and b) you can't dispatch on anonymous functions.

Also, for changing the proposals to pass the whole NestedModel- could you just pass a prior_and_loglike function instead? This might make tests easier since you don't have to construct a NestedModel just to run the MCMC proposal. The current way the proposals are written is pretty agnostic of the NestedModel and Nested structs.

Agreed. I actually had it implemented this way initially, but then I changed it to passing the model around. Now I'm thinking I should revert back to passing around the prior_transform_and_loglikelihood for the proposals.

EDIT: Oh I remember. The idea was that you could use other types of AbstractModel as long as they had an implementation of prior_transform and loglikelihood (we should rename this IMO, e.g. loglikelihood_from_uniform).

EDIT 2: And the current implementation also supports prior_transform_and_loglikelihood since it also implements prior_transform and loglikelihood:)

Also, the tests are failing because you need to update these tests to use the new proposal interface:

NestedSamplers.jl/test/proposals/proposals.jl

Lines 16 to 17 in 2b3423c

	logl(X) = -sum(x->x^2, X)
	prior(u) = 2u .- 1 # Uniform -1, to 1

Yeah I realized, just hadn't gotten around to it yet, but thanks 👍

[codecov]

Codecov Report

Merging #76 (bfb05c8) into main (ef42bd4) will decrease coverage by 1.92%.
The diff coverage is 71.69%.

@@            Coverage Diff             @@
##             main      #76      +/-   ##
==========================================
- Coverage   90.55%   88.62%   -1.93%     
==========================================
  Files          13       13              
  Lines         540      554      +14     
==========================================
+ Hits          489      491       +2     
- Misses         51       63      +12     

Impacted Files	Coverage Δ
src/NestedSamplers.jl	100.00% <ø> (ø)
src/model.jl	54.16% <47.61%> (-45.84%)	⬇️
src/step.jl	93.65% <69.23%> (-0.71%)	⬇️
src/proposals/Proposals.jl	92.98% <100.00%> (-0.32%)	⬇️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 2b3423c...bfb05c8. Read the comment docs.

[ParadaCarleton]

Well... that is definitely the weirdest mix of test failures and successes I have ever seen. Fails on 1.3 Windows and 1.7 Mac, but totally fine on 1.3 Mac and 1.7 Windows. And then somehow, everything passes on Julia-latest, which is weird in and of itself.

[torfjelde]

Well... that is definitely the weirdest mix of test failures and successes I have ever seen. Fails on 1.3 Windows and 1.7 Mac, but totally fine on 1.3 Mac and 1.7 Windows. And then somehow, everything passes on Julia-latest, which is weird in and of itself.

If you look at the logs, you can see that it's a numerical errors, i.e. that the approximation isn't good enough. Hence the "weird" pattern of behavior is likely just due to difference in RNG on the different architectures.

[ParadaCarleton]

Well... that is definitely the weirdest mix of test failures and successes I have ever seen. Fails on 1.3 Windows and 1.7 Mac, but totally fine on 1.3 Mac and 1.7 Windows. And then somehow, everything passes on Julia-latest, which is weird in and of itself.

If you look at the logs, you can see that it's a numerical errors, i.e. that the approximation isn't good enough. Hence the "weird" pattern of behavior is likely just due to difference in RNG on the different architectures.

I thought that might be the case, but everything is supposed to be using StableRNG with the same seed, so it should be the same regardless of the OS.

That being said I believe I found the error, in models.jl:
chain_res = sample(chain, Weights(vec(chain[:weights])), length(chain))
Forgot to specify an RNG here.

[torfjelde]

Ah, nice catch 👍

Trying with the global rng (also found some additional sample statements for which it was missing).

[ParadaCarleton]

Ok, now Windows nightly is failing some of the random numeric tests, which is very weird given everything I see has StableRNG attached to it. Literally went through every call to rand, randn, and sample to make sure nothing was missing.

[ParadaCarleton]

An idea: We can create a struct DummyRNG <: Random.AbstractRNG end that just throws an error whenever you call it, then set that as the global RNG. If there's anything still calling the global RNG, we'll see where the error is.

[mileslucas]

thanks for cleaning up the rng in the testing @ParadaCarleton. Overall the PR looks good, if there's just this test failing I'm happy to merge and I'll deal with that myself.

[ParadaCarleton]

thanks for cleaning up the rng in the testing @ParadaCarleton. Overall the PR looks good, if there's just this test failing I'm happy to merge and I'll deal with that myself.

That does seem to be the only problem, yeah.

[ParadaCarleton]

@mileslucas want to merge this, or should I?

[mileslucas]

Thanks @torfjelde @ParadaCarleton !