Replication section in docs

docs/source/background/literature.rst

@@ -10,3 +10,5 @@ Below you find related literature to provide a background of the `dc-egm` algori
 - Iskhakov, Jørgensen, 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*
 
 - Loretti I. Dobrescu & Akshay Shanker (2022). `Fast Upper-Envelope Scan for Discrete-Continuous Dynamic Programming <https://dx.doi.org/10.2139/ssrn.4181302>`_.
+
+- Fedor Iskhakov & Michael Keane (2021). `Effects of Taxes and Safety Net Pensions on Life-Cycle Labor Supply, Savings and Human Capital: The Case of Australia <https://doi.org/10.1016/j.jeconom.2020.01.017>`_. *Journal of Econometrics*. See :ref:`replications` for a `dcegm` replication.

docs/source/index.rst

@@ -37,6 +37,14 @@ Check out our :ref:`guides<guides/index.rst>` to find information on getting sta
 
 
 
+.. toctree::
+   :maxdepth: 1
+   :caption: Replications
+   :hidden:
+
+   replications/index
+
+
 .. toctree::
    :maxdepth: 2
    :caption: Background

docs/source/replications/index.rst

@@ -0,0 +1,22 @@
+.. _replications:
+
+Replications
+============
+
+This section walks through published studies re-implemented with `dcegm`. Unlike
+the :doc:`guides <../guides/practitioner_guide>`, which teach the interface
+through small worked examples, replications show the package applied to a
+full-scale, published life-cycle model.
+
+A replication in these docs is a **calibrated structural replication**: we
+implement the paper's model faithfully and use its published parameter
+estimates, but we do not re-run the paper's own estimation procedure (which
+typically requires restricted-access microdata we don't have). The goal is
+to demonstrate that `dcegm` reproduces the paper's *mechanism* and
+*qualitative* implications, not to match its estimated moments point for
+point. Each notebook states its simplifications explicitly.
+
+.. toctree::
+   :maxdepth: 1
+
+   iskhakov_keane_2021.ipynb

docs/source/replications/params.yaml

@@ -0,0 +1,107 @@
+---
+# Calibration for the Iskhakov & Keane (2021) replication notebook.
+#
+# All monetary values are in $1000 AUD, matching the paper. Values are the
+# paper's own published point estimates (not re-estimated here):
+#   - Table 5 (preference parameters), Table 6 (human capital), Table 7 (misc)
+#     of the supplementary material.
+#   - Eq. (27) tax function, main paper page 20.
+#   - Eq. (2)/(3) survival + pension function, supplementary material page 15-16
+#     (we use the post-2010 pension regime constant throughout; see notebook).
+#   - Table 1 (main paper) bottom panel for fixed/calibrated constants.
+#
+# Education groups are indexed 0=dropout, 1=highschool, 2=college throughout
+# the notebook, matching the key order below.
+education_groups:
+  - dropout
+  - highschool
+  - college
+hours_by_choice:
+  - 0.0
+  - 1000.0
+  - 2000.0
+  - 2250.0
+  - 2500.0
+  - 3000.0
+t0: 19  # model period 0 = age 19 for everyone; college students are
+# choice-restricted to h=0 (in school) until age 23, see notebook.
+t_retire: 85  # compulsory retirement age (last period agents may work)
+college_start_age: 23
+# --- Table 5: preference parameters ---
+preferences:
+  zeta: 0.79488  # CRRA coefficient in consumption
+  gamma:  # disutility by hours level (gamma_0=0, unused)
+    - 0.0
+    - 1.4139
+    - 2.0088
+    - 2.9213
+    - 2.8639
+    - 3.8775
+  kappa_1: 0.50321  # low-type disutility correction
+  kappa_2: 0.00008  # quadratic age term (older workers)
+  kappa_3: 0.05083  # linear age term (younger workers)
+  xi: 0.48834  # CRRA coefficient of bequest
+  b_scale: 0.68659  # bequest scale
+  taste_shock_scale: 0.29950  # lambda
+  beta_by_education:
+    dropout: 0.96806
+    highschool: 0.96732
+    college: 0.96963
+# --- Table 6: human capital production function, eq. (4) ---
+human_capital:
+  eta0_high_type: 0.39311
+  eta3: 0.02676
+  eta4: -0.00076
+  eta0_by_education:
+    dropout: 2.45647
+    highschool: 2.56761
+    college: 2.78766
+  eta1_by_education:
+    dropout: 0.01974
+    highschool: 0.02164
+    college: 0.03041
+  eta2_by_education:
+    dropout: 0.00000
+    highschool: -0.00002
+    college: -0.00017
+# --- Table 7: misc structural parameters ---
+misc:
+  sigma0: 0.24485  # wage shock std: constant
+  sigma1: 0.00421  # wage shock std: age slope
+  tr: 5.51308  # parental transfer ($1000/year, up to age 23)
+  rho_super_by_education:
+    dropout: 6.47838
+    highschool: 5.43473
+    college: 6.30347
+  high_type_share_by_education:  # only used to draw simulated agents' types
+    dropout: 0.69306
+    highschool: 0.80130
+    college: 0.90089
+# --- Eq. (27): income tax. NB the printed additive constant for the top
+# bracket (rate1*thld1) creates a small downward jump in tax liability right
+# at the second threshold; we use the continuity-preserving constant
+# rate1*(thld2-thld1) instead, as is standard for bracket-style tax rules. ---
+tax:
+  thld1: 17.39184
+  thld2: 73.17661
+  rate1: 0.29907
+  rate2: 0.37930
+# --- Supplementary eq. (3): pension function. We fix the post-2010 regime
+# constant throughout (see notebook scope discussion). ---
+pension:
+  benefit_max: 12.60665  # 10.75973 + 1.84692
+  income_taper: 0.27794
+  asset_taper: 0.00499
+  asset_threshold: 117.08260
+  pension_age: 65
+# --- Supplementary eq. (2): survival function ---
+survival:
+  age_threshold: 40
+  a: 0.0006569
+  b: 0.1078507
+# --- Table 1 bottom panel: fixed/calibrated ---
+fixed:
+  credit_constraint: 20.0  # a0, in $1000
+  interest_rate: 0.04
+  consumption_floor: 0.05  # numerical safety floor, not in the paper
+  superannuation_age: 65

