eryn 1.2.0__py3-none-any.whl

eryn/utils/__init__.py

@@ -0,0 +1,10 @@
+# -*- coding: utf-8 -*-
+
+# from .plot import PlotContainer
+from .utility import *
+from .periodic import *
+from .transform import *
+from .updates import *
+from .stopping import *
+
+__all__ = ["PlotContainer", "PeriodicContainer", "TransformContainer"]

eryn/utils/periodic.py

@@ -0,0 +1,134 @@
+from ast import Import
+import numpy as np
+
+try:
+    import cupy as xp
+
+except (ModuleNotFoundError, ImportError) as e:
+    pass
+
+
+class PeriodicContainer:
+    """Perform operations for periodic parameters
+
+    Args:
+        periodic_in (dict): Keys are ``branch_names``. Values are
+            dictionaries. These dictionaries have keys as the parameter
+            indexes and values their associated period.
+
+    """
+
+    def __init__(self, periodic):
+
+        # store all the information
+        self.periodic = periodic
+        self.inds_periodic = {
+            key: np.asarray([i for i in periodic[key].keys()]) for key in periodic
+        }
+        self.periods = {
+            key: np.asarray([i for i in periodic[key].values()]) for key in periodic
+        }
+
+    def distance(self, p1, p2, xp=None):
+        """Move from p1 to p2 with periodic distance control
+
+        Args:
+            p1 (dict): If dict, keys are ``branch_names``
+                and values are positions with parameters along the final dimension.
+            p2 (dict): If dict, keys are ``branch_names``
+                and values are positions with parameters along the final dimension.
+            xp (object, optional): ``numpy`` or ``cupy``. If ``None``, use ``numpy``.
+                (default: ``None``)
+
+        Returns:
+            dict: Distances accounting for periodicity.
+                    Keys are branch names and values are distance arrays.
+
+        """
+
+        # cupy or numpy
+        if xp is None:
+            xp = np
+
+        # make sure both have same branches
+        assert list(p1.keys()) == list(p2.keys())
+
+        names = list(p1.keys())
+
+        # prepare output
+        out_diff = {}
+        for key in names:
+
+            # get basic distance
+            diff = p2[key] - p1[key]
+
+            # no periodic parameters for this key
+            if key not in self.periods:
+                out_diff[key] = diff
+                continue
+
+            # get period info
+            periods = xp.asarray(self.periods[key])
+            inds_periodic = xp.asarray(self.inds_periodic[key])
+
+            if len(self.periods[key]) > 0:
+                # get specific periodic parameterss
+                diff_periodic = diff[:, :, inds_periodic]
+
+                # fix when the distance is over 1/2 period away
+                inds_fix = (
+                    xp.abs(diff_periodic) > periods[xp.newaxis, xp.newaxis, :] / 2.0
+                )
+
+                # wrap back to make proper periodic distance
+                new_s = -(
+                    periods[xp.newaxis, xp.newaxis, :] - p1[key][:, :, inds_periodic]
+                ) * (diff_periodic < 0.0) + (
+                    periods[xp.newaxis, xp.newaxis, :] + p1[key][:, :, inds_periodic]
+                ) * (
+                    diff_periodic >= 0.0
+                )
+
+                # fill new information
+                diff_periodic[inds_fix] = (
+                    p2[key][:, :, inds_periodic][inds_fix] - new_s[inds_fix]
+                )
+                diff[:, :, inds_periodic] = diff_periodic
+
+            out_diff[key] = diff
+
+        return out_diff
+
+    def wrap(self, p, xp=None):
+        """Wrap p with periodic distance control
+
+        Args:
+            p (dict): If dict, keys are ``branch_names``
+                and values are positions with parameters along the final dimension.
+            xp (object, optional): ``numpy`` or ``cupy``. If ``None``, use ``numpy``.
+                (default: ``None``)
+
+        """
+
+        # cupy or numpy
+        if xp is None:
+            xp = np
+
+        names = list(p.keys())
+        # wrap for each branch
+        for key in names:
+            pos = p[key]
+
+            if len(self.periods[key]) > 0:
+                # get periodic information
+                periods = xp.asarray(self.periods[key])
+                inds_periodic = xp.asarray(self.inds_periodic[key])
+                # wrap
+                pos[:, :, inds_periodic] = (
+                    pos[:, :, inds_periodic] % periods[xp.newaxis, xp.newaxis, :]
+                )
+
+            # fill new info
+            p[key] = pos
+
+        return p

