eryn 1.2.0__py3-none-any.whl

eryn/git_version.py.in

@@ -0,0 +1,7 @@
+"""Metadata deduced from git at build time."""
+
+id: str
+short_id: str
+
+id = "@ERYN_GIT_COMMIT_ID@"
+short_id = "@ERYN_GIT_COMMIT_SHORT_ID@"

eryn/model.py

@@ -0,0 +1,18 @@
+# -*- coding: utf-8 -*-
+
+from collections import namedtuple
+
+__all__ = ["Model"]
+
+
+Model = namedtuple(
+    "Model",
+    (
+        "log_like_fn",
+        "compute_log_like_fn",
+        "compute_log_prior_fn",
+        "temperature_control",
+        "map_fn",
+        "random",
+    ),
+)

eryn/moves/__init__.py

@@ -0,0 +1,42 @@
+# -*- coding: utf-8 -*-
+
+# from .de import DEMove
+# from .de_snooker import DESnookerMove
+from .gaussian import GaussianMove
+
+# from .kde import KDEMove
+from .mh import MHMove
+from .move import Move
+from .red_blue import RedBlueMove
+from .stretch import StretchMove
+
+# from .walk import WalkMove
+from .tempering import TemperatureControl
+from .rj import ReversibleJumpMove
+from .distgenrj import DistributionGenerateRJ
+from .distgen import DistributionGenerate
+from .multipletry import MultipleTryMove
+from .group import GroupMove
+from .groupstretch import GroupStretchMove
+from .combine import CombineMove
+
+# from .basicmodelswaprj import BasicSymmetricModelSwapRJMove
+from .mtdistgen import MTDistGenMove
+from .mtdistgenrj import MTDistGenMoveRJ
+from .multipletry import MultipleTryMove
+
+__all__ = [
+    "Move",
+    "MHMove",
+    "GaussianMove",
+    "RedBlueMove",
+    "StretchMove",
+    "DistributionGenerateRJ",
+    "DistributionGenerate",
+    "TemperatureControl",
+    "ReversibleJumpMove",
+    "MultipleTryMove",
+    "GroupMove",
+    "GroupStretchMove",
+    "CombineMove",
+]

eryn/moves/combine.py

@@ -0,0 +1,135 @@
+# -*- coding: utf-8 -*-
+
+from ..state import BranchSupplemental
+from . import Move
+import numpy as np
+import tqdm
+
+__all__ = ["CombineMove"]
+
+
+class CombineMove(Move):
+    """Move that combines specific moves in order.
+
+    Args:
+        moves (list): List of moves, similar to how ``moves`` is submitted
+            to :class:`eryn.ensemble.EnsembleSampler`. If weights are provided,
+            they will be ignored.
+        *args (tuple, optional): args to be passed to :class:`Move`.
+        verbose (bool, optional): If ``True``, use ``tqdm`` to show progress throught steps.
+            This can be very helpful when debugging.
+        **kwargs (dict, optional): kwargs to be passed to :class:`Move`.
+
+
+    """
+
+    def __init__(self, moves, *args, verbose=False, **kwargs):
+        # store moves
+        self.moves = moves
+        self.verbose = verbose
+        Move.__init__(self, *args, **kwargs)
+
+    @property
+    def accepted(self):
+        """Accepted counts for each move."""
+        if self._accepted is None:
+            raise ValueError(
+                "accepted must be inititalized with the init_accepted function if you want to use it."
+            )
+        # this retrieves the accepted arrays from the individual moves
+        accepted_out = [move.accepted for move in self.moves]
+        return accepted_out
+
+    @accepted.setter
+    def accepted(self, accepted):
+        # set the accepted arrays for all moves
+        assert isinstance(accepted, np.ndarray)
+        for move in self.moves:
+            move.accepted = accepted.copy()
+
+    @property
+    def acceptance_fraction(self):
+        """get acceptance fraction averaged over all moves"""
+        acceptance_fraction_out = np.mean(
+            [move.acceptance_fraction for move in self.moves], axis=0
+        )
+        return acceptance_fraction_out
+
+    @property
+    def acceptance_fraction_separate(self):
+        """get acceptance fraction from each move"""
+        acceptance_fraction_out = [move.acceptance_fraction for move in self.moves]
+        return acceptance_fraction_out
+
+    @property
+    def temperature_control(self):
+        """temperature controller"""
+        return self._temperature_control
+
+    @temperature_control.setter
+    def temperature_control(self, temperature_control):
+        # when setting the temperature control object
+        # need to apply it to each move
+        for i, move in enumerate(self.moves):
+            # if weights were provided with moves, remove move class
+            if isinstance(move, tuple):
+                move = move[0]
+            # set temperature control for each move
+            move.temperature_control = temperature_control
+
+        # main temperature control here for reference
+        self._temperature_control = temperature_control
+
+    @property
+    def periodic(self):
+        """periodic parameter information"""
+        return self._periodic
+
+    @periodic.setter
+    def periodic(self, periodic):
+        # when setting the periodic parameters
+        # need to apply it to each move
+        for i, move in enumerate(self.moves):
+            # if weights were provided with moves, remove move class
+            if isinstance(move, tuple):
+                move = move[0]
+            move.periodic = periodic
+        self._periodic = periodic
+
+    def propose(self, model, state):
+        """Propose a combined move.
+
+        Args:
+            model (:class:`eryn.model.Model`): Carrier of sampler information.
+            state (:class:`State`): Current state of the sampler.
+
+        Returns:
+            tuple: (state, accepted)
+                The first return is the state of the sampler after the move.
+                The second return value is the accepted count array for each walker
+                counting for all proposals.
+
+        """
+
+        # prepare for verbosity if needed
+        iterator = enumerate(self.moves)
+        if self.verbose:
+            iterator = tqdm.tqdm(iterator)
+
+        # we will set this inside the loop during the first iteration
+        accepted_out = None
+        for i, move in iterator:
+            # get move out of tuple
+            if isinstance(move, tuple):
+                move = move[0]
+
+            # run move
+            state, accepted = move.propose(model, state)
+
+            # set (first iteration) or add (after first iteration) to accepted_out
+            if accepted_out is None:
+                accepted_out = accepted.copy()
+            else:
+                accepted_out += accepted
+
+        return state, accepted_out

