eryn 1.2.4__tar.gz → 1.2.5__tar.gz

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

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

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

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

{eryn-1.2.4 → 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
 

{eryn-1.2.4 → 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.4 → 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

{eryn-1.2.4 → 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.4 → 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.4 → 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