src/dcegm/egm/interpolate_marginal_utility.py

@@ -64,7 +64,13 @@ def interpolate_value_and_marg_util(
     ]
     compute_marginal_utility = model_funcs["compute_marginal_utility"]
     compute_utility = model_funcs["compute_utility"]
-    discount_factor = model_funcs["read_funcs"]["discount_factor"](params)
+    # `state_choice_vec` here is the *full batch*, not a single state-choice
+    # (each branch below only reduces it to scalar via its own internal
+    # vmap), so we pass the (unevaluated) read function through and let it
+    # be resolved deep inside each branch, at the point state_choice_vec is
+    # actually scalar/consumed -- see interp1d.py, interp1d_dj.py,
+    # interp2d_irregular.py and interpnd_regular.py.
+    read_discount_factor = model_funcs["read_funcs"]["discount_factor"]
 
     # Check if interpolation needs to be multidimensional and irregular
     multi_dim = continuous_grids_info["has_additional_continuous_state"]
@@ -83,7 +89,7 @@ def interpolate_value_and_marg_util(
             policy_child_state_choice=policy_child_state_choice,
             value_child_state_choice=value_child_state_choice,
             params=params,
-            discount_factor=discount_factor,
+            read_discount_factor=read_discount_factor,
         )
 
     elif multi_dim & (not irregular):
@@ -99,7 +105,7 @@ def interpolate_value_and_marg_util(
             policy_child_state_choice=policy_child_state_choice,
             value_child_state_choice=value_child_state_choice,
             params=params,
-            discount_factor=discount_factor,
+            read_discount_factor=read_discount_factor,
         )
     else:
         # Selects inside if jorgensen_druedahl or fues (different treatment of budget constraint)
@@ -117,7 +123,7 @@ def interpolate_value_and_marg_util(
             policy_child_state_choice,
             value_child_state_choice,
             params,
-            discount_factor,
+            read_discount_factor,
             upper_envelope_method == "druedahl_jorgensen",
         )
 
