From c1b7781a32d97073600b1227e314e1b9cfefcaf2 Mon Sep 17 00:00:00 2001 From: Mateo VG Date: Tue, 4 Jul 2023 20:04:41 -0400 Subject: [PATCH 1/2] Diagnostic test --- HARK/distribution.py | 2 ++ HARK/tests/test_distribution.py | 34 +++++++++++++++++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/HARK/distribution.py b/HARK/distribution.py index 37be02226..e9c810e1b 100644 --- a/HARK/distribution.py +++ b/HARK/distribution.py @@ -1235,6 +1235,8 @@ def func_wrapper(x, *args): if len(kwargs): f_query = func(self.dataset, *args, **kwargs) + # TODO: seems like the issue is that self.probability has to be a dataarray with + # dimension atom. But it is sometimes a numpy array. Need to figure out why. ldd = DiscreteDistributionLabeled.from_dataset(f_query, self.probability) return ldd._weighted.mean("atom") diff --git a/HARK/tests/test_distribution.py b/HARK/tests/test_distribution.py index ced58fbf5..635bd9d37 100644 --- a/HARK/tests/test_distribution.py +++ b/HARK/tests/test_distribution.py @@ -603,3 +603,37 @@ def test_combine_labeled_dist(self): np.concatenate([de.expected(), abc.expected()]), ) ) + + +import xarray as xr + +class labeled_transition_tests(unittest.TestCase): + + def setUp(self) -> None: + return super().setUp() + + def test_mytest(self): + + # Create a basic labeled distribution + base_dist = DiscreteDistributionLabeled( + pmv=np.array([0.5, 0.5]), atoms=np.array([[1.0, 2.0], [3.0, 4.0]]), var_names=["a", "b"] + ) + + # Define a transition function + def transition(shocks, state): + state_new = {} + state_new["m"] = state["m"]*shocks["a"] + state_new["n"] = state["n"]*shocks["b"] + return state_new + + m = xr.DataArray(np.linspace(0, 10, 11), name="m", dims=("grid",)) + n = xr.DataArray(np.linspace(0, -10, 11), name="n", dims=("grid",)) + state_grid = xr.Dataset({"m": m, "n": n}) + + # Evaluate labeled transformation + + # This works + new_state_dstn = base_dist.expected(transition, state=state_grid) + # This does not work + new_state_dstn = base_dist.dist_of_func(transition, state=state_grid) + new_state_dstn.expected() From 9bac21ecbe00a075ffe0977656000e33ffbd244d Mon Sep 17 00:00:00 2001 From: MateoVG Date: Wed, 5 Jul 2023 10:16:41 -0400 Subject: [PATCH 2/2] Fix and make test pass --- HARK/distribution.py | 4 +--- HARK/tests/test_distribution.py | 29 +++++++++++++++++------------ 2 files changed, 18 insertions(+), 15 deletions(-) diff --git a/HARK/distribution.py b/HARK/distribution.py index e9c810e1b..ea808bbf4 100644 --- a/HARK/distribution.py +++ b/HARK/distribution.py @@ -1181,7 +1181,7 @@ def func_wrapper(x: np.ndarray, *args: Any) -> np.ndarray: if len(kwargs): f_query = func(self.dataset, **kwargs) - ldd = DiscreteDistributionLabeled.from_dataset(f_query, self.pmv) + ldd = DiscreteDistributionLabeled.from_dataset(f_query, self.probability) return ldd @@ -1235,8 +1235,6 @@ def func_wrapper(x, *args): if len(kwargs): f_query = func(self.dataset, *args, **kwargs) - # TODO: seems like the issue is that self.probability has to be a dataarray with - # dimension atom. But it is sometimes a numpy array. Need to figure out why. ldd = DiscreteDistributionLabeled.from_dataset(f_query, self.probability) return ldd._weighted.mean("atom") diff --git a/HARK/tests/test_distribution.py b/HARK/tests/test_distribution.py index 635bd9d37..ec685eaa3 100644 --- a/HARK/tests/test_distribution.py +++ b/HARK/tests/test_distribution.py @@ -1,6 +1,7 @@ import unittest import numpy as np +import xarray as xr from HARK.distribution import ( Bernoulli, @@ -605,25 +606,23 @@ def test_combine_labeled_dist(self): ) -import xarray as xr - class labeled_transition_tests(unittest.TestCase): - def setUp(self) -> None: return super().setUp() - - def test_mytest(self): + def test_expectation_transformation(self): # Create a basic labeled distribution base_dist = DiscreteDistributionLabeled( - pmv=np.array([0.5, 0.5]), atoms=np.array([[1.0, 2.0], [3.0, 4.0]]), var_names=["a", "b"] + pmv=np.array([0.5, 0.5]), + atoms=np.array([[1.0, 2.0], [3.0, 4.0]]), + var_names=["a", "b"], ) # Define a transition function def transition(shocks, state): state_new = {} - state_new["m"] = state["m"]*shocks["a"] - state_new["n"] = state["n"]*shocks["b"] + state_new["m"] = state["m"] * shocks["a"] + state_new["n"] = state["n"] * shocks["b"] return state_new m = xr.DataArray(np.linspace(0, 10, 11), name="m", dims=("grid",)) @@ -632,8 +631,14 @@ def transition(shocks, state): # Evaluate labeled transformation - # This works - new_state_dstn = base_dist.expected(transition, state=state_grid) - # This does not work + # Direct expectation + exp1 = base_dist.expected(transition, state=state_grid) + # Expectation after transformation new_state_dstn = base_dist.dist_of_func(transition, state=state_grid) - new_state_dstn.expected() + # TODO: needs a cluncky identity function with an extra argument because + # DDL.expected() behavior is very different with and without kwargs. + # Fix! + exp2 = new_state_dstn.expected(lambda x, unused: x, unused=0) + + assert np.all(exp1["m"] == exp2["m"]).item() + assert np.all(exp1["n"] == exp2["n"]).item()