mikekatz04 / Eryn

An omni-sampler for Bayesian MCMC
45 stars 10 forks source link

`eryn.moves.DistributionGenerate` isn't working #8

Closed rjrosati closed 1 year ago

rjrosati commented 1 year ago

Hi, I've been trying to add in proposals drawn from my prior, but have been having some trouble. I think there are actually two issues here.

The first one is that DistributionGenerate always expects a dictionary of ProbDistContainer objects, unlike EnsembleSampler which can handle a raw ProbDistContainer. The docs also suggest that a raw ProbDistContainer should be fine.

Here's an example that fails:

from eryn.prior import ProbDistContainer, uniform_dist
from eryn.moves import DistributionGenerate

priors = ProbDistContainer({0: uniform_dist(-1.0,1.0)})
DistributionGenerate(priors)

# fails with:
#Traceback (most recent call last):
#  File "/Users/rjrosati/code/sgwb-data-analysis/issue_eryn.py", line 5, in <module>
#    DistributionGenerate(priors)
#  File "/Users/rjrosati/opt/anaconda3/lib/python3.9/site-packages/eryn/moves/distgen.py", line 21, in __init__
#    for key in generate_dist:
#TypeError: 'ProbDistContainer' object is not iterable

Putting in the dictionary structure manually, it looks like DistributionGenerate still isn't actually able to generate any proposals. There's an extra argument on DistributionGenerate.get_proposal that isn't added in the sampling. Here's another minimal example and the errors:

from eryn.prior import ProbDistContainer, uniform_dist
from eryn.moves import DistributionGenerate,StretchMove
from eryn.ensemble import EnsembleSampler
import numpy as np
ndim = 5
nwalkers = 100

def log_like_fn(x, mu, invcov):
    diff = x - mu
    return -0.5 * (diff * np.dot(invcov, diff.T).T).sum()

means = np.zeros(ndim)
cov = np.diag(np.ones(ndim))
invcov = np.linalg.inv(cov)

lims = 5.0
priors_in = {i: uniform_dist(-lims+means[i], lims+means[i]) for i in range(ndim)}
priors = {'model_0': ProbDistContainer(priors_in)}

ensemble = EnsembleSampler(nwalkers,ndim,log_like_fn,priors,args=[means,invcov],
                           moves = [
                               (StretchMove(),0.5),
                               (DistributionGenerate(priors),0.5),
                               ])

theta0 = priors['model_0'].rvs(size=(nwalkers,))
result = ensemble.run_mcmc(theta0,1000,burn=200,progress=True, thin_by=2)

#Traceback (most recent call last):
#  File "/Users/rjrosati/code/sgwb-data-analysis/issue_eryn.py", line 28, in <module>
#    result = ensemble.run_mcmc(theta0,1000,burn=200,progress=True, thin_by=2)
# File "/Users/rjrosati/opt/anaconda3/lib/python3.9/site-packages/eryn/ensemble.py", line 1002, in run_mcmc
#    for results in self.sample(initial_state, iterations=burn, **burn_kwargs):
#  File "/Users/rjrosati/opt/anaconda3/lib/python3.9/site-packages/eryn/ensemble.py", line 892, in sample
#   state, accepted_out = move.propose(model, state)
#  File "/Users/rjrosati/opt/anaconda3/lib/python3.9/site-packages/eryn/moves/mh.py", line 110, in propose
#    q, factors = self.get_proposal(
#TypeError: get_proposal() missing 1 required positional argument: 'random'

The extra argument is never used in the function body as far as I can tell, so this one is probably a pretty easy fix.

mikekatz04 commented 1 year ago

I think I fixed both of these. There was an update needed for the calls in the distribution generate class. However, it now gives an error when entering a non-dictionary. So the code as before will now give an error. It needs to be like this because in the get_proposal function it uses the model name to find the proper distribution.