Consolidate distribution classes

HARK/Calibration/Income/IncomeProcesses.py

@@ -11,7 +11,6 @@
     DiscreteDistributionLabeled,
     IndexDistribution,
     MeanOneLogNormal,
-    TimeVaryingDiscreteDistribution,
     Lognormal,
     Uniform,
 )
@@ -813,14 +812,14 @@ def get_PermShkDstn_from_IncShkDstn(IncShkDstn, RNG):
     PermShkDstn = [
         this.make_univariate(0, seed=RNG.integers(0, 2**31 - 1)) for this in IncShkDstn
     ]
-    return TimeVaryingDiscreteDistribution(PermShkDstn, seed=RNG.integers(0, 2**31 - 1))
+    return IndexDistribution(distributions=PermShkDstn, seed=RNG.integers(0, 2**31 - 1))
 
 
 def get_TranShkDstn_from_IncShkDstn(IncShkDstn, RNG):
     TranShkDstn = [
         this.make_univariate(1, seed=RNG.integers(0, 2**31 - 1)) for this in IncShkDstn
     ]
-    return TimeVaryingDiscreteDistribution(TranShkDstn, seed=RNG.integers(0, 2**31 - 1))
+    return IndexDistribution(distributions=TranShkDstn, seed=RNG.integers(0, 2**31 - 1))
 
 
 def get_PermShkDstn_from_IncShkDstn_markov(IncShkDstn, RNG):

HARK/core.py

@@ -25,7 +25,6 @@
 from HARK.distributions import (
     Distribution,
     IndexDistribution,
-    TimeVaryingDiscreteDistribution,
     combine_indep_dstns,
 )
 from HARK.parallel import multi_thread_commands, multi_thread_commands_fake