eryn/utils/stopping.py

@@ -0,0 +1,164 @@
+# -*- coding: utf-8 -*-
+
+from abc import ABC
+
+import numpy as np
+
+
+class Stopping(ABC, object):
+    """Base class for stopping.
+    
+    Stopping checks are only performed every ``thin_by`` iterations.
+    
+    """
+
+    @classmethod
+    def __call__(self, iter, last_sample, sampler):
+        """Call update function.
+
+        Args:
+            iter (int): Iteration of the sampler.
+            last_sample (obj): Last state of sampler (:class:`eryn.state.State`).
+            sampler (obj): Full sampler oject (:class:`eryn.ensemble.EnsembleSampler`).
+
+        Returns:
+            bool: Value of ``stop``. If ``True``, stop sampling.
+            
+        """
+        raise NotImplementedError
+
+
+class SearchConvergeStopping(Stopping):
+    """Stopping function based on a convergence to a maximunm Likelihood.
+
+    Stopping checks are only performed every ``thin_by`` iterations.
+    Therefore, the iterations of stopping checks are really every
+    ``sampler iterations * thin_by``.  
+
+    All arguments are stored as attributes.
+
+    Args:
+        n_iters (int, optional): Number of iterative stopping checks that need to pass
+            in order to stop the sampler. (default: ``30``)
+        diff (float, optional): Change in the Likelihood needed to fail the stopping check. In other words,
+            if the new maximum Likelihood is more than ``diff`` greater than the old, all iterative checks 
+            reset. (default: 0.1). 
+        start_iteration (int, optional): Iteration of sampler to start checking to stop. (default: 0)
+        verbose (bool, optional): If ``True``, print information. (default: ``False``)
+
+    Attributes:
+        iters_consecutive (int): Number of consecutive passes of the stopping check.
+        past_like_best (float): Previous best Likelihood. The initial value is ``-np.inf``.
+    
+    """
+
+    def __init__(self, n_iters=30, diff=0.1, start_iteration=0, verbose=False):
+
+        # store all the relevant information
+        self.n_iters = n_iters
+
+        self.diff = diff
+        self.verbose = verbose
+        self.start_iteration = start_iteration
+
+        # initialize important info
+        self.iters_consecutive = 0
+        self.past_like_best = -np.inf
+
+    def __call__(self, iter, sample, sampler):
+        """Call update function.
+
+        Args:
+            iter (int): Iteration of the sampler.
+            last_sample (obj): Last state of sampler (:class:`eryn.state.State`).
+            sampler (obj): Full sampler oject (:class:`eryn.ensemble.EnsembleSampler`).
+
+        Returns:
+            bool: Value of ``stop``. If ``True``, stop sampling.
+            
+        """
+
+        # if we have not reached the start iteration return
+        if iter < self.start_iteration:
+            return False
+
+        # get best Likelihood so far
+        like_best = sampler.get_log_like(discard=self.start_iteration).max()
+
+        # compare to last
+        # if it is less than diff change it passes
+        if np.abs(like_best - self.past_like_best) < self.diff:
+            self.iters_consecutive += 1
+
+        else:
+            # if it fails reset iters consecutive
+            self.iters_consecutive = 0
+
+            # store new best
+            self.past_like_best = like_best
+
+        # print information
+        if self.verbose:
+            print(
+                f"\nITERS CONSECUTIVE: {self.iters_consecutive}",
+                f"Previous best LL: {self.past_like_best}",
+                f"Current best LL: {like_best}\n",
+            )
+
+        if self.iters_consecutive >= self.n_iters:
+            # if we have passes the number of iters necessary, return True and reset
+            self.iters_consecutive = 0
+            return True
+
+        else:
+            return False
+
+
+"""
+class AutoCorrelationStop(Stopping):
+    # TODO: check and doc this
+    def __init__(self, autocorr_multiplier=50, verbose=False):
+        self.autocorr_multiplier = autocorr_multiplier
+        self.verbose = verbose
+
+        self.time = 0
+
+    def __call__(self, iter, last_sample, sampler):
+
+        tau = sampler.backend.get_autocorr_time(multiply_thin=False)
+
+        if self.time > 0:
+            # backend iteration
+            iteration = sampler.backend.iteration
+
+            finish = []
+
+            for name, values in tau.items():
+                converged = np.all(tau[name] * self.autocorr_multiplier < iteration)
+                converged &= np.all(
+                    np.abs(self.old_tau[name] - tau[name]) / tau[name] < 0.01
+                )
+
+                finish.append(converged)
+
+            stop = True if np.all(finish) else False
+            if self.verbose:
+                print(
+                    "\ntau:",
+                    tau,
+                    "\nIteration:",
+                    iteration,
+                    "\nAutocorrelation multiplier:",
+                    self.autocorr_multiplier,
+                    "\nStopping:",
+                    stop,
+                    "\n",
+                )
+
+        else:
+            stop = False
+
+        self.old_tau = tau
+        self.time += 1
+        return stop
+"""

