eryn 1.2.4__py3-none-any.whl → 1.2.6__py3-none-any.whl

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/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/ensemble.py

@@ -13,7 +13,7 @@ from .pbar import get_progress_bar
 from .state import State
 from .prior import ProbDistContainer
 
-# from .utils import PlotContainer
+from .utils import PlotContainer
 from .utils import PeriodicContainer
 from .utils.utility import groups_from_inds
 
@@ -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.
@@ -232,7 +233,7 @@ class EnsembleSampler(object):
         blobs_dtype=None,  # TODO check this
         plot_iterations=-1,  # TODO: do plot stuff?
         plot_generator=None,
-        plot_name=None,
+        plot_folder=None,
         periodic=None,
         update_fn=None,
         update_iterations=-1,
@@ -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 = (
@@ -657,18 +662,15 @@ class EnsembleSampler(object):
         self.plot_iterations = plot_iterations
 
         if plot_generator is None and self.plot_iterations > 0:
-            raise NotImplementedError
             # set to default if not provided
-            if plot_name is not None:
-                name = plot_name
-            else:
-                name = "output"
+            if plot_folder is  None:
+                plot_folder = "./runtime_plots"
             self.plot_generator = PlotContainer(
-                fp=name, backend=self.backend, thin_chain_by_ac=True
+                backend=self.backend, plots=['base', 'rj'], parent_folder=plot_folder, discard=0.2
             )
         elif self.plot_iterations > 0:
-            raise NotImplementedError
             self.plot_generator = plot_generator
+            self.plot_generator.backend = self.backend # make sure backend is correctly set
 
         # prepare stopping functions
         self.stopping_fn = stopping_fn
@@ -689,7 +691,7 @@ class EnsembleSampler(object):
 
         """
         return self._random.get_state()
-
+    
     @random_state.setter  # NOQA
     def random_state(self, state):
         """
@@ -751,13 +753,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 +768,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
@@ -1101,7 +1104,7 @@ class EnsembleSampler(object):
             # diagnostic plots
             # TODO: adjust diagnostic plots
             if self.plot_iterations > 0 and (i + 1) % (self.plot_iterations) == 0:
-                self.plot_generator.generate_plot_info()  # TODO: remove defaults
+                self.plot_generator.produce_plots(sampler=self)  # TODO: remove defaults
 
             # check for stopping before updating
             if (

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/utils/__init__.py

@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 
-# from .plot import PlotContainer
+from .plot import PlotContainer
 from .utility import *
 from .periodic import *
 from .transform import *

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):