eryn 1.2.0__py3-none-any.whl

eryn/moves/groupstretch.py

@@ -0,0 +1,120 @@
+# -*- coding: utf-8 -*-
+try:
+    import cupy as xp
+except (ModuleNotFoundError, ImportError):
+    pass
+
+import numpy as np
+
+from .group import GroupMove
+from .stretch import StretchMove
+
+__all__ = ["GroupStretchMove"]
+
+
+class GroupStretchMove(GroupMove, StretchMove):
+    """Proposal like stretch with stationary compliment.
+
+    This move uses the stretch proposal method and math, but the compliment
+    of walkers used to propose a new point is chosen from a stationary group
+    rather than the current walkers in the ensemble.
+
+    This move allows for "stretch"-like proposal to be used in Reversible Jump MCMC.
+
+    Args:
+        **kwargs (dict, optional): Keyword arguments passed to :class:`GroupMove` and
+            :class:`StretchMove`.
+
+    """
+
+    def __init__(self, **kwargs):
+        GroupMove.__init__(self, **kwargs)
+        StretchMove.__init__(self, **kwargs)
+
+    def get_proposal(
+        self,
+        s_all,
+        random,
+        gibbs_ndim=None,
+        s_inds_all=None,
+        branch_supps=None,
+        **kwargs
+    ):
+        """Generate group stretch proposal coordinates
+
+        Args:
+            s_all (dict): Keys are ``branch_names`` and values are coordinates
+                for which a proposal is to be generated.
+            random (object): Random state object.
+            gibbs_ndim (int or np.ndarray, optional): If Gibbs sampling, this indicates
+                the true dimension. If given as an array, must have shape ``(ntemps, nwalkers)``.
+                See the tutorial for more information.
+                (default: ``None``)
+            s_inds_all (dict, optional): Keys are ``branch_names`` and values are
+                ``inds`` arrays indicating which leaves are currently used. (default: ``None``)
+            branch_supps (dict, optional): Keys are ``branch_names`` and values are
+                :class:`BranchSupplemental` objects. For the group stretch,
+                ``branch_supps`` are the best device for passing and tracking useful
+                information. (default: ``None``)
+
+        Returns:
+            tuple: First entry is new positions. Second entry is detailed balance factors.
+
+        Raises:
+            ValueError: Issues with dimensionality.
+
+        """
+        # needs to be set before we reach the end
+        self.zz = None
+        random_number_generator = random if not self.use_gpu else self.xp.random
+        newpos = {}
+
+        # iterate over branches
+        for i, name in enumerate(s_all):
+            # get points to move
+            s = self.xp.asarray(s_all[name])
+
+            Ns = s.shape[1]
+
+            if s_inds_all is not None:
+                s_inds = self.xp.asarray(s_inds_all[name])
+            else:
+                s_inds = None
+
+            ntemps, nwalkers, nleaves_max, ndim_here = s.shape
+
+            # gets rid of any values of exactly zero
+            ndim_temp = nleaves_max * ndim_here
+
+            # need to properly handle ndim
+            if i == 0:
+                ndim = ndim_temp
+                Ns_check = Ns
+
+            else:
+                ndim += ndim_temp
+                if Ns_check != Ns:
+                    raise ValueError("Different number of walkers across models.")
+
+            Ns = nwalkers
+
+            # get actual compliment values
+            c_temp = self.choose_c_vals(
+                name, s, s_inds=s_inds, branch_supps=branch_supps
+            )
+
+            # use stretch to get new proposals
+            newpos[name] = self.get_new_points(
+                name, s, c_temp, Ns, s.shape, i, random_number_generator
+            )
+
+        # proper factors
+        factors = (ndim - 1.0) * self.xp.log(self.zz)
+        if self.use_gpu and not self.return_gpu:
+            factors = factors.get()
+
+        if gibbs_ndim is not None:
+            # adjust factors in place
+            self.adjust_factors(factors, ndim, gibbs_ndim)
+
+        return newpos, factors

eryn/moves/mh.py