@@ -131,7 +137,7 @@ def interp1d_value_and_marg_util_for_state_choice(
     policy_child_state_choice: jnp.ndarray,
     value_child_state_choice: jnp.ndarray,
     params: Dict[str, float],
-    discount_factor: float,
+    read_discount_factor: Callable,
     use_dj_interpolation: bool,
 ) -> Tuple[jnp.ndarray, jnp.ndarray]:
     """Interpolate value and policy for given child state and compute marginal utility.
@@ -182,7 +188,7 @@ def interp_on_single_wealth_point(wealth_point):
                 compute_utility=compute_utility,
                 state_choice_vec=state_choice_vec,
                 params=params,
-                discount_factor=discount_factor,
+                read_discount_factor=read_discount_factor,
             )
         else:
             policy_interp, value_interp = interp1d_policy_and_value_on_wealth(
@@ -193,7 +199,7 @@ def interp_on_single_wealth_point(wealth_point):
                 compute_utility=compute_utility,
                 state_choice_vec=state_choice_vec,
                 params=params,
-                discount_factor=discount_factor,
+                read_discount_factor=read_discount_factor,
             )
         marg_util_interp = compute_marginal_utility(
             consumption=policy_interp, params=params, **state_choice_vec
@@ -229,7 +235,7 @@ def _interpolate_value_and_marg_util_2d_irregular(
     policy_child_state_choice: jnp.ndarray,
     value_child_state_choice: jnp.ndarray,
     params: Dict[str, float],
-    discount_factor: float,
+    read_discount_factor: Callable,
 ) -> Tuple[jnp.ndarray, jnp.ndarray]:
     """Interpolate value and marginal utility on the irregular FUES 2D grid.
 
@@ -276,7 +282,7 @@ def _interpolate_value_and_marg_util_2d_irregular(
         policy_child_state_choice,
         value_child_state_choice,
         params,
-        discount_factor,
+        read_discount_factor,
     )
 
 
@@ -292,7 +298,7 @@ def _interpolate_value_and_marg_util_nd_regular(
     policy_child_state_choice: jnp.ndarray,
     value_child_state_choice: jnp.ndarray,
     params: Dict[str, float],
-    discount_factor: float,
+    read_discount_factor: Callable,
 ) -> Tuple[jnp.ndarray, jnp.ndarray]:
     """Interpolate value and marginal utility on the regular n-D grid.
 
@@ -320,7 +326,7 @@ def _interpolate_value_and_marg_util_nd_regular(
             state_choice_child_states=state_choice_vec,
             compute_utility=compute_utility,
             params=params,
-            discount_factor=discount_factor,
+            read_discount_factor=read_discount_factor,
         )
     )
 
@@ -405,7 +411,7 @@ def interp2d_value_and_marg_util_for_state_choice(
     policy_child_state_choice: jnp.ndarray,
     value_child_state_choice: jnp.ndarray,
     params: Dict[str, float],
-    discount_factor: float,
+    read_discount_factor: Callable,
 ) -> Tuple[jnp.ndarray, jnp.ndarray]:
     """Interpolate value and policy for given child state and compute marginal utility.
 
@@ -458,7 +464,7 @@ def interp_on_single_wealth_point(wealth_point, second_cont_grid_point):
                 compute_utility=compute_utility,
                 state_choice_vec=state_choice_vec,
                 params=params,