eryn/moves/delayedrejection.py

@@ -0,0 +1,229 @@
+# -*- coding: utf-8 -*-
+
+import numpy as np
+from copy import deepcopy
+
+from ..state import State
+from .move import Move
+from ..state import BranchSupplemental
+
+__all__ = ["MHMove"]
+
+
+class DelayedRejectionContainer:
+    def __init__(self, **kwargs):
+        for key, item in kwargs.items():
+            setattr(self, key, item)
+
+        # Initialize
+        self.coords = []
+        self.log_prob = []
+        self.log_prior = []
+        self.alpha = []
+
+    def append(self, new_coords, new_log_prob, new_log_prior, new_alpha):
+        self.coords.append(new_coords)
+        self.log_prob.append(new_log_prob)
+        self.log_prior.append(new_log_prior)
+        self.alpha.append(new_alpha)
+
+
+class DelayedRejection(Move):
+    r"""
+    Delayed Rejection scheme assuming symmetric and non-adaptive proposal distribution.
+    We apply the DR algorithm only on the cases where we have rejected a +1 proposal for
+    a given Reversible Jump proposal and branch.
+
+    Refernces:
+
+    Tierney L and Mira A, Stat. Med. 18 2507 (1999)
+    Haario et al, Stat. Comput. 16:339-354 (2006)
+    Mira A, Metron - International Journal of Statistics, vol. LIX, issue 3-4, 231-241 (2001)
+    M. Trias, et al, https://arxiv.org/abs/0904.2207
+
+    """
+
+    def __init__(self, proposal, max_iter=10, **kwargs):
+        self.proposal = proposal
+        self.max_iter = max_iter
+        self.dr_container = None
+        super(DelayedRejection, self).__init__(**kwargs)
+
+    def dr_scheme(
+        self,
+        state,
+        new_state,
+        keep_rejected,
+        model,
+        ntemps,
+        nwalkers,
+        inds_for_change,
+        inds=None,
+        dr_iter=0,
+    ):
+        """Calcuate the delayed rejection acceptace ratio.
+
+        Args:
+            stateslist (:class:`State`): a python list containing the proposed states
+
+        Returns:
+            logalpha: a numpy array containing the acceptance ratios per temperature and walker.
+        """
+        # Draw a uniform random for the previously rejected points
+        randU = model.random.rand(
+            ntemps, nwalkers
+        )  # We draw for all temps x walkers but we ignore
+        # previously accepted points by setting prior[rej] = - inf
+
+        old_new_state = State(new_state, copy=True)
+
+        # Propose a new point
+        new_state, log_proposal_ratio = self.get_new_state(
+            model, new_state, keep_rejected
+        )  # Get a new state
+
+        # Compute log-likelihood and log-prior
+        logp = new_state.log_prior
+        logl = new_state.log_like
+
+        # Compute the logposterior for all
+        logP = self.compute_log_posterior(logl, logp)
+
+        # Compute log-likelihood and log-prior
+        prev_logp = old_new_state.log_prior
+        prev_logl = old_new_state.log_like
+
+        # Compute the logposterior for all
+        prev_logP = self.compute_log_posterior(prev_logl, prev_logp)
+
+        # Compute the acceptance ratio
+        lndiff = logP - prev_logP + log_proposal_ratio
+        alpha_1 = np.exp(lndiff)
+        alpha_1[alpha_1 > 1.0] = 1.0  # np.min((1, alpha))
+
+        # update delayed rejection alpha
+        dr_alpha = np.exp(
+            lndiff
+            + np.log(1.0 - alpha_1)
+            - np.log(1.0 - old_new_state.supplemental[:]["past_alpha"])
+        )
+        dr_alpha[dr_alpha > 1.0] = 1.0  # np.min((1., dr_alpha ))
+        dr_alpha = np.nan_to_num(dr_alpha)  # Automatically reject NaNs
+
+        new_state.supplemental[:] = {"alpha": dr_alpha}  # Replace current dr alpha
+
+        new_accepted = np.logical_or(
+            dr_alpha >= 1.0, randU < dr_alpha
+        )  # Decide on accepted points
+
+        # Update state with the new accepted points
+        state = self.update(state, new_state, new_accepted)
+
+        return state, new_accepted, new_state
+
+    def get_new_state(self, model, state, keep):
+        """A utility function to propose new points"""
+        qn, factors = self.proposal.get_proposal(
+            state.branches_coords, state.branches_inds, model.random
+        )
+
+        # Compute prior of the proposed position
+        logp = model.compute_log_prior_fn(qn, inds=state.branches_inds)
+        logp[~keep] = -np.inf  # This trick help us compute only the indeces of interest
+
+        # Compute the lnprobs of the proposed position.
+        logl, new_blobs = model.compute_log_like_fn(
+            qn, inds=state.branches_inds, logp=logp
+        )
+
+        # Update the parameters, update the state. TODO: Fix blobs?
+        new_state = State(
+            qn,
+            log_like=logl,
+            log_prior=logp,
+            blobs=new_blobs,
+            inds=state.branches_inds,
+            supplemental=state.supplemental,
+        )  # I create a new initial state that all are accepted
+        return new_state, factors
+
+    def propose(
+        self,
+        log_diff_0,
+        accepted,
+        model,
+        state,
+        new_state,
+        inds,
+        inds_for_change,
+        factors,
+    ):
+        """Use the move to generate a proposal and compute the acceptance
+
+        Args:
+            accepted ():
+            model (:class:`eryn.model.Model`): Carrier of sampler information.
+            state (:class:`State`): Current state of the sampler.
+            rj_inds (): Dictionary containing the indices where the Reversible Jump
+            move proposed "birth" of a model. Will we operate a Delayed Rejection type
+            of move only on those cases. The keys of the dictionary are the names of the
+            models.
+
+        Returns:
+            :class:`State`: State of sampler after proposal is complete.
+
+        """
+
+        # Check to make sure that the dimensions match.
+        ntemps, nwalkers, _, _ = state.branches[list(state.branches.keys())[0]].shape
+
+        alpha_0 = np.exp(log_diff_0)
+        alpha_0[alpha_0 > 1.0] = 1.0  # np.min((1.0, alpha_0))
+        new_state.supplemental = BranchSupplemental(
+            {"past_alpha": alpha_0}, base_shape=(ntemps, nwalkers)
+        )
+
+        # Check to make sure that the dimensions match.
+        ntemps, nwalkers, _, _ = state.branches[list(state.branches.keys())[0]].shape
+
+        dr_iter = 0  # Initialize
+
+        # Begin main DR loop. Stop when we exceed the maximum iterations, or (extreme case) all proposals are accepted
+        while dr_iter <= self.max_iter and not np.all(accepted):
+            rejected = ~accepted  # Get rejected points
+
+            # Get the +1 proposals that got previously rejected
+            plus_one_rej_inds = {}
+            for name in inds_for_change:
+                plus_one_inds = inds_for_change[name]["+1"][:, :2]
+            plus_one_rej_inds[name] = plus_one_inds[
+                rejected[(plus_one_inds[:, 0], plus_one_inds[:, 1])]
+            ]
+
+            # Generate the indeces of the proposals that got rejected
+            keep_rejected = np.unique(
+                np.concatenate(list(plus_one_rej_inds.values())), axis=0
+            )
+            run_dr = np.zeros_like(rejected)
+            run_dr[tuple(keep_rejected.T)] = True
+
+            # Pass into the DR scheme
+            state, new_accepted, new_state = self.dr_scheme(
+                state,
+                new_state,
+                run_dr,
+                model,
+                ntemps,
+                nwalkers,
+                inds_for_change,
+                inds=inds,
+            )
+
+            # Update the accepted, increment current iteration
+            accepted += new_accepted
+            dr_iter += 1
+
+        if self.temperature_control is not None:
+            state, accepted = self.temperature_control.temper_comps(state, accepted)
+
+        return state, accepted

