Multivariate normal with discrete approximation

Documentation/CHANGELOG.md

@@ -17,6 +17,7 @@ Release Data: TBD
 * Adds a constructor for LogNormal distributions from mean and standard deviation [#891](https://github.com/econ-ark/HARK/pull/891/)
 * Uses new LogNormal constructor in ConsPortfolioModel [#891](https://github.com/econ-ark/HARK/pull/891/)
 * calcExpectations method for taking the expectation of a distribution over a function [#884](https://github.com/econ-ark/HARK/pull/884/] (#897)[https://github.com/econ-ark/HARK/pull/897/)
+* Implements the multivariate normal as a supported distribution, with a discretization method. See [#948](https://github.com/econ-ark/HARK/pull/948).
 * Centralizes the definition of value, marginal value, and marginal marginal value functions that use inverse-space
 interpolation for problems with CRRA utility. See [#888](https://github.com/econ-ark/HARK/pull/888).
 * MarkovProcess class used in ConsMarkovModel, ConsRepAgentModel, ConsAggShockModel [#902](https://github.com/econ-ark/HARK/pull/902) [#929](https://github.com/econ-ark/HARK/pull/929)

HARK/distribution.py

@@ -1,4 +1,5 @@
 from HARK.utilities import memoize
+from itertools import product
 import math
 import numpy as np
 from scipy.special import erf, erfc
@@ -298,7 +299,103 @@ def approx(self, N):
         return DiscreteDistribution(
             pmf, X, seed=self.RNG.randint(0, 2 ** 31 - 1, dtype="int32")
         )
+
+    def approx_equiprobable(self, N):
+
+        CDF = np.linspace(0,1,N+1)
+        lims = stats.norm.ppf(CDF)
+        scores = (lims - self.mu)/self.sigma
+        pdf = stats.norm.pdf(scores)
+
+        # Find conditional means using Mills's ratio
+        pmf = np.diff(CDF)
+        X = self.mu - np.diff(pdf)/pmf
 
+        return DiscreteDistribution(
+            pmf, X, seed=self.RNG.randint(0, 2 ** 31 - 1, dtype="int32")
+        )
+
+
+class MVNormal(Distribution):
+    """
+    A Multivariate Normal distribution.
+
+    Parameters
+    ----------
+    mu : numpy array
+        Mean vector.
+    Sigma : 2-d numpy array. Each dimension must have length equal to that of
+            mu.
+        Variance-covariance matrix.
+    seed : int
+        Seed for random number generator.
+    """
+
+    mu = None
+    Sigma = None
+
+    def __init__(self, mu = np.array([1,1]), Sigma = np.array([[1,0],[0,1]]), seed=0):
+        self.mu = mu
+        self.Sigma = Sigma
+        self.M = len(self.mu)
+        super().__init__(seed)
+
+    def draw(self, N):
+        """
+        Generate an array of multivariate normal draws.
+
+        Parameters
+        ----------
+        N : int
+            Number of multivariate draws.
+
+        Returns
+        -------
+        draws : np.array
+            Array of dimensions N x M containing the random draws, where M is
+            the dimension of the multivariate normal and N is the number of
+            draws. Each row represents a draw.
+        """
+        draws = self.RNG.multivariate_normal(self.mu, self.Sigma, N)
+
+        return draws
+
+    def approx(self, N):
+        """
+        Returns a discrete approximation of this distribution.
+
+        The discretization will have N**M points, where M is the dimension of
+        the multivariate normal.
+
+        It uses the fact that:
+            - Being positive definite, Sigma can be factorized as Sigma = QVQ',
+              with V diagonal. So, letting A=Q*sqrt(V), Sigma = A*A'.
+            - If Z is an N-dimensional multivariate standard normal, then
+              A*Z ~ N(0,A*A' = Sigma).
+
+        The idea therefore is to construct an equiprobable grid for a standard
+        normal and multiply it by matrix A.
+        """
+
+        # Start by computing matrix A.
+        v, Q = np.linalg.eig(self.Sigma)
+        sqrtV = np.diag(np.sqrt(v))
+        A = np.matmul(Q,sqrtV) 
+
+        # Now find a discretization for a univariate standard normal.
+        z_approx = Normal().approx_equiprobable(N)
+
+        # Now create the multivariate grid and pmf
+        Z = np.array(list(product(*[z_approx.X]*self.M)))
+        pmf = np.prod(np.array(list(product(*[z_approx.pmf]*self.M))),axis = 1)
+
+        # Apply mean and standard deviation to the Z grid
+        X = np.tile(np.reshape(self.mu, (1, self.M)), (N**self.M,1)) + np.matmul(Z, A.T)
+
+        # Construct and return discrete distribution
+        return DiscreteDistribution(
+            pmf, X, seed=self.RNG.randint(0, 2 ** 31 - 1, dtype="int32")
+        )
 
 class Weibull(Distribution):
     """

HARK/tests/test_approxDstns.py

@@ -33,3 +33,43 @@ def test_mu_lognormal_from_normal(self):
                     )
                     < 1e-12
                 )
+
+
+class test_MVNormalApprox(unittest.TestCase):
+    def setUp(self):
+
+        N = 5
+
+        # 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_approx = self.dist2D.approx(N)
+
+        # 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_approx = self.dist3D.approx(N)
+
+    def test_means(self):
+
+        mu_2D = distribution.calcExpectation(self.dist2D_approx)
+        self.assertTrue(np.allclose(mu_2D, self.mu2, rtol=1e-5))
+
+        mu_3D = distribution.calcExpectation(self.dist3D_approx)
+        self.assertTrue(np.allclose(mu_3D, self.mu3, rtol=1e-5))
+
+    def test_VCOV(self):
+
+        vcov_fun = lambda X, mu: np.outer(X - mu, X - mu)
+
+        Sig_2D = distribution.calcExpectation(self.dist2D_approx, vcov_fun, self.mu2)[
+            :, :, 0
+        ]
+        self.assertTrue(np.allclose(Sig_2D, self.Sigma2, rtol=1e-5))
+
+        Sig_3D = distribution.calcExpectation(self.dist3D_approx, vcov_fun, self.mu3)[
+            :, :, 0
+        ]
+        self.assertTrue(np.allclose(Sig_3D, self.Sigma3, rtol=1e-5))

HARK/tests/test_distribution.py

@@ -9,10 +9,11 @@
     Lognormal,
     MeanOneLogNormal,
     Normal,
+    MVNormal,
     Uniform,
     Weibull,
     calcExpectation,
-    combineIndepDstns
+    combineIndepDstns,
 )
 
 
@@ -24,18 +25,13 @@ class DiscreteDistributionTests(unittest.TestCase):
 
     def test_draw(self):
         self.assertEqual(
-            DiscreteDistribution(
-                np.ones(1),
-                np.zeros(1)).draw(1)[
-                0
-            ],
-            0,
+            DiscreteDistribution(np.ones(1), np.zeros(1)).draw(1)[0], 0,
         )
 
     def test_calcExpectation(self):
         dd_0_1_20 = Normal().approx(20)
-        dd_1_1_40 = Normal(mu = 1).approx(40)
-        dd_10_10_100 = Normal(mu = 10, sigma = 10).approx(100)
+        dd_1_1_40 = Normal(mu=1).approx(40)
+        dd_10_10_100 = Normal(mu=10, sigma=10).approx(100)
 
         ce1 = calcExpectation(dd_0_1_20)
         ce2 = calcExpectation(dd_1_1_40)
@@ -45,32 +41,20 @@ def test_calcExpectation(self):
         self.assertAlmostEqual(ce2, 1.0)
         self.assertAlmostEqual(ce3, 10.0)
 
-        ce4= calcExpectation(
-            dd_0_1_20,
-            lambda x: 2 ** x
-        )
+        ce4 = calcExpectation(dd_0_1_20, lambda x: 2 ** x)
 
         self.assertAlmostEqual(ce4, 1.27153712)
 
-        ce5 = calcExpectation(
-            dd_1_1_40,
-            lambda x: 2 * x
-        )
+        ce5 = calcExpectation(dd_1_1_40, lambda x: 2 * x)
 
         self.assertAlmostEqual(ce5, 2.0)
 
-        ce6 = calcExpectation(
-            dd_10_10_100,
-            lambda x, y: 2 * x + y,
-            20
-        )
+        ce6 = calcExpectation(dd_10_10_100, lambda x, y: 2 * x + y, 20)
 
         self.assertAlmostEqual(ce6, 40.0)
 
         ce7 = calcExpectation(
-            dd_0_1_20,
-            lambda x, y: x + y,
-            np.hstack(np.array([0,1,2,3,4,5]))
+            dd_0_1_20, lambda x, y: x + y, np.hstack(np.array([0, 1, 2, 3, 4, 5]))
         )
 
         self.assertAlmostEqual(ce7.flat[3], 3.0)
@@ -79,24 +63,19 @@ def test_calcExpectation(self):
         TranShkDstn = MeanOneLogNormal().approx(200)
         IncShkDstn = combineIndepDstns(PermShkDstn, TranShkDstn)
 
-        ce8 = calcExpectation(
-            IncShkDstn,
-            lambda X: X[0] + X[1]
-        )
+        ce8 = calcExpectation(IncShkDstn, lambda X: X[0] + X[1])
 
         self.assertAlmostEqual(ce8, 2.0)
 
         ce9 = calcExpectation(
             IncShkDstn,
             lambda X, a, r: r / X[0] * a + X[1],
-            np.array([0,1,2,3,4,5]), # an aNrmNow grid?
-            1.05 # an interest rate?
+            np.array([0, 1, 2, 3, 4, 5]),  # an aNrmNow grid?
+            1.05,  # an interest rate?
         )
 
-        self.assertAlmostEqual(
-            ce9[3],
-            9.518015322143837
-        )
+        self.assertAlmostEqual(ce9[3], 9.518015322143837)
+
 
 class DistributionClassTests(unittest.TestCase):
     """
@@ -128,27 +107,32 @@ def test_Normal(self):
 
         self.assertEqual(dist.draw(1)[0], 1.764052345967664)
 
+    def test_MVNormal(self):
+        dist = MVNormal()
+
+        self.assertTrue(
+            np.allclose(dist.draw(1)[0], np.array([2.76405235, 1.40015721]))
+        )
+
+        dist.draw(100)
+        dist.reset()
+
+        self.assertTrue(
+            np.allclose(dist.draw(1)[0], np.array([2.76405235, 1.40015721]))
+        )
+
     def test_Weibull(self):
-        self.assertEqual(
-            Weibull().draw(1)[0],
-            0.79587450816311)
+        self.assertEqual(Weibull().draw(1)[0], 0.79587450816311)
 
     def test_Uniform(self):
         uni = Uniform()
 
-        self.assertEqual(
-            Uniform().draw(1)[0],
-            0.5488135039273248)
+        self.assertEqual(Uniform().draw(1)[0], 0.5488135039273248)
 
-        self.assertEqual(
-            calcExpectation(uni.approx(10)),
-            0.5
-        )
+        self.assertEqual(calcExpectation(uni.approx(10)), 0.5)
 
     def test_Bernoulli(self):
-        self.assertEqual(
-            Bernoulli().draw(1)[0], False
-        )
+        self.assertEqual(Bernoulli().draw(1)[0], False)
 
 
 class MarkovProcessTests(unittest.TestCase):
@@ -158,9 +142,7 @@ class MarkovProcessTests(unittest.TestCase):
 
     def test_draw(self):
 
-        mrkv_array = np.array(
-            [[.75, .25],[0.1, 0.9]]
-        )
+        mrkv_array = np.array([[0.75, 0.25], [0.1, 0.9]])
 
         mp = distribution.MarkovProcess(mrkv_array)