Merge remote-tracking branch 'upstream/master' into cism_Rfree

Documentation/CHANGELOG.md

@@ -31,6 +31,7 @@ Release Date: TBD
 * Option to have newborn IndShockConsumerType agents with a transitory income shock in the first period. Default is false, meaning they only have a permanent income shock in period 1 and permanent AND transitory in the following ones. [#1126](https://github.com/econ-ark/HARK/pull/1126)
 * Adds `benchmark` utility to profile the performance of `HARK` solvers. [#1131](https://github.com/econ-ark/HARK/pull/1131)
 * Fixes scaling bug in Normal equiprobable approximation method. [1139](https://github.com/econ-ark/HARK/pull/1139)
+* Removes the extra-dimension that was returned by `calc_expectations` in some instances. [#1149](https://github.com/econ-ark/HARK/pull/1149)
 
 ### 0.12.0
 

HARK/ConsumptionSaving/ConsGenIncProcessModel.py

@@ -354,7 +354,7 @@ def vp_next(shocks, a_lvl, p_lvl):
             * self.Rfree
             * calc_expectation(self.IncShkDstn, vp_next, self.aLvlNow, self.pLvlNow)
         )
-        EndOfPrdvP = EndOfPrdvP[:, :, 0]
+
         return EndOfPrdvP
 
     def make_EndOfPrdvFunc(self, EndOfPrdvP):
@@ -382,7 +382,7 @@ def v_lvl_next(shocks, a_lvl, p_lvl):
         vLvlNext = calc_expectation(
             self.IncShkDstn, v_lvl_next, self.aLvlNow, self.pLvlNow
         )
-        vLvlNext = vLvlNext[:, :, 0]
+
         # expected value, averaging across states
         EndOfPrdv = self.DiscFacEff * vLvlNext
         # value transformed through inverse utility
@@ -719,7 +719,7 @@ def vpp_next(shocks, a_lvl, p_lvl):
             * self.Rfree
             * calc_expectation(self.IncShkDstn, vpp_next, self.aLvlNow, self.pLvlNow)
         )
-        EndOfPrdvPP = EndOfPrdvPP[:, :, 0]
+
 
         dcda = EndOfPrdvPP / self.uPP(np.array(cLvl[1:, 1:]))
         MPC = dcda / (dcda + 1.0)

HARK/ConsumptionSaving/ConsIndShockModel.py

@@ -2512,7 +2512,7 @@ def check_FVAC(self, verbose=None):
         """
         EpShkuInv = calc_expectation(
             self.PermShkDstn[0], lambda x: x ** (1 - self.CRRA)
-        )
+        )[0]
 
         if self.CRRA != 1.0:
             uInvEpShkuInv = EpShkuInv ** (
@@ -2576,7 +2576,7 @@ def check_conditions(self, verbose=None):
         # would be referenced below as:
         # [url]/#Uncertainty-Modified-Conditions
 
-        self.Ex_PermShkInv = calc_expectation(self.PermShkDstn[0], lambda x: 1 / x)
+        self.Ex_PermShkInv = calc_expectation(self.PermShkDstn[0], lambda x: 1 / x)[0]
         # $\Ex_{t}[\psi^{-1}_{t+1}]$ (in first eqn in sec)
 
         # [url]/#Pat, adjusted to include mortality

HARK/ConsumptionSaving/ConsPortfolioModel.py

@@ -658,9 +658,7 @@ def dvds_dist(shocks, b_nrm, Share_next):
         dvdb_intermed = calc_expectation(
             self.IncShkDstn, dvdb_dist, self.bNrmNext, self.ShareNext
         )
-        # calc_expectation returns one additional "empty" dimension, remove it
-        # this line can be deleted when calc_expectation is fixed
-        dvdb_intermed = dvdb_intermed[:, :, 0]
+
         dvdbNvrs_intermed = self.uPinv(dvdb_intermed)
         dvdbNvrsFunc_intermed = BilinearInterp(
             dvdbNvrs_intermed, self.bNrmGrid, self.ShareGrid
@@ -671,9 +669,7 @@ def dvds_dist(shocks, b_nrm, Share_next):
         dvds_intermed = calc_expectation(
             self.IncShkDstn, dvds_dist, self.bNrmNext, self.ShareNext
         )
-        # calc_expectation returns one additional "empty" dimension, remove it
-        # this line can be deleted when calc_expectation is fixed
-        dvds_intermed = dvds_intermed[:, :, 0]
+
         dvdsFunc_intermed = BilinearInterp(dvds_intermed, self.bNrmGrid, self.ShareGrid)
 
         # Make tiled arrays to calculate future realizations of bNrm and Share when integrating over RiskyDstn
@@ -711,9 +707,7 @@ def EndOfPrddvds_dist(shock, a_nrm, Share_next):
                 self.RiskyDstn, EndOfPrddvda_dist, self.aNrm_tiled, self.ShareNext
             )
         )
-        # calc_expectation returns one additional "empty" dimension, remove it
-        # this line can be deleted when calc_expectation is fixed
-        self.EndOfPrddvda = self.EndOfPrddvda[:, :, 0]
+
         self.EndOfPrddvdaNvrs = self.uPinv(self.EndOfPrddvda)
 
         # Calculate end-of-period marginal value of risky portfolio share by taking expectations
@@ -724,9 +718,6 @@ def EndOfPrddvds_dist(shock, a_nrm, Share_next):
                 self.RiskyDstn, EndOfPrddvds_dist, self.aNrm_tiled, self.ShareNext
             )
         )
-        # calc_expectation returns one additional "empty" dimension, remove it
-        # this line can be deleted when calc_expectation is fixed
-        self.EndOfPrddvds = self.EndOfPrddvds[:, :, 0]
 
     def optimize_share(self):
         """
@@ -877,9 +868,7 @@ def v_intermed_dist(shocks, b_nrm, Share_next):
         v_intermed = calc_expectation(
             self.IncShkDstn, v_intermed_dist, self.bNrmNext, self.ShareNext
         )
-        # calc_expectation returns one additional "empty" dimension, remove it
-        # this line can be deleted when calc_expectation is fixed
-        v_intermed = v_intermed[:, :, 0]
+
         vNvrs_intermed = self.uinv(v_intermed)
         vNvrsFunc_intermed = BilinearInterp(
             vNvrs_intermed, self.bNrmGrid, self.ShareGrid
@@ -903,9 +892,7 @@ def EndOfPrdv_dist(shock, a_nrm, Share_next):
                 self.RiskyDstn, EndOfPrdv_dist, self.aNrm_tiled, self.ShareNext
             )
         )
-        # calc_expectation returns one additional "empty" dimension, remove it
-        # this line can be deleted when calc_expectation is fixed
-        self.EndOfPrdv = self.EndOfPrdv[:, :, 0]
+
         self.EndOfPrdvNvrs = self.uinv(self.EndOfPrdv)
 
     def make_vFunc(self):
@@ -1218,7 +1205,7 @@ def EndOfPrddvds_dist(shocks, a_nrm, shares):
                 self.ShockDstn, EndOfPrddvda_dists, self.aNrm_tiled, self.Share_tiled
             )
         )
