Stochastic transitions

atomica/demos.py

@@ -11,10 +11,30 @@
 from .scenarios import BudgetScenario
 from .utils import TimeSeries
 
-__all__ = ["demo", "make_demo_scenarios"]
-
-
-def demo(which: str = None, do_run: bool = True, addprogs: bool = True) -> Project:
+__all__ = ["demo", "demo_projects", "make_demo_scenarios"]
+
+
+demo_projects = [
+    "udt",
+    "udt_dyn",
+    "usdt",
+    "cervicalcancer",
+    "sir",
+    "diabetes",
+    "combined",
+    "service",
+    "hypertension",
+    "hypertension_dyn",
+    "hiv",
+    "hiv_dyn",
+    "tb_simple",
+    "tb_simple_dyn",
+    "tb",
+    "dt",
+]
+
+
+def demo(which: str = "sir", do_run: bool = True, addprogs: bool = True) -> Project:
     """
     Return a demo project
 
@@ -25,30 +45,11 @@
 
     """
 
-    options = [
-        "udt",
-        "udt_dyn",
-        "usdt",
-        "cervicalcancer",
-        "sir",
-        "diabetes",
-        "combined",
-        # 'service',
-        "hypertension",
-        "hypertension_dyn",
-        "hiv",
-        "hiv_dyn",
-        "tb_simple",
-        "tb_simple_dyn",
-        "environment",
-        "tb",
-    ]
-
-    dtdict = sc.odict.fromkeys(options, 1.0)
+    dtdict = sc.odict.fromkeys(demo_projects, 1.0)
     dtdict["tb"] = 0.5
 
-    if which is None or which not in options:
-        raise Exception("Supported project types are:\n%s" % ("\n".join(options)))
+    if which not in demo_projects:
+        raise Exception("Supported demonstration projects are:\n%s" % ("\n".join(demos)))
 
     framework = LIBRARY_PATH / f"{which}_framework.xlsx"
     databook = LIBRARY_PATH / f"{which}_databook.xlsx"
@@ -61,7 +62,7 @@
     if do_run:
         P.run_sim(store_results=True)
 
-    if addprogs:
+    if addprogs and progbook.exists():
         logger.debug("Loading progbook")
         P.load_progbook(progbook)
 

atomica/excel.py

@@ -448,7 +448,8 @@
             raise Exception('Unknown TimeDependentConnections type - must be "transfer" or "interaction"')
 
     def __repr__(self):
-        return '<TDC %s "%s">' % (self.type.title(), self.code_name)
+        # return '<TDC %s "%s">' % (self.type.title(), self.code_name)
+        return sc.prepr(self)
 
     @classmethod
     def from_tables(cls, tables: list, interaction_type):
@@ -859,13 +860,13 @@
             from_pop, to_pop = interaction
             if not self.enable_diagonal and from_pop == to_pop:
                 raise Exception("Trying to write a diagonal entry to a table that is not allowed to contain diagonal terms")  # This is because data loss will occur if the user adds entries on the diagonal, then writes the table, and then reads it back in
             from_idx = self.from_pops.index(from_pop)
             to_idx = self.to_pops.index(to_pop)
             if boolean_choice:
-                value = "Y" if value else "N"
+                value = "Y" if (value.vals or value.assumption) else "N"
             content[from_idx, to_idx] = value
 
         # Write the content
         for from_idx in range(0, len(self.from_pops)):
             for to_idx in range(0, len(self.to_pops)):
                 row = start_row + 1 + from_idx

atomica/migration.py

@@ -868,4 +868,5 @@ def _progset_nontargetable_flag(progset):
     for d in progset.comps.values():
         if 'non_targetable' not in d:
             d['non_targetable'] = False
-    return progset
+    return progset
+

atomica/parameters.py

@@ -96,7 +96,7 @@ def interpolate(self, tvec, pop_name: str) -> np.array:
 
         return self.ts[pop_name].interpolate(tvec, method=self._interpolation_method)
 
