eryn 1.2.0__py3-none-any.whl

eryn/moves/multipletry.py

@@ -0,0 +1,776 @@
+from multiprocessing.sharedctypes import Value
+import numpy as np
+import warnings
+from copy import deepcopy
+from abc import ABC
+
+# from scipy.special import logsumexp
+
+try:
+    import cupy as cp
+
+    gpu_available = True
+except (ModuleNotFoundError, ImportError):
+    import numpy as cp
+
+    gpu_available = False
+
+from .rj import ReversibleJumpMove
+from ..prior import ProbDistContainer
+from ..utils.utility import groups_from_inds
+
+___ = ["MultipleTryMove"]
+
+
+def logsumexp(a, axis=None, xp=None):
+    if xp is None:
+        xp = np
+
+    max = xp.max(a, axis=axis)
+    ds = a - max[:, None]
+
+    sum_of_exp = xp.exp(ds).sum(axis=axis)
+    return max + xp.log(sum_of_exp)
+
+
+def get_mt_computations(logP, log_proposal_pdf, symmetric=False, xp=None):
+
+    if xp is None:
+        xp = np
+
+        # set weights based on if symmetric
+    if symmetric:
+        log_importance_weights = logP
+    else:
+        log_importance_weights = logP - log_proposal_pdf
+
+    # get the sum of weights
+    log_sum_weights = logsumexp(log_importance_weights, axis=-1, xp=xp)
+
+    # probs = wi / sum(wi)
+    log_of_probs = log_importance_weights - log_sum_weights[:, None]
+
+    # probabilities to choose try
+    probs = xp.exp(log_of_probs)
+
+    # draw based on likelihood
+    inds_keep = (probs.cumsum(1) > xp.random.rand(probs.shape[0])[:, None]).argmax(1)
+
+    return log_importance_weights, log_sum_weights, inds_keep
+
+
+class MultipleTryMove(ABC):
+    """Generate multiple proposal tries.
+
+    This class should be inherited by another proposal class
+    with the ``@classmethods`` overwritten. See :class:`eryn.moves.MTDistGenMove`
+    and :class:`MTDistGenRJ` for examples.
+
+    Args:
+        num_try (int, optional): Number of tries. (default: 1)
+        independent (bool, optional): Set to ``True`` if the proposal is independent of the current points.
+            (default: ``False``).
+        symmetric (bool, optional): Set to ``True`` if the proposal is symmetric.
+            (default: ``False``).
+        rj (bool, optional): Set to ``True`` if this is a nested reversible jump proposal.
+            (default: ``False``).
+        **kwargs (dict, optional): for compatibility with other proposals.
+
+        Raises:
+            ValueError: Input issues.
+
+    """
+
+    def __init__(
+        self,
+        num_try=1,
+        independent=False,
+        symmetric=False,
+        rj=False,
+        use_gpu=None,
+        **kwargs
+    ):
+        self.num_try = num_try
+
+        self.independent = independent
+        self.symmetric = symmetric
+        self.rj = rj
+
+        if self.rj:
+            if self.symmetric or self.independent:
+                raise ValueError(
+                    "If rj==True, symmetric and independt must both be False."
+                )
+
+        self.use_gpu = use_gpu
+
+    @property
+    def xp(self):
+        xp = cp if self.use_gpu else np
+        return xp
+
+    @classmethod
+    def special_like_func(self, generated_coords, *args, inds_leaves_rj=None, **kwargs):
+        """Calculate the Likelihood for sampled points.
+
+        Args:
+            generated_coords (np.ndarray): Generated coordinates with shape ``(number of independent walkers, num_try)``.
+            *args (tuple, optional): additional arguments passed by overwriting the
+                ``get_proposal`` function and passing ``args_like`` keyword argument.
+            inds_leaves_rj (np.ndarray): Index into each individual walker giving the
+                leaf index associated with this proposal. Should only be used if ``self.rj is True``. (default: ``None``)
+            **kwargs (tuple, optional): additional keyword arguments passed by overwriting the
+                ``get_proposal`` function and passing ``kwargs_like`` keyword argument.
+
+        Returns:
+            np.ndarray: Likelihood values with shape ``(generated_coords.shape[0], num_try).``
+
+        Raises:
+            NotImplementedError: Function not included.
+
+        """
+        raise NotImplementedError
+
+    @classmethod
+    def special_prior_func(self, generated_coords, *args, **kwargs):
+        """Calculate the Prior for sampled points.
+
+        Args:
+            generated_coords (np.ndarray): Generated coordinates with shape ``(number of independent walkers, num_try)``.
+            *args (tuple, optional): additional arguments passed by overwriting the
+                ``get_proposal`` function and passing ``args_prior`` keyword argument.
+            inds_leaves_rj (np.ndarray): Index into each individual walker giving the
+                leaf index associated with this proposal. Should only be used if ``self.rj is True``. (default: ``None``)
+            **kwargs (tuple, optional): additional keyword arguments passed by overwriting the
+                ``get_proposal`` function and passing ``kwargs_prior`` keyword argument.
+
+        Returns:
+            np.ndarray: Prior values with shape ``(generated_coords.shape[0], num_try).``
+
+        Raises:
+            NotImplementedError: Function not included.
+
+        """
+        raise NotImplementedError
+
+    @classmethod
+    def special_generate_func(
+        coords, random, size=1, *args, fill_tuple=None, fill_values=None, **kwargs
+    ):
+        """Generate samples and calculate the logpdf of their proposal function.
+
+        Args:
+            coords (np.ndarray): Current coordinates of walkers.
+            random (obj): Random generator.
+            *args (tuple, optional): additional arguments passed by overwriting the
+                ``get_proposal`` function and passing ``args_generate`` keyword argument.
+            size (int, optional): Number of tries to generate.
+            fill_tuple (tuple, optional): Length 2 tuple with the indexing of which values to fill
+                when generating. Can be used for auxillary proposals or reverse RJ proposals. First index is the index into walkers and the second index is
+                the index into the number of tries. (default: ``None``)
+            fill_values (np.ndarray): values to fill associated with ``fill_tuple``. Should
+                have size ``(len(fill_tuple[0]), ndim)``. (default: ``None``).
+            **kwargs (tuple, optional): additional keyword arguments passed by overwriting the
+                ``get_proposal`` function and passing ``kwargs_generate`` keyword argument.
+
+        Returns:
+            tuple: (generated points, logpdf of generated points).
+
+        Raises:
+            NotImplementedError: Function not included.
+
+        """
+        raise NotImplementedError
+
+    @classmethod
+    def special_generate_logpdf(self, coords):
+        """Get logpdf of generated coordinates.
+
+        Args:
+            coords (np.ndarray): Current coordinates of walkers.
+
+        Returns:
+            np.ndarray: logpdf of generated points.
+
+        Raises:
+            NotImplementedError: Function not included.
+        """
+        raise NotImplementedError
+
+    def get_mt_log_posterior(self, ll, lp, betas=None):
+        """Calculate the log of the posterior for all tries.
+
+        Args:
+            ll (np.ndarray): Log Likelihood values with shape ``(nwalkers, num_tries)``.
+            lp (np.ndarray): Log Prior values with shape ``(nwalkers, num_tries)``.
+            betas (np.ndarray, optional): Inverse temperatures to include in log Posterior computation.
+                (default: ``None``)
+
+        Returns:
+            np.ndarray: Log of the Posterior with shape ``(nwalkers, num_tries)``.
+
+        """
+        if betas is None:
+            ll_temp = ll.copy()
+        else:
+            assert isinstance(betas, self.xp.ndarray)
+            if ll.ndim > 1:
+                betas_tmp = self.xp.expand_dims(betas, ll.ndim - 1)
+            else:
+                betas_tmp = betas
+            ll_temp = betas_tmp * ll
+
+        return ll_temp + lp
+
+    def readout_adjustment(self, out_vals, all_vals_prop, aux_all_vals):
+        """Read out values from the proposal.
+
+        Allows the user to read out any values from the proposal that may be needed elsewhere. This function must be overwritten.
+
+        Args:
+            out_vals (list): ``[logP_out, ll_out, lp_out, log_proposal_pdf_out, log_sum_weights]``.
+            all_vals_prop (list): ``[logP, ll, lp, log_proposal_pdf, log_sum_weights]``.
+            aux_all_vals (list): ``[aux_logP, aux_ll, aux_lp, aux_log_proposal_pdf, aux_log_sum_weights]``.
+
+        """
+        pass
+
+    def get_mt_proposal(
+        self,
+        coords,
+        random,
+        args_generate=(),
+        kwargs_generate={},
+        args_like=(),
+        kwargs_like={},
+        args_prior=(),
+        kwargs_prior={},
+        betas=None,
+        ll_in=None,
+        lp_in=None,
+        inds_leaves_rj=None,
+        inds_reverse_rj=None,
+    ):
+        """Make a multiple-try proposal
+
+        Here, ``nwalkers`` refers to all independent walkers which generally
+        will mean ``nwalkers * ntemps`` in terms of the rest of the sampler.
+
+        Args:
+            coords (np.ndarray): Current coordinates of walkers.
+            random (obj): Random generator.
+            args_generate (tuple, optional): Additional ``*args`` to pass to generate function.
+                Must overwrite ``get_proposal`` function to use these.
+                (default: ``()``)
+            kwargs_generate (dict, optional): Additional ``**kwargs`` to pass to generate function.
+                (default: ``{}``)
+                Must overwrite ``get_proposal`` function to use these.
+            args_like (tuple, optional): Additional ``*args`` to pass to Likelihood function.
+                Must overwrite ``get_proposal`` function to use these.
+                (default: ``()``)
+            kwargs_like (dict, optional): Additional ``**kwargs`` to pass to Likelihood function.
+                Must overwrite ``get_proposal`` function to use these.
+                (default: ``{}``)
+            args_prior (tuple, optional): Additional ``*args`` to pass to Prior function.
+                Must overwrite ``get_proposal`` function to use these.
+                (default: ``()``)
+            kwargs_prior (dict, optional): Additional ``**kwargs`` to pass to Prior function.
+                Must overwrite ``get_proposal`` function to use these.
+                (default: ``{}``)
+            betas (np.ndarray, optional): Inverse temperatures passes to the proposal with shape ``(nwalkers,)``.
+            ll_in (np.ndarray, optional): Log Likelihood values coming in for current coordinates. Must be provided
+                if ``self.rj is True``. If ``self.rj is True``, must be nested.
+                Also, for all proposed removals, this value must be the Likelihood with the binary
+                removed so all proposals are pretending to add a binary.
+                Useful if ``self.independent is True``. (default: ``None``)
+            lp_in (np.ndarray, optional): Log Prior values coming in for current coordinates. Must be provided
+                if ``self.rj is True``. If ``self.rj is True``, must be nested.
+                Also, for all proposed removals, this value must be the Likelihood with the binary
+                removed so all proposals are pretending to add a binary.
+                Useful if ``self.independent is True``. (default: ``None``)
+            inds_leaves_rj (np.ndarray, optional): Array giving the leaf index of each incoming walker.
+                Must be provided if ``self.rj is True``. (default: ``None``)
+            inds_reverse_rj (np.ndarray, optional): Array giving the walker index for which proposals are
+                reverse proposal removing a leaf.
+                Must be provided if ``self.rj is True``. (default: ``None``)
+
+        Returns:
+            tuple: (generated points, factors).
+
+        Raises:
+            ValueError: Inputs are incorrect.
+
+        """
+
+        # check if rj and make sure we have all the information in that case
+        if self.rj:
+            try:
+                assert ll_in is not None and lp_in is not None
+                assert inds_leaves_rj is not None and inds_reverse_rj is not None
+            except AssertionError:
+                raise ValueError(
+                    "If using rj, must provide ll_in, lp_in, inds_leaves_rj, and inds_reverse_rj."
+                )
+
+            # if using reversible jump, fill first spot with values that are proposed to remove
+            fill_tuple = (inds_reverse_rj, np.zeros_like(inds_reverse_rj))
+            fill_values = coords[inds_reverse_rj]
+        else:
+            fill_tuple = None
+            fill_values = None
+
+        # generate new points and get log of the proposal probability
+        generated_points, log_proposal_pdf = self.special_generate_func(
+            coords,
+            random,
+            *args_generate,
+            size=self.num_try,
+            fill_values=fill_values,
+            fill_tuple=fill_tuple,
+            **kwargs_generate
+        )
+
+        # compute the Likelihood functions
+        ll = self.special_like_func(
+            generated_points, *args_like, inds_leaves_rj=inds_leaves_rj, **kwargs_like
+        )
+
+        # check for nans
+        if self.xp.any(self.xp.isnan(ll)):
+            warnings.warn("Getting nans for ll in multiple try.")
+            ll[self.xp.isnan(ll)] = -1e300
+
+        # compute the Prior functions
+        lp = self.special_prior_func(
+            generated_points, *args_prior, inds_leaves_rj=inds_leaves_rj, **kwargs_prior
+        )
+
+        # if rj, make proposal distribution for all other leaves the prior value
+        # this will properly cancel the prior with the proposal for leaves that already exists
+        if self.rj:
+            log_proposal_pdf += lp_in[:, None]
+
+        # get posterior distribution including tempering
+        logP = self.get_mt_log_posterior(ll, lp, betas=betas)
+
+        log_importance_weights, log_sum_weights, inds_keep = get_mt_computations(
+            logP, log_proposal_pdf, symmetric=self.symmetric, xp=self.xp
+        )
+
+        # tuple of index arrays of which try chosen per walker
+        inds_tuple = (self.xp.arange(len(inds_keep)), inds_keep)
+
+        if self.rj:
+            # this just ensures the cancellation of logP and aux_logP outside of proposal
+            inds_tuple[1][inds_reverse_rj] = 0
+
+        # get chosen prior, Likelihood, posterior information
+        lp_out = lp[inds_tuple]
+        ll_out = ll[inds_tuple]
+        logP_out = logP[inds_tuple]
+
+        # store this information for access outside of multiple try part
+        self.mt_lp = lp_out
+        self.mt_ll = ll_out
+
+        # choose points and get the log of the proposal for storage
+        generated_points_out = generated_points[inds_tuple].copy()  # theta^j
+        log_proposal_pdf_out = log_proposal_pdf[inds_tuple]
+
+        # prepare auxillary information based on if it is nested rj, independent, or not
+        if self.independent:
+            # if independent, all the tries can be repeated for the auxillary draws
+            aux_ll = ll.copy()
+            aux_lp = lp.copy()
+
+            # sub in the generation pdf for the current coordinates
+            aux_log_proposal_pdf_sub = self.special_generate_logpdf(coords)
+
+            # set sub ll based on if it is provided
+            if ll_in is None:
+                aux_ll_sub = self.special_generate_like(coords)
+
+            else:
+                assert ll_in.shape[0] == coords.shape[0]
+                aux_ll_sub = ll_in
+
+            # set sub lp based on if it is provided
+            if lp_in is None:
+                aux_lp_sub = self.special_generate_prior(coords)
+
+            else:
+                assert lp_in.shape[0] == coords.shape[0]
+                aux_lp_sub = lp_in
+
+            # sub in this information from the current coordinates
+            aux_ll[inds_tuple] = aux_ll_sub
+            aux_lp[inds_tuple] = aux_lp_sub
+
+            # get auxillary posterior
+            aux_logP = self.get_mt_log_posterior(aux_ll, aux_lp, betas=betas)
+
+            # get aux_log_proposal_pdf information
+            aux_log_proposal_pdf = log_proposal_pdf.copy()
+            aux_log_proposal_pdf[inds_tuple] = aux_log_proposal_pdf_sub
+
+            # set auxillary weights
+            aux_log_importance_weights = aux_logP - aux_log_proposal_pdf
+
+        elif self.rj:
+            # in rj, set aux_ll and aux_lp to be repeats of the model with one less leaf
+            aux_ll = np.repeat(ll_in[:, None], self.num_try, axis=-1)
+            aux_lp = np.repeat(lp_in[:, None], self.num_try, axis=-1)
+
+            # probability is the prior for the existing points
+            aux_log_proposal_pdf = aux_lp.copy()
+
+            # get log posterior
+            aux_logP = self.get_mt_log_posterior(aux_ll, aux_lp, betas=betas)
+
+            # get importance weights
+            aux_log_importance_weights = aux_logP - aux_log_proposal_pdf
+
+        else:
+            # generate auxillary points based on chosen new points
+            aux_generated_points, aux_log_proposal_pdf = self.special_generate_func(
+                generated_points_out,
+                random,
+                *args_generate,
+                size=self.num_try,
+                fill_tuple=inds_tuple,
+                fill_values=generated_points_out,
+                **kwargs_generate
+            )
+
+            # get ll, lp, and lP
+            aux_ll = self.special_like_func(
+                aux_generated_points, *args_like, **kwargs_like
+            )
+
+            aux_lp = self.special_prior_func(aux_generated_points)
+
+            aux_logP = self.get_mt_log_posterior(aux_ll, aux_lp, betas=betas)
+
+            # set auxillary weights
+            if not self.symmetric:
+                aux_log_importance_weights = aux_logP - aux_log_proposal_pdf_sub
+            else:
+                aux_log_importance_weights = aux_logP
+
+        # chosen output old Posteriors
+        aux_logP_out = aux_logP[inds_tuple]
+        # get sum of log weights
+        aux_log_sum_weights = logsumexp(aux_log_importance_weights, axis=-1, xp=self.xp)
+
+        aux_log_proposal_pdf_out = aux_log_proposal_pdf[inds_tuple]
+        # this is setup to make clear with the math.
+        # setting up factors properly means the
+        # final lnpdiff will be effectively be the ratio of the sums
+        # of the weights
+
+        # IMPORTANT: logP_out must be subtracted against log_sum_weights before anything else due to -1e300s.
+        factors = (
+            (aux_logP_out - aux_log_sum_weights)
+            - aux_log_proposal_pdf_out
+            + aux_log_proposal_pdf_out
+        ) - ((logP_out - log_sum_weights) - log_proposal_pdf_out + log_proposal_pdf_out)
+
+        if self.rj:
+            # adjust all information for reverese rj proposals
+            factors[inds_reverse_rj] *= -1
+            self.mt_ll[inds_reverse_rj] = ll_in[inds_reverse_rj]
+            self.mt_lp[inds_reverse_rj] = lp_in[inds_reverse_rj]
+
+        # store output information
+        self.aux_logP_out = aux_logP_out
+        self.logP_out = logP_out
+        self.aux_ll = aux_ll
+        self.aux_lp = aux_lp
+
+        self.log_sum_weights = log_sum_weights
+        self.aux_log_sum_weights = aux_log_sum_weights
+
+        if self.rj:
+            self.inds_reverse_rj = inds_reverse_rj
+            self.inds_forward_rj = np.delete(
+                np.arange(coords.shape[0]), inds_reverse_rj
+            )
+
+        # prepare to readout any information the user would like in readout_adjustment
+        out_vals = [logP_out, ll_out, lp_out, log_proposal_pdf_out, log_sum_weights]
+        all_vals_prop = [logP, ll, lp, log_proposal_pdf, log_sum_weights]
+        aux_all_vals = [
+            aux_logP,
+            aux_ll,
+            aux_lp,
+            aux_log_proposal_pdf,
+            aux_log_sum_weights,
+        ]
+        self.readout_adjustment(out_vals, all_vals_prop, aux_all_vals)
+
+        return (
+            generated_points_out,
+            factors,
+        )
+
+    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:
+            ValueError: Input issues.
+
+        """
+
+        # mutliple try is only made for one branch here
+        if len(list(branches_coords.keys())) > 1:
+            raise ValueError("Can only propose change to one model at a time with MT.")
+
+        # get main key
+        key_in = list(branches_coords.keys())[0]
+        self.key_in = key_in
+
+        # get inds information
+        if branches_inds is None:
+            branches_inds = {}
+            branches_inds[key_in] = np.ones(
+                branches_coords[key_in].shape[:-1], dtype=bool
+            )
+
+        # Make sure for base proposals that there is only one leaf
+        if np.any(branches_inds[key_in].sum(axis=-1) > 1):
+            raise ValueError
+
+        ntemps, nwalkers, _, _ = branches_coords[key_in].shape
+
+        # get temperature information
+        betas_here = np.repeat(
+            self.temperature_control.betas[:, None],
+            np.prod(branches_coords[key_in].shape[1:-1]),
+        ).reshape(branches_inds[key_in].shape)[branches_inds[key_in]]
+
+        # previous Likelihoods in case proposal is independent
+        ll_here = np.repeat(
+            self.current_state.log_like[:, :, None],
+            branches_coords[key_in].shape[2],
+            axis=-1,
+        ).reshape(branches_inds[key_in].shape)[branches_inds[key_in]]
+
+        # previous Priors in case proposal is independent
+        lp_here = np.repeat(
+            self.current_state.log_prior[:, :, None],
+            branches_coords[key_in].shape[2],
+            axis=-1,
+        ).reshape(branches_inds[key_in].shape)[branches_inds[key_in]]
+
+        # get mt proposal
+        generated_points, factors = self.get_mt_proposal(
+            branches_coords[key_in][branches_inds[key_in]],
+            random,
+            betas=betas_here,
+            ll_in=ll_here,
+            lp_in=lp_here,
+        )
+
+        # store this information for access outside
+        self.mt_ll = self.mt_ll.reshape(ntemps, nwalkers)
+        self.mt_lp = self.mt_lp.reshape(ntemps, nwalkers)
+
+        return (
+            {key_in: generated_points.reshape(ntemps, nwalkers, 1, -1)},
+            factors.reshape(ntemps, nwalkers),
+        )
+
+
+class MultipleTryMoveRJ(MultipleTryMove):
+    def get_proposal(
+        self,
+        branches_coords,
+        branches_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 (list): Minimum values of leaf ount for each model. Must have same order as ``all_cords``.
+            nleaves_max_all (list): 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.
+
+        """
+
+        if len(list(branches_coords.keys())) > 1:
+            raise ValueError("Can only propose change to one model at a time with MT.")
+
+        # get main key
+        key_in = list(branches_coords.keys())[0]
+        self.key_in = key_in
+
+        if branches_inds is None:
+            raise ValueError("In MT RJ proposal, branches_inds cannot be None.")
+
+        ntemps, nwalkers, nleaves_max, ndim = branches_coords[key_in].shape
+
+        # get temperature information
+        betas_here = np.repeat(
+            self.temperature_control.betas[:, None], nwalkers, axis=-1
+        ).flatten()
+
+        # current Likelihood and prior information
+        ll_here = self.current_state.log_like.flatten()
+        lp_here = self.current_state.log_prior.flatten()
+
+        # do rj setup
+        assert len(nleaves_min_all) == 1 and len(nleaves_max_all) == 1
+        nleaves_min = nleaves_min_all[key_in]
+        nleaves_max = nleaves_max_all[key_in]
+
+        if nleaves_min == nleaves_max:
+            raise ValueError("MT RJ proposal requires that nleaves_min != nleaves_max.")
+        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 = self.get_model_change_proposal(
+            branches_inds[key_in], random, nleaves_min, nleaves_max
+        )
+
+        # preparing leaf information for going into the proposal
+        inds_leaves_rj = np.zeros(ntemps * nwalkers, dtype=int)
+        coords_in = np.zeros((ntemps * nwalkers, ndim))
+        inds_reverse_rj = None
+
+        # prepare proposal dictionaries
+        new_inds = deepcopy(branches_inds)
+        q = deepcopy(branches_coords)
+        for change in all_inds_for_change.keys():
+            if change not in ["+1", "-1"]:
+                raise ValueError("MT RJ is only implemented for +1/-1 moves.")
+
+            # get indicies of changing leaves
+            temp_inds = all_inds_for_change[change][:, 0]
+            walker_inds = all_inds_for_change[change][:, 1]
+            leaf_inds = all_inds_for_change[change][:, 2]
+
+            # leaf index to change
+            inds_leaves_rj[temp_inds * nwalkers + walker_inds] = leaf_inds
+            coords_in[temp_inds * nwalkers + walker_inds] = branches_coords[key_in][
+                (temp_inds, walker_inds, leaf_inds)
+            ]
+
+            # adjustment of indices
+            new_val = {"+1": True, "-1": False}[change]
+
+            # adjust indices
+            new_inds[key_in][(temp_inds, walker_inds, leaf_inds)] = new_val
+
+            if change == "-1":
+                # which walkers are removing
+                inds_reverse_rj = temp_inds * nwalkers + walker_inds
+
+        if inds_reverse_rj is not None:
+            # setup reversal coords and inds
+            # need to determine Likelihood and prior of removed binaries.
+            # this goes into the multiple try proposal as previous ll and lp
+            temp_reverse_coords = {}
+            temp_reverse_inds = {}
+
+            for key in self.current_state.branches:
+                (
+                    ntemps_tmp,
+                    nwalkers_tmp,
+                    nleaves_max_tmp,
+                    ndim_tmp,
+                ) = self.current_state.branches[key].shape
+
+                # coords from reversal
+                temp_reverse_coords[key] = self.current_state.branches[
+                    key
+                ].coords.reshape(ntemps_tmp * nwalkers_tmp, nleaves_max_tmp, ndim_tmp)[
+                    inds_reverse_rj
+                ][
+                    None, :
+                ]
+
+                # which inds array to use
+                inds_tmp_here = (
+                    new_inds[key]
+                    if key == key_in
+                    else self.current_state.branches[key].inds
+                )
+                temp_reverse_inds[key] = inds_tmp_here.reshape(
+                    ntemps * nwalkers, nleaves_max_tmp
+                )[inds_reverse_rj][None, :]
+
+            # calculate information for the reverse
+            lp_reverse_here = self.current_model.compute_log_prior_fn(
+                temp_reverse_coords, inds=temp_reverse_inds
+            )[0]
+            ll_reverse_here = self.current_model.compute_log_like_fn(
+                temp_reverse_coords, inds=temp_reverse_inds, logp=lp_here
+            )[0]
+
+            # fill the here values
+            ll_here[inds_reverse_rj] = ll_reverse_here
+            lp_here[inds_reverse_rj] = lp_reverse_here
+
+        # get mt proposal
+        generated_points, factors = self.get_mt_proposal(
+            coords_in,
+            random,
+            betas=betas_here,
+            ll_in=ll_here,
+            lp_in=lp_here,
+            inds_leaves_rj=inds_leaves_rj,
+            inds_reverse_rj=inds_reverse_rj,
+        )
+
+        # for reading outside
+        self.mt_ll = self.mt_ll.reshape(ntemps, nwalkers)
+        self.mt_lp = self.mt_lp.reshape(ntemps, nwalkers)
+
+        # which walkers have information added
+        inds_forward_rj = np.delete(np.arange(coords_in.shape[0]), inds_reverse_rj)
+
+        # updated the coordinates
+        temp_inds = all_inds_for_change["+1"][:, 0]
+        walker_inds = all_inds_for_change["+1"][:, 1]
+        leaf_inds = all_inds_for_change["+1"][:, 2]
+        q[key_in][(temp_inds, walker_inds, leaf_inds)] = generated_points[
+            inds_forward_rj
+        ]
+
+        return q, new_inds, factors.reshape(ntemps, nwalkers)