-        self.EndOfPrddvda = self.EndOfPrddvda[:, :, 0]
+
         self.EndOfPrddvdaNvrs = self.uPinv(self.EndOfPrddvda)
 
         # Calculate end-of-period marginal value of risky portfolio share by taking expectations
@@ -1229,7 +1216,6 @@ def EndOfPrddvds_dist(shocks, a_nrm, shares):
                 self.ShockDstn, EndOfPrddvds_dist, self.aNrm_tiled, self.Share_tiled
             )
         )
-        self.EndOfPrddvds = self.EndOfPrddvds[:, :, 0]
 
     def make_EndOfPrdvFunc(self):
         """
@@ -1257,7 +1243,7 @@ def v_dist(shocks, a_nrm, shares):
                 self.ShockDstn, v_dist, self.aNrm_tiled, self.Share_tiled
             )
         )
-        self.EndOfPrdv = self.EndOfPrdv[:, :, 0]
+
         self.EndOfPrdvNvrs = self.uinv(self.EndOfPrdv)
 
     def solve(self):

HARK/ConsumptionSaving/ConsRiskyContribModel.py

@@ -1217,7 +1217,7 @@ def post_return_derivs(inc_shocks, b_aux, g_aux, s):
         # next period's derivatives and value.
         pr_derivs = calc_expectation(
             IncShkDstn, post_return_derivs, b_aux_tiled, g_aux_tiled, Share_tiled
-        )[:, :, :, :, 0]
+        )
 
         # Unpack results and create interpolators
         pr_dvdb_func = MargValueFuncCRRA(
@@ -1363,7 +1363,7 @@ def end_of_period_derivs(shocks, a, nTil, s):
         aNrm_tiled,
         nNrm_tiled,
         Share_tiled,
-    )[:, :, :, :, 0]
+    )
 
     # Unpack results
     eop_dvdaNvrs = uPinv(eop_derivs[0])

HARK/distribution.py

@@ -1292,7 +1292,8 @@ def calc_expectation(dstn, func=lambda x: x, *args):
         The distribution over which the function is to be evaluated.
     func : function
         The function to be evaluated.
-        This function should take an array of shape dstn.dim().
+        This function should take an array of shape dstn.dim() and return
+        either arrays of arbitrary shape or scalars.
         It may also take other arguments *args.
     *args :
         Other inputs for func, representing the non-stochastic arguments.
@@ -1311,17 +1312,12 @@ def calc_expectation(dstn, func=lambda x: x, *args):
 
     f_query = np.stack(f_query, axis=-1)
 
-    f_exp = np.dot(f_query, np.vstack(dstn.pmf))
-
-    # TODO: f_exp will have an extra dimension of length 1 at the end
-    # because of the way numpy handles dot products. It would be good
-    # either keep or eliminate that dimension _always_ and not only
-    # when the result has some particular shape as is done below. Keeping
-    # it as is because some models have come to expect it.
-    if f_exp.size == 1:
-        f_exp = f_exp.flat[0]
-    elif f_exp.shape[0] == f_exp.size:
-        f_exp = f_exp.flatten()
+    # From the numpy.dot documentation:
+    # If both a and b are 1-D arrays, it is inner product of vectors (without complex conjugation).
+    # If a is an N-D array and b is a 1-D array, it is a sum product over the last axis of a and b.
+    # Thus, if func returns scalars, f_exp will be a scalar and if it returns arrays f_exp
+    # will be an array of the same shape.
+    f_exp = np.dot(f_query, dstn.pmf)
 
     return f_exp
 

HARK/tests/test_approxDstns.py

@@ -66,10 +66,10 @@ def test_VCOV(self):
 
         Sig_2D = distribution.calc_expectation(
             self.dist2D_approx, vcov_fun, self.mu2
-        )[...,0]
+        )
         self.assertTrue(np.allclose(Sig_2D, self.Sigma2, rtol=1e-5))
 
         Sig_3D = distribution.calc_expectation(
             self.dist3D_approx, vcov_fun, self.mu3
-        )[...,0]
+        )
         self.assertTrue(np.allclose(Sig_3D, self.Sigma3, rtol=1e-5))

HARK/tests/test_distribution.py

@@ -60,7 +60,7 @@ def test_distr_of_function(self):
         binorm = combine_indep_dstns(norm_a, norm_b)
         mysum = distr_of_function(binorm, lambda x: np.sum(x))
         exp = calc_expectation(mysum)
-        self.assertAlmostEqual(exp, mu_a + mu_b)
+        self.assertAlmostEqual(exp[0], mu_a + mu_b)
 
         # Function n -> m
         # Mean and variance of two normals
@@ -83,21 +83,21 @@ def test_calc_expectation(self):
         ce2 = calc_expectation(dd_1_1_40)
         ce3 = calc_expectation(dd_10_10_100)
 
-        self.assertAlmostEqual(ce1, 0.0)
-        self.assertAlmostEqual(ce2, 1.0)
-        self.assertAlmostEqual(ce3, 10.0)
+        self.assertAlmostEqual(ce1[0], 0.0)
+        self.assertAlmostEqual(ce2[0], 1.0)
+        self.assertAlmostEqual(ce3[0], 10.0)
 
         ce4 = calc_expectation(dd_0_1_20, lambda x: 2 ** x)
 
-        self.assertAlmostEqual(ce4, 1.27153712)
+        self.assertAlmostEqual(ce4[0], 1.27153712)
 
         ce5 = calc_expectation(dd_1_1_40, lambda x: 2 * x)
 
-        self.assertAlmostEqual(ce5, 2.0)
+        self.assertAlmostEqual(ce5[0], 2.0)
 
         ce6 = calc_expectation(dd_10_10_100, lambda x, y: 2 * x + y, 20)
 
-        self.assertAlmostEqual(ce6, 40.0)
+        self.assertAlmostEqual(ce6[0], 40.0)
 
         ce7 = calc_expectation(
             dd_0_1_20, lambda x, y: x + y, np.hstack(np.array([0, 1, 2, 3, 4, 5]))