eryn 1.2.3__tar.gz → 1.2.5__tar.gz

{eryn-1.2.3 → eryn-1.2.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: eryn
-Version: 1.2.3
+Version: 1.2.5
 Summary: Eryn: an omni-MCMC sampling package.
 Author: Michael Katz
 Author-email: Michael Katz <mikekatz04@gmail.com>

{eryn-1.2.3 → eryn-1.2.5}/pyproject.toml

@@ -8,7 +8,7 @@ requires = [
 [project]
 name = "eryn" #@NAMESUFFIX@
 
-version = "1.2.3"
+version = "1.2.5"
 
 description = "Eryn: an omni-MCMC sampling package."
 

{eryn-1.2.3 → eryn-1.2.5}/src/eryn/backends/backend.py

@@ -83,6 +83,7 @@ class Backend(object):
         nbranches=1,
         rj=False,
         moves=None,
+        key_order=None,
         **info,
     ):
         """Clear the state of the chain and empty the backend
@@ -106,6 +107,9 @@ class Backend(object):
                 (default: ``False``)
             moves (list, optional): List of all of the move classes input into the sampler.
                 (default: ``None``)
+            key_order (dict, optional): Keys are ``branch_names`` and values are lists of key ordering for each
+                branch. For example, ``{"model_0": ["x1", "x2", "x3"]}``. 
+                (default: ``None``)
             **info (dict, optional): Any other key-value pairs to be added
                 as attributes to the backend.
 
@@ -118,6 +122,7 @@ class Backend(object):
             branch_names=branch_names,
             rj=rj,
             moves=moves,
+            key_order=key_order,
             info=info,
         )
 
@@ -184,6 +189,7 @@ class Backend(object):
         self.branch_names = branch_names
         self.ndims = ndims
         self.nleaves_max = nleaves_max
+        self.key_order = key_order 
 
         self.iteration = 0
 
@@ -703,8 +709,11 @@ class Backend(object):
             "thermo",
             "ti",
         ]:
-            logls = np.mean(logls_all, axis=(0, -1))
-            logZ, dlogZ = thermodynamic_integration_log_evidence(betas, logls)
+            logls = logls_all.copy()
+            logls[~np.isfinite(logls)] = np.nan
+            meanlogls = np.nanmean(logls, axis=(0, -1))
+            logZ, dlogZ = thermodynamic_integration_log_evidence(betas, meanlogls)
+            
         elif method.lower() in [
             "stepping stone",
             "ss",

{eryn-1.2.3 → eryn-1.2.5}/src/eryn/backends/hdfbackend.py

@@ -176,6 +176,7 @@ class HDFBackend(Backend):
         nbranches=1,
         rj=False,
         moves=None,
+        key_order=None,
         **info,
     ):
         """Clear the state of the chain and empty the backend
@@ -199,6 +200,9 @@ class HDFBackend(Backend):
                 (default: ``False``)
             moves (list, optional): List of all of the move classes input into the sampler.
                 (default: ``None``)
+            key_order (dict, optional): Keys are ``branch_names`` and values are lists of key ordering for each
+                branch. For example, ``{"model_0": ["x1", "x2", "x3"]}``. 
+                (default: ``None``)
             **info (dict, optional): Any other key-value pairs to be added
                 as attributes to the backend. These are also added to the HDF5 file.
 
@@ -343,6 +347,7 @@ class HDFBackend(Backend):
 
             chain = g.create_group("chain")
             inds = g.create_group("inds")
+            k_o_g = g.create_group("key_order")
 
             for name in branch_names:
                 nleaves = self.nleaves_max[name]
@@ -365,6 +370,9 @@ class HDFBackend(Backend):
                     compression_opts=self.compression_opts,
                 )
 
+                if key_order is not None:
+                    k_o_g.attrs[name] = key_order[name] 
+
             # store move specific information
             if moves is not None:
                 move_group = g.create_group("moves")
@@ -388,6 +396,12 @@ class HDFBackend(Backend):
 
             self.blobs = None
 
+    @property
+    def key_order(self):
+        """Key order of parameters for each model."""
+        with self.open() as f:
+            return {key: value for key, value in f[self.name]["key_order"].attrs.items()}
+
     @property
     def nwalkers(self):
         """Get nwalkers from h5 file."""
@@ -456,6 +470,7 @@ class HDFBackend(Backend):
             branch_names=self.branch_names,
             rj=self.rj,
             moves=self.moves,