@@ -0,0 +1,193 @@
+# -*- coding: utf-8 -*-
+
+import numpy as np
+from copy import deepcopy
+from ..state import State
+from .move import Move
+
+__all__ = ["MHMove"]
+
+
+class MHMove(Move):
+    r"""A general Metropolis-Hastings proposal
+
+    Concrete implementations can be made by providing a ``get_proposal`` method.
+    For standard Gaussian Metropolis moves, :class:`moves.GaussianMove` can be used.
+
+    """
+
+    def __init__(self, **kwargs):
+
+        Move.__init__(self, **kwargs)
+
+    def get_proposal(self, branches_coords, random, branches_inds=None, **kwargs):
+        """Get 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: (Proposed coordinates, factors) -> (dict, np.ndarray)
+
+        Raises:
+            NotImplementedError: If proposal is not implemented in a subclass.
+
+        """
+
+        raise NotImplementedError("The proposal must be implemented by " "subclasses")
+
+    def setup(self, branches_coords):
+        """Any setup for the proposal.
+
+        Args:
+            branches_coords (dict): Keys are ``branch_names``. Values are
+                np.ndarray[ntemps, nwalkers, nleaves_max, ndim]. These are the curent
+                coordinates for all the walkers.
+
+        """
+
+    def propose(self, model, state):
+        """Use the move to generate a proposal and compute the acceptance
+
+        Args:
+            model (:class:`eryn.model.Model`): Carrier of sampler information.
+            state (:class:`State`): Current state of the sampler.
+
+        Returns:
+            :class:`State`: State of sampler after proposal is complete.
+
+        """
+
+        self.setup(state.branches_coords)
+
+        # get all branch names for gibbs setup
+        all_branch_names = list(state.branches.keys())
+
+        # get initial shape information
+        ntemps, nwalkers, _, _ = state.branches[all_branch_names[0]].shape
+
+        # in case there are no leaves yet
+        accepted = np.zeros((ntemps, nwalkers), dtype=bool)
+
+        # iterate through gibbs setup
+        for branch_names_run, inds_run in self.gibbs_sampling_setup_iterator(
+            all_branch_names
+        ):
+            # setup supplemental information
+            if not np.all(
+                np.asarray(list(state.branches_supplemental.values())) == None
+            ):
+                new_branch_supps = deepcopy(state.branches_supplemental)
+            else:
+                new_branch_supps = None
+
+            if state.supplemental is not None:
+                new_supps = deepcopy(state.supplemental)
+            else:
+                new_supps = None
+
+            # setup information according to gibbs info
+            (
+                coords_going_for_proposal,
+                inds_going_for_proposal,
+                at_least_one_proposal,
+            ) = self.setup_proposals(
+                branch_names_run, inds_run, state.branches_coords, state.branches_inds
+            )
+
+            # if no walkers are actually being proposed
+            if not at_least_one_proposal:
+                continue
+
+            self.current_model = model
+            self.current_state = state
+
+            # Get the move-specific proposal.
+            q, factors = self.get_proposal(
+                coords_going_for_proposal,
+                model.random,
+                branches_inds=inds_going_for_proposal,
+                supps=new_supps,
+                branch_supps=new_branch_supps,
+            )
+
+            # account for gibbs sampling
+            self.cleanup_proposals_gibbs(
+                branch_names_run, inds_run, q, state.branches_coords
+            )
+
+            # order everything properly
+            q, _, new_branch_supps = self.ensure_ordering(
+                list(state.branches.keys()), q, state.branches_inds, new_branch_supps
+            )
+
+            # if not wrapping with mutliple try (normal route)
+            if not hasattr(self, "mt_ll") or not hasattr(self, "mt_lp"):
+                # Compute prior of the proposed position
+                logp = model.compute_log_prior_fn(q, inds=state.branches_inds)
+
+                self.fix_logp_gibbs(
+                    branch_names_run, inds_run, logp, state.branches_inds
+                )
+
+                # Compute the lnprobs of the proposed position.
+                # Can adjust supplementals in place
+                logl, new_blobs = model.compute_log_like_fn(
+                    q,
+                    inds=state.branches_inds,
+                    logp=logp,
+                    supps=new_supps,
+                    branch_supps=new_branch_supps,
+                )
+
+            else:
+                # if already computed in multiple try
+                logl = self.mt_ll
+                logp = self.mt_lp
+                new_blobs = None
+
+            # get log posterior
+            logP = self.compute_log_posterior(logl, logp)
+
+            # get previous information
+            prev_logl = state.log_like
+
+            prev_logp = state.log_prior
+
+            # takes care of tempering
+            prev_logP = self.compute_log_posterior(prev_logl, prev_logp)
+
+            # difference
+            lnpdiff = factors + logP - prev_logP
+
+            # draw against acceptance fraction
+            accepted = lnpdiff > np.log(model.random.rand(ntemps, nwalkers))
+
+            # Update the parameters
+            new_state = State(
+                q,
+                log_like=logl,
+                log_prior=logp,
+                blobs=new_blobs,
+                inds=state.branches_inds,
+                supplemental=new_supps,
+                branch_supplemental=new_branch_supps,
+            )
+            state = self.update(state, new_state, accepted)
+
+            # add to move-specific accepted information
+            self.accepted += accepted
+            self.num_proposals += 1
+
+        # temperature swaps
+        if self.temperature_control is not None:
+            state = self.temperature_control.temper_comps(state)
+
+        return state, accepted