-                discount_factor=discount_factor,
+                read_discount_factor=read_discount_factor,
             )
         )
         marg_util_interp = compute_marginal_utility(

src/dcegm/egm/solve_euler_equation.py

@@ -95,7 +95,9 @@ def compute_optimal_policy_and_value(
     compute_utility = model_funcs["compute_utility"]
     compute_stochastic_transition_vec = model_funcs["compute_stochastic_transition_vec"]
 
-    discount_factor = model_funcs["read_funcs"]["discount_factor"](params)
+    discount_factor = model_funcs["read_funcs"]["discount_factor"](
+        params=params, **state_choice_vec
+    )
     interest_rate = model_funcs["read_funcs"]["interest_rate"](params)
 
     policy, expected_value = solve_euler_equation(

src/dcegm/interpolation/interp1d.py

@@ -49,7 +49,7 @@ def interp1d_policy_and_value_on_wealth(
     compute_utility: Callable,
     state_choice_vec: Dict[str, int],
     params: Dict[str, float],
-    discount_factor,
+    read_discount_factor: Callable,
 ) -> Tuple[float, float]:
     """Interpolate policy and value function given a single wealth grid point.
 
@@ -90,7 +90,7 @@ def interp1d_policy_and_value_on_wealth(
         value_at_zero_wealth=value_grid[0],
         state_choice_vec=state_choice_vec,
         params=params,
-        discount_factor=discount_factor,
+        read_discount_factor=read_discount_factor,
     )
 
     return policy_interp, value_interp
@@ -103,7 +103,7 @@ def interp_value_on_wealth(
     compute_utility: Callable,
     state_choice_vec: Dict[str, int],
     params: Dict[str, float],
-    discount_factor: float,
+    read_discount_factor: Callable,
 ) -> jnp.ndarray | float:
     """Interpolate value function on a single wealth point.
 
@@ -131,7 +131,7 @@ def interp_value_on_wealth(
         value_at_zero_wealth=value[0],
         state_choice_vec=state_choice_vec,
         params=params,
-        discount_factor=discount_factor,
+        read_discount_factor=read_discount_factor,
     )
 
     return value_interp
@@ -179,7 +179,7 @@ def interp_value_and_check_creditconstraint(
     value_at_zero_wealth: float | jnp.ndarray,
     state_choice_vec: Dict[str, int],
     params: Dict[str, float],
-    discount_factor: float,
+    read_discount_factor: Callable,
 ) -> float | jnp.ndarray:
     """Calculate the interpolated value with accounting for a possible credit
     constrained solution.
@@ -223,6 +223,7 @@ def interp_value_and_check_creditconstraint(
         params=params,
         **state_choice_vec,
     )
+    discount_factor = read_discount_factor(params=params, **state_choice_vec)
     value_interp_closed_form = utility + discount_factor * value_at_zero_wealth
 
     # Check if we are in the credit constrained region

src/dcegm/interpolation/interp1d_dj.py

@@ -16,7 +16,7 @@ def interp1d_policy_and_value_on_wealth_dj(
     compute_utility: Callable,
     state_choice_vec: Dict[str, int],
     params: Dict[str, float],
-    discount_factor: float,
+    read_discount_factor: Callable,
 ) -> Tuple[jnp.ndarray | float, jnp.ndarray | float]:
     """1D interpolation for DJ with consume-all overwrite for policy and value."""
     ind_high, ind_low = get_index_high_and_low(x=wealth_grid, x_new=wealth)
@@ -42,7 +42,7 @@ def interp1d_policy_and_value_on_wealth_dj(
         compute_utility=compute_utility,
         state_choice_vec=state_choice_vec,
         params=params,
-        discount_factor=discount_factor,
+        read_discount_factor=read_discount_factor,
     )
     overwrite_mask = consume_all_value > value_interp_on_grid
     policy = jnp.where(overwrite_mask, wealth, policy_interp)
@@ -57,7 +57,7 @@ def interp1d_value_on_wealth_dj(
     compute_utility: Callable,
     state_choice_vec: Dict[str, int],
     params: Dict[str, float],
-    discount_factor: float,
+    read_discount_factor: Callable,
 ) -> jnp.ndarray | float:
     """1D value interpolation for DJ with consume-all overwrite."""
     _, value = interp1d_policy_and_value_on_wealth_dj(
@@ -68,7 +68,7 @@ def interp1d_value_on_wealth_dj(
         compute_utility=compute_utility,
         state_choice_vec=state_choice_vec,
         params=params,
-        discount_factor=discount_factor,
+        read_discount_factor=read_discount_factor,
     )
     return value
 
@@ -79,9 +79,10 @@ def _consume_all_value(
     compute_utility: Callable,
     state_choice_vec: Dict[str, int],
     params: Dict[str, float],
-    discount_factor: float,
+    read_discount_factor: Callable,
 ) -> jnp.ndarray:
     util = compute_utility(consumption=wealth, params=params, **state_choice_vec)
     if isinstance(util, tuple):
         util = util[0]
+    discount_factor = read_discount_factor(params=params, **state_choice_vec)
     return jnp.asarray(util) + discount_factor * value_at_zero_wealth