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

Refactor solution and simulation #108

Open
wants to merge 20 commits into
base: main
Choose a base branch
from

Conversation

timmens
Copy link
Member

@timmens timmens commented Feb 26, 2025

In this PR, I try to make clear where simulation and solution code differs, and where it does not. I also perform some additional refactoring along the way.

Specifically, I

  • Refactored the next state module, to showcase the difference between simulation and solution functionality
  • Rename the model_functions module to utility_and_feasibility and rewrite for clarity
  • Use productmap instead of spacemap in the solution
  • Rename spacemap and adjust it so that it's usage during the simulation is clear
  • Integrate comments from the previous PRs
    • I remove the default random seed
  • Create the conditional_continuation module
  • Start refactoring the entry point.
  • Fix stochastic test test_get_lcm_function_with_simulate_target (test did not account for stochastics...)
  • Implement solve_and_simulate function
  • Refactor simulation
  • Add continuous choice grids to state-choice-space
  • Combine both state-choice-space modules

Copy link

Check out this pull request on  ReviewNB

See visual diffs & provide feedback on Jupyter Notebooks.


Powered by ReviewNB

@timmens timmens changed the title Split solution and simulation code and refactor Refactor solution and simulation Feb 27, 2025
@timmens timmens requested a review from hmgaudecker February 28, 2025 15:32
Copy link
Member

@hmgaudecker hmgaudecker left a comment

Choose a reason for hiding this comment

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

Very nice! Only managed to look at the files I have comments in so far, will try to do more on the weekend but no guarantees.

@@ -30,7 +34,7 @@ def get_lcm_function(
targets: Literal["solve", "simulate", "solve_and_simulate"],
debug_mode: bool = True,
jit: bool = True,
) -> tuple[Callable[..., list[Array] | pd.DataFrame], ParamsDict]:
) -> tuple[Callable[..., dict[int, Array] | pd.DataFrame], ParamsDict]:
Copy link
Member

Choose a reason for hiding this comment

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

I don't see from the docstring how "requested targets" are reflected in the output type. Could you change that? Also, in the summary line we don't need need to repeat entry_point. Maybe that could be "Return the solve/simulate functions and a template for the model's parameters." ?

)

target_func: Callable[..., list[Array] | pd.DataFrame]
target_func: Callable[..., dict[int, Array] | pd.DataFrame]
Copy link
Member

Choose a reason for hiding this comment

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

Maybe add a comment that this is just repeating the signature to get rid of mypy complaints? (I suppose)

)

solve_and_simulate_model = partial(
solve_and_simulate,
Copy link
Member

Choose a reason for hiding this comment

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

I would define the function solve_and_simulate here, I think. Or in a separate module of the same name. But I would not look for it in the simulation.simulate module...

If there was a clever way of telling simulate_model that the partialled argument to simulate(vf_arr_dict) will be the output of solve_model (2nd order partialling, I guess), that would be even better...

logger=logger,
)

_next_state_simulate = get_next_state_function(
Copy link
Member

Choose a reason for hiding this comment

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

Move below solve_model = .... I appreciate having both jit calls next to each other but having a "solve" and a "simulate" block would be even better.

period=period,
is_last_period=is_last_period,
)

compute_ccv = create_compute_conditional_continuation_value(
compute_ccv = get_compute_conditional_continuation_value(
Copy link
Member

Choose a reason for hiding this comment

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

Suggested change
compute_ccv = get_compute_conditional_continuation_value(
compute_cc_val = get_compute_conditional_continuation_value(

Symmetry because I'd like ccp to be cc_pol


compute_ccv_argmax = create_compute_conditional_continuation_policy(
compute_ccp = get_compute_conditional_continuation_policy(
Copy link
Member

Choose a reason for hiding this comment

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

Suggested change
compute_ccp = get_compute_conditional_continuation_policy(
compute_cc_pol = get_compute_conditional_continuation_policy(

Since we have CCP estimators and there P stands for probability, this is too confusing IMO.

compute_ccv_functions=compute_ccv_functions,
emax_calculators=emax_calculators,
emax_calculators=solve_discrete_problem_functions,
Copy link
Member

Choose a reason for hiding this comment

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

Would there be a value in renaming the emax_calculators argument, too?

@@ -6,6 +7,7 @@

from lcm.grids import ContinuousGrid, DiscreteGrid, Grid
from lcm.typing import InternalUserFunction, ParamsDict, ShockType
from lcm.utils import first_non_none


@dataclass(frozen=True)
Copy link
Member

Choose a reason for hiding this comment

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

Does this make the mutable attributes frozen, too? Else, would there be a value in using FrozenDicts or the like?

*,
is_last_period: bool,
) -> StateSpaceInfo:
"""Create a state-space information for the model solution.
Copy link
Member

Choose a reason for hiding this comment

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

Suggested change
"""Create a state-space information for the model solution.
"""Collect information on the state space for the model solution.

)


def _validate_initial_states(
Copy link
Member

Choose a reason for hiding this comment

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

Suggested change
def _validate_initial_states(
def _validate_initial_state_names(

This does not do validation of values, right? With the earlier name I'd expect that kind of checks.

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.

2 participants