eryn 1.2.0__py3-none-any.whl

eryn/moves/distgenrj.py

@@ -0,0 +1,222 @@
+# -*- coding: utf-8 -*-
+
+import numpy as np
+
+from .rj import ReversibleJumpMove
+from ..prior import ProbDistContainer
+
+__all__ = ["DistributionGenerateRJ"]
+
+
+class DistributionGenerateRJ(ReversibleJumpMove):
+    """Generate Revesible-Jump proposals from a distribution.
+
+    The prior can be entered as the ``generate_dist`` to generate proposals directly from the prior.
+
+    Args:
+        generate_dist (dict): Keys are branch names and values are :class:`ProbDistContainer` objects 
+            that have ``logpdf`` and ``rvs`` methods. If you 
+        *args (tuple, optional): Additional arguments to pass to parent classes.
+        **kwargs (dict, optional): Keyword arguments passed to parent classes.
+
+    """
+
+    def __init__(self, generate_dist, *args, **kwargs):
+
+        # make sure all inputs are distribution Containers
+        for key in generate_dist:
+            if not isinstance(generate_dist[key], ProbDistContainer):
+                raise ValueError(
+                    "Distributions need to be eryn.prior.ProbDistContiner object."
+                )
+        self.generate_dist = generate_dist
+        super(DistributionGenerateRJ, self).__init__(*args, **kwargs)
+
+    def get_model_change_proposal(self, inds, random, nleaves_min, nleaves_max):
+        """Helper function for changing the model count by 1.
+        
+        This helper function works with nested models where you want to add or remove
+        one leaf at a time. 
+
+        Args:
+            inds (np.ndarray): ``inds`` values for this specific branch with shape 
+                ``(ntemps, nwalkers, nleaves_max)``.
+            random (object): Current random state of the sampler.
+            nleaves_min (int): Minimum allowable leaf count for this branch.
+            nleaves_max (int): Maximum allowable leaf count for this branch.
+
+        Returns:
+            dict: Keys are ``"+1"`` and ``"-1"``. Values are indexing information.
+                    ``"+1"`` and ``"-1"`` indicate if a source is being added or removed, respectively.
+                    The indexing information is a 2D array with shape ``(number changing, 3)``.
+                    The length 3 is the index into each of the ``(ntemps, nwalkers, nleaves_max)``.
+        
+        """
+
+        ntemps, nwalkers, _ = inds.shape
+
+        nleaves = inds.sum(axis=-1)
+
+        # choose whether to add or remove
+        if self.fix_change is None:
+            change = random.choice([-1, +1], size=nleaves.shape)
+        else:
+            change = np.full(nleaves.shape, self.fix_change)
+
+        # fix edge cases
+        change = (
+            change * ((nleaves != nleaves_min) & (nleaves != nleaves_max))
+            + (+1) * (nleaves == nleaves_min)
+            + (-1) * (nleaves == nleaves_max)
+        )
+
+        # setup storage for this information
+        inds_for_change = {}
+        num_increases = np.sum(change == +1)
+        inds_for_change["+1"] = np.zeros((num_increases, 3), dtype=int)
+        num_decreases = np.sum(change == -1)
+        inds_for_change["-1"] = np.zeros((num_decreases, 3), dtype=int)
+
+        # TODO: not loop ? Is it necessary?
+        # TODO: might be able to subtract new inds from old inds type of thing
+        # fill the inds_for_change
+        increase_i = 0
+        decrease_i = 0
+        for t in range(ntemps):
+            for w in range(nwalkers):
+                # check if add or remove
+                change_tw = change[t][w]
+                # inds array from specific walker
+                inds_tw = inds[t][w]
+
+                # adding
+                if change_tw == +1:
+                    # find where leaves are not currently used
+                    inds_false = np.where(inds_tw == False)[0]
+                    # decide which spot to add
+                    ind_change = random.choice(inds_false)
+                    # put in the indexes into inds arrays
+                    inds_for_change["+1"][increase_i] = np.array(
+                        [t, w, ind_change], dtype=int
+                    )
+                    # count increases
+                    increase_i += 1
+
+                # removing
+                elif change_tw == -1:
+                    # change_tw == -1
+                    # find which leavs are used
+                    inds_true = np.where(inds_tw == True)[0]
+                    # choose which to remove
+                    ind_change = random.choice(inds_true)
+                    # add indexes into inds
+                    if inds_for_change["-1"].shape[0] > 0:
+                        inds_for_change["-1"][decrease_i] = np.array(
+                            [t, w, ind_change], dtype=int
+                        )
+                        decrease_i += 1
+                    # do not care currently about what we do with discarded coords, they just sit in the state
+                # model component number not changing
+                else:
+                    pass
+        return inds_for_change
+
+    def get_proposal(
+        self, all_coords, all_inds, nleaves_min_all, nleaves_max_all, random, **kwargs
+    ):
+        """Make a proposal
+
+        Args:
+            all_coords (dict): Keys are ``branch_names``. Values are
+                np.ndarray[ntemps, nwalkers, nleaves_max, ndim]. These are the curent
+                coordinates for all the walkers.
+            all_inds (dict): Keys are ``branch_names``. Values are
+                np.ndarray[ntemps, nwalkers, nleaves_max]. These are the boolean
+                arrays marking which leaves are currently used within each walker.
+            nleaves_min_all (dict): Minimum values of leaf ount for each model. Must have same order as ``all_cords``. 
+            nleaves_max_all (dict): Maximum values of leaf ount for each model. Must have same order as ``all_cords``. 
+            random (object): Current random state of the sampler.
+            **kwargs (ignored): For modularity. 
+
+        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 new ``inds`` array with
+                boolean values flipped for added or removed sources. Third 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.
+
+        """
+        # prepare the output dictionaries
+        q = {}
+        new_inds = {}
+        all_inds_for_change = {}
+
+        # loop over the models included here
+        assert len(nleaves_min_all)
+        assert len(all_coords.keys()) == len(nleaves_max_all.keys())
+        for (name, inds) in all_inds.items():
+            nleaves_max = nleaves_max_all[name]
+            nleaves_min = nleaves_min_all[name]
+            if nleaves_min == nleaves_max:
+                continue
+            elif nleaves_min > nleaves_max:
+                raise ValueError("nleaves_min is greater than nleaves_max. Not allowed.")
+
+            # get the inds adjustment information
+            all_inds_for_change[name] = self.get_model_change_proposal(
+                inds, random, nleaves_min, nleaves_max
+            )
+
+        # loop through branches and propose new points from the prio
+        for i, (name, coords, inds) in enumerate(
+            zip(all_coords.keys(), all_coords.values(), all_inds.values(),)
+        ):
+            # put in base information
+            ntemps, nwalkers, nleaves_max, ndim = coords.shape
+            new_inds[name] = inds.copy()
+            q[name] = coords.copy()
+
+            if i == 0:
+                factors = np.zeros((ntemps, nwalkers))
+
+            # if not included
+            if name not in all_inds_for_change:
+                continue
+
+            # inds changing for this branch
+            inds_for_change = all_inds_for_change[name]
+
+            # adjust inds
+
+            # adjust deaths from True -> False
+            inds_here = tuple(inds_for_change["-1"].T)
+            new_inds[name][inds_here] = False
+
+            # factor is +log q()
+            current_generate_dist = self.generate_dist[name]
+            factors[inds_here[:2]] += +1 * current_generate_dist.logpdf(
+                q[name][inds_here]
+            )
+
+            # adjust births from False -> True
+            inds_here = tuple(inds_for_change["+1"].T)
+            new_inds[name][inds_here] = True
+
+            # add coordinates for new leaves
+            current_generate_dist = self.generate_dist[name]
+            inds_here = tuple(inds_for_change["+1"].T)
+            num_inds_change = len(inds_here[0])
+
+            q[name][inds_here] = current_generate_dist.rvs(size=num_inds_change)
+
+            # factor is -log q()
+            factors[inds_here[:2]] += -1 * current_generate_dist.logpdf(
+                q[name][inds_here]
+            )
+
+        return q, new_inds, factors

