[WIP] Class representing discretized transitions - #1548
Conversation
|
This is a work in progress, do not merge yet. I need to, for example, add tests. |
|
I finally had a chance to look at this. Main comments are related to both prongs of For the infinite horizon version, you don't need to loop on For the lifecycle model, your steady state age distribution is wrong. You can't just take the cumulative survival probabilities and normalize them, because that ignores death and replacement. Instead, make a square array A of "age transitions", with two non-zero elements in each row: A_{i,i+1} = LivPrb_i and A_{i,0} = 1-LivPrb_i. Just like above, take the transpose of A, then find the eigenvector associated with its largest eigenvalue and normalize it. That's the steady state age distribution. Your survival-conditional state distributions are correct. See L168-202 of estimation.py in the DistributionOfWealthMPC REMARK for where this was done in cstwMPC, but note that when I wrote the function there, I didn't know that the largest eigenvalue is always 1 and/or EDIT: Ignore that last paragraph, it's wrong. |
|
Addendum to the above: The lifecycle method assumes that you want a constant population, with newborns exactly replacing decedents. If you instead want a constant growth rate of the population, the math is slightly different (but you can probably work it out). |
|
Hit pause on my comment about the steady state age distribution. Normalized cumulative survival probabilities and the first normalized eigenvector might be equivalent. I'm on a train right now and can't check. |
|
Yeah nevermind, the methods are equivalent. |
There was a problem hiding this comment.
Pull Request Overview
This PR introduces the DiscreteTransitions class for simulating population transitions in discretized state spaces for both life-cycle and infinite-horizon models. Key changes include the addition of a new dataclass with methods to iterate distributions forward, calculate conditional age distributions, and compute steady-state distributions.
| ) | ||
| else: | ||
| return _iterate_dstn_forward_ih( | ||
| dstn_init[0], self.living_tmat[0], self.surv_prob[0], self.newborn_dstn |
There was a problem hiding this comment.
In infinite-horizon mode, the properties 'self.living_tmat' and 'self.surv_prob' are referenced, but the class defines 'living_tmats' and 'surv_probs'. Please update the variable names to ensure consistency.
| dstn_init[0], self.living_tmat[0], self.surv_prob[0], self.newborn_dstn | |
| dstn_init[0], self.living_tmats[0], self.surv_probs[0], self.newborn_dstn |
| The class assumes that: | ||
| - Death is exogenous and independent of every state. | ||
| - Agents that die are replaced by newborns. | ||
| - Newborns draw their state from a distribution that is constatn over time. |
There was a problem hiding this comment.
Correct the typo 'constatn' to 'constant'.
| - Newborns draw their state from a distribution that is constatn over time. | |
| - Newborns draw their state from a distribution that is constant over time. |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 1 out of 1 changed files in this pull request and generated 1 comment.
Suppressed comments (8)
HARK/mat_methods.py:17
- Spelling typo in the class docstring: "constatn" should be "constant".
- Newborns draw their state from a distribution that is constatn over time.
HARK/mat_methods.py:70
- The infinite-horizon branch of
iterate_dstn_forwardreferences non-existent attributes (self.living_tmat,self.surv_prob) and assumesdstn_initis indexable. This will raiseAttributeErrorand is inconsistent with the docstring (which says a single array is allowed).
else:
return _iterate_dstn_forward_ih(
dstn_init[0], self.living_tmat[0], self.surv_prob[0], self.newborn_dstn
)
HARK/mat_methods.py:144
- If the public method is renamed to
find_conditional_age_dstn, the helper should also be renamed and all internal references updated to avoidNameError.
def _find_conditional_age_dsnt(dstn_init, living_tmats):
dstns = [dstn_init]
for tmat in living_tmats:
dstns.append(dstns[-1] @ tmat)
return dstns
HARK/mat_methods.py:46
- In infinite-horizon mode,
__post_init__doesn't validate thatliving_tmats/surv_probshave length 1, but later methods index[0]. This can lead to IndexError with unclear provenance; validate lengths up-front for a clearer error.
else:
self.T = 1
HARK/mat_methods.py:92
- The public API name
find_conditional_age_dsntlooks like a typo ("dsnt" vs the consistently-used abbreviation "dstn"). Because this is a new public method, fixing the name now will avoid locking in a confusing API.
This issue also appears on line 140 of the same file.
def find_conditional_age_dsnt(self, dstn_init):
"""
Given a distribution of agents over states for the first period of life,
find the distribution of agents over states in every age conditional on
their survival.
HARK/mat_methods.py:129
_iterate_dstn_forward_lcstoresnewborn_dstndirectly innew_dstnand then scales it in-place (new_dstn[0] *= dead_mass), which mutates the caller-providednewborn_dstnarray. This hidden side effect is surprising and can corrupt subsequent uses of the samenewborn_dstn.
new_dstn = [newborn_dstn]
HARK/mat_methods.py:195
_find_steady_state_dstn_ihsilently stops whenmax_iteris exceeded and also allowsnormalize_every=0(which would raise a modulo-by-zero error). For a new public solver, it's safer to (1) validatenormalize_every >= 1, and (2) raise an explicit error when iteration doesn't converge withinmax_iter.
if dstn_init is None:
dstn = newborn_dstn
else:
dstn = dstn_init
go = True
HARK/mat_methods.py:13
- New
DiscreteTransitionsfunctionality is introduced inHARK/mat_methods.py, but the existing test suite for this module (tests/test_mat_methods.py) doesn't cover any of these new behaviors (life-cycle vs infinite-horizon forward iteration, steady-state solver, etc.). Adding unit tests now would help prevent regressions as this API evolves.
class DiscreteTransitions:
"""
Class to facilitate simulating transitions of populations in discretized state spaces,
supporting both life-cycle and infinite-horizon models.
| living_tmats: list | ||
| surv_probs: list | ||
| life_cycle: bool = False | ||
| newborn_dstn: np.array |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 1 out of 1 changed files in this pull request and generated no new comments.
Suppressed comments (10)
HARK/mat_methods.py:35
- Dataclass fields are ordered with a defaulted field (life_cycle) before a non-default field (newborn_dstn), which will raise
TypeError: non-default argument 'newborn_dstn' follows default argumentat class creation.
living_tmats: list
surv_probs: list
life_cycle: bool = False
newborn_dstn: np.array
HARK/mat_methods.py:70
- Infinite-horizon branch references undefined attributes (
living_tmat,surv_prob) and indexesdstn_init[0]even though the docstring saysdstn_initis a single array in infinite-horizon mode. This will raise AttributeError/TypeError at runtime.
return _iterate_dstn_forward_ih(
dstn_init[0], self.living_tmat[0], self.surv_prob[0], self.newborn_dstn
)
HARK/mat_methods.py:140
- Helper name
_find_conditional_age_dsntappears to be a typo ("dsnt" vs "dstn"); consider renaming to avoid cementing the misspelling in internal APIs.
def _find_conditional_age_dsnt(dstn_init, living_tmats):
HARK/mat_methods.py:151
- This call site uses
_find_conditional_age_dsnt; if you rename the helper to_find_conditional_age_dstn, this needs to be updated too to avoid NameError.
age_dstns = _find_conditional_age_dsnt(newborn_dstn, living_tmats)
HARK/mat_methods.py:17
- Spelling: "constatn" -> "constant" in the class docstring.
- Newborns draw their state from a distribution that is constatn over time.
HARK/mat_methods.py:72
- Method name
find_conditional_age_dsntappears to be a typo ("dsnt" vs "dstn"); this looks like an accidental public API name that will be hard to support long-term.
This issue also appears on line 140 of the same file.
def find_conditional_age_dsnt(self, dstn_init):
HARK/mat_methods.py:129
new_dstn = [newborn_dstn]followed bynew_dstn[0] *= dead_massmutates the input arraynewborn_dstnin-place. When called withself.newborn_dstn, this permanently scales the stored newborn distribution, breaking subsequent iterations.
new_dstn = [newborn_dstn]
HARK/mat_methods.py:160
- In infinite-horizon iteration,
dead_massis computed as1 - surv_probregardless of the total mass indstn_init. Ifdstn_initis not normalized to sum to 1, births will be under/over-scaled and total mass will drift incorrectly.
dead_mass = 1.0 - surv_prob
new_dstn = surv_prob * dstn_init @ living_tmat
new_dstn += dead_mass * newborn_dstn
HARK/mat_methods.py:10
- This PR introduces a new public class (
DiscreteTransitions) with non-trivial numerical behavior (LC vs IH iteration, steady state).HARK/mat_methods.pyalready has unit tests, but there are no tests added for these new APIs/edge cases (mass conservation, normalization, LC vs IH consistency, convergence/max_iterbehavior).
class DiscreteTransitions:
HARK/mat_methods.py:91
find_conditional_age_dsntreturnsdstn_initdirectly in infinite-horizon mode, which is inconsistent with the documented return type (list of distributions) and with the life-cycle branch (returns a list). This inconsistency can break callers expecting a uniform return shape.
This issue also appears on line 151 of the same file.
if self.life_cycle:
return _find_conditional_age_dsnt(dstn_init, self.living_tmats)
else:
return dstn_init
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 1 out of 1 changed files in this pull request and generated 3 comments.
Suppressed comments (6)
HARK/mat_methods.py:72
- The name
dsntappears to be a typo/inconsistent abbreviation for 'distribution' (dstnis used elsewhere in this file). Consider renaming both the public method and helper tofind_conditional_age_dstn/_find_conditional_age_dstnfor consistency and discoverability.
def find_conditional_age_dsnt(self, dstn_init):
HARK/mat_methods.py:89
- The name
dsntappears to be a typo/inconsistent abbreviation for 'distribution' (dstnis used elsewhere in this file). Consider renaming both the public method and helper tofind_conditional_age_dstn/_find_conditional_age_dstnfor consistency and discoverability.
return _find_conditional_age_dsnt(dstn_init, self.living_tmats)
HARK/mat_methods.py:140
- The name
dsntappears to be a typo/inconsistent abbreviation for 'distribution' (dstnis used elsewhere in this file). Consider renaming both the public method and helper tofind_conditional_age_dstn/_find_conditional_age_dstnfor consistency and discoverability.
def _find_conditional_age_dsnt(dstn_init, living_tmats):
HARK/mat_methods.py:17
- Correct typo: 'constatn' → 'constant'.
- Newborns draw their state from a distribution that is constatn over time.
HARK/mat_methods.py:195
- The comment says the function returns 'as list', but it actually returns a single array (
dstn). Either update the comment to match the behavior, or change the return value to a list to match the stated compatibility goal.
# Return as list just for compatibility with LC methods that return
# a list of age dstns
return dstn
HARK/mat_methods.py:191
normalize_every=0(or other non-positive values) will raiseZeroDivisionErrorini % normalize_every. If this is user-configurable via**kwargs, add a guard (e.g., requirenormalize_every >= 1, or interpretNone/0as 'do not renormalize').
# Renormalize every given number of iterations
if i % normalize_every == 0:
dstn /= np.sum(dstn)
| living_tmats: list | ||
| surv_probs: list | ||
| life_cycle: bool = False | ||
| newborn_dstn: np.array |
| if self.life_cycle: | ||
| return _iterate_dstn_forward_lc( | ||
| dstn_init, self.living_tmats, self.surv_probs, self.newborn_dstn | ||
| ) | ||
| else: | ||
| return _iterate_dstn_forward_ih( | ||
| dstn_init[0], self.living_tmat[0], self.surv_prob[0], self.newborn_dstn | ||
| ) |
| @njit | ||
| def _iterate_dstn_forward_lc(dstn_init, living_tmats, surv_probs, newborn_dstn): | ||
| new_dstn = [newborn_dstn] | ||
| dead_mass = 0.0 | ||
| for i, (d0, tmat) in enumerate(zip(dstn_init, living_tmats)): | ||
| new_dstn.append((surv_probs[i] * d0) @ tmat) | ||
| dead_mass += (1.0 - surv_probs[i]) * np.sum(d0) | ||
| dead_mass += np.sum(dstn_init[-1]) | ||
| new_dstn[0] *= dead_mass | ||
|
|
||
| return new_dstn |
I'm making public versions of a few methods that I have used to work for simulations with discretized/matrix transitions.
The idea is to have a simple class that holds all the ingredients needed to, say, find the steady state distribution of an infinite horizon or life-cycle model, or to iterate a distribution forward a single time period.
Ideally, people will use this and make it more efficient and find useful sparse representations and who knows what else.