Add infimum and supremum information to limit dicts

The limit attribute of a DiscreteDistribution has data about the continuous distribution from which it was derived. This dictionary now contains infimum and supremum fields, which are also derived from the original distribution. It's present for *at least* the income distributions used by ConsIndShock.

HARK/Calibration/Income/IncomeProcesses.py

@@ -75,14 +75,27 @@ class LognormPermIncShk(DiscreteDistribution):
 
     def __init__(self, sigma, n_approx, neutral_measure=False, seed=0):
         # Construct an auxiliary discretized normal
-        logn_approx = MeanOneLogNormal(sigma).discretize(
+        lognormal_dstn = MeanOneLogNormal(sigma)
+        logn_approx = lognormal_dstn.discretize(
             n_approx if sigma > 0.0 else 1, method="equiprobable", tail_N=0
         )
+
+        limit = {
+            "dist": lognormal_dstn,
+            "method": "equiprobable",
+            "N": n_approx,
+            "endpoints": False,
+            "infimum": logn_approx.limit["infimum"],
+            "supremum": logn_approx.limit["supremum"],
+        }
+
         # Change the pmv if necessary
         if neutral_measure:
             logn_approx.pmv = (logn_approx.atoms * logn_approx.pmv).flatten()
 
-        super().__init__(pmv=logn_approx.pmv, atoms=logn_approx.atoms, seed=seed)
+        super().__init__(
+            pmv=logn_approx.pmv, atoms=logn_approx.atoms, limit=limit, seed=seed
+        )
 
 
 class MixtureTranIncShk(DiscreteDistribution):
@@ -111,15 +124,22 @@ class MixtureTranIncShk(DiscreteDistribution):
     """
 
     def __init__(self, sigma, UnempPrb, IncUnemp, n_approx, seed=0):
-        dstn_approx = MeanOneLogNormal(sigma).discretize(
+        dstn_orig = MeanOneLogNormal(sigma)
+        dstn_approx = dstn_orig.discretize(
             n_approx if sigma > 0.0 else 1, method="equiprobable", tail_N=0
         )
+
         if UnempPrb > 0.0:
             dstn_approx = add_discrete_outcome_constant_mean(
                 dstn_approx, p=UnempPrb, x=IncUnemp
             )
 
-        super().__init__(pmv=dstn_approx.pmv, atoms=dstn_approx.atoms, seed=seed)
+        super().__init__(
+            pmv=dstn_approx.pmv,
+            atoms=dstn_approx.atoms,
+            limit=dstn_approx.limit,
+            seed=seed,
+        )
 
 
 class MixtureTranIncShk_HANK(DiscreteDistribution):
@@ -240,6 +260,7 @@ def __init__(
             var_names=["PermShk", "TranShk"],
             pmv=joint_dstn.pmv,
             atoms=joint_dstn.atoms,
+            limit=joint_dstn.limit,
             seed=seed,
         )
 

HARK/distribution.py

@@ -53,6 +53,10 @@ def __init__(self, seed: Optional[int] = 0) -> None:
         self._seed: int = _seed
         self._rng: random.Generator = random.default_rng(self._seed)
 
+        # Bounds of distribution support should be overwritten by subclasses
+        self.infimum = np.array([])
+        self.supremum = np.array([])
+
     @property
     def seed(self) -> int:
         """
@@ -139,7 +143,7 @@ def discretize(
 
         Returns
         -------
-        DiscreteDistribution
+        discretized_dstn : DiscreteDistribution
             Discretized distribution.
 
         Raises
@@ -158,7 +162,10 @@ def discretize(
             )
 
         approx = getattr(self, approx_method)
-        return approx(N, endpoints, **kwds)
+        discretized_dstn = approx(N, endpoints, **kwds)
+        discretized_dstn.limit["infimum"] = self.infimum.copy()
+        discretized_dstn.limit["supremum"] = self.supremum.copy()
+        return discretized_dstn
 
 
 # CONTINUOUS DISTRIBUTIONS
@@ -214,6 +221,9 @@ def __init__(self, mu=0.0, sigma=1.0, seed=0):
 
         super().__init__(stats.norm, loc=mu, scale=sigma, seed=seed)
 
+        self.infimum = -np.inf * np.ones(self.mu.size)
+        self.supremum = np.inf * np.ones(self.mu.size)
+
     def discretize(self, N, method="hermite", endpoints=False):
         """
         For normal distributions, the Gauss-Hermite quadrature rule is
@@ -227,8 +237,6 @@ def _approx_hermite(self, N, endpoints=False):
         Returns a discrete approximation of this distribution
         using the Gauss-Hermite quadrature rule.
 
-        TODO: add endpoints option
-
         Parameters
         ----------
         N : int
@@ -259,8 +267,6 @@ def _approx_equiprobable(self, N, endpoints=False):
         """
         Returns a discrete equiprobable approximation of this distribution.
 
-        TODO: add endpoints option
-
         Parameters
         ----------
         N : int
@@ -371,6 +377,9 @@ def __init__(
             stats.lognorm, s=self.sigma, scale=np.exp(self.mu), loc=0, seed=seed
         )
 
+        self.infimum = np.array([0.0])
+        self.supremum = np.array([np.inf])
+
     def _approx_equiprobable(
         self, N, endpoints=False, tail_N=0, tail_bound=None, tail_order=np.e
     ):
@@ -594,6 +603,9 @@ def __init__(self, bot=0.0, top=1.0, seed=0):
             stats.uniform, loc=self.bot, scale=self.top - self.bot, seed=seed
         )
 
+        self.infimum = self.bot
+        self.supremum = self.top
+
     def _approx_equiprobable(self, N, endpoints=False):
         """
         Makes a discrete approximation to this uniform distribution.
@@ -653,9 +665,13 @@ class Weibull(ContinuousFrozenDistribution):
     def __init__(self, scale=1.0, shape=1.0, seed=0):
         self.scale = np.asarray(scale)
         self.shape = np.asarray(shape)
+
         # Set up the RNG
         super().__init__(stats.weibull_min, c=shape, scale=scale, seed=seed)
 
+        self.infimum = np.array([0.0])
+        self.supremum = np.array([np.inf])
+
 
 # MULTIVARIATE DISTRIBUTIONS
 
@@ -682,6 +698,9 @@ def __init__(self, mu=[1, 1], Sigma=[[1, 0], [0, 1]], seed=0):
         multivariate_normal_frozen.__init__(self, mean=self.mu, cov=self.Sigma)
         Distribution.__init__(self, seed=seed)
 
+        self.infimum = -np.inf * np.ones(self.M)
+        self.supremum = np.inf * np.ones(self.M)
+
     def discretize(self, N, method="hermite", endpoints=False):
         """
         For multivariate normal distributions, the Gauss-Hermite
@@ -786,6 +805,8 @@ def __init__(self, p=0.5, seed=0):
         self.pmv = [1 - self.p, self.p]
         self.atoms = [0, 1]
         self.limit = {"dist": self}
+        self.infimum = np.array([0.0])
+        self.supremum = np.array([1.0])
 
 
 class DiscreteDistribution(Distribution):
@@ -801,6 +822,10 @@ class DiscreteDistribution(Distribution):
         For multivariate distributions, the last dimension of atoms must index
         "atom" or the random realization. For instance, if atoms.shape == (2,6,4),
         the random variable has 4 possible realizations and each of them has shape (2,6).
+    limit : dict
+        Dictionary with information about the continuous distribution from which
+        this distribution was generated. The reference distribution is in the entry
+        called 'dist'.
     seed : int
         Seed for random number generator.
     """
@@ -1526,6 +1551,7 @@ def draw(self, condition):
         return draws
 
 
+# TODO: This function does not generate the limit attribute
 def approx_lognormal_gauss_hermite(N, mu=0.0, sigma=1.0, seed=0):
     d = Normal(mu, sigma).discretize(N, method="hermite")
     return DiscreteDistribution(d.pmv, np.exp(d.atoms), seed=seed)
@@ -1763,8 +1789,17 @@ 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:
+        # apply recursively on all the internal distributions
+        return TimeVaryingDiscreteDistribution(
+            [
+                add_discrete_outcome_constant_mean(d, x, p)
+                for d in distribution.distributions
+            ],
+            seed=distribution.seed,
+        )
 
-    if type(distribution) != TimeVaryingDiscreteDistribution:
+    else:
         atoms = np.append(x, distribution.atoms * (1 - p * x) / (1 - p))
         pmv = np.append(p, distribution.pmv * (1 - p))
 
@@ -1773,16 +1808,37 @@ def add_discrete_outcome_constant_mean(distribution, x, p, sort=False):
             atoms = atoms[indices]
             pmv = pmv[indices]
 
-        return DiscreteDistribution(pmv, atoms)
-    elif type(distribution) == TimeVaryingDiscreteDistribution:
-        # apply recursively on all the internal distributions
-        return TimeVaryingDiscreteDistribution(
-            [
-                add_discrete_outcome_constant_mean(d, x, p)
-                for d in distribution.distributions
-            ],
-            seed=distribution.seed,
-        )
+        # Update infimum and supremum
+        temp_x = np.array(x, ndmin=1)
+        try:
+            infimum = np.array(
+                [
+                    np.minimum(temp_x[i], distribution.limit["infimum"][i])
+                    for i in range(temp_x.size)
+                ]
+            )
+        except:
+            infimum = np.min(atoms, axis=-1, keepdims=True)
+        try:
+            supremum = np.array(
+                [
+                    np.maximum(temp_x[i], distribution.limit["supremum"][i])
+                    for i in range(temp_x.size)
+                ]
+            )
+        except:
+            supremum = np.max(atoms, axis=-1, keepdims=True)
+
+        limit = {
+            "dist": distribution,
+            "method": "add_discrete_outcome_constant_mean",
+            "x": x,
+            "p": p,
+            "infimum": infimum,
+            "supremum": supremum,
+        }
+
+        return DiscreteDistribution(pmv, atoms, seed=distribution.seed, limit=limit)
 
 
 def add_discrete_outcome(distribution, x, p, sort=False):
@@ -1814,7 +1870,38 @@ def add_discrete_outcome(distribution, x, p, sort=False):
         atoms = atoms[indices]
         pmv = pmv[indices]
 
-    return DiscreteDistribution(pmv, atoms)
+    # Update infimum and supremum
+    # Update infimum and supremum
+    temp_x = np.array(x, ndmin=1)
+    try:
+        infimum = np.array(
+            [
+                np.minimum(temp_x[i], distribution.limit["infimum"][i])
+                for i in range(temp_x.size)
+            ]
+        )
+    except:
+        infimum = np.min(atoms, axis=-1, keepdims=True)
+    try:
+        supremum = np.array(
+            [
+                np.maximum(temp_x[i], distribution.limit["supremum"][i])
+                for i in range(temp_x.size)
+            ]
+        )
+    except:
+        supremum = np.max(atoms, axis=-1, keepdims=True)
+
+    limit = {
+        "dist": distribution,
+        "method": "add_discrete_outcome",
+        "x": x,
+        "p": p,
+        "infimum": infimum,
+        "supremum": supremum,
+    }
+
+    return DiscreteDistribution(pmv, atoms, seed=distribution.seed, limit=limit)
 
 
 def combine_indep_dstns(*distributions, seed=0):
@@ -1827,6 +1914,8 @@ def combine_indep_dstns(*distributions, seed=0):
     distributions : DiscreteDistribution
         Arbitrary number of discrete distributionss to combine. Their realizations must be
         vector-valued (for each D in distributions, it must be the case that len(D.dim())==1).
+    seed : int, optional
+        Value to use as the RNG seed for the combined distribution, default is 0.
 
     Returns
     -------
@@ -1854,7 +1943,7 @@ def combine_indep_dstns(*distributions, seed=0):
         else:
             var_labels += tuple([""] * dist.dim()[0])
 
-    number_of_distributions = len(distributions)
+    dstn_count = len(distributions)
 
     all_labeled = all(dist_is_labeled)
     labels_are_unique = len(var_labels) == len(set(var_labels))
@@ -1878,19 +1967,36 @@ def combine_indep_dstns(*distributions, seed=0):
 
     assert np.isclose(np.sum(P_out), 1), "Probabilities do not sum to 1!"
 
+    # Make the limit dictionary
+    infimum = np.concatenate(
+        [distributions[i].limit["infimum"] for i in range(dstn_count)]
+    )
+    supremum = np.concatenate(
+        [distributions[i].limit["supremum"] for i in range(dstn_count)]
+    )
+    limit = {
+        "dist": distributions,
+        "method": "combine_indep_dstns",
+        "infimum": infimum,
+        "supremum": supremum,
+    }
+    # except:
+    #    limit=None
+
     if all_labeled and labels_are_unique:
         combined_dstn = DiscreteDistributionLabeled(
             pmv=P_out,
             atoms=atoms_out,
             var_names=var_labels,
+            limit=limit,
             seed=seed,
         )
     else:
         if all_labeled and not labels_are_unique:
             warn(
                 "There are duplicated labels in the provided distributions. Returning a non-labeled combination"
             )
-        combined_dstn = DiscreteDistribution(P_out, atoms_out, seed=seed)
+        combined_dstn = DiscreteDistribution(P_out, atoms_out, limit=limit, seed=seed)
 
     return combined_dstn