eryn/moves/gaussian.py

@@ -0,0 +1,190 @@
+# -*- coding: utf-8 -*-
+
+import numpy as np
+
+from .mh import MHMove
+
+__all__ = ["GaussianMove"]
+
+
+class GaussianMove(MHMove):
+    """A Metropolis step with a Gaussian proposal function.
+
+    This class is heavily based on the same class in ``emcee``.
+
+    Args:
+        cov (dict): The covariance of the proposal function. The keys are branch names and the
+            values are covariance information. This information can be provided as a scalar,
+            vector, or matrix and the proposal will be assumed isotropic,
+            axis-aligned, or general, respectively.
+        mode (str, optional): Select the method used for updating parameters. This
+            can be one of ``"vector"``, ``"random"``, or ``"sequential"``. The
+            ``"vector"`` mode updates all dimensions simultaneously,
+            ``"random"`` randomly selects a dimension and only updates that
+            one, and ``"sequential"`` loops over dimensions and updates each
+            one in turn. (default: ``"vector"``)
+        factor (float, optional): If provided the proposal will be made with a
+            standard deviation uniformly selected from the range
+            ``exp(U(-log(factor), log(factor))) * cov``. This is invalid for
+            the ``"vector"`` mode. (default: ``None``)
+        **kwargs (dict, optional): Kwargs for parent classes. (default: ``{}``)
+
+    Raises:
+        ValueError: If the proposal dimensions are invalid or if any of any of
+            the other arguments are inconsistent.
+
+    """
+
+    def __init__(self, cov_all, mode="vector", factor=None, **kwargs):
+        self.all_proposal = {}
+        for name, cov in cov_all.items():
+            # Parse the proposal type.
+            try:
+                float(cov)
+
+            except TypeError:
+                cov = np.atleast_1d(cov)
+                if len(cov.shape) == 1:
+                    # A diagonal proposal was given.
+                    ndim = len(cov)
+                    proposal = _diagonal_proposal(np.sqrt(cov), factor, mode)
+
+                elif len(cov.shape) == 2 and cov.shape[0] == cov.shape[1]:
+                    # The full, square covariance matrix was given.
+                    ndim = cov.shape[0]
+                    proposal = _proposal(cov, factor, mode)
+
+                else:
+                    raise ValueError("Invalid proposal scale dimensions")
+
+            else:
+                # This was a scalar proposal.
+                ndim = None
+                proposal = _isotropic_proposal(np.sqrt(cov), factor, mode)
+            self.all_proposal[name] = proposal
+
+        super(GaussianMove, self).__init__(**kwargs)
+
+    def get_proposal(self, branches_coords, random, branches_inds=None, **kwargs):
+        """Get proposal from Gaussian distribution
+
+        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)
+
+        """
+
+        # initialize ouput
+        q = {}
+        for name, coords in zip(branches_coords.keys(), branches_coords.values()):
+            ntemps, nwalkers, nleaves_max, ndim = coords.shape
+
+            # setup inds accordingly
+            if branches_inds is None:
+                inds = np.ones((ntemps, nwalkers, nleaves_max), dtype=bool)
+            else:
+                inds = branches_inds[name]
+
+            # get the proposal for this branch
+            proposal_fn = self.all_proposal[name]
+            inds_here = np.where(inds == True)
+
+            # copy coords
+            q[name] = coords.copy()
+
+            # get new points
+            new_coords, _ = proposal_fn(coords[inds_here], random)
+
+            # put into coords in proper location
+            q[name][inds_here] = new_coords.copy()
+
+        # handle periodic parameters
+        if self.periodic is not None:
+            q = self.periodic.wrap(
+                {
+                    name: tmp.reshape((ntemps * nwalkers,) + tmp.shape[-2:])
+                    for name, tmp in q.items()
+                },
+                xp=self.xp,
+            )
+
+            q = {
+                name: tmp.reshape(
+                    (
+                        ntemps,
+                        nwalkers,
+                    )
+                    + tmp.shape[-2:]
+                )
+                for name, tmp in q.items()
+            }
+
+        return q, np.zeros((ntemps, nwalkers))
+
+
+class _isotropic_proposal(object):
+    allowed_modes = ["vector", "random", "sequential"]
+
+    def __init__(self, scale, factor, mode):
+        self.index = 0
+        self.scale = scale
+        self.invscale = np.linalg.inv(np.linalg.cholesky(scale))
+        if factor is None:
+            self._log_factor = None
+        else:
+            if factor < 1.0:
+                raise ValueError("'factor' must be >= 1.0")
+            self._log_factor = np.log(factor)
+
+        if mode not in self.allowed_modes:
+            raise ValueError(
+                ("'{0}' is not a recognized mode. " "Please select from: {1}").format(
+                    mode, self.allowed_modes
+                )
+            )
+        self.mode = mode
+
+    def get_factor(self, rng):
+        if self._log_factor is None:
+            return 1.0
+        return np.exp(rng.uniform(-self._log_factor, self._log_factor))
+
+    def get_updated_vector(self, rng, x0):
+        return x0 + self.get_factor(rng) * self.scale * rng.randn(*(x0.shape))
+
+    def __call__(self, x0, rng):
+        nw, nd = x0.shape
+        xnew = self.get_updated_vector(rng, x0)
+        if self.mode == "random":
+            m = (range(nw), rng.randint(x0.shape[-1], size=nw))
+        elif self.mode == "sequential":
+            m = (range(nw), self.index % nd + np.zeros(nw, dtype=int))
+            self.index = (self.index + 1) % nd
+        else:
+            return xnew, np.zeros(nw)
+        x = np.array(x0)
+        x[m] = xnew[m]
+        return x, np.zeros(nw)
+
+
+class _diagonal_proposal(_isotropic_proposal):
+    def get_updated_vector(self, rng, x0):
+        return x0 + self.get_factor(rng) * self.scale * rng.randn(*(x0.shape))
+
+
+class _proposal(_isotropic_proposal):
+    allowed_modes = ["vector"]
+
+    def get_updated_vector(self, rng, x0):
+        return x0 + self.get_factor(rng) * rng.multivariate_normal(
+            np.zeros(len(self.scale)), self.scale, size=len(x0)
+        )