eryn/utils/transform.py

@@ -0,0 +1,226 @@
+try:
+    import cupy as xp
+
+except (ModuleNotFoundError, ImportError) as e:
+    pass
+
+import numpy as np
+
+
+class TransformContainer:
+    """Container for helpful transformations
+
+    Args:
+        parameter_transforms (dict, optional): Keys are ``int`` or ``tuple``
+            of ``int`` that contain the indexes into the parameters
+            that correspond to the transformation added as the Values to the
+            dict. If using ``fill_values``, you must be careful with
+            making sure parameter transforms properly comes before or after
+            filling values. ``int`` indicate single parameter transforms. These
+            are performed first. ``tuple`` of ``int`` indicates multiple
+            parameter transforms. These are performed after single-parameter transforms. 
+            (default: ``None``)
+        fill_dict (dict, optional): Keys must contain ``'ndim_full'``, ``'fill_inds'``,
+            and ``'fill_values'``. ``'ndim_full'`` is the full last dimension of the final
+            array after fill_values are added. 'fill_inds' and 'fill_values' are
+            np.ndarray[number of fill values] that contain the indexes and corresponding values
+            for filling. (default: ``None``)
+
+    Raises:
+        ValueError: Input information is not correct.
+
+    """
+
+    def __init__(self, parameter_transforms=None, fill_dict=None):
+
+        # store originals
+        self.original_parameter_transforms = parameter_transforms
+        if parameter_transforms is not None:
+            # differentiate between single and multi parameter transformations
+            self.base_transforms = {"single_param": {}, "mult_param": {}}
+
+            # iterate through transforms and setup single and multiparameter transforms
+            for key, item in parameter_transforms.items():
+                if isinstance(key, int):
+                    self.base_transforms["single_param"][key] = item
+                elif isinstance(key, tuple):
+                    self.base_transforms["mult_param"][key] = item
+                else:
+                    raise ValueError(
+                        "Parameter transform keys must be int or tuple of ints. {} is neither.".format(
+                            key
+                        )
+                    )
+        else:
+            self.base_transforms = None
+
+        if fill_dict is not None:
+            if not isinstance(fill_dict, dict):
+                raise ValueError("fill_dict must be a dictionary.")
+
+            self.fill_dict = fill_dict
+            fill_dict_keys = list(self.fill_dict.keys())
+            for key in ["ndim_full", "fill_inds", "fill_values"]:
+                # check to make sure it has all necessary pieces
+                if key not in fill_dict_keys:
+                    raise ValueError(
+                        f"If providing fill_inds, dictionary must have {key} as a key."
+                    )
+            # check all the inputs
+            if not isinstance(fill_dict["ndim_full"], int):
+                raise ValueError("fill_dict['ndim_full'] must be an int.")
+            if not isinstance(fill_dict["fill_inds"], np.ndarray):
+                raise ValueError("fill_dict['fill_inds'] must be an np.ndarray.")
+            if not isinstance(fill_dict["fill_values"], np.ndarray):
+                raise ValueError("fill_dict['fill_values'] must be an np.ndarray.")
+
+            # set up test_inds accordingly
+            self.fill_dict["test_inds"] = np.delete(
+                np.arange(self.fill_dict["ndim_full"]), self.fill_dict["fill_inds"]
+            )
+
+        else:
+            self.fill_dict = None
+
+    def transform_base_parameters(
+        self, params, copy=True, return_transpose=False, xp=None
+    ):
+        """Transform the base parameters
+
+        Args:
+            params (np.ndarray[..., ndim]): Array with coordinates. This array is
+                transformed according to the ``self.base_transforms`` dictionary.
+            copy (bool, optional): If True, copy the input array.
+                (default: ``True``)
+            return_transpose (bool, optional): If True, return the transpose of the
+                array. (default: ``False``)
+            xp (object, optional): ``numpy`` or ``cupy``. If ``None``, use ``numpy``.
+                (default: ``None``) 
+
+        Returns:
+            np.ndarray[..., ndim]: Transformed ``params`` array.
+
+        """
+
+        # cupy or numpy
+        if xp is None:
+            xp = np
+
+        if self.base_transforms is not None:
+            params_temp = params.copy() if copy else params
+            params_temp = params_temp.T
+            # regular transforms
+            for ind, trans_fn in self.base_transforms["single_param"].items():
+                params_temp[ind] = trans_fn(params_temp[ind])
+
+            # multi parameter transforms
+            for inds, trans_fn in self.base_transforms["mult_param"].items():
+                temp = trans_fn(*[params_temp[i] for i in inds])
+                for j, i in enumerate(inds):
+                    params_temp[i] = temp[j]
+
+            # its actually the opposite now
+            if return_transpose:
+                return params_temp
+            else:
+                return params_temp.T
+
+        else:
+            if return_transpose:
+                return params.T
+            else:
+                return params
+
+    def fill_values(self, params, xp=None):
+        """fill fixed parameters
+
+        Args:
+            params (np.ndarray[..., ndim]): Array with coordinates. This array is
+                filled with values according to the ``self.fill_dict`` dictionary.
+            xp (object, optional): ``numpy`` or ``cupy``. If ``None``, use ``numpy``.
+                (default: ``None``) 
+
+        Returns:
+            np.ndarray[..., ndim_full]: Filled ``params`` array.
+
+        """
+        if self.fill_dict is not None:
+            if xp is None:
+                xp = np
+
+            # get shape
+            shape = params.shape
+
+            # setup new array to fill
+            params_filled = xp.zeros(shape[:-1] + (self.fill_dict["ndim_full"],))
+            test_inds = xp.asarray(self.fill_dict["test_inds"])
+            # special indexing to properly fill array with params
+            indexing_test_inds = tuple([slice(0, temp) for temp in shape[:-1]]) + (
+                test_inds,
+            )
+
+            # fill values directly from params array
+            params_filled[indexing_test_inds] = params
+
+            fill_inds = xp.asarray(self.fill_dict["fill_inds"])
+            # special indexing to fill fill_values
+            indexing_fill_inds = tuple([slice(0, temp) for temp in shape[:-1]]) + (
+                fill_inds,
+            )
+
+            # add fill_values at fill_inds
+            params_filled[indexing_fill_inds] = xp.asarray(
+                self.fill_dict["fill_values"]
+            )
+
+            return params_filled
+
+        else:
+            return params
+
+    def both_transforms(
+        self, params, copy=True, return_transpose=False, reverse=False, xp=None
+    ):
+        """Transform the parameters and fill fixed parameters
+
+        This fills the fixed parameters and then transforms all of them. Therefore, the user
+        must be careful with the indexes input. 
+
+        This is generally the direction recommended because fixed parameters may change
+        non-fixed parameters during parameter transformations. This can be reversed
+        with the ``reverse`` kwarg.
+
+        Args:
+            params (np.ndarray[..., ndim]): Array with coordinates. This array is
+                transformed according to the ``self.base_transforms`` dictionary.
+            copy (bool, optional): If True, copy the input array.
+                (default: ``True``)
+            return_transpose (bool, optional): If ``True``, return the transpose of the
+                array. (default: ``False``)
+            reverse (bool, optional): If ``True`` perform the filling after the transforms. This makes
+                indexing easier, but removes the ability of fixed parameters to affect transforms. 
+                (default: ``False``)
+            xp (object, optional): ``numpy`` or ``cupy``. If ``None``, use ``numpy``.
+                (default: ``None``) 
+
+        Returns:
+            np.ndarray[..., ndim]: Transformed and filleds ``params`` array.
+
+        """
+        # numpy or cupy
+        if xp is None:
+            xp = np
+
+        # run transforms first
+        if reverse:
+            temp = self.transform_base_parameters(
+                params, copy=copy, return_transpose=return_transpose, xp=xp
+            )
+            temp = self.fill_values(temp, xp=xp)
+
+        else:
+            temp = self.fill_values(params, xp=xp)
+            temp = self.transform_base_parameters(
+                temp, copy=copy, return_transpose=return_transpose, xp=xp
+            )
+        return temp

