Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[WIP] Transition matrices (general, life-cycle) #1286

Closed
wants to merge 32 commits into from

Conversation

Mv77
Copy link
Contributor

@Mv77 Mv77 commented Jun 15, 2023

This PR contains preliminary work that seeks to construct state-transition matrices for any model in HARK, including life-cycle versions.

I am making the raw code public now because I expect to collaborate around these ideas with @wdu9, and in case other people are interested in contributing.

I think this is an important piece of functionality in our goal of communicating with the sequence-space-jacobian methods and tools.

Even though the two examples that I have included are infinite-horizon, the method works for LC parametrizations. It creates a list of transition matrices. One for every age

There are many things to be done:

  • Clean up the code and make the methods more readable and user-friendly (using labeled distributions and xarrays, for instance).
  • Add an example showing that it works with LC models (we really need an init_portfolio_lifecyle param set).
  • Incorporate death (the transition matrices currently are conditional on survival).
  • Add tests that compare with the existing model-specific methods for correctness.
  • Add the required methods (state-to-state transition, distribution_engine, etc) to more of the existing models.

And here are the ususal to-dos.

  • Tests for new functionality/models or Tests to reproduce the bug-fix in code.
  • Updated documentation of features that add new functionality.
  • Update CHANGELOG.md with major/minor changes.

@alanlujan91
Copy link
Member

great ideas! @AMonninger and I just came back from NBER ssj workshop and are blown away by their tools. Anything we can do to more tightly integrate with their framework will be very beneficial

@Mv77
Copy link
Contributor Author

Mv77 commented Jun 15, 2023

The main idea that the PR works off of is that once you have solved a model, the only things that you need to produce a transition matrix that represents it are:

  • A transition equation: which takes the state today, policy functions, and shock realizations tomorrow; and transforms them into the state tomorrow.
  • The joint distribution of all the shocks in the model.

I have used the ConsPortfolioModel as a working example, and I added methods that construct both objects for it:

State transition equation:
https://github.com/Mv77/HARK/blob/c4c5064bea6665d096886fe049fc5ec3922c27fc/HARK/ConsumptionSaving/ConsPortfolioModel.py#L382-L425

Shock-distribution engine:
https://github.com/Mv77/HARK/blob/c4c5064bea6665d096886fe049fc5ec3922c27fc/HARK/ConsumptionSaving/ConsPortfolioModel.py#L324C14-L348

@Mv77
Copy link
Contributor Author

Mv77 commented Jun 15, 2023

Once we have the transition equation, shock distribution, and the grids we want to use for the state-space, we can use the tools that

  • Find the distribution of a function of a matrix-valued random variable.
  • Clip points to a general grid using what the SSJ people call the "lottery method."

to produce the transition matrix. It's easy!:
https://github.com/Mv77/HARK/blob/c4c5064bea6665d096886fe049fc5ec3922c27fc/HARK/core.py#L941-L959

@Mv77
Copy link
Contributor Author

Mv77 commented Jun 15, 2023

So, in principle, this means that if you write a state_to_state_transition method for your model and a method to produce the joint distribution of all shocks, I can get you the transition matrices.

@codecov
Copy link

codecov bot commented Jun 15, 2023

Codecov Report

Patch coverage: 95.98% and project coverage change: +0.48% 🎉

Comparison is base (4f644bf) 72.64% compared to head (41f6ee2) 73.12%.
Report is 10 commits behind head on master.

❗ Current head 41f6ee2 differs from pull request most recent head 497a7f2. Consider uploading reports for the commit 497a7f2 to get more accurate results

Additional details and impacted files
@@            Coverage Diff             @@
##           master    #1286      +/-   ##
==========================================
+ Coverage   72.64%   73.12%   +0.48%     
==========================================
  Files          78       78              
  Lines       13028    13302     +274     
==========================================
+ Hits         9464     9727     +263     
- Misses       3564     3575      +11     
Files Changed Coverage Δ
HARK/core.py 88.44% <88.75%> (+0.04%) ⬆️
HARK/ConsumptionSaving/ConsIndShockModel.py 87.15% <94.11%> (+0.10%) ⬆️
HARK/ConsumptionSaving/ConsPortfolioModel.py 90.59% <97.87%> (+0.77%) ⬆️
...ConsumptionSaving/tests/test_ConsPortfolioModel.py 100.00% <100.00%> (ø)
...nsumptionSaving/tests/test_IndShockConsumerType.py 80.70% <100.00%> (+1.68%) ⬆️
HARK/mat_methods.py 58.19% <100.00%> (+31.05%) ⬆️

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

@Mv77
Copy link
Contributor Author

Mv77 commented Jun 15, 2023

