DiscreteDistribution keeps link to parent ContinuousDistribution

Documentation/CHANGELOG.md

@@ -55,6 +55,7 @@ interpolation for problems with CRRA utility. See [#888](https://github.com/econ
 * remove Model.__call__; use Model init in Market and AgentType init to standardize on parameters dictionary [#947](https://github.com/econ-ark/HARK/issues/947)
 * Moves state MrkvNow to shocks['Mrkv'] in AggShockMarkov and KrusellSmith models [#935](https://github.com/econ-ark/HARK/pull/935)
 * Replaces `ConsIndShock`'s `init_lifecycle` with an actual life-cycle calibration [#951](https://github.com/econ-ark/HARK/pull/951).
+* New ContinuousDistribution class; DiscreteDistribution now tracks Continuous parent and approximation parameters [#954](https://github.com/econ-ark/HARK/pull/954)
 
 #### Minor Changes
 

HARK/ConsumptionSaving/ConsAggShockModel.py

@@ -1986,10 +1986,10 @@ def make_AggShkDstn(self):
         None
         """
         self.TranShkAggDstn = MeanOneLogNormal(sigma=self.TranShkAggStd).approx(
-            N=self.TranShkAggCount
+            n=self.TranShkAggCount
         )
         self.PermShkAggDstn = MeanOneLogNormal(sigma=self.PermShkAggStd).approx(
-            N=self.PermShkAggCount
+            n=self.PermShkAggCount
         )
         self.AggShkDstn = combine_indep_dstns(self.PermShkAggDstn, self.TranShkAggDstn)
 
@@ -2240,10 +2240,10 @@ def make_AggShkDstn(self):
         None
         """
         self.TranShkAggDstn = MeanOneLogNormal(sigma=self.TranShkAggStd).approx(
-            N=self.TranShkAggCount
+            n=self.TranShkAggCount
         )
         self.PermShkAggDstn = MeanOneLogNormal(sigma=self.PermShkAggStd).approx(
-            N=self.PermShkAggCount
+            n=self.PermShkAggCount
         )
         self.AggShkDstn = combine_indep_dstns(self.PermShkAggDstn, self.TranShkAggDstn)
 
@@ -2492,12 +2492,12 @@ def make_AggShkDstn(self):
         for i in range(StateCount):
             TranShkAggDstn.append(
                 MeanOneLogNormal(sigma=self.TranShkAggStd[i]).approx(
-                    N=self.TranShkAggCount
+                    n=self.TranShkAggCount
                 )
             )
             PermShkAggDstn.append(
                 MeanOneLogNormal(sigma=self.PermShkAggStd[i]).approx(
-                    N=self.PermShkAggCount
+                    n=self.PermShkAggCount
                 )
             )
             AggShkDstn.append(combine_indep_dstns(PermShkAggDstn[-1], TranShkAggDstn[-1]))

HARK/ConsumptionSaving/ConsMedModel.py

@@ -624,7 +624,7 @@ def update_med_shock_process(self):
             MedShkDstnNow = Lognormal(
                 mu=np.log(MedShkAvgNow) - 0.5 * MedShkStdNow ** 2, sigma=MedShkStdNow
             ).approx(
-                N=self.MedShkCount, tail_N=self.MedShkCountTail, tail_bound=[0, 0.9]
+                n=self.MedShkCount, tail_N=self.MedShkCountTail, tail_bound=[0, 0.9]
             )
             MedShkDstnNow = add_discrete_outcome_constant_mean(
                 MedShkDstnNow, 0.0, 0.0, sort=True

HARK/ConsumptionSaving/ConsPrefShockModel.py

@@ -124,7 +124,7 @@ def update_pref_shock_process(self):
             PrefShkStd = self.PrefShkStd[t]
             new_dstn = MeanOneLogNormal(
                 sigma=PrefShkStd, seed=self.RNG.randint(0, 2 ** 31 - 1)
-            ).approx(N=self.PrefShkCount, tail_N=self.PrefShk_tail_N,)
+            ).approx(n=self.PrefShkCount, tail_N=self.PrefShk_tail_N,)
             PrefShkDstn.append(new_dstn)
 
         # Store the preference shocks in self (time-varying) and restore time flow

HARK/core.py

@@ -1439,7 +1439,7 @@ def distribute_params(agent, param_name, param_count, distribution):
         will be split between the agents of the returned
         list in proportion to the given distribution.
     """
-    param_dist = distribution.approx(N=param_count)
+    param_dist = distribution.approx(n=param_count)
 
     agent_set = [deepcopy(agent) for i in range(param_count)]
 

HARK/distribution.py

@@ -174,8 +174,61 @@ def draw(self, condition):
 
 ### CONTINUOUS DISTRIBUTIONS
 
+class ContinuousDistribution(Distribution):
+    """
+    A continuous distribution.
 
-class Lognormal(Distribution):
+    Parameters
+    ----------
+    seed : int
+        Seed for random number generator.
+    """
+    def __init__(self, seed=0):
+        super().__init__(seed)
+
+    def pmf_X(self, n, **kwargs):
+        """
+        The point mass function and values for
+        the discretized version of this distribution.
+
+        This is a private method accessed by the public method
+        approx.
+
+        It should be overwritten for each subclass.
+
+        Parameters
+        ----------
+        n : int
+            Number of 
+
+        Returns
+        ----------
+        pmf : np.array
+            An array of floats representing a probability mass function.
+        X : np.array or [np.array]
+            Discrete point values for each probability mass.
+            May be multivariate (list of arrays). 
+        """
+        return None, None
+
+    def approx(self, n, **kwargs):
+        """
+        Returns a DiscreteDistribution that is an approximation
+        of this distribution.
+
+        n: int
+            Number of discrete points in the "main part" of the approximation.
+        """
+        pmf, X = self.pmf_X(n, **kwargs)
+        return DiscreteDistribution(
+            pmf,
+            X,
+            parent=self,
+            approx_args=(n, kwargs),
+            seed = self.RNG.randint(0, 2 ** 31 - 1, dtype="int32") # or 0? why is the seed 0 here?,
+        )
+
+class Lognormal(ContinuousDistribution):
     """
     A Lognormal distribution
 
@@ -238,7 +291,7 @@ def draw(self, N):
         # TODO: change return type to np.array?
         return draws[0] if len(draws) == 1 else draws
 
-    def approx(self, N, tail_N=0, tail_bound=None, tail_order=np.e):
+    def pmf_X(self, n, tail_N=0, tail_bound=None, tail_order=np.e):
         """
         Construct a discrete approximation to a lognormal distribution with underlying
         normal distribution N(mu,sigma).  Makes an equiprobable distribution by
@@ -247,7 +300,7 @@ def approx(self, N, tail_N=0, tail_bound=None, tail_order=np.e):
 
         Parameters
         ----------
-        N: int
+        n: int
             Number of discrete points in the "main part" of the approximation.
         tail_N: int
             Number of points in each "tail part" of the approximation; 0 = no tail.
@@ -261,9 +314,11 @@ def approx(self, N, tail_N=0, tail_bound=None, tail_order=np.e):
 
         Returns
         -------
-        d : DiscreteDistribution
-            Probability associated with each point in array of discrete
-            points for discrete probability mass function.
+        pmf : np.array
+            An array of floats representing a probability mass function.
+        X : np.array or [np.array]
+            Discrete point values for each probability mass.
+            May be multivariate (list of arrays).
         """
         tail_bound = tail_bound if tail_bound is not None else [0.02, 0.98]
         # Find the CDF boundaries of each segment
@@ -276,7 +331,7 @@ def approx(self, N, tail_N=0, tail_bound=None, tail_order=np.e):
                 hi_cut = 1.0
             inner_size = hi_cut - lo_cut
             inner_CDF_vals = [
-                lo_cut + x * N ** (-1.0) * inner_size for x in range(1, N)
+                lo_cut + x * n ** (-1.0) * inner_size for x in range(1, n)
             ]
             if inner_size < 1.0:
                 scale = 1.0 / tail_order
@@ -335,11 +390,10 @@ def approx(self, N, tail_N=0, tail_bound=None, tail_order=np.e):
                     )
 
         else:
-            pmf = np.ones(N) / N
-            X = np.exp(self.mu) * np.ones(N)
-        return DiscreteDistribution(
-            pmf, X, seed=self.RNG.randint(0, 2 ** 31 - 1, dtype="int32")
-        )
+            pmf = np.ones(n) / n
+            X = np.exp(self.mu) * np.ones(n)
+
+        return pmf, X
 
     @classmethod
     def from_mean_std(cls, mean, std, seed = 0):
@@ -380,7 +434,7 @@ def __init__(self, sigma=1.0, seed=0):
         mu = -0.5 * sigma ** 2
         super().__init__(mu=mu, sigma=sigma, seed=seed)
 
-class Normal(Distribution):
+class Normal(ContinuousDistribution):
     """
     A Normal distribution.
 
@@ -429,15 +483,24 @@ def draw(self, N):
 
         return draws
 
-    def approx(self, N):
+    def pmf_X(self, n, **kwargs):
         """
         Returns a discrete approximation of this distribution.
+
+        Returns
+        -------
+        pmf : np.array
+            An array of floats representing a probability mass function.
+        X : np.array or [np.array]
+            Discrete point values for each probability mass.
+            May be multivariate (list of arrays).
         """
-        x, w = np.polynomial.hermite.hermgauss(N)
+        x, w = np.polynomial.hermite.hermgauss(n)
         # normalize w
         pmf = w * np.pi ** -0.5
         # correct x
         X = math.sqrt(2.0) * self.sigma * x + self.mu
+
         return DiscreteDistribution(
             pmf, X, seed=self.RNG.randint(0, 2 ** 31 - 1, dtype="int32")
         )
@@ -542,7 +605,7 @@ def approx(self, N, equiprobable = False):
             pmf, X, seed=self.RNG.randint(0, 2 ** 31 - 1, dtype="int32")
         )
 
-class Weibull(Distribution):
+class Weibull(ContinuousDistribution):
     """
     A Weibull distribution.
 
@@ -600,7 +663,7 @@ def draw(self, N):
         return draws[0] if len(draws) == 1 else draws
 
 
-class Uniform(Distribution):
+class Uniform(ContinuousDistribution):
     """
     A Uniform distribution.
 
@@ -654,30 +717,30 @@ def draw(self, N):
             )
         return draws[0] if len(draws) == 1 else draws
 
-    def approx(self, N):
+    def pmf_X(self, n, **kwargs):
         """
         Makes a discrete approximation to this uniform distribution.
 
         Parameters
         ----------
-        N : int
+        n : int
             The number of points in the discrete approximation
 
         Returns
         -------
-        d : DiscreteDistribution
-            Probability associated with each point in array of discrete
-            points for discrete probability mass function.
+        pmf : np.array
+            An array of floats representing a probability mass function.
+        X : np.array or [np.array]
+            Discrete point values for each probability mass.
+            May be multivariate (list of arrays).
         """
-        pmf = np.ones(N) / float(N)
+        pmf = np.ones(n) / float(n)
         center = (self.top + self.bot) / 2.0
         width = (self.top - self.bot) / 2.0
-        X = center + width * np.linspace(-(N - 1.0) / 2.0, (N - 1.0) / 2.0, N) / (
-            N / 2.0
-        )
-        return DiscreteDistribution(
-            pmf, X, seed=self.RNG.randint(0, 2 ** 31 - 1, dtype="int32")
+        X = center + width * np.linspace(-(n - 1.0) / 2.0, (n - 1.0) / 2.0, n) / (
+            n / 2.0
         )
+        return pmf, X
 
 ### DISCRETE DISTRIBUTIONS
 
@@ -743,10 +806,14 @@ class DiscreteDistribution(Distribution):
 
     pmf = None
     X = None
+    parent = None
+    approx_args = None
 
-    def __init__(self, pmf, X, seed=0):
+    def __init__(self, pmf, X, parent=None, approx_args=None,seed=0):
         self.pmf = pmf
         self.X = X
+        self.parent = parent
+        self.approx_args = approx_args
         # Set up the RNG
         super().__init__(seed)
 

HARK/tests/test_distribution.py

@@ -113,6 +113,16 @@ def test_Lognormal(self):
 
         self.assertEqual(dist.draw(1)[0], 5.836039190663969)
 
+        self.assertEqual(
+            dist.approx(5).approx_args[0],
+            5
+        )
+
+        self.assertEqual(
+            dist.approx(10).parent.sigma,
+            1.0
+        )
+
     def test_Normal(self):
         dist = Normal()
 
@@ -146,7 +156,18 @@ def test_Uniform(self):
         self.assertEqual(Uniform().draw(1)[0], 0.5488135039273248)
 
         self.assertEqual(
-            calc_expectation(uni.approx(10)),
+            uni.approx(10).approx_args[0],
+            10
+        )
+
+        self.assertEqual(
+            uni.approx(10).parent.top,
+            1.0
+        )
+
+        self.assertEqual(
+            calcExpectation(uni.approx(10)),
+
             0.5
         )