eryn/moves/group.py

@@ -0,0 +1,281 @@
+# -*- coding: utf-8 -*-
+from abc import ABC
+from copy import deepcopy
+import numpy as np
+import warnings
+
+from ..state import BranchSupplemental, State
+from .move import Move
+
+
+__all__ = ["GroupMove"]
+
+
+class GroupMove(Move, ABC):
+    """
+    A "group" ensemble move based on the :class:`eryn.moves.RedBlueMove`.
+
+    In moves like the :class:`eryn.moves.StretchMove`, the complimentary
+    group for which the proposal is used is chosen from the current points in
+    the ensemble. In "group" moves the complimentary group is a stationary group
+    that is updated every `n_iter_update` iterations. This update is performed with the
+    last set of coordinates to maintain detailed balance.
+
+    Args:
+        nfriends (int, optional): The number of friends to draw from as the complimentary
+            ensemble. This group is determined from the stationary group. If ``None``, it will
+            be set to the number of walkers. (default: ``None``)
+        n_iter_update (int, optional): Number of iterations to run before updating the
+            stationary distribution. (default: 100).
+        live_dangerously (bool, optional): If ``True``, allow for ``n_iter_update == 1``.
+            (deafault: ``False``)
+
+    ``kwargs`` are passed to :class:`Move` class.
+
+    """
+
+    def __init__(
+        self, nfriends=None, n_iter_update=100, live_dangerously=False, **kwargs
+    ):
+
+        Move.__init__(self, **kwargs)
+        self.nfriends = int(nfriends)
+        self.n_iter_update = n_iter_update
+
+        if self.n_iter_update <= 1 and not live_dangerously:
+            raise ValueError("n_iter_update must be greather than or equal to 2.")
+
+        self.iter = 0
+
+    def find_friends(self, name, s, s_inds=None, branch_supps=None):
+        """Function for finding friends.
+
+        Args:
+            name (str): Branch name for proposal coordinates.
+            s (np.ndarray): Coordinates array for the points to be moved.
+            s_inds (np.ndarray, optional): ``inds`` arrays that represent which leaves are present.
+                (default: ``None``)
+            branch_supps (dict, optional): Keys are ``branch_names`` and values are
+                :class:`BranchSupplemental` objects. For group proposals,
+                ``branch_supps`` are the best device for passing and tracking useful
+                information. (default: ``None``)
+
+        Return:
+            np.ndarray: Complimentary values.
+
+        """
+        raise NotImplementedError
+
+    def choose_c_vals(self, name, s, s_inds=None, branch_supps=None):
+        """Get the complimentary values."""
+        return self.find_friends(name, s, s_inds=s_inds, branch_supps=branch_supps)
+
+    def setup(self, branches):
+        """Any setup necessary for the proposal"""
+        pass
+
+    def setup_friends(self, branches):
+        """Setup anything for finding friends.
+
+        Args:
+            branches (dict): Dictionary with all the current branches in the sampler.
+
+        """
+        raise NotImplementedError
+
+    def fix_friends(self, branches):
+        """Fix any friends that were born through RJ.
+
+        This function is not required. If not implemented, it will just return immediately.
+
+        Args:
+            branches (dict): Dictionary with all the current branches in the sampler.
+
+        """
+        return
+
+    @classmethod
+    def get_proposal(self, s_all, random, gibbs_ndim=None, s_inds_all=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``)
+
+        Returns:
+            tuple: First entry is new positions. Second entry is detailed balance factors.
+
+        Raises:
+            ValueError: Issues with dimensionality.
+
+        """
+
+        raise NotImplementedError("The proposal must be implemented by " "subclasses")
+
+    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:
+            tuple: (state, accepted)
+                The first return is the state of the sampler after the move.
+                The second return value is the accepted count array.
+
+        """
+
+        # Check that the dimensions are compatible.
+        ndim_total = 0
+        for branch in state.branches.values():
+            ntemps, nwalkers, nleaves_, ndim_ = branch.shape
+            ndim_total += ndim_ * nleaves_
+
+        if self.nfriends is None:
+            self.nfriends = nwalkers
+
+        # Run any move-specific setup.
+        self.setup(state.branches)
+
+        if self.iter == 0 or self.iter % self.n_iter_update == 0:
+            self.setup_friends(state.branches)
+
+        if self.iter != 0 and self.iter % self.n_iter_update == 0:
+            # store old values to maintain detailed balance when updating
+            old_branches = deepcopy(state.branches)
+
+        # fix any friends that may have come through rj
+        if self.iter != 0 and self.iter % self.n_iter_update != 0:
+            self.fix_friends(state.branches)
+
+        # Split the ensemble in half and iterate over these two halves.
+        accepted = np.zeros((ntemps, nwalkers), dtype=bool)
+
+        all_branch_names = list(state.branches.keys())
+
+        # get gibbs sampling information
+        for branch_names_run, inds_run in self.gibbs_sampling_setup_iterator(
+            all_branch_names
+        ):
+
+            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 proposals based on Gibbs sampling
+            (
+                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 not at_least_one_proposal:
+                continue
+
+            # need to trick stretch proposal into using the dimenionality associated
+            # with Gibbs sampling if it is being used
+            gibbs_ndim = 0
+            for brn, ir in zip(branch_names_run, inds_run):
+                if ir is not None:
+                    gibbs_ndim += ir.sum()
+                else:
+                    # nleaves * ndim
+                    gibbs_ndim += np.prod(state.branches[brn].shape[-2:])
+
+            self.current_model = model
+            self.current_state = state
+            # Get the move-specific proposal.
+            q, factors = self.get_proposal(
+                coords_going_for_proposal,
+                model.random,
+                gibbs_ndim=gibbs_ndim,
+                s_inds_all=inds_going_for_proposal,
+                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
+            )
+
+            # Compute prior of the proposed position
+            # new_inds_prior is adjusted if product-space is used
+            logp = model.compute_log_prior_fn(q, inds=state.branches_inds)
+
+            self.fix_logp_gibbs(branch_names_run, inds_run, logp, state.branches_inds)
+
+            # 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,
+            )
+
+            # 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
+
+        if self.temperature_control is not None:
+            state = self.temperature_control.temper_comps(state)
+
+        if self.iter != 0 and self.iter % self.n_iter_update == 0:
+            # use old values to maintain detailed balance when updating
+            # nfriends
+            self.setup_friends(old_branches)
+
+        self.iter += 1
+        return state, accepted