Finish utility kwargs (#6)

.gitignore

@@ -139,4 +139,4 @@ venv.bak/
 
 
 # Debugging files should be ignored
-*.png
+*.png

README.md

@@ -9,5 +9,10 @@ Extension of the Fast Upper-Envelope Scan (FUES) for solving discrete-continuous
 
 ## References
 
+1. Iskhakov, Jorgensen, Rust, & Schjerning (2017).
+   [The Endogenous Grid Method for Discrete-Continuous Dynamic Choice Models with (or without) Taste Shocks](http://onlinelibrary.wiley.com/doi/10.3982/QE643/full).
+   *Quantitative Economics*
+
+
 1. Loretti I. Dobrescu & Akshay Shanker (2022).
-   [Fast Upper-Envelope Scan for Discrete-Continuous Dynamic Programming](https://dx.doi.org/10.2139/ssrn.4181302).
+   [Fast Upper-Envelope Scan for Discrete-Continuous Dynamic Programming](https://dx.doi.org/10.2139/ssrn.4181302).

src/upper_envelope/upper_envelope_jax.py

@@ -26,9 +26,9 @@ def fast_upper_envelope_wrapper(
     policy: jnp.ndarray,
     value: jnp.ndarray,
     expected_value_zero_savings: float,
-    state_choice_vec: jnp.ndarray,
-    params: Dict[str, float],
-    compute_utility: Callable,
+    utility_function: Callable,
+    utility_kwargs: Dict,
+    disc_factor: float,
 ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]:
     """Drop suboptimal points and refines the endogenous grid, policy, and value.
 
@@ -62,9 +62,10 @@ def fast_upper_envelope_wrapper(
             containing the current state- and choice-specific value function.
         expected_value_zero_savings (float): The agent's expected value given that she
             saves zero.
-        choice (int): The current choice.
-        compute_value (callable): Function to compute the agent's utility.
-        params (dict): Dictionary containing the model parameters.
+        utility_function (callable): The utility function. The first argument is
+            assumed to be consumption.
+        utility_kwargs (dict): The keyword arguments to be passed to the utility
+            function.
 
     Returns:
         tuple:
@@ -98,9 +99,9 @@ def fast_upper_envelope_wrapper(
     values_to_add = vmap(_compute_value, in_axes=(0, None, None, None, None))(
         grid_points_to_add,
         expected_value_zero_savings,
-        state_choice_vec,
-        params,
-        compute_utility,
+        utility_function,
+        utility_kwargs,
+        disc_factor,
     )
 
     grid_augmented = jnp.append(grid_points_to_add, endog_grid)
@@ -1359,11 +1360,10 @@ def create_indicator_if_value_function_is_switched(
 
 
 def _compute_value(
-    consumption, next_period_value, state_choice_vec, params, compute_utility
+    consumption, next_period_value, utility_function, utility_kwargs, discount_factor
 ):
-    utility = compute_utility(
-        consumption=consumption,
-        params=params,
-        **state_choice_vec,
+    utility = utility_function(
+        consumption,
+        **utility_kwargs,
     )
-    return utility + params["beta"] * next_period_value
+    return utility + discount_factor * next_period_value

src/upper_envelope/upper_envelope_numba.py

@@ -20,9 +20,9 @@ def fast_upper_envelope_wrapper(
     value: np.ndarray,
     exog_grid: np.ndarray,
     expected_value_zero_savings: float,
-    state_choice_vec: np.ndarray,
-    params: Dict[str, float],
-    compute_utility: Callable,
+    utility_function: Callable,
+    utility_kwargs: Dict,
+    discount_factor: float,
 ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
     """Drop suboptimal points and refine the endogenous grid, policy, and value.
 
@@ -89,12 +89,12 @@ def fast_upper_envelope_wrapper(
             endog_grid=endog_grid,
             value=value,
             policy=policy,
-            state_choice_vec=state_choice_vec,
             expected_value_zero_savings=expected_value_zero_savings,
             min_wealth_grid=min_wealth_grid,
             n_grid_wealth=n_grid_wealth,
-            params=params,
-            compute_utility=compute_utility,
+            utility_function=utility_function,
+            utility_kwargs=utility_kwargs,
+            discount_factor=discount_factor,
         )
         exog_grid = np.append(np.zeros(n_grid_wealth // 10 - 1), exog_grid)
 
@@ -732,12 +732,12 @@ def _augment_grids(
     endog_grid: np.ndarray,
     value: np.ndarray,
     policy: np.ndarray,
-    state_choice_vec: np.ndarray,
     expected_value_zero_savings: np.ndarray,
     min_wealth_grid: float,
     n_grid_wealth: int,
-    compute_utility: Callable,
-    params: Dict[str, float],
+    utility_function: Callable,
+    utility_kwargs: Dict[str, float],
+    discount_factor: float,
 ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
     """Extends the endogenous wealth grid, value, and policy functions to the left.
 
@@ -779,12 +779,11 @@ def _augment_grids(
         min_wealth_grid, endog_grid[0], n_grid_wealth // 10
     )[:-1]
 
-    utility = compute_utility(
-        consumption=grid_points_to_add,
-        params=params,
-        **state_choice_vec,
+    utility = utility_function(
+        grid_points_to_add,
+        **utility_kwargs,
     )
-    values_to_add = utility + params["beta"] * expected_value_zero_savings
+    values_to_add = utility + discount_factor * expected_value_zero_savings
 
     grid_augmented = np.append(grid_points_to_add, endog_grid)
     value_augmented = np.append(values_to_add, value)

tests/test_upper_envelope_jax.py

@@ -6,7 +6,6 @@
 import numpy as np
 import pytest
 from numpy.testing import assert_array_almost_equal as aaae
-from upper_envelope.shared import process_function_args_to_kwargs
 from upper_envelope.upper_envelope_jax import fast_upper_envelope
 from upper_envelope.upper_envelope_jax import (
     fast_upper_envelope_wrapper,
@@ -81,9 +80,7 @@ def setup_model():
 
     options["state_space"]["exogenous_states"] = {"exog_state": [0]}
 
-    compute_utility = process_function_args_to_kwargs(utility_crra)
-
-    return params, exog_savings_grid, state_choice_vars, compute_utility
+    return params, exog_savings_grid, state_choice_vars
 
 
 @pytest.mark.parametrize("period", [2, 4, 9, 10, 18])
@@ -112,8 +109,12 @@ def test_fast_upper_envelope_wrapper(period, setup_model):
         ~np.isnan(value_refined_fedor).any(axis=0),
     ]
 
-    params, _exog_savings_grid, state_choice_vars, compute_utility = setup_model
+    params, _exog_savings_grid, state_choice_vars = setup_model
 
+    utility_kwargs = {
+        "choice": state_choice_vars["choice"],
+        "params": params,
+    }
     (
         endog_grid_refined,
         policy_left_refined,
@@ -124,9 +125,9 @@ def test_fast_upper_envelope_wrapper(period, setup_model):
         policy=policy_egm[1, 1:],
         value=value_egm[1, 1:],
         expected_value_zero_savings=value_egm[1, 0],
-        state_choice_vec=state_choice_vars,
-        params=params,
-        compute_utility=compute_utility,
+        utility_function=utility_crra,
+        utility_kwargs=utility_kwargs,
+        disc_factor=params["beta"],
     )
 
     wealth_max_to_test = np.max(endog_grid_refined[~np.isnan(endog_grid_refined)]) + 100
@@ -164,7 +165,7 @@ def test_fast_upper_envelope_against_org_fues(setup_model):
     value_egm = np.genfromtxt(
         TEST_RESOURCES_DIR / "upper_envelope_period_tests/val10.csv", delimiter=","
     )
-    _params, exog_savings_grid, state_choice_vars, compute_utility = setup_model
+    _params, exog_savings_grid, state_choice_vars = setup_model
 
     (
         endog_grid_refined,
@@ -185,7 +186,7 @@ def test_fast_upper_envelope_against_org_fues(setup_model):
         value=value_egm[1],
         exog_grid=exog_savings_grid,
         choice=state_choice_vars["choice"],
-        compute_utility=compute_utility,
+        compute_utility=utility_crra,
     )
 
     endog_grid_expected = endog_grid_org[~np.isnan(endog_grid_org)]
@@ -208,22 +209,27 @@ def test_fast_upper_envelope_against_fedor(period, setup_model):
         delimiter=",",
     )
 
-    params, exog_savings_grid, state_choice_vec, compute_utility = setup_model
+    params, exog_savings_grid, state_choice_vec = setup_model
 
     _policy_fedor, _value_fedor = upper_envelope(
         policy=policy_egm,
         value=value_egm,
         exog_grid=exog_savings_grid,
-        state_choice_vec=state_choice_vec,
+        state_choice_vec={"choice": state_choice_vec["choice"]},
         params=params,
-        compute_utility=compute_utility,
+        compute_utility=utility_crra,
     )
     policy_expected = _policy_fedor[:, ~np.isnan(_policy_fedor).any(axis=0)]
     value_expected = _value_fedor[
         :,
         ~np.isnan(_value_fedor).any(axis=0),
     ]
 
+    utility_kwargs = {
+        "choice": state_choice_vec["choice"],
+        "params": params,
+    }
+
     (
         endog_grid_fues,
         policy_fues_left,
@@ -234,9 +240,9 @@ def test_fast_upper_envelope_against_fedor(period, setup_model):
         policy=policy_egm[1, 1:],
         value=value_egm[1, 1:],
         expected_value_zero_savings=value_egm[1, 0],
-        state_choice_vec=state_choice_vec,
-        params=params,
-        compute_utility=compute_utility,
+        utility_function=utility_crra,
+        utility_kwargs=utility_kwargs,
+        disc_factor=params["beta"],
     )
 
     wealth_max_to_test = np.max(endog_grid_fues[~np.isnan(endog_grid_fues)]) + 100

tests/test_upper_envelope_numba.py

@@ -4,7 +4,6 @@
 import numpy as np
 import pytest
 from numpy.testing import assert_array_almost_equal as aaae
-from upper_envelope.shared import process_function_args_to_kwargs
 from upper_envelope.upper_envelope_numba import fast_upper_envelope
 from upper_envelope.upper_envelope_numba import fast_upper_envelope_wrapper
 
@@ -59,9 +58,7 @@ def setup_model():
 
     state_choice_vec = {"choice": 0, "lagged_choice": 0}
 
-    compute_utility = process_function_args_to_kwargs(utility_crra)
-
-    return params, state_choice_vec, exog_savings_grid, compute_utility
+    return params, state_choice_vec, exog_savings_grid
 
 
 @pytest.mark.parametrize("period", [2, 4, 9, 10, 18])
@@ -90,17 +87,22 @@ def test_fast_upper_envelope_wrapper(period, setup_model):
         ~np.isnan(value_refined_fedor).any(axis=0),
     ]
 
-    params, state_choice_vec, _exog_savings_grid, compute_utility = setup_model
+    params, state_choice_vec, _exog_savings_grid = setup_model
+
+    utility_kwargs = {
+        "choice": state_choice_vec["choice"],
+        "params": params,
+    }
 
     endog_grid_refined, policy_refined, value_refined = fast_upper_envelope_wrapper(
         endog_grid=policy_egm[0, 1:],
         policy=policy_egm[1, 1:],
         value=value_egm[1, 1:],
         expected_value_zero_savings=value_egm[1, 0],
         exog_grid=_exog_savings_grid,
-        state_choice_vec=state_choice_vec,
-        params=params,
-        compute_utility=compute_utility,
+        utility_function=utility_crra,
+        utility_kwargs=utility_kwargs,
+        discount_factor=params["beta"],
     )
 
     wealth_max_to_test = np.max(endog_grid_refined[~np.isnan(endog_grid_refined)]) + 100
@@ -138,7 +140,7 @@ def test_fast_upper_envelope_against_org_fues(setup_model):
         TEST_RESOURCES_DIR / "upper_envelope_period_tests/val10.csv", delimiter=","
     )
 
-    _params, state_choice_vec, exog_savings_grid, compute_utility = setup_model
+    _params, state_choice_vec, exog_savings_grid = setup_model
 
     endog_grid_refined, value_refined, policy_refined = fast_upper_envelope(
         endog_grid=policy_egm[0],
@@ -153,7 +155,7 @@ def test_fast_upper_envelope_against_org_fues(setup_model):
         value=value_egm[1],
         exog_grid=exog_savings_grid,
         choice=state_choice_vec["choice"],
-        compute_utility=compute_utility,
+        compute_utility=utility_crra,
     )
 
     endog_grid_expected = endog_grid_org[~np.isnan(endog_grid_org)]
@@ -176,31 +178,35 @@ def test_fast_upper_envelope_against_fedor(period, setup_model):
         delimiter=",",
     )
 
-    params, state_choice_vec, exog_savings_grid, compute_utility = setup_model
+    params, state_choice_vec, exog_savings_grid = setup_model
 
     _policy_fedor, _value_fedor = upper_envelope(
         policy=policy_egm,
         value=value_egm,
         exog_grid=exog_savings_grid,
-        state_choice_vec=state_choice_vec,
+        state_choice_vec={"choice": state_choice_vec["choice"]},
         params=params,
-        compute_utility=compute_utility,
+        compute_utility=utility_crra,
     )
     policy_expected = _policy_fedor[:, ~np.isnan(_policy_fedor).any(axis=0)]
     value_expected = _value_fedor[
         :,
         ~np.isnan(_value_fedor).any(axis=0),
     ]
+    utility_kwargs = {
+        "choice": state_choice_vec["choice"],
+        "params": params,
+    }
 
     endog_grid_fues, policy_fues, value_fues = fast_upper_envelope_wrapper(
         endog_grid=policy_egm[0, 1:],
         policy=policy_egm[1, 1:],
         value=value_egm[1, 1:],
         exog_grid=np.append(0, exog_savings_grid),
         expected_value_zero_savings=value_egm[1, 0],
-        state_choice_vec=state_choice_vec,
-        params=params,
-        compute_utility=compute_utility,
+        utility_function=utility_crra,
+        utility_kwargs=utility_kwargs,
+        discount_factor=params["beta"],
     )
 
     wealth_max_to_test = np.max(endog_grid_fues[~np.isnan(endog_grid_fues)]) + 100