@@ -1051,7 +1050,7 @@ def check_elements_of_time_vary_are_lists(self):
                 continue
             if not isinstance(
                 getattr(self, param),
-                (TimeVaryingDiscreteDistribution, IndexDistribution),
+                (IndexDistribution,),
             ):
                 assert type(getattr(self, param)) == list, (
                     param

HARK/distributions/__init__.py

@@ -3,7 +3,6 @@
     "DiscreteDistributionLabeled",
     "Distribution",
     "IndexDistribution",
-    "TimeVaryingDiscreteDistribution",
     "Lognormal",
     "MeanOneLogNormal",
     "Normal",
@@ -29,7 +28,6 @@
     Distribution,
     IndexDistribution,
     MarkovProcess,
-    TimeVaryingDiscreteDistribution,
 )
 from HARK.distributions.continuous import (
     Lognormal,

HARK/distributions/base.py

@@ -219,9 +219,8 @@ class IndexDistribution(Distribution):
     class (such as Bernoulli, LogNormal, etc.) with information
     about the conditions on the parameters of the distribution.
 
-    For example, an IndexDistribution can be defined as
-    a Bernoulli distribution whose parameter p is a function of
-    a different input parameter.
+    It can also wrap a list of pre-discretized distributions (previously
+    provided by TimeVaryingDiscreteDistribution) and provide the same API.
 
     Parameters
     ----------
@@ -235,14 +234,19 @@ class (such as Bernoulli, LogNormal, etc.) with information
         Keys should match the arguments to the engine class
         constructor.
 
+    distributions: [DiscreteDistribution]
+        Optional. A list of discrete distributions to wrap directly.
+
     seed : int
         Seed for random number generator.
     """
 
     conditional = None
     engine = None
 
-    def __init__(self, engine, conditional, RNG=None, seed=0):
+    def __init__(
+        self, engine=None, conditional=None, distributions=None, RNG=None, seed=0
+    ):
         if RNG is None:
             # Set up the RNG
             super().__init__(seed)
@@ -255,11 +259,24 @@ def __init__(self, engine, conditional, RNG=None, seed=0):
             # and create a new one.
             self.seed = seed
 
-        self.conditional = conditional
+        # Mode 1: wrapping a list of discrete distributions
+        if distributions is not None:
+            self.distributions = distributions
+            self.engine = None
+            self.conditional = None
+            self.dstns = []
+            return
+
+        # Mode 2: engine + conditional parameters (original IndexDistribution)
+        self.conditional = conditional if conditional is not None else {}
         self.engine = engine
 
         self.dstns = []
 
+        # If no engine/conditional were provided, remain empty (should not happen in normal use)
+        if self.engine is None and not self.conditional:
+            return
+
         # Test one item to determine case handling
         item0 = list(self.conditional.values())[0]
 
@@ -273,7 +290,7 @@ def __init__(self, engine, conditional, RNG=None, seed=0):
 
         elif type(item0) is float:
             self.dstns = [
-                self.engine(seed=self._rng.integers(0, 2**31 - 1), **conditional)
+                self.engine(seed=self._rng.integers(0, 2**31 - 1), **self.conditional)
             ]
 
         else:
@@ -284,6 +301,9 @@ def __init__(self, engine, conditional, RNG=None, seed=0):
             )
 
     def __getitem__(self, y):
+        # Prefer discrete list mode if present
+        if hasattr(self, "distributions") and self.distributions:
+            return self.distributions[y]
         return self.dstns[y]
 
     def discretize(self, N, **kwds):
@@ -302,16 +322,16 @@ def discretize(self, N, **kwds):
 
         Returns:
         ------------
-        dists : [DiscreteDistribution]
-            A list of DiscreteDistributions that are the
-            approximation of engine distribution under each condition.
-
-            TODO: It would be better if there were a conditional discrete
-            distribution representation. But that integrates with the
-            solution code. This implementation will return the list of
-            distributions representations expected by the solution code.
+        dists : [DiscreteDistribution] or IndexDistribution
+            If parameterization is constant, returns a single DiscreteDistribution.
+            If parameterization varies with index, returns an IndexDistribution in
+            discrete-list mode, wrapping the corresponding discrete distributions.
         """
 
+        # If already in discrete list mode, return self (already discretized)
+        if hasattr(self, "distributions") and self.distributions:
+            return self
+
         # test one item to determine case handling
         item0 = list(self.conditional.values())[0]
 
@@ -320,8 +340,12 @@ def discretize(self, N, **kwds):
             return self.dstns[0].discretize(N, **kwds)
 
         if type(item0) is list:
-            return TimeVaryingDiscreteDistribution(
-                [self[i].discretize(N, **kwds) for i, _ in enumerate(item0)]
+            # Return an IndexDistribution wrapping a list of discrete distributions
+            return IndexDistribution(
+                distributions=[
+                    self[i].discretize(N, **kwds) for i, _ in enumerate(item0)
+                ],
+                seed=self.seed,
             )
 
     def draw(self, condition):
@@ -345,6 +369,15 @@ def draw(self, condition):
         # are of the same type.
         # this matches the HARK 'time-varying' model architecture.
 
+        # If wrapping discrete distributions, draw from those
+        if hasattr(self, "distributions") and self.distributions:
+            draws = np.zeros(condition.size)
+            for c in np.unique(condition):
+                these = c == condition
+                N = np.sum(these)
+                draws[these] = self.distributions[c].draw(N)
+            return draws
+
         # test one item to determine case handling
         item0 = list(self.conditional.values())[0]
 
@@ -367,70 +400,6 @@ def draw(self, condition):
                 these = c == condition
                 N = np.sum(these)
 
-                cond = {key: val[c] for (key, val) in self.conditional.items()}
                 draws[these] = self[c].draw(N)
 
             return draws
-
-
-class TimeVaryingDiscreteDistribution(Distribution):
-    """
-    This class provides a way to define a discrete distribution that
-    is conditional on an index.
-
-    Wraps a list of discrete distributions.
-
-    Parameters
-    ----------
-
-    distributions : [DiscreteDistribution]
-        A list of discrete distributions
-
-    seed : int
-        Seed for random number generator.
-    """
-
-    distributions = []
-
-    def __init__(self, distributions, seed=0):
-        # Set up the RNG
-        super().__init__(seed)
-
-        self.distributions = distributions
-
-    def __getitem__(self, y):
-        return self.distributions[y]
-
-    def draw(self, condition):
-        """
-        Generate arrays of draws.
-        The input is an array containing the conditions.
-        The output is an array of the same length (axis 1 dimension)
-        as the conditions containing random draws of the conditional
-        distribution.
-
-        Parameters
-        ----------
-        condition : np.array
-            The input conditions to the distribution.
-
-        Returns:
-        ------------
-        draws : np.array
-        """
-        # for now, assume that all the conditionals
-        # are of the same type.
-        # this matches the HARK 'time-varying' model architecture.
-
-        # conditions are indices into list
-        # somewhat convoluted sampling strategy retained
-        # for test backwards compatibility
-        draws = np.zeros(condition.size)
-
-        for c in np.unique(condition):
-            these = c == condition
-            N = np.sum(these)
-
-            draws[these] = self.distributions[c].draw(N)
-
-        return draws

HARK/distributions/utils.py

@@ -4,7 +4,7 @@
 import numpy as np
 from scipy import stats
 
-from HARK.distributions.base import TimeVaryingDiscreteDistribution
+from HARK.distributions.base import IndexDistribution
 from HARK.distributions.discrete import (
     DiscreteDistribution,
     DiscreteDistributionLabeled,
@@ -265,10 +265,14 @@ def add_discrete_outcome_constant_mean(distribution, x, p, sort=False):
         Probability associated with each point in array of discrete
         points for discrete probability mass function.
     """
-    if type(distribution) == TimeVaryingDiscreteDistribution:
+    if (
+        isinstance(distribution, IndexDistribution)
+        and hasattr(distribution, "distributions")
+        and distribution.distributions
+    ):
         # apply recursively on all the internal distributions
-        return TimeVaryingDiscreteDistribution(
-            [
+        return IndexDistribution(
+            distributions=[
                 add_discrete_outcome_constant_mean(d, x, p)
                 for d in distribution.distributions
             ],

HARK/simulation/monte_carlo.py

@@ -10,7 +10,6 @@
 from HARK.distributions import (
     Distribution,
     IndexDistribution,
-    TimeVaryingDiscreteDistribution,
 )
 from HARK.model import Aggregate
 from HARK.model import DBlock
@@ -47,10 +46,7 @@ def draw_shocks(shocks: Mapping[str, Distribution], conditions: Sequence[int]):
             draws[shock_var] = np.ones(len(conditions)) * shock
         elif isinstance(shock, Aggregate):
             draws[shock_var] = shock.dist.draw(1)[0]
-        elif isinstance(shock, IndexDistribution) or isinstance(
-            shock, TimeVaryingDiscreteDistribution
-        ):
-            ## TODO  his type test is awkward. They should share a superclass.
+        elif isinstance(shock, IndexDistribution):
             draws[shock_var] = shock.draw(conditions)
         else:
             draws[shock_var] = shock.draw(len(conditions))

README.md

@@ -237,6 +237,12 @@ To install for development see the [Quickstart Guide](https://docs.econ-ark.org/
 For more information on contributing to HARK please see [the contributing guide](https://docs.econ-ark.org/docs/guides/contributing.html).
 This is the guide that collaborators follow in maintaining the Econ-ARK project.
 
+## Migration notes
+
+- Distributions: `TimeVaryingDiscreteDistribution` has been consolidated into `IndexDistribution`.
+  - If you previously constructed a time-indexed list of discrete distributions via `TimeVaryingDiscreteDistribution([...])`, use `IndexDistribution(distributions=[...])` instead.
+  - Existing `IndexDistribution(engine=..., conditional=...)` usage is unchanged.
+
 ## Disclaimer
 
 This is a beta version of HARK. The code has not been extensively tested as it should be. We hope it is useful, but there are absolutely no guarantees (expressed or implied) that it works or will do what you want. Use at your own risk. And please, let us know if you find bugs by posting an issue to [the GitHub page](https://github.com/econ-ark/HARK)!

docs/CHANGELOG.md

@@ -34,6 +34,7 @@ See documentation notebooks in /examples/SequenceSpaceJacobians/ . The capabilit
 - Documentation files have been moved from /Documentation/ to /docs/ [#1579](https://github.com/econ-ark/HARK/pull/1579)
 - All tests have been consolidated into a single directory, rather than being scattered about. [#1578](https://github.com/econ-ark/HARK/pull/1578)
 - Add a special README so that the robots know we're on their side when the singularity arrives. [#1577](https://github.com/econ-ark/HARK/pull/1577)
+- Consolidates `TimeVaryingDiscreteDistribution` into `IndexDistribution`. For time-varying discrete behavior, use `IndexDistribution(distributions=[...])`. [#1592](https://github.com/econ-ark/HARK/pull/1592)
 
 
 ### 0.16.0