+            key_order=self.key_order,
         )
 
     @property

{eryn-1.2.3 → eryn-1.2.5}/src/eryn/ensemble.py

@@ -198,6 +198,7 @@ class EnsembleSampler(object):
             if it is changed. In this case, the user should declare a new backend and use the last
             state from the previous backend. **Warning**: If the order of moves of the same move class
             is changed, the check may not catch it, so the tracking may mix move acceptance fractions together.
+            (default: ``True'')
         info (dict, optional): Key and value pairs reprenting any information
             the user wants to add to the backend if the user is not inputing
             their own backend.
@@ -597,6 +598,7 @@ class EnsembleSampler(object):
                 nleaves_max=nleaves_max,
                 rj=self.has_reversible_jump,
                 moves=move_keys,
+                key_order=self.key_order,
                 **info,
             )
             state = np.random.get_state()
@@ -615,6 +617,9 @@ class EnsembleSampler(object):
                         "Configuration of moves has changed. Cannot use the same backend. Declare a new backend and start from the previous state. If you would prefer not to track move acceptance fraction, set track_moves to False in the EnsembleSampler."
                     )
 
+            if self.key_order != self.backend.key_order:
+                raise ValueError("Input key order from priors does not match backend.")
+            
             # Check the backend shape
             for i, (name, shape) in enumerate(self.backend.shape.items()):
                 test_shape = (
@@ -689,7 +694,7 @@ class EnsembleSampler(object):
 
         """
         return self._random.get_state()
-
+    
     @random_state.setter  # NOQA
     def random_state(self, state):
         """
@@ -751,13 +756,14 @@ class EnsembleSampler(object):
         else:
             raise ValueError("Priors must be a dictionary.")
 
+        self.key_order = {key: value.key_order for key, value in self._priors.items()}
         return
 
     @property
     def iteration(self):
         return self.backend.iteration
 
-    def reset(self, **info):
+    def reset(self, **kwargs):
         """
         Reset the backend.
 
@@ -765,7 +771,7 @@ class EnsembleSampler(object):
             **info (dict, optional): information to pass to backend reset method.
 
         """
-        self.backend.reset(self.nwalkers, self.ndims, **info)
+        self.backend.reset(self.nwalkers, self.ndims, **kwargs)
 
     def __getstate__(self):
         # In order to be generally picklable, we need to discard the pool
@@ -1208,6 +1214,9 @@ class EnsembleSampler(object):
                 # vectorized because everything is rectangular (no groups to indicate model difference)
                 prior_out += prior_out_temp.sum(axis=-1)
 
+        if np.any(np.isnan(prior_out)):
+            raise ValueError("The prior function is returning Nan.")
+        
         return prior_out
 
     def compute_log_like(
@@ -1493,8 +1502,9 @@ class EnsembleSampler(object):
             ll[inds_fix_zeros] = self.fill_zero_leaves_val
 
             # deal with blobs
-            blobs_out = np.zeros((nwalkers_all, results.shape[1] - 1))
-            blobs_out[unique_groups] = results[:, 1:]
+            _blobs_out = np.zeros((nwalkers_all, results.shape[1] - 1))
+            _blobs_out[unique_groups] = results[:, 1:]
+            blobs_out = _blobs_out.reshape(ntemps, nwalkers)
 
         elif results.dtype == "object":
             # TODO: check blobs and add this capability
@@ -1531,6 +1541,9 @@ class EnsembleSampler(object):
                                 for key in branch_supps_in_2[name_i]
                             }
 
+        if np.any(np.isnan(ll)):
+            raise ValueError("The likelihood function is returning Nan.")
+        
         # return Likelihood and blobs
         return ll.reshape(ntemps, nwalkers), blobs_out
 

{eryn-1.2.3 → eryn-1.2.5}/src/eryn/moves/gaussian.py

@@ -137,7 +137,12 @@ class _isotropic_proposal(object):
     def __init__(self, scale, factor, mode):
         self.index = 0
         self.scale = scale
-        self.invscale = np.linalg.inv(np.linalg.cholesky(scale))
+        
+        if isinstance(scale, float):
+            self.invscale = 1. / scale
+        else:
+            self.invscale = np.linalg.inv(np.linalg.cholesky(scale))
+
         if factor is None:
             self._log_factor = None
         else:

{eryn-1.2.3 → eryn-1.2.5}/src/eryn/prior.py

@@ -248,24 +248,69 @@ class ProbDistContainer:
         # to separate out in list form
         self.priors = []
 
+        self.has_strings = False
+        self.has_ints = False
+
+        # this is for the strings (for the ints it just counts them)
+        current_ind = 0
+        key_order = []
+
         # setup lists
         temp_inds = []
         for inds, dist in priors_in.items():
             # multiple index
             if isinstance(inds, tuple):
-                inds_in = np.asarray(inds)
+                inds_tmp = []
+                for i in range(len(inds)):
+                    if isinstance(inds[i], str):
+                        assert not self.has_ints
+                        self.has_strings = True
+                        inds_tmp.append(current_ind)
+                        key_order.append(inds[i])
+
+                    elif isinstance(inds[i], int):
+                        assert not self.has_strings
+                        self.has_ints = True
+                        inds_tmp.append(i)
+
+                    else:
+                        raise ValueError("Index in tuple must be int or str and all be the same type.")
+
+                    current_ind += 1
+
+                inds_in = np.asarray(inds_tmp)
                 self.priors.append([inds_in, dist])
 
             # single index
             elif isinstance(inds, int):
+                self.has_ints = True
+                assert not self.has_strings
                 inds_in = np.array([inds])
                 self.priors.append([inds_in, dist])
+                current_ind += 1
+
+            elif isinstance(inds, str):
+                assert not self.has_ints
+                self.has_strings = True
+                key_order.append(inds)
+                inds_in = np.array([current_ind])
+                current_ind += 1
+                self.priors.append([inds_in, dist])
 
             else:
                 raise ValueError(
-                    "Keys for prior dictionary must be an integer or tuple."
+                    "Keys for prior dictionary must be an integer, string, or tuple."
                 )
 
+            if self.has_strings:
+                assert not self.has_ints
+                # key order is already set
+                self.key_order = key_order
+
+            if self.has_ints:
+                self.key_order = [i for i in range(current_ind)]  # here current_ind is the total count
+                assert not self.has_strings
+             
             temp_inds.append(np.asarray([inds_in]))
 
         uni_inds = np.unique(np.concatenate(temp_inds, axis=1).flatten())

{eryn-1.2.3 → eryn-1.2.5}/src/eryn/utils/periodic.py

@@ -18,15 +18,31 @@ class PeriodicContainer:
 
     """
 
-    def __init__(self, periodic):
+    def __init__(self, periodic, key_order=None):
 
         # store all the information
         self.periodic = periodic
+        inds_periodic = {}
+        periods = {}
+        for key in periodic:
+            if periodic[key] is None:
+                continue
+            inds_periodic[key] = []
+            periods[key] = []
+            for var, period in periodic[key].items():
+                if isinstance(var, str):
+                    if key_order is None:
+                        raise ValueError(f"If providing str values for the variable names, must provide key_order argument.")
+
+                    index = key_order[key].index(var)
+                    inds_periodic[key].append(index)
+                    periods[key].append(period)
+
         self.inds_periodic = {
-            key: np.asarray([i for i in periodic[key].keys()]) for key in periodic
+            key: np.asarray(tmp) for key, tmp in inds_periodic.items()
         }
         self.periods = {
-            key: np.asarray([i for i in periodic[key].values()]) for key in periodic
+            key: np.asarray(tmp) for key, tmp in periods.items()
         }
 
     def distance(self, p1, p2, xp=None):

{eryn-1.2.3 → eryn-1.2.5}/src/eryn/utils/transform.py

@@ -11,6 +11,10 @@ class TransformContainer:
     """Container for helpful transformations
 
     Args:
+        input_basis (list): List of integers or strings representing each 
+            basis element from the input basis.
+        output_basis (list): List of integers or strings representing each 
+            basis element for the output basis.
         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
@@ -31,53 +35,70 @@ class TransformContainer:
 
     """
 
-    def __init__(self, parameter_transforms=None, fill_dict=None):
+    def __init__(self, input_basis=None, output_basis=None, parameter_transforms=None, fill_dict=None, key_map={}):
 
+        
         # store originals
         self.original_parameter_transforms = parameter_transforms
+        self.ndim_full = len(output_basis)
+        self.ndim = len(input_basis)
+
+        self.input_basis, self.output_basis = input_basis, output_basis
+
+        test_inds = []
+        for key in input_basis:
+            if key not in output_basis and key not in key_map:
+                raise ValueError("All keys in input_basis must be present in output basis, or you must provide a key_map")
+            key_in = key if key not in key_map else key_map[key]
+            test_inds.append(output_basis.index(key_in))
+
+        self.test_inds = test_inds = np.asarray(test_inds)
         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
+                if isinstance(key, str) or isinstance(key, int):
+                    if key not in output_basis:
+                        assert key in key_map
+                        key = key_map[key]
+                    key_in = output_basis.index(key)
+                    self.base_transforms["single_param"][key_in] = item
                 elif isinstance(key, tuple):
-                    self.base_transforms["mult_param"][key] = item
+                    _tmp = []
+                    for i in range(len(key)):
+                        key_tmp = key[i]
+                        if  key_tmp not in output_basis:
+                            assert key_tmp in key_map
+                            key_tmp = key_map[key_tmp]
+                        _tmp.append(output_basis.index(key_tmp))
+                    self.base_transforms["mult_param"][tuple(_tmp)] = item
                 else:
                     raise ValueError(
-                        "Parameter transform keys must be int or tuple of ints. {} is neither.".format(
+                        "Parameter transform keys must be str (or int) or tuple of strs (or ints). {} is neither.".format(
                             key
                         )
                     )
         else:
             self.base_transforms = None
 
+        self.original_fill_dict = fill_dict
         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.")
+            self.fill_dict = {}
+            self.fill_dict["fill_inds"] = []
+            self.fill_dict["fill_values"] = []
+            for key in fill_dict.keys():
+                self.fill_dict["fill_inds"].append(output_basis.index(key))
+                self.fill_dict["fill_values"].append(fill_dict[key])
 
             # set up test_inds accordingly
-            self.fill_dict["test_inds"] = np.delete(
-                np.arange(self.fill_dict["ndim_full"]), self.fill_dict["fill_inds"]
-            )
+            self.fill_dict["test_inds"] = test_inds
+            self.fill_dict["fill_inds"] = np.asarray(self.fill_dict["fill_inds"])
+            self.fill_dict["fill_values"] = np.asarray(self.fill_dict["fill_values"])
 
         else:
             self.fill_dict = None
@@ -134,6 +155,8 @@ class TransformContainer:
     def fill_values(self, params, xp=None):
         """fill fixed parameters
 
+        This also adjusts parameter order as needed between the two bases. 
+
         Args:
             params (np.ndarray[..., ndim]): Array with coordinates. This array is
                 filled with values according to the ``self.fill_dict`` dictionary.
@@ -152,7 +175,7 @@ class TransformContainer:
             shape = params.shape
 
             # setup new array to fill
-            params_filled = xp.zeros(shape[:-1] + (self.fill_dict["ndim_full"],))
+            params_filled = xp.zeros(shape[:-1] + (self.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]]) + (
@@ -179,7 +202,7 @@ class TransformContainer:
             return params
 
     def both_transforms(
-        self, params, copy=True, return_transpose=False, reverse=False, xp=None
+        self, params, copy=True, return_transpose=False, xp=None
     ):
         """Transform the parameters and fill fixed parameters
 
@@ -197,9 +220,6 @@ class TransformContainer:
                 (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``) 
 
@@ -212,15 +232,8 @@ class TransformContainer:
             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
-            )
+        temp = self.fill_values(params, xp=xp)
+        temp = self.transform_base_parameters(
+            temp, copy=copy, return_transpose=return_transpose, xp=xp
+        )
         return temp

{eryn-1.2.3 → eryn-1.2.5}/src/eryn/utils/utility.py

@@ -237,11 +237,12 @@ def stepping_stone_log_evidence(betas, logls, block_len=50, repeats=100):
 
     def calculate_stepping_stone(betas, logls):
         n = logls.shape[0]
-        delta_betas = betas[1:] - betas[:-1]
         n_T = betas.shape[0]
-        log_ratio = logsumexp(delta_betas * logls[:, :-1], axis=0) - np.log(n)
+        delta_betas = betas[1:] - betas[:-1]
+        throwaways = np.any(~np.isfinite(logls), axis=1) # a safeguard against non-finite entries
+        log_ratio = logsumexp(delta_betas * logls[~throwaways, :-1], axis=0) - (n_T - 1.0)*np.log(n - np.sum(throwaways))
         return np.sum(log_ratio), log_ratio
-
+    
     # make sure they are the same length
     if len(betas) != logls.shape[1]:
         raise ValueError(