Coverage for bilby/core/sampler/dynesty_utils.py: 90%
305 statements
« prev ^ index » next coverage.py v7.6.1, created at 2025-05-06 04:57 +0000
1import warnings
3import numpy as np
4from dynesty.nestedsamplers import MultiEllipsoidSampler, UnitCubeSampler
5from dynesty.utils import apply_reflect, get_random_generator
7from ...bilby_mcmc.chain import calculate_tau
8from ..utils.log import logger
9from .base_sampler import _SamplingContainer
12class LivePointSampler(UnitCubeSampler):
13 """
14 Modified version of dynesty UnitCubeSampler that adapts the MCMC
15 length in addition to the proposal scale, this corresponds to
16 :code:`bound=live`.
18 In order to support live-point based proposals, e.g., differential
19 evolution (:code:`diff`), the live points are added to the
20 :code:`kwargs` passed to the evolve method.
22 Note that this does not perform ellipsoid clustering as with the
23 :code:`bound=multi` option, if ellipsoid-based proposals are used, e.g.,
24 :code:`volumetric`, consider using the
25 :code:`MultiEllipsoidLivePointSampler` (:code:`sample=live-multi`).
26 """
28 rebuild = False
30 def update_user(self, blob, update=True):
31 """
32 Update the proposal parameters based on the number of accepted steps
33 and MCMC chain length.
35 There are a number of logical checks performed:
36 - if the ACT tracking rwalk method is being used and any parallel
37 process has an empty cache, set the :code:`rebuild` flag to force
38 the cache to rebuild at the next call. This improves the efficiency
39 when using parallelisation.
40 - update the :code:`walks` parameter to asymptotically approach the
41 desired number of accepted steps for the :code:`FixedRWalk` proposal.
42 - update the ellipsoid scale if the ellipsoid proposals are being used.
43 """
44 # do we need to trigger rebuilding the cache
45 if blob.get("remaining", 0) == 1:
46 self.rebuild = True
47 if update:
48 self.kwargs["rebuild"] = self.rebuild
49 self.rebuild = False
51 # update walks to match target naccept
52 accept_prob = max(0.5, blob["accept"]) / self.kwargs["walks"]
53 delay = self.nlive // 10 - 1
54 n_target = getattr(_SamplingContainer, "naccept", 60)
55 self.walks = (self.walks * delay + n_target / accept_prob) / (delay + 1)
56 self.kwargs["walks"] = min(int(np.ceil(self.walks)), _SamplingContainer.maxmcmc)
58 self.scale = blob["accept"]
60 update_rwalk = update_user
62 def propose_live(self, *args):
63 """
64 We need to make sure the live points are passed to the proposal
65 function if we are using live point-based proposals.
66 """
67 self.kwargs["nlive"] = self.nlive
68 self.kwargs["live"] = self.live_u
69 i = self.rstate.integers(self.nlive)
70 u = self.live_u[i, :]
71 return u, np.identity(self.ncdim)
74class MultiEllipsoidLivePointSampler(MultiEllipsoidSampler):
75 """
76 Modified version of dynesty MultiEllipsoidSampler that adapts the MCMC
77 length in addition to the proposal scale, this corresponds to
78 :code:`bound=live-multi`.
80 Additionally, in order to support live point-based proposals, e.g.,
81 differential evolution (:code:`diff`), the live points are added to the
82 :code:`kwargs` passed to the evolve method.
84 When just using the :code:`diff` proposal method, consider using the
85 :code:`LivePointSampler` (:code:`sample=live`).
86 """
88 rebuild = False
90 def update_user(self, blob, update=True):
91 LivePointSampler.update_user(self, blob=blob, update=update)
92 super(MultiEllipsoidLivePointSampler, self).update_rwalk(
93 blob=blob, update=update
94 )
96 update_rwalk = update_user
98 def propose_live(self, *args):
99 """
100 We need to make sure the live points are passed to the proposal
101 function if we are using ensemble proposals.
102 """
103 self.kwargs["nlive"] = self.nlive
104 self.kwargs["live"] = self.live_u
105 return super(MultiEllipsoidLivePointSampler, self).propose_live(*args)
108class FixedRWalk:
109 """
110 Run the MCMC walk for a fixed length. This is nearly equivalent to
111 :code:`bilby.sampling.sample_rwalk` except that different proposal
112 distributions can be used.
113 """
115 def __call__(self, args):
116 current_u = args.u
117 naccept = 0
118 ncall = 0
120 periodic = args.kwargs["periodic"]
121 reflective = args.kwargs["reflective"]
122 boundary_kwargs = dict(periodic=periodic, reflective=reflective)
124 proposals, common_kwargs, proposal_kwargs = _get_proposal_kwargs(args)
125 walks = len(proposals)
127 accepted = list()
129 for prop in proposals:
130 u_prop = proposal_funcs[prop](
131 u=current_u, **common_kwargs, **proposal_kwargs[prop]
132 )
133 u_prop = apply_boundaries_(u_prop=u_prop, **boundary_kwargs)
134 if u_prop is None:
135 accepted.append(0)
136 continue
138 v_prop = args.prior_transform(u_prop)
139 logl_prop = args.loglikelihood(v_prop)
140 ncall += 1
142 if logl_prop > args.loglstar:
143 current_u = u_prop
144 current_v = v_prop
145 logl = logl_prop
146 naccept += 1
147 accepted.append(1)
148 else:
149 accepted.append(0)
151 if naccept == 0:
152 logger.debug(
153 "Unable to find a new point using walk: returning a random point. "
154 "If this warning occurs often, increase naccept."
155 )
156 # Technically we can find out the likelihood value
157 # stored somewhere
158 # But I'm currently recomputing it
159 current_u = common_kwargs["rstate"].uniform(0, 1, len(current_u))
160 current_v = args.prior_transform(current_u)
161 logl = args.loglikelihood(current_v)
163 blob = {
164 "accept": naccept,
165 "reject": walks - naccept,
166 "scale": args.scale,
167 }
169 return current_u, current_v, logl, ncall, blob
172class ACTTrackingRWalk:
173 """
174 Run the MCMC sampler for many iterations in order to reliably estimate
175 the autocorrelation time.
177 This builds a cache of :code:`nact / thin` points consistent with the
178 likelihood constraint.
179 While this approach is the most robust, it is not well optimized for
180 parallelised sampling as the length of the MCMC will be different for each
181 parallel process.
182 """
184 # the _cache is a class level variable to avoid being forgotten at every
185 # iteration when using multiprocessing
186 _cache = list()
188 def __init__(self):
189 self.act = 1
190 self.thin = getattr(_SamplingContainer, "nact", 2)
191 self.maxmcmc = getattr(_SamplingContainer, "maxmcmc", 5000) * 50
193 def __call__(self, args):
194 self.args = args
195 if args.kwargs.get("rebuild", False):
196 logger.debug("Force rebuilding cache")
197 self.build_cache()
198 while self.cache[0][2] < args.loglstar:
199 self.cache.pop(0)
200 current_u, current_v, logl, ncall, blob = self.cache.pop(0)
201 blob["remaining"] = len(self.cache)
202 return current_u, current_v, logl, ncall, blob
204 @property
205 def cache(self):
206 if len(self._cache) == 0:
207 self.build_cache()
208 else:
209 logger.debug(f"Not rebuilding cache, remaining size {len(self._cache)}")
210 return self._cache
212 def build_cache(self):
213 args = self.args
214 # Bounds
215 periodic = args.kwargs.get("periodic", None)
216 reflective = args.kwargs.get("reflective", None)
217 boundary_kwargs = dict(periodic=periodic, reflective=reflective)
219 proposals, common_kwargs, proposal_kwargs = _get_proposal_kwargs(args)
221 # Setup
222 current_u = args.u
223 check_interval = int(np.ceil(self.act))
224 target_nact = 50
225 next_check = check_interval
226 n_checks = 0
228 # Initialize internal variables
229 current_v = args.prior_transform(np.array(current_u))
230 logl = args.loglikelihood(np.array(current_v))
231 accept = 0
232 reject = 0
233 nfail = 0
234 ncall = 0
235 act = np.inf
236 u_list = list()
237 v_list = list()
238 logl_list = list()
239 most_failures = 0
240 current_failures = 0
242 iteration = 0
243 while iteration < min(target_nact * act, self.maxmcmc):
244 iteration += 1
246 prop = proposals[iteration % len(proposals)]
247 u_prop = proposal_funcs[prop](
248 u=current_u, **common_kwargs, **proposal_kwargs[prop]
249 )
250 u_prop = apply_boundaries_(u_prop=u_prop, **boundary_kwargs)
251 success = False
252 if u_prop is not None:
253 v_prop = args.prior_transform(np.array(u_prop))
254 logl_prop = args.loglikelihood(np.array(v_prop))
255 ncall += 1
256 if logl_prop > args.loglstar:
257 success = True
258 current_u = u_prop
259 current_v = v_prop
260 logl = logl_prop
262 u_list.append(current_u)
263 v_list.append(current_v)
264 logl_list.append(logl)
266 if success:
267 accept += 1
268 most_failures = max(most_failures, current_failures)
269 current_failures = 0
270 else:
271 nfail += 1
272 current_failures += 1
274 # If we've taken the minimum number of steps, calculate the ACT
275 if iteration > next_check and accept > target_nact:
276 n_checks += 1
277 most_failures = max(most_failures, current_failures)
278 act = self._calculate_act(
279 accept=accept,
280 iteration=iteration,
281 samples=np.array(u_list),
282 most_failures=most_failures,
283 )
284 to_next_update = (act * target_nact - iteration) // 2
285 to_next_update = max(to_next_update, iteration // 100)
286 next_check += to_next_update
287 logger.debug(
288 f"it: {iteration}, accept: {accept}, reject: {reject}, "
289 f"fail: {nfail}, act: {act:.2f}, nact: {iteration / act:.2f} "
290 )
291 elif iteration > next_check:
292 next_check += check_interval
294 most_failures = max(most_failures, current_failures)
295 self.act = self._calculate_act(
296 accept=accept,
297 iteration=iteration,
298 samples=np.array(u_list),
299 most_failures=most_failures,
300 )
301 reject += nfail
302 blob = {"accept": accept, "reject": reject, "scale": args.scale}
303 iact = int(np.ceil(self.act))
304 thin = self.thin * iact
306 if accept == 0:
307 logger.debug(
308 "Unable to find a new point using walk: returning a random point"
309 )
310 u = common_kwargs["rstate"].uniform(size=len(current_u))
311 v = args.prior_transform(u)
312 logl = args.loglikelihood(v)
313 self._cache.append((u, v, logl, ncall, blob))
314 elif not np.isfinite(act):
315 logger.warning(
316 "Unable to find a new point using walk: try increasing maxmcmc"
317 )
318 self._cache.append((current_u, current_v, logl, ncall, blob))
319 elif (self.thin == -1) or (len(u_list) <= thin):
320 self._cache.append((current_u, current_v, logl, ncall, blob))
321 else:
322 u_list = u_list[thin::thin]
323 v_list = v_list[thin::thin]
324 logl_list = logl_list[thin::thin]
325 n_found = len(u_list)
326 accept = max(accept // n_found, 1)
327 reject //= n_found
328 nfail //= n_found
329 ncall_list = [ncall // n_found] * n_found
330 blob_list = [
331 dict(accept=accept, reject=reject, fail=nfail, scale=args.scale)
332 ] * n_found
333 self._cache.extend(zip(u_list, v_list, logl_list, ncall_list, blob_list))
334 logger.debug(
335 f"act: {self.act:.2f}, max failures: {most_failures}, thin: {thin}, "
336 f"iteration: {iteration}, n_found: {n_found}"
337 )
338 logger.debug(
339 f"Finished building cache with length {len(self._cache)} after "
340 f"{iteration} iterations with {ncall} likelihood calls and ACT={self.act:.2f}"
341 )
343 @staticmethod
344 def _calculate_act(accept, iteration, samples, most_failures):
345 """
346 Take the maximum of three ACT estimates, leading to a conservative estimate:
348 - a full ACT estimate as done in :code:`bilby_mcmc`. This is almost always the
349 longest estimator, and is the most computationally expensive. The other
350 methods mostly catch cases where the estimated ACT is very small.
351 - the naive ACT used for the acceptance tracking walk.
352 - the most failed proposals between any pair of accepted steps. This is a strong
353 lower bound, because we know that if we thin by less than this, there will be
354 duplicate points.
355 """
356 if accept > 0:
357 naive_act = 2 / accept * iteration - 1
358 else:
359 return np.inf
360 return max(calculate_tau(samples), naive_act, most_failures)
363class AcceptanceTrackingRWalk:
364 """
365 This is a modified version of dynesty.sampling.sample_rwalk that runs the
366 MCMC random walk for a user-specified number of a crude approximation to
367 the autocorrelation time.
369 This is the proposal method used by default for :code:`Bilby<2` and
370 corresponds to specifying :code:`sample="rwalk"`
371 """
373 # to retain state between calls to pool.Map, this needs to be a class
374 # level attribute
375 old_act = None
377 def __init__(self, old_act=None):
378 self.maxmcmc = getattr(_SamplingContainer, "maxmcmc", 5000)
379 self.nact = getattr(_SamplingContainer, "nact", 40)
381 def __call__(self, args):
382 rstate = get_random_generator(args.rseed)
384 periodic = args.kwargs.get("periodic", None)
385 reflective = args.kwargs.get("reflective", None)
386 boundary_kwargs = dict(periodic=periodic, reflective=reflective)
388 u = args.u
389 nlive = args.kwargs.get("nlive", args.kwargs.get("walks", 100))
391 proposals, common_kwargs, proposal_kwargs = _get_proposal_kwargs(args)
393 accept = 0
394 reject = 0
395 nfail = 0
396 act = np.inf
398 iteration = 0
399 while iteration < self.nact * act:
400 iteration += 1
402 prop = proposals[iteration % len(proposals)]
403 u_prop = proposal_funcs[prop](u, **common_kwargs, **proposal_kwargs[prop])
404 u_prop = apply_boundaries_(u_prop, **boundary_kwargs)
406 if u_prop is None:
407 nfail += 1
408 continue
410 # Check proposed point.
411 v_prop = args.prior_transform(np.array(u_prop))
412 logl_prop = args.loglikelihood(np.array(v_prop))
413 if logl_prop > args.loglstar:
414 u = u_prop
415 v = v_prop
416 logl = logl_prop
417 accept += 1
418 else:
419 reject += 1
421 # If we've taken the minimum number of steps, calculate the ACT
422 if iteration > self.nact:
423 act = self.estimate_nmcmc(
424 accept_ratio=accept / (accept + reject + nfail),
425 safety=1,
426 tau=nlive,
427 )
429 # If we've taken too many likelihood evaluations then break
430 if accept + reject > self.maxmcmc:
431 warnings.warn(
432 f"Hit maximum number of walks {self.maxmcmc} with accept={accept},"
433 f" reject={reject}, and nfail={nfail} try increasing maxmcmc"
434 )
435 break
437 if not (np.isfinite(act) and accept > 0):
438 logger.debug(
439 "Unable to find a new point using walk: returning a random point"
440 )
441 u = rstate.uniform(size=len(u))
442 v = args.prior_transform(u)
443 logl = args.loglikelihood(v)
445 blob = {"accept": accept, "reject": reject + nfail, "scale": args.scale}
446 AcceptanceTrackingRWalk.old_act = act
448 ncall = accept + reject
449 return u, v, logl, ncall, blob
451 def estimate_nmcmc(self, accept_ratio, safety=5, tau=None):
452 """Estimate autocorrelation length of chain using acceptance fraction
454 Using ACL = (2/acc) - 1 multiplied by a safety margin. Code adapted from:
456 - https://github.com/farr/Ensemble.jl
457 - https://github.com/johnveitch/cpnest/blob/master/cpnest/sampler.py
459 Parameters
460 ==========
461 accept_ratio: float [0, 1]
462 Ratio of the number of accepted points to the total number of points
463 old_act: int
464 The ACT of the last iteration
465 maxmcmc: int
466 The maximum length of the MCMC chain to use
467 safety: int
468 A safety factor applied in the calculation
469 tau: int (optional)
470 The ACT, if given, otherwise estimated.
472 Notes
473 =====
474 This method does not compute a reliable estimate of the autocorrelation
475 length for our proposal distributions.
476 """
477 if tau is None:
478 tau = self.maxmcmc / safety
480 if accept_ratio == 0.0:
481 if self.old_act is None:
482 Nmcmc_exact = np.inf
483 else:
484 Nmcmc_exact = (1 + 1 / tau) * self.old_act
485 else:
486 estimated_act = 2 / accept_ratio - 1
487 Nmcmc_exact = safety * estimated_act
488 if self.old_act is not None:
489 Nmcmc_exact = (1 - 1 / tau) * self.old_act + Nmcmc_exact / tau
490 Nmcmc_exact = float(min(Nmcmc_exact, self.maxmcmc))
491 return max(safety, Nmcmc_exact)
494def _get_proposal_kwargs(args):
495 """
496 Resolve the proposal cycle from the provided keyword arguments.
498 The steps involved are:
500 - extract the requested proposal types from the :code:`_SamplingContainer`.
501 If none are specified, only differential evolution will be used.
502 - differential evolution requires the live points to be passed. If they are
503 not present, raise an error.
504 - if a dictionary, e.g., :code:`dict(diff=5, volumetric=1)` is specified,
505 the keys will be used weighted by the values, e.g., 5:1 differential
506 evolution to volumetric.
507 - each proposal needs different keyword arguments, see the specific functions
508 for what requires which.
511 Parameters
512 ==========
513 args: dynesty.sampler.SamplerArgument
514 Object that carries around various pieces of information about the
515 analysis.
516 """
517 rstate = get_random_generator(args.rseed)
518 walks = args.kwargs.get("walks", 100)
519 current_u = args.u
520 n_cluster = args.axes.shape[0]
522 proposals = getattr(_SamplingContainer, "proposals", None)
523 if proposals is None:
524 proposals = ["diff"]
525 if "diff" in proposals:
526 live = args.kwargs.get("live", None)
527 if live is None:
528 raise ValueError(
529 "Live points not passed for differential evolution, specify "
530 "bound='live' to use differential evolution proposals."
531 )
532 live = np.unique(live, axis=0)
533 matches = np.where(np.equal(current_u, live).all(axis=1))[0]
534 np.delete(live, matches, 0)
536 if isinstance(proposals, (list, set, tuple)):
537 proposals = rstate.choice(proposals, int(walks))
538 elif isinstance(proposals, dict):
539 props, weights = zip(*proposals.items())
540 weights = np.array(weights) / sum(weights)
541 proposals = rstate.choice(list(props), int(walks), p=weights)
543 common_kwargs = dict(
544 n=len(current_u),
545 n_cluster=n_cluster,
546 rstate=rstate,
547 )
548 proposal_kwargs = dict()
549 if "diff" in proposals:
550 proposal_kwargs["diff"] = dict(
551 live=live[:, :n_cluster],
552 mix=0.5,
553 scale=2.38 / (2 * n_cluster) ** 0.5,
554 )
555 if "volumetric" in proposals:
556 proposal_kwargs["volumetric"] = dict(
557 axes=args.axes,
558 scale=args.scale,
559 )
560 return proposals, common_kwargs, proposal_kwargs
563def propose_differetial_evolution(
564 u,
565 live,
566 n,
567 n_cluster,
568 rstate,
569 mix=0.5,
570 scale=1,
571):
572 r"""
573 Propose a new point using ensemble differential evolution
574 (`ter Braak + (2006) <https://doi.org/10.1007/s11222-006-8769-1>`_).
576 .. math::
578 u_{\rm prop} = u + \gamma (v_{a} - v_{b})
580 We consider two choices for :math:`\gamma`: weighted by :code:`mix`.
582 - :math:`\gamma = 1`: this is a mode-hopping mode for efficiently
583 exploring multi-modal spaces
584 - :math:`\gamma \sim \Gamma\left(\gamma; k=4, \theta=\frac{\kappa}{4}\right)`
586 Here :math:`\kappa = 2.38 / \sqrt{2 n}` unless specified by the user and
587 we scale by a random draw from a Gamma distribution. The specific
588 distribution was chosen somewhat arbitrarily to have mean and mode close to
589 :math:`\kappa` and give good acceptance and autocorrelation times on a subset
590 of problems.
592 Parameters
593 ----------
594 u: np.ndarray
595 The current point.
596 live: np.ndarray
597 The ensemble of live points to select :math:`v` from.
598 n: int
599 The number of dimensions being explored
600 n_cluster: int
601 The number of dimensions to run the differential evolution over, the
602 first :code:`n_cluster` dimensions are used. The rest are randomly
603 sampled from the prior.
604 rstate: np.random.RandomState
605 The random state to use to generate random numbers.
606 mix: float
607 The fraction of proposed points that should follow the specified scale
608 rather than mode hopping. :code:`default=0.5`
609 scale: float
610 The amount to scale the difference vector by.
611 :code:`default = 2.38 / (2 * n_cluster)**0.5)`
613 Returns
614 -------
615 u_prop: np.ndarray
616 The proposed point.
617 """
618 delta = np.diff(rstate.choice(live, 2, replace=False), axis=0)[0]
619 if rstate.uniform(0, 1) < mix:
620 if scale is None:
621 scale = 2.38 / (2 * n_cluster) ** 0.5
622 scale *= rstate.gamma(4, 0.25)
623 else:
624 scale = 1
625 u_prop = u.copy()
626 u_prop[:n_cluster] += delta * scale
627 u_prop[n_cluster:] = rstate.uniform(0, 1, n - n_cluster)
628 return u_prop
631def propose_volumetric(
632 u,
633 axes,
634 scale,
635 n,
636 n_cluster,
637 rstate,
638):
639 """
640 Propose a new point using the default :code:`dynesty` proposal.
642 The new difference vector is scaled by a vector isotropically drawn
643 from an ellipsoid centered on zero with covariance given by the
644 provided axis. Note that the magnitude of this proposal is heavily
645 skewed to the size of the ellipsoid.
647 Parameters
648 ----------
649 u: np.ndarray
650 The current point.
651 n: int
652 The number of dimensions being explored.
653 scale: float
654 The amount to scale the proposed point by.
655 n_cluster: int
656 The number of dimensions to run the differential evolution over, the
657 first :code:`n_cluster` dimensions are used. The rest are randomly
658 sampled from the prior.
659 rstate: np.random.RandomState
660 The random state to use to generate random numbers.
662 Returns
663 -------
664 u_prop: np.ndarray
665 The proposed point.
666 """
667 # Propose a direction on the unit n-sphere.
668 drhat = rstate.normal(0, 1, n_cluster)
669 drhat /= np.linalg.norm(drhat)
671 # Scale based on dimensionality.
672 dr = drhat * rstate.uniform(0, 1) ** (1.0 / n_cluster)
674 # Transform to proposal distribution.
675 delta = scale * np.dot(axes, dr)
676 u_prop = u.copy()
677 u_prop[:n_cluster] += delta
678 u_prop[n_cluster:] = rstate.uniform(0, 1, n - n_cluster)
679 return u_prop
682def apply_boundaries_(u_prop, periodic, reflective):
683 """
684 Apply the periodic and reflective boundaries and test if we are inside the
685 unit cube.
687 Parameters
688 ----------
689 u_prop: np.ndarray
690 The proposed point in the unit hypercube space.
691 periodic: np.ndarray
692 Indices of the parameters with periodic boundaries.
693 reflective: np.ndarray
694 Indices of the parameters with reflective boundaries.
696 Returns
697 =======
698 [np.ndarray, None]:
699 Either the remapped proposed point, or None if the proposed point
700 lies outside the unit cube.
701 """
702 # Wrap periodic parameters
703 if periodic is not None:
704 u_prop[periodic] = np.mod(u_prop[periodic], 1)
706 # Reflect
707 if reflective is not None:
708 u_prop[reflective] = apply_reflect(u_prop[reflective])
710 if u_prop.min() < 0 or u_prop.max() > 1:
711 return None
712 else:
713 return u_prop
716proposal_funcs = dict(diff=propose_differetial_evolution, volumetric=propose_volumetric)