salilab / hmc Goto Github PK

Mici provides an implementation of dynamic HMC not tied to a specific probabilistic programming language. It is implemented in Python. It would be a bit cleaner to use it as a dependency for IMP.hmc instead of having the Julia language as a dependency; I'm not certain how its feature list or efficiency compares, but it's worth looking into.

When variables are marked as optimized and sampled by HMC, the assumption is that any constraints that act on them during Model::update() only do so via translation, as these transformations leave the momenta unchanged. Any other constraint will need to also transform the momentum, which would require all constraints to be HMC-aware. As an aside, such constraints also change the effective posterior distribution, so in general, how to handle these will need to be considered on a case-by-case basis.

At the moment, this is a silent failure. We can help the user by providing an internal check for debug builds to ValueGradientInterface::set_values that sets the coordinates, calls Model::update(), gets the coordinates, and errors if they are not equal.

This package uses AdvancedHMC.jl as the backend because currently there is not a general-purpose dynamic HMC sampler in Python. That's will soon change with BlackJAX, which will implement various HMC samplers in Python to be jit-compiled to JAX. BlackJAX will be the sampling backend for Numpyro, PyMC, and MCX.

When BlackJAX is finished, I intend to switch over to using it as the backend. Having a native Python package as the backend will make it easier to install the module and maintain it. This will essentially require a rewrite of this module, but in the end it will probably result in deletion of most of the code.

Supersedes #6

As raised in salilab/imp#1020, IMP currently silently fails when a derivative is requested from a non-differentiable restraint. That's not good in general, but that's catastrophic for HMC. In addition to the check in #1, we should add a check that uses central difference approximations to validate the gradient at the current model configuration. It's ad hoc and local but should identify if there are any major issues with the gradient.

This and #1 can be bundled into a single function that locally validates that things are working correctly for IMP.hmc and that the user would be recommended to run on a new model to diagnose issues before using IMP.hmc.

@ichem001 pointed out that we're hitting some deprecation warnings in AdvancedHMC:

┌ Warning: `find_good_eps` is deprecated, use `find_good_stepsize` instead.
│   caller = invokelatest(::Any, ::Any, ::Vararg{Any,N} where N; kwargs::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}) at essentials.jl:712
└ @ Base ./essentials.jl:712
┌ Warning: `Preconditioner(args...)` is deprecated, use `MassMatrixAdaptor(args...)` instead.
│   caller = #invokelatest#1 at essentials.jl:712 [inlined]
└ @ Core ./essentials.jl:712
┌ Warning: `StanHMCAdaptor(n_adapts, pc, ssa)` is deprecated, use `initialize!(StanHMCAdaptor(pc, ssa), n_adapts)` instead.
│   caller = ip:0x0
└ @ Core :-1

Not certain if these need to be fixed here or in HMCUtilities.jl. The environment was

julia> using Pkg

julia> Pkg.status()
Status `~/.julia/environments/v1.4/Project.toml`
  [0bf59076] AdvancedHMC v0.2.24
  [26cc04aa] FiniteDifferences v0.10.0 #master (https://github.com/JuliaDiff/FiniteDifferences.jl)
  [2ce65f13] HMCUtilities v0.2.1 #v0.2.1 (https://github.com/salilab/HMCUtilities.jl)
  [438e738f] PyCall v1.91.4