-    def sample(self, constant: bool) -> None:
+    def sample(self, constant: bool, rng_sampler=None) -> None:
         """
         Perturb parameter based on uncertainties
 
@@ -105,11 +105,11 @@ def sample(self, constant: bool) -> None:
 
         :param constant: If True, time series will be perturbed by a single constant offset. If False,
                          a different perturbation will be applied to each time specific value independently.
-
+        :param rng_sampler: Optional random number generator that may have been seeded to generate consistent results
         """
 
         for k, ts in self.ts.items():
-            self.ts[k] = ts.sample(constant)
+            self.ts[k] = ts.sample(constant, rng_sampler)
 
     def smooth(self, tvec, method="smoothinterp", pop_names=None, **kwargs):
         """
@@ -530,19 +530,20 @@ def get_par(self, name: str, pop: str = None) -> Parameter:
 
         raise KeyError(f'Parameter "{name}" not found')
 
-    def sample(self, constant=True):
+    def sample(self, constant: bool = True, rng_sampler=None):
         """
         Return a sampled copy of the ParameterSet
 
         :param constant: If True, time series will be perturbed by a single constant offset. If False,
                          a different perturbation will be applied to each time specific value independently.
+        :param rng_sampler: Optional random number generator that may have been seeded to generate consistent results
         :return: A new :class:`ParameterSet` with perturbed values
 
         """
-
         new = sc.dcp(self)
-        for par in new.all_pars():
-            par.sample(constant)
+
+        for i, par in enumerate(new.all_pars()):
+            par.sample(constant=constant, rng_sampler=rng_sampler)
         return new
 
     def make_constant(self, year: float):

atomica/plotting.py

@@ -516,6 +516,7 @@ def placeholder_pop():
                             numerator = sum(aggregated_outputs[x][output_name] * weights[x] for x in pop_labels)  # Add together all the outputs
                             denominator = sum([weights[x] for x in pop_labels])
                             vals = np.divide(numerator, denominator, out=np.full(numerator.shape, np.nan, dtype=float), where=numerator != 0)
+
                         else:
                             raise Exception(f'Unknown population aggregation method "{aggregation}"')
 
@@ -1132,7 +1133,9 @@ def set_colors(self, colors=None, results="all", pops="all", outputs="all", over
         elif isinstance(colors, list):
             assert len(colors) == len(targets), "Number of colors must either be a string, or a list with as many elements as colors to set"
             colors = colors
-        elif colors.startswith("#") or colors not in [m for m in plt.cm.datad if not m.endswith("_r")]:
+        elif isinstance(colors, tuple) and colors[0].startswith("#"): # tuple with (color, opacity)
+            colors = [colors for _ in range(len(targets))]  # Apply color to all requested outputs
+        elif isinstance(colors, str) and (colors.startswith("#") or colors not in [m for m in plt.cm.datad if not m.endswith("_r")]):
             colors = [colors for _ in range(len(targets))]  # Apply color to all requested outputs
         else:
             color_norm = matplotlib_colors.Normalize(vmin=-1, vmax=len(targets))
@@ -1658,7 +1661,8 @@ def process_input_stacks(input_stacks, available_items):
     return figs
 
 
-def plot_series(plotdata, plot_type="line", axis=None, data=None, legend_mode=None, lw=None, n_cols: int = None) -> list:
+def plot_series(plotdata, plot_type="line", axis=None, data=None, legend_mode=None, lw=None, n_cols: int = None,
+                colors=None) -> list:
     """
     Produce a time series plot
 
