Remove TimeVaryingDiscreteDistribution, extend IndexDistribution functionality

Co-authored-by: alanlujan91 <[email protected]>

Author: cursoragent

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,
@@ -50,7 +48,7 @@
     approx_lognormal_gauss_hermite,
     calc_expectation,
     calc_lognormal_style_pars_from_normal_pars,
-    calc_normal_style_pars_from_lognormal_pars,
+    calc_lognormal_style_pars_from_lognormal_pars,
     combine_indep_dstns,
     distr_of_function,
     expected,

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,17 @@ 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 +257,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 +288,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 +299,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 +320,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 +338,10 @@ 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 +365,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 +396,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,10 @@ 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))