You can tell the methods what "states" to keep track off. This is important because:

  • Sometimes we do not want to keep track of some states (e.g. we are using Harmenberg or we do not care about PLvl).
  • Some of the models have versions that augment the state-space. E.g., if in the portfolio model AdjPrb<1.0, the share and the Calvo shock become states.

Here are tests that show that the method can accommodate both versions of the portfolio model.
https://github.com/Mv77/HARK/blob/c4c5064bea6665d096886fe049fc5ec3922c27fc/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py#L310-L345

@Mv77
Copy link
Contributor Author

Mv77 commented Jun 15, 2023

This will benefit from some ongoing work:

  • There is a lot of ugly code that finds the right parameters to use given the age of the agent and time_vary and time_inv. E.g., https://github.com/Mv77/HARK/blob/c4c5064bea6665d096886fe049fc5ec3922c27fc/HARK/core.py#L930-L937 . If @alanlujan91 gets his Parameters class working and merged in, this could be a lot cleaner, like Parameters[t].
  • This is a step towards thinking of models using stochastic transition equations---that is one of the things that needs to be defined for the method to work. I know @sbenthall has been giving some thought to model-representation.
  • There is also ugly code in packing and unpacking matrices that represent states. Labeled xarrays and functions that work with them would make the code more readable. For instance, having a dist_of_func method with labeled distributions as both inputs and outputs would avoid a lot of unpacking.

@Mv77 Mv77 mentioned this pull request Jun 16, 2023
3 tasks
@Mv77
Copy link
Contributor Author

Mv77 commented Jun 16, 2023

@alanlujan91 has pointed me to some work that might help with the cleaning up of stochastic transition equations using DiscreteDistributionLabeled and xarray. Linking for future self-reference:

#1283 (comment)

I think what I want might be possible using DiscreteDistributionLabeled.dist_of_func or lines like these:

HARK/HARK/distribution.py

Lines 1237 to 1238 in 37134b9

f_query = func(self.dataset, *args, **kwargs)
ldd = DiscreteDistributionLabeled.from_dataset(f_query, self.probability)

Comment on lines +410 to +427
def state_to_state_trans(self, shocks_next, solution, state, PermGroFac, Rfree):
# Consumption
cNrm = solution.cFunc(state["mNrm"], state["Share"], state["Adjust"])
# Savings
aNrm = state["mNrm"] - cNrm
# Share
Share_next = solution.ShareFunc(state["mNrm"], state["Share"], state["Adjust"])
# Shock transition
state_next = post_state_transition(
shocks_next,
state["PLvl"],
aNrm,
Share_next,
PermGroFac,
Rfree,
)

return state_next
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just updated the methods to use xarrays and labeled distributions. They now take more expressive transition equations like these.

Comment on lines 931 to 959
def test_compare_will_mateo(self):
# No deaths for now in Mateo's method
# but Will's requires LivPrb < 1
params = deepcopy(dict_harmenberg)
params["LivPrb"] = [0.999]

# Create and solve agent
agent = IndShockConsumerType(**params)
agent.cycles = 0
agent.solve()
# Activate harmenberg
agent.neutral_measure = True
agent.update_income_process()

m_grid = np.linspace(0, 20, 10)

# Will's methods
agent.define_distribution_grid(dist_mGrid=m_grid)
agent.calc_transition_matrix()
tm_will = agent.tran_matrix

# Mateo's methods
agent.make_state_grid(mNrmGrid=m_grid)
agent.full_shock_dstns = agent.IncShkDstn
agent.find_transition_matrices()
tm_mateo = agent.trans_mats[0]

# Compare
self.assertTrue(np.allclose(tm_will, tm_mateo.T, atol=5e-3))
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just added a test that compares the tools I am working on with @wdu9 's methods for IndShockConsumerType.

We are getting the same! (except I am not dealing with death/newborns yet).

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The cool feature that my tools are adding is that, to get this, I only had to define the model's state-to-state transition function.

    def state_to_state_trans(self, shocks_next, solution, state, PermGroFac, Rfree):


        # Consumption
        cNrm = solution.cFunc(state["mNrm"])
        # Savings
        aNrm = state["mNrm"] - cNrm


        # Shock realization
        PermGroShk = shocks_next["PermShk"] * PermGroFac
        PLvl_next = state["PLvl"] * PermGroShk
        mNrm_next = aNrm * Rfree / PermGroShk + shocks_next["TranShk"]
        
        return {
            "PLvl": PLvl_next,
            "mNrm": mNrm_next,
        }

This equation is now pretty expressive and could be used to simulate the model.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just added a way to handle newborns and the methods exactly match.

@sbenthall
Copy link
Contributor

sbenthall commented Jul 21, 2023

Just to make sure it's on your radar....

I see you're contributing this as additions to the ConsIndShock and ConsPortfolio models, and that this involves yet another representation of the model equations to be hard coded into the model.

