Recreate PR #1507

HARK/distributions/__init__.py

@@ -9,8 +9,8 @@
     "Normal",
     "Weibull",
     "Bernoulli",
-    "MVLogNormal",
-    "MVNormal",
+    "MultivariateLogNormal",
+    "MultivariateNormal",
     "approx_beta",
     "approx_lognormal_gauss_hermite",
     "calc_expectation",
@@ -43,7 +43,7 @@
     DiscreteDistribution,
     DiscreteDistributionLabeled,
 )
-from HARK.distributions.multivariate import MVLogNormal, MVNormal
+from HARK.distributions.multivariate import MultivariateLogNormal, MultivariateNormal
 from HARK.distributions.utils import (
     add_discrete_outcome_constant_mean,
     approx_beta,

HARK/distributions/multivariate.py

@@ -13,7 +13,7 @@
 # MULTIVARIATE DISTRIBUTIONS
 
 
-class MVNormal(multivariate_normal_frozen, Distribution):
+class MultivariateNormal(multivariate_normal_frozen, Distribution):
     """
     A Multivariate Normal distribution.
 
@@ -92,7 +92,7 @@ def _approx(self, N, method="hermite", endpoints=False):
         )
 
 
-class MVLogNormal(multi_rv_frozen, Distribution):
+class MultivariateLogNormal(multi_rv_frozen, Distribution):
     """
     A Multivariate Lognormal distribution.
 
@@ -161,7 +161,7 @@ def _cdf(self, x: Union[list, np.ndarray]):
         if (x.shape != self.M) & (x.shape[1] != self.M):
             raise ValueError(f"x must be and {self.M}-dimensional input")
 
-        return MVNormal(mu=self.mu, Sigma=self.Sigma).cdf(np.log(x))
+        return MultivariateNormal(mu=self.mu, Sigma=self.Sigma).cdf(np.log(x))
 
     def _pdf(self, x: Union[list, np.ndarray]):
         """
@@ -267,7 +267,7 @@ def rvs(self, size: int = 1, random_state=None):
             Random sample from the distribution.
         """
 
-        Z = MVNormal(mu=self.mu, Sigma=self.Sigma)
+        Z = MultivariateNormal(mu=self.mu, Sigma=self.Sigma)
 
         return np.exp(Z.rvs(size, random_state=random_state))
 
@@ -310,11 +310,11 @@ def _approx_equiprobable(
             )
 
         if np.array_equal(self.Sigma, np.diag(np.diag(self.Sigma))):
-            ind_atoms = np.empty((self.M, N))
+            ind_atoms = np.empty((self.M, N + 2 * tail_N))
 
             for i in range(self.M):
                 if self.Sigma[i, i] == 0.0:
-                    x_atoms = np.repeat(np.exp(self.mu[i]), N)
+                    x_atoms = np.repeat(np.exp(self.mu[i]), N + 2 * tail_N)
                     ind_atoms[i] = x_atoms
                 else:
                     x_atoms = (
@@ -330,7 +330,22 @@ def _approx_equiprobable(
             atoms = np.stack(
                 [ar.flatten() for ar in list(np.meshgrid(*atoms_list))], axis=1
             ).T
-            pmv = np.repeat(1 / (N**self.M), N**self.M)
+
+            interiors = np.empty([self.M, (N + 2 * tail_N) ** (self.M)])
+
+            inners = np.zeros(N + 2 * tail_N)
+
+            if tail_N > 0:
+                inners[:tail_N] = [(tail_N - i) for i in range(tail_N)]
+                inners[-tail_N:] = [(i + 1) for i in range(tail_N)]
+
+            for i in range(self.M):
+                inners_i = [inners for _ in range((N + 2 * tail_N) ** i)]
+
+                interiors[i] = np.repeat(
+                    [*inners_i], (N + 2 * tail_N) ** (self.M - (i + 1))
+                )
+
         else:
             if tail_bound is not None:
                 if type(tail_bound) is float:
@@ -371,10 +386,10 @@ def eval(params, z):
                 excl = []
 
                 for j in range(len(z)):
-                    if z[j, 0] != z[j, 1]:
-                        inds.append(j)
-                    else:
+                    if z[j, 0] == z[j, 1]:
                         excl.append(j)
+                    elif params[j] != 0.0:
+                        inds.append(j)
 
                 dim = len(inds)
 
@@ -458,21 +473,21 @@ def eval(params, z):
 
                     atoms[i] = xi_atoms
 
-            max_locs = np.argmax(np.abs(interiors), axis=0)
+        max_locs = np.argmax(np.abs(interiors), axis=0)
 
-            max_inds = np.stack([max_locs, np.arange(len(max_locs))], axis=1)
+        max_inds = np.stack([max_locs, np.arange(len(max_locs))], axis=1)
 
-            prob_locs = interiors[max_inds[:, 0], max_inds[:, 1]]
+        prob_locs = interiors[max_inds[:, 0], max_inds[:, 1]]
 
-            def prob_assign(x):
-                if x == 0:
-                    return 1 / (N**self.M)
-                else:
-                    return 0.0
+        def prob_assign(x):
+            if x == 0:
+                return 1 / (N**self.M)
+            else:
+                return 0.0
 
-            prob_vec = np.vectorize(prob_assign)
+        prob_vec = np.vectorize(prob_assign)
 
-            pmv = prob_vec(prob_locs)
+        pmv = prob_vec(prob_locs)
 
         limit = {
             "dist": self,

HARK/tests/test_approxDstns.py

@@ -46,13 +46,13 @@ def setUp(self):
         # 2-D distribution
         self.mu2 = np.array([5, -10])
         self.Sigma2 = np.array([[2, -0.6], [-0.6, 1]])
-        self.dist2D = distribution.MVNormal(self.mu2, self.Sigma2)
+        self.dist2D = distribution.MultivariateNormal(self.mu2, self.Sigma2)
         self.dist2D_approx = self.dist2D.discretize(N, method="hermite")
 
         # 3-D Distribution
         self.mu3 = np.array([5, -10, 0])
         self.Sigma3 = np.array([[2, -0.6, 0.1], [-0.6, 1, 0.2], [0.1, 0.2, 3]])
-        self.dist3D = distribution.MVNormal(self.mu3, self.Sigma3)
+        self.dist3D = distribution.MultivariateNormal(self.mu3, self.Sigma3)
         self.dist3D_approx = self.dist3D.discretize(N, method="hermite")
 
     def test_means(self):

HARK/tests/test_distribution.py

@@ -11,8 +11,6 @@
     Lognormal,
     MarkovProcess,
     MeanOneLogNormal,
-    MVNormal,
-    MVLogNormal,
     Normal,
     Uniform,
     Weibull,
@@ -23,6 +21,8 @@
     distr_of_function,
     expected,
     make_tauchen_ar1,
+    MultivariateNormal,
+    MultivariateLogNormal,
 )
 from HARK.tests import HARK_PRECISION
 
@@ -306,9 +306,9 @@ def test_Normal(self):
 
         dist.draw(1)[0]
 
-    def test_MVNormal(self):
+    def test_MultivariateNormal(self):
         # Are these tests generator/backend specific?
-        dist = MVNormal()
+        dist = MultivariateNormal()
 
         # self.assertTrue(
         #    np.allclose(dist.draw(1)[0], np.array([2.76405, 1.40016]))
@@ -321,8 +321,8 @@ def test_MVNormal(self):
         #    np.allclose(dist.draw(1)[0], np.array([2.76405, 1.40016]))
         # )
 
-    def test_MVLogNormal(self):
-        dist = MVLogNormal()
+    def test_MultivariateLogNormal(self):
+        dist = MultivariateLogNormal()
 
         dist.draw(100)
         dist.reset()