eryn/utils/updates.py

@@ -0,0 +1,69 @@
+# -*- coding: utf-8 -*-
+
+from abc import ABC
+
+import numpy as np
+
+
+class Update(ABC, object):
+    """Update the sampler."""
+
+    @classmethod
+    def __call__(self, iter, last_sample, sampler):
+        """Call update function.
+
+        Args:
+            iter (int): Iteration of the sampler.
+            last_sample (obj): Last state of sampler (:class:`eryn.state.State`).
+            sampler (obj): Full sampler oject (:class:`eryn.ensemble.EnsembleSampler`).
+
+        """
+        raise NotImplementedError
+
+
+class AdjustStretchProposalScale(Update):
+    def __init__(
+        self,
+        target_acceptance=0.22,
+        supression_factor=0.1,
+        max_change=0.5,
+        verbose=False,
+    ):
+        """Adjusted scale for stretch proposal based on cold chain acceptance rate"""
+        self.target_acceptance = target_acceptance
+        self.verbose = verbose
+        self.max_change, self.supression_factor = max_change, supression_factor
+
+        self.time = 0
+
+    def __call__(self, iter, last_sample, sampler):
+
+        mean_af = 0.0
+        change = 1.0
+        if self.time > 0:
+            # cold chain -> 0
+            mean_af = np.mean(
+                (sampler.backend.accepted[:, 0] - self.previously_accepted)
+                / (sampler.backend.iteration - self.previous_iter)
+            )
+
+            if mean_af > self.target_acceptance:
+                factor = self.supression_factor * (mean_af / self.target_acceptance)
+                if factor > self.max_change:
+                    factor = self.max_change
+                change = 1 + self.supression_factor * factor
+
+            else:
+                factor = self.supression_factor * (self.target_acceptance / mean_af)
+                if factor > self.max_change:
+                    factor = self.max_change
+                change = 1 - factor
+
+            sampler._moves[0].a *= change
+
+        self.previously_accepted = sampler.backend.accepted[:, 0].copy()
+        print(
+            self.previously_accepted, "\n", mean_af, change, "\n", sampler._moves[0].a
+        )
+        self.previous_iter = sampler.backend.iteration
+        self.time += 1