Transition matrix methods are awesome, but it would be even more awesome if they were implemented generically, so that model-specific information could be passed to them.

What I've been working on in #1296 is a standalone engine for the Monte Carlo simulation that is currently bundled in AgentType. The idea is you pass it the equations and shock definitions you want, and it does the work.

Can something similar be done for transition matrices?

@Mv77
Copy link
Contributor Author

Mv77 commented Jul 21, 2023

Hey @sbenthall,

Yeah, one of the main points of this PR is to have the methods work for any model, receiving their transition equation as a parameter. The method find_transition_matrix belongs to AgentType and takes self.state_to_state_transition as given---whatever it is.

It moves us away from model-speficit methods like

HARK/HARK/utilities.py

Lines 653 to 655 in b123b0e

def gen_tran_matrix_1D(
dist_mGrid, bNext, shk_prbs, perm_shks, tran_shks, LivPrb, NewBornDist
):

There is not a current transition equation definition that is used across models in HARK, so I came up with my own placeholder and showed that it can be defined for multiple models easily. Once there is an agreed-upon representation, this code will be easy to modify to take that one and not the one I am using.

We agree with the goal of having a consistent representation of models that general methods can receive and work with. I am coding with that in mind.

@Mv77
Copy link
Contributor Author

Mv77 commented Jul 21, 2023

To be clear, notice that the main work is being done in HARK.core.AgentType. That is where the "methods" are being added.

But there is no current representation of transitions that is common to ConsIndShock and ConsPortfolio. The additions to those files that you see are just a makeshift common representation for me to work with. Once there is a hark-wide standard, I'll work with it.

But yeah, if you check out the methods, they don't assume much about the model. They work in n-dimensions, you can tell them what subset of states to consider, they work in infinite horizon and life-cycle versions ... they just need a consistent representation for models. I know you are working towards that and I'll be happy to use it once it is agreed upon and merged in.

@sbenthall
Copy link
Contributor

Super, @Mv77 .

@mnwhite I wonder if you ultimately want this sort of thing to be in AgentType (which I agree is the natural place to put it given the current architecture) or if it would make more sense to put it in something like HARK.simulation.transition_matrices. My preference is for the latter.... That could happen even if there isn't a HARK-wide standard for model equations. (It should be relatively easy to shift between ad hoc model equation formulations, since they do very similar things).

I don't mean for this to block the PR, which obviously does a lot of good as is.

@llorracc
Copy link
Collaborator

One generic comment on all of this is to remember that we ultimately want to be specifying our models at a higher level of abstraction in which, for example, in the case of mapping a continuous space into a continuous space, the "transition equation" is described symbolically and with abstract information like the range of of the initial states and the domain of the post-states. Using that as an input, a specific solution of the model would then take that generic description and specialize it. For example, in the solution stage if using a transition matrix method, the specialization would need to invoke the transition matrix representation, and specify the details of the transition matrix being constructed (including, for example, the method of assigning probability weights to future gridpoints). Then in the simulation stage we could choose ueither to use the already-computed transition matrix (guaranteeing compatibility with the solution numerics), or to say use a Monte Carlo, or whether to construct a different "simulation" transition matrix than was used at the solution stage.

Point is to think ahead about this eventuality as we implement these more specialized things here.


# Construct transition matrix from the object above
tmat = np.zeros((mesh_size, mesh_size))
for i in range(mesh_size):
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  • Vectorize

@Mv77
Copy link
Contributor Author

Mv77 commented Aug 11, 2023

I believe I have achieved what I wanted out of this PR in terms of scope. I now need to document things and clean up.

The PR uses representations of transition equations, which I know are in the works. I've introduced my own placeholder for two models.

I'd like to decide if I should wait for the canonical representation of transition equations that I know is in the works, or push towards merging this in with the placeholder representations.

@sbenthall @alanlujan91 @llorracc , could you offer any thoughts or updates?

@sbenthall
Copy link
Contributor

My opinion is that you should go ahead with the clean up and we should merge this in as soon as it's ready.

This PR is good progress, and we can always do the work of aligning it with new transition equation representations later.

If you don't mind, I might go to you first to ask you to adapt this code once the new way of doing equations is settled. But in any case, this is important work to get into HARK!

@llorracc
Copy link
Collaborator

I agree with Seb, that you should go ahead, and we should merge when it is ready.

My point is just to always have people thinking about the longer-term goal so that we do as much as we can to avoid building things in a way that, if we had only done it slightly differently, it would be much easier to upgrade to the new paradigm.

@sbenthall sbenthall mentioned this pull request Aug 29, 2023
3 tasks
@Mv77 Mv77 closed this Apr 23, 2024
@Mv77 Mv77 deleted the lc-transition branch July 12, 2024 23:25
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants