Resumed pooled dynamic nested sampler will not cache prior transform?

[OGdodge]

Dynesty version Specify the exact version, and how you installed it (pip/git) 1.2.2 thru git clone Describe the bug [A clear and concise description of what the bug is.] Resuming a dynamic nested sampler in a dynesty pool does not correctly recognise the provided prior transform function, leading to an error on resumption. I think I've got the setup pretty much as described in the dynamic sampler tutorial on read the docs. Setup I am runningthe following dynesty setup:

import dynesty
...
def dynamic_dynesty(self, name):
        if not HAS_DYNESTY:
           print("Error: Cannot run nested sampler "
                 "wihout the dynesty package installed!")
           sys.exit(1)
        resume_flag = False
        self._prepare_MCMC()
        resume_file = glob.glob(f'{name}.sampler') #finds resume file 
        if len(resume_file) == 1:
           print(f'Saved run found! Resuming w/ {name}...')
           resume_flag = True
        with open(f'{name}.dyMN', 'wb') as f:
            pk.dump(self, f)
        if self.use_pool:
            loglike = self.LogLikelihood
            prior_transform = self.inv_CDF_prior_transform
            if resume_flag:
                with dynesty.pool.Pool(self.npool, loglike, prior_transform) as respool: #resuming here
                    sampler = dynesty.DynamicNestedSampler.restore(resume_file[0], pool=respool)
                    result = sampler.run_nested(resume=True)
                    return result
            with dynesty.pool.Pool(self.npool, loglike, prior_transform) as pool: #where sampler is first created
               sampler = dynesty.DynamicNestedSampler(pool.loglike, pool.prior_transform,
                                              ndim=self.ndim,
                                              bound='multi',
                                              update_interval=10,
                                              pool=pool,
                                              #use_pool={'prior_transform':False,
                                              #          'loglikelihood':False}
                                              )
               result = sampler.run_nested(dlogz_init=0.05, 
                                           nlive_init=self.nlive,
                                           nlive_batch=int(self.nlive/5),
                                           checkpoint_file=f'{name}.sampler')
              return result

Dynesty output

Saved run found! Resuming w/ MN_pickles/run4...
0it [00:00, ?it/s]Exception while calling prior_transform function:
  params: [0.99492014 0.14897805 0.9476684  0.8725722  0.79028606 0.97062346
 0.97538978]
  args: []
  kwargs: {}
  exception:

Bug I find the prior transform function is not found by the resumed sampler which gives the following error: Traceback (most recent call last):

  File "/scratch/nas_spiders/olidodge/PSRJ1810+1744/icarus_dynesty_modelling/phot_only/localdynesty/dynesty.py", line 1098, in __call__
    return self.func(np.asarray(x).copy(), *self.args, **self.kwargs)
  File "/scratch/nas_spiders/olidodge/PSRJ1810+1744/icarus_dynesty_modelling/phot_only/localdynesty/pool.py", line 46, in prior_transform_cache
    return FunctionCache.prior_transform(x, *FunctionCache.ptform_args,
AttributeError: type object 'FunctionCache' has no attribute 'prior_transform'
0it [00:00, ?it/s]
Traceback (most recent call last):
  File "j1810_dynamic.py", line 226, in <module>
    results = dyMN.dynamic_dynesty(name=f'MN_pickles/run{options.nrun}')
  File "/scratch/nas_spiders/olidodge/icarus-local/dynesty_utils.py", line 1834, in dynamic_dynesty
    result = sampler.run_nested(resume=True)
  File "/scratch/nas_spiders/olidodge/PSRJ1810+1744/icarus_dynesty_modelling/phot_only/localdynesty/dynamicsampler.py", line 1813, in run_nested
    for results in self.sample_initial(
  File "/scratch/nas_spiders/olidodge/PSRJ1810+1744/icarus_dynesty_modelling/phot_only/localdynesty/dynamicsampler.py", line 914, in sample_initial
    for it, results in enumerate(
  File "/scratch/nas_spiders/olidodge/PSRJ1810+1744/icarus_dynesty_modelling/phot_only/localdynesty/sampler.py", line 771, in sample
    u, v, logl, nc = self._new_point(loglstar_new)
  File "/scratch/nas_spiders/olidodge/PSRJ1810+1744/icarus_dynesty_modelling/phot_only/localdynesty/sampler.py", line 382, in _new_point
    u, v, logl, nc, blob = self._get_point_value(loglstar)
  File "/scratch/nas_spiders/olidodge/PSRJ1810+1744/icarus_dynesty_modelling/phot_only/localdynesty/sampler.py", line 365, in _get_point_value
    self._fill_queue(loglstar)
  File "/scratch/nas_spiders/olidodge/PSRJ1810+1744/icarus_dynesty_modelling/phot_only/localdynesty/sampler.py", line 358, in _fill_queue
    self.queue = list(mapper(evolve_point, args))
  File "/scratch/nas_spiders/olidodge/PSRJ1810+1744/icarus_dynesty_modelling/phot_only/localdynesty/sampling.py", line 87, in sample_unif
    v = args.prior_transform(np.asarray(args.u))
  File "/scratch/nas_spiders/olidodge/PSRJ1810+1744/icarus_dynesty_modelling/phot_only/localdynesty/dynesty.py", line 1098, in __call__
    return self.func(np.asarray(x).copy(), *self.args, **self.kwargs)
  File "/scratch/nas_spiders/olidodge/PSRJ1810+1744/icarus_dynesty_modelling/phot_only/localdynesty/pool.py", line 46, in prior_transform_cache
    return FunctionCache.prior_transform(x, *FunctionCache.ptform_args,
AttributeError: type object 'FunctionCache' has no attribute 'prior_transform'

The behaviour I expected was a simple resumption of the nested sampling. My setup works without a pool. Apologies if this isn't at all helpful, my first bug report and tbh I'm not even sure it's a bug, I could be doing something stupid wrong! Just figured I'd submit it as a bug as the question section said if you're not sure to submit it as a bug :) Thanks for dynesty, it's a wonderfully simple sampler overall and has made my life a lot easier!

[segasai]

Hi,

Thanks for the report. I can reproduce the issue with this script (it first has to be run with resume=False, interrupted then resumed):

import numpy as np
import dynesty

import time

class CL:

    def __init__(self):
        pass

    def __call__(self, x):
        time.sleep(.01)
        return -0.5 * np.sum((x - .5)**2)

    def pt(self, x):
        return x

if __name__ == '__main__':
    C = CL()
    resume = True
    if not resume:
        with dynesty.pool.Pool(4, C, C.pt) as pool:
            dns = dynesty.DynamicNestedSampler(pool.loglike,
                                               pool.prior_transform,
                                               4,
                                               pool=pool)
            dns.run_nested(checkpoint_file='xx.pkl', checkpoint_every=10)
    else:
        with dynesty.pool.Pool(4, C, C.pt) as pool:
            dns = dynesty.DynamicNestedSampler.restore('xx.pkl', pool=pool)
            dns.run_nested(resume=True)

I haven't previously tested the dynesty.pool with the resume.

If you use the standard multiprocessing pool it will certainly work without issues.

But with the dynesty.pool, I need to think about it. it may be not trivial (to be honest it may even be impossible)

[segasai]

I think I've fixed this issue with https://github.com/joshspeagle/dynesty/commit/3411eecfab468130d290b07a6db4b6091c91890c please check if it works for you. The fix is somewhat tentative.

[OGdodge]

Works perfectly, thanks for the fix much appreciated! :)