eryn/moves/distgen.py

@@ -0,0 +1,104 @@
+# -*- coding: utf-8 -*-
+
+import numpy as np
+from .mh import MHMove
+from ..prior import ProbDistContainer
+
+__all__ = ["DistributionGenerate"]
+
+
+class DistributionGenerate(MHMove):
+    """Generate proposals from a distribution
+
+    Args:
+        generate_dist (dict): Dictionary with keys as branch names and items as
+            :class:`ProbDistContainer` objects that have ``logpdf``
+            and ``rvs`` methods.
+
+    """
+
+    def __init__(self, generate_dist, *args, **kwargs):
+        if not isinstance(generate_dist, dict):
+            raise ValueError(
+                "When entering directly into the DistributionGenerate class, generate_dist must be a dictionary. The keys are branch names and the items are ProbDistContainer objects."
+            )
+
+        for key in generate_dist:
+            if not isinstance(generate_dist[key], ProbDistContainer):
+                raise ValueError(
+                    "Distributions need to be eryn.prior.ProbDistContainer object."
+                )
+        self.generate_dist = generate_dist
+        super(DistributionGenerate, self).__init__(*args, **kwargs)
+
+    def get_proposal(self, branches_coords, random, branches_inds=None, **kwargs):
+        """Make a proposal
+
+        Args:
+            branches_coords (dict): Keys are ``branch_names`` and values are
+                np.ndarray[ntemps, nwalkers, nleaves_max, ndim] representing
+                coordinates for walkers.
+            random (object): Current random state object.
+            branches_inds (dict, optional): Keys are ``branch_names`` and values are
+                np.ndarray[ntemps, nwalkers, nleaves_max] representing which
+                leaves are currently being used. (default: ``None``)
+            **kwargs (ignored): This is added for compatibility. It is ignored in this function.
+
+        Returns:
+            tuple: Tuple containing proposal information.
+                First entry is the new coordinates as a dictionary with keys
+                as ``branch_names`` and values as
+                ``double `` np.ndarray[ntemps, nwalkers, nleaves_max, ndim] containing
+                proposed coordinates. Second entry
+                is the factors associated with the
+                proposal necessary for detailed balance. This is effectively
+                any term in the detailed balance fraction. +log of factors if
+                in the numerator. -log of factors if in the denominator.
+
+        """
+
+        # set up all initial holders
+        q = {}
+        factors = {}
+        new_inds = {}
+        if branches_inds is None:
+            branches_inds = {
+                name: np.ones(coords.shape[:-1], dtype=bool)
+                for name, coords in branches_coords
+            }
+
+        # iterate through branches and propose new points where inds == True
+        for i, (name, coords, inds) in enumerate(
+            zip(
+                branches_coords.keys(),
+                branches_coords.values(),
+                branches_inds.values(),
+            )
+        ):
+            # copy over previous info
+            ntemps, nwalkers, _, _ = coords.shape
+            q[name] = coords.copy()
+            new_inds[name] = inds.copy()
+
+            if i == 0:
+                factors = np.zeros((ntemps, nwalkers))
+
+            # add coordinates for new leaves
+            current_generate_dist = self.generate_dist[name]
+            inds_here = np.where(inds == True)
+            num_inds_change = len(inds_here[0])
+
+            old_points = coords[inds_here]
+
+            # old points so + log(qold)
+            factors[inds_here[:2]] += +1 * current_generate_dist.logpdf(old_points)
+
+            # Draw
+            new_points = current_generate_dist.rvs(size=num_inds_change)
+
+            # new point, so -log(qnew)
+            factors[inds_here[:2]] += -1 * current_generate_dist.logpdf(new_points)
+
+            q[name][inds_here] = new_points
+
+        return q, factors