@@ -1671,6 +1675,7 @@ def plot_series(plotdata, plot_type="line", axis=None, data=None, legend_mode=No
     :param lw: override the default line width
     :param n_cols: If None (default), separate figures will be created for each axis. If provided, axes will be tiled as subplots in a single figure
                    window with the requested number of columns
+    :param colors: Colors to be passed to plotdata.set_colors
     :return: A list of newly created Figures
 
     """
@@ -1724,7 +1729,7 @@ def _prepare_figures(dim1, dim2, n_cols):
         logger.warning("At least one Series has only one timepoint. Series must have at least 2 time points to be rendered as a line - `plot_bars` may be more suitable for such data")
 
     if axis == "results":
-        plotdata.set_colors(results=plotdata.results.keys())
+        plotdata.set_colors(colors=colors, results=plotdata.results.keys())
 
         figs, axes = _prepare_figures(plotdata.pops, plotdata.outputs, n_cols)
 
@@ -1759,7 +1764,7 @@ def _prepare_figures(dim1, dim2, n_cols):
                 _render_legend(ax, plot_type)
 
     elif axis == "pops":
-        plotdata.set_colors(pops=plotdata.pops.keys())
+        plotdata.set_colors(colors=colors, pops=plotdata.pops.keys())
 
         figs, axes = _prepare_figures(plotdata.results, plotdata.outputs, n_cols)
 
@@ -1792,7 +1797,7 @@ def _prepare_figures(dim1, dim2, n_cols):
                 _render_legend(ax, plot_type)
 
     elif axis == "outputs":
-        plotdata.set_colors(outputs=plotdata.outputs.keys())
+        plotdata.set_colors(colors=colors, outputs=plotdata.outputs.keys())
 
         figs, axes = _prepare_figures(plotdata.results, plotdata.pops, n_cols)
 

atomica/programs.py

@@ -1141,7 +1141,7 @@
 
         return {(covout.par, covout.pop): covout.get_outcome(prop_coverage) for covout in self.covouts.values()}
 
-    def sample(self, constant: bool = True):
+    def sample(self, constant: bool = True, rng_sampler=None):
         """
         Perturb programs based on uncertainties
 
@@ -1158,9 +1158,9 @@
 
         new = sc.dcp(self)
         for prog in new.programs.values():
-            prog.sample(constant)
+            prog.sample(constant, rng_sampler=rng_sampler)
         for covout in new.covouts.values():
-            covout.sample()
+            covout.sample(rng_sampler=rng_sampler)
         return new
 
 
@@ -1221,24 +1221,25 @@
         """
         return "/year" not in self.unit_cost.units
 
-    def sample(self, constant: bool) -> None:
+    def sample(self, constant: bool, rng_sampler=None) -> None:
         """
         Perturb program values based on uncertainties
 
         Calling this function will perturb the original values based on their uncertainties. The
         values will change in-place. Normally, this method would be called by
         :meth:`ProgramSet.sample()` which will copy the ``Program`` instance first.
 
         :param constant: If True, time series will be perturbed by a single constant offset. If False,
                          an different perturbation will be applied to each time specific value independently.
+        :param rng_sampler: Optional random number generator that may have been seeded to generate consistent results
         """
 
         self.spend_data = self.spend_data.sample(constant)
         self.unit_cost = self.unit_cost.sample(constant)
-        self.capacity_constraint = self.capacity_constraint.sample(constant)
+        self.capacity_constraint = self.capacity_constraint.sample(constant, rng_sampler)
         self.saturation = self.saturation.sample(constant)
         self.coverage = self.coverage.sample(constant)
 
     def __repr__(self):
         output = sc.prepr(self)
         output += "          Program name: %s\n" % self.name
@@ -1397,25 +1398,28 @@
 
         return len(self.progs)
 
-    def sample(self) -> None:
+    def sample(self, rng_sampler=None) -> None:
         """
         Perturb the values entered in the databook
 
         The :class:`Covout` instance is modified in-place. Note that the program outcomes are scalars
         that do not vary over time - therefore, :meth:`Covout.sample()` does not have a ``constant``
         argument.
-
+        :param rng_sampler: Optional random number generator that may have been seeded to generate consistent results
         """
 
         if self.sigma is None:
             return
 
+        if rng_sampler is None:
+            rng_sampler = np.random
+
         for k, v in self.progs.items():
-            self.progs[k] = v + self.sigma * np.random.randn(1)[0]
+            self.progs[k] = v + self.sigma * rng_sampler.standard_normal(1)[0]
         # Perturb the interactions
         if self._interactions:
             for k, v in self.interactions.items():
-                self.interactions[k] = v + self.sigma * np.random.randn(1)[0]
+                self.interactions[k] = v + self.sigma * rng_sampler.standard_normal(1)[0]
             tokens = ["%s=%.4f" % ("+".join(k), v) for k, v in self.interactions.items()]
             self.imp_interaction = ",".join(tokens)