[pre-commit.ci] pre-commit autoupdate

.pre-commit-config.yaml

@@ -48,7 +48,7 @@ repos:
   #       args:
   #         - --py37-plus
   - repo: https://github.com/pycqa/isort
-    rev: 8.0.1
+    rev: 9.0.0a3
     hooks:
       - id: isort
         name: isort
@@ -59,7 +59,7 @@ repos:
 #    hooks:
 #      - id: setup-cfg-fmt
   - repo: https://github.com/psf/black-pre-commit-mirror
-    rev: 26.1.0
+    rev: 26.3.1
     hooks:
       - id: black
         language_version: python3.13
@@ -73,7 +73,7 @@ repos:
 #            docs/source/conf.py|
 #          )$
   - repo: https://github.com/PyCQA/docformatter
-    rev: v1.7.7
+    rev: v1.7.8
     hooks:
       - id: docformatter
         args:

docs/source/conf.py

@@ -86,11 +86,20 @@
         {
             "name": "GitHub",
             "url": "https://github.com/OpenSourceEconomics/dc-egm",
-            "html": """
-                <svg stroke="currentColor" fill="currentColor" stroke-width="0" viewBox="0 0 16 16">
-                    <path fill-rule="evenodd" d="M8 0C3.58 0 0 3.58 0 8c0 3.54 2.29 6.53 5.47 7.59.4.07.55-.17.55-.38 0-.19-.01-.82-.01-1.49-2.01.37-2.53-.49-2.69-.94-.09-.23-.48-.94-.82-1.13-.28-.15-.68-.52-.01-.53.63-.01 1.08.58 1.23.82.72 1.21 1.87.87 2.33.66.07-.52.28-.87.51-1.07-1.78-.2-3.64-.89-3.64-3.95 0-.87.31-1.59.82-2.15-.08-.2-.36-1.02.08-2.12 0 0 .67-.21 2.2.82.64-.18 1.32-.27 2-.27.68 0 1.36.09 2 .27 1.53-1.04 2.2-.82 2.2-.82.44 1.1.16 1.92.08 2.12.51.56.82 1.27.82 2.15 0 3.07-1.87 3.75-3.65 3.95.29.25.54.73.54 1.48 0 1.07-.01 1.93-.01 2.2 0 .21.15.46.55.38A8.013 8.013 0 0 0 16 8c0-4.42-3.58-8-8-8z"></path>
-                </svg>
-            """,
+            "html": """<svg stroke="currentColor" fill="currentColor" stroke-width="0"
+                    viewBox="0 0 16 16"> <path fill-rule="evenodd" d="M8 0C3.58 0 0 3.58
+                    0 8c0 3.54 2.29 6.53 5.47 7.59.4.07.55-.17.55-.38 0-.19-.01-.82-.01-
+                    1.49-2.01.37-2.53-.49-2.69-.94-.09-.23-.48-.94-.82-1.13-.28-.15-.68-
+                    .52-.01-.53.63-.01 1.08.58 1.23.82.72 1.21 1.87.87
+                    2.33.66.07-.52.28-.87.51-1.07-1.78-.2-3.64-.89-3.64-3.95
+                    0-.87.31-1.59.82-2.15-.08-.2-.36-1.02.08-2.12 0 0 .67-.21
+                    2.2.82.64-.18 1.32-.27 2-.27.68 0 1.36.09 2 .27 1.53-1.04 2.2-.82
+                    2.2-.82.44 1.1.16 1.92.08 2.12.51.56.82 1.27.82 2.15 0 3.07-1.87
+                    3.75-3.65 3.95.29.25.54.73.54 1.48 0 1.07-.01 1.93-.01 2.2 0
+                    .21.15.46.55.38A8.013 8.013 0 0 0 16 8c0-4.42-3.58-8-8-8z"></path>
+                    </svg>"""
+
+               ,
             "class": "",
         },
     ],

src/dcegm/egm/solve_euler_equation.py

@@ -17,7 +17,6 @@ def calculate_candidate_solutions_from_euler_equation(
     params: Dict[str, float],
 ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]:
     """Calculate candidates for the optimal policy and value function."""
-
     feasible_marg_utils_child = jnp.take(
         marg_util_next, idx_post_decision_child_states, axis=0
     )
@@ -208,7 +207,6 @@ def solve_euler_equation(
             choice-specific expected value. Has shape (n_grid_wealth,).
 
     """
-
     transition_vec = compute_stochastic_transition_vec(
         params=params, **state_choice_vec
     )

src/dcegm/interfaces/jit_large_arrays.py

@@ -2,7 +2,6 @@ def split_structure_and_batch_info(model_structure, batch_info):
     """Splits the model structure and batch info into static parts, which we can not jit
     compile and (large) arrays that we want to include in the function call for
     jitting."""
-
     struct_keys_not_for_jit = [
         "discrete_states_names",
         "state_names_without_stochastic",

src/dcegm/interfaces/model_class.py

@@ -62,7 +62,6 @@ def __init__(
             use_stochastic_sparsity (bool, optional): EXPERIMENTAL: Use stochastic transition sparsity.
 
         """
-
         if model_load_path is not None:
             model_dict = load_model_dict(
                 model_config=model_config,
@@ -152,7 +151,6 @@ def solve(self, params, load_sol_path=None, save_sol_path=None):
                 state a transition matrix vector.
 
         """
-
         params_processed = process_params(
             params, params_check_info=self.params_check_info
         )
@@ -249,7 +247,6 @@ def solve_and_simulate(
 
     def get_solve_func(self):
         """Create a fast function for solving that is jit compiled in the first call."""
-
         (
             model_structure_for_jit,
             batch_info_for_jit,
@@ -307,7 +304,6 @@ def get_solve_and_simulate_func(
     ):
         """Create a fast function for solving and simulation that is jit compiled in the
         first call."""
-
         # Fix everything except params, solution of the model and model_structure which contains large arrays.
         sim_func = lambda params, value, policy, endog_gid, model_structure: simulate_all_periods(
             states_initial=states_initial,

src/dcegm/interfaces/sol_interface.py

@@ -121,7 +121,6 @@ def value_for_states_and_choices(self, states, choices):
             The value for the given state and choice.
 
         """
-
         return value_for_state_and_choice(
             states=states,
             choices=choices,
@@ -144,7 +143,6 @@ def policy_for_states_and_choices(self, states, choices):
             The policy for the given state and choice.
 
         """
-
         return policy_for_state_choice_vec(
             states=states,
             choices=choices,

src/dcegm/interpolation/interp1d.py

@@ -63,7 +63,6 @@ def interp1d_policy_and_value_on_wealth(
         - value_interp (float): Interpolated value for wealth.
 
     """
-
     # For all choices, the wealth is the same in the solution
     ind_high, ind_low = get_index_high_and_low(x=wealth_grid, x_new=wealth)
 
@@ -114,7 +113,6 @@ def interp_value_on_wealth(
         jnp.ndarray | float: Interpolated value for wealth.
 
     """
-
     ind_high, ind_low = get_index_high_and_low(x=wealth_grid, x_new=wealth)
 
     value_interp = interp_value_and_check_creditconstraint(
@@ -152,7 +150,6 @@ def interp_policy_on_wealth(
         float: Interpolated value for wealth.
 
     """
-
     ind_high, ind_low = get_index_high_and_low(x=endog_grid, x_new=wealth)
 
     policy_interp = linear_interpolation_formula(
@@ -207,7 +204,6 @@ def interp_value_and_check_creditconstraint(
         - value_interp (float): Interpolated value function.
 
     """
-
     value_interp_on_grid = linear_interpolation_formula(
         y_high=value_high,
         y_low=value_low,

src/dcegm/interpolation/interp2d.py

@@ -62,7 +62,6 @@ def interp2d_policy_and_value_on_wealth_and_regular_grid(
             value function.
 
     """
-
     regular_points, wealth_points, coords_idxs = find_grid_coords_for_interp(
         regular_grid=regular_grid,
         wealth_grid=wealth_grid,
@@ -139,7 +138,6 @@ def interp2d_value_on_wealth_and_regular_grid(
             (regular, wealth) point.
 
     """
-
     regular_points, wealth_points, coords_idxs = find_grid_coords_for_interp(
         regular_grid=regular_grid,
         wealth_grid=wealth_grid,
@@ -197,7 +195,6 @@ def interp2d_policy_on_wealth_and_regular_grid(
             (regular, wealth) point.
 
     """
-
     regular_points, wealth_points, coords_idxs = find_grid_coords_for_interp(
         regular_grid=regular_grid,
         wealth_grid=wealth_grid,
@@ -249,7 +246,6 @@ def interp2d_policy(
             (regular, wealth) point.
 
     """
-
     policy_known = policy_grid[coords_idxs[:, 0], coords_idxs[:, 1]]
 
     policy_interp = interp2d(
@@ -402,7 +398,6 @@ def find_grid_coords_for_interp(
         - The fourth element is the upper left point.
 
     """
-
     # Determine the closest points in the regular direction
     regular_idx_right, regular_idx_left = get_index_high_and_low(
         regular_grid, regular_point_to_interp
@@ -459,7 +454,6 @@ def interp2d(x_coords, y_coords, z_vals, x_new, y_new):
         float: The interpolated value at the point (x_new, y_new).
 
     """
-
     # Map the irregular quadrilateral onto a canonical unit square
     x_vec, y_vec = TRANSFORMATION_MAT @ x_coords, TRANSFORMATION_MAT @ y_coords
 
@@ -496,7 +490,6 @@ def determine_coordinates_in_unit_square(x, y, x_vec, y_vec):
             the x-axis, and `y_rel` is the relative position along the y-axis.
 
     """
-
     x_rel = (x - x_vec[0]) / x_vec[1]
     y_rel = (y - y_vec[0] - y_vec[1] * x_rel) / (y_vec[2] + y_vec[3] * x_rel)
 
@@ -521,5 +514,4 @@ def compute_vertex_weights(x, y):
             unit square, respectively.
 
     """
-
     return jnp.array([(1 - x) * (1 - y), x * (1 - y), x * y, (1 - x) * y])

src/dcegm/likelihood.py

@@ -392,7 +392,6 @@ def calc_choice_prob_for_state_choices(
     and then interpolates the wealth at the beginning of period on them.
 
     """
-
     choice_prob_across_choices = calc_choice_probs_for_states(
         value_solved=value_solved,
         endog_grid_solved=endog_grid_solved,

src/dcegm/pre_processing/alternative_sim_functions.py

@@ -48,7 +48,6 @@ def generate_alternative_sim_functions(
         budget_constraint (Callable): User supplied budget constraint.
 
     """
-
     model_config = check_model_config_and_process(model_config)
 
     model_funcs, _ = process_alternative_sim_functions(

src/dcegm/pre_processing/batches/batch_creation.py

@@ -45,7 +45,6 @@ def create_batches_and_information(
             batch arrays from the step before.
 
     """
-
     last_two_period_info = add_last_two_period_information(
         n_periods=n_periods,
         model_structure=model_structure,

src/dcegm/pre_processing/batches/single_segment.py

@@ -9,7 +9,6 @@ def create_single_segment_of_batches(bool_state_choices_to_batch, model_structur
     If the last batch is not evenly we correct it.
 
     """
-
     state_choice_space = model_structure["state_choice_space"]
     state_choice_space_dict = model_structure["state_choice_space_dict"]
 

src/dcegm/pre_processing/check_model_config.py

@@ -4,7 +4,6 @@
 
 def check_model_config_and_process(model_config):
     """Check if options are valid and set defaults."""
-
     processed_model_config = {}
 
     if not isinstance(model_config, dict):

src/dcegm/pre_processing/model_structure/deterministic_states.py

@@ -8,7 +8,6 @@
 
 def process_endog_state_specifications(state_space_options):
     """Get number of endog states which we loop over when creating the state space."""
-
     if state_space_options.get("deterministic_states"):
 
         endog_states_names = list(state_space_options["deterministic_states"].keys())

src/dcegm/pre_processing/model_structure/shared.py

@@ -5,7 +5,6 @@
 
 def create_indexer_for_space(space, max_var_values=None):
     """Create indexer for space."""
-
     # Indexer has always unsigned data type with integers starting at zero
     # Leave one additional value for the invalid number
     data_type = get_smallest_int_type(space.shape[0] + 1)

src/dcegm/pre_processing/model_structure/state_choice_space.py

@@ -59,7 +59,6 @@ def create_state_choice_space_and_child_state_mapping(
             (ii) to expand state level arrays to the state-choice level.
 
     """
-
     state_names_without_stochastic = state_space_arrays[
         "state_names_without_stochastic"
     ]

src/dcegm/pre_processing/model_structure/stochastic_states.py

@@ -52,7 +52,6 @@ def process_stochastic_transitions(
         options (dict): Options dictionary.
 
     """
-
     func_list = []
     func_dict = {}
 

src/dcegm/pre_processing/setup_model.py

@@ -166,7 +166,6 @@ def create_model_dict_and_save(
     than recreating the model from scratch.
 
     """
-
     model_dict = create_model_dict(
         model_config=model_config,
         model_specs=model_specs,
@@ -202,7 +201,6 @@ def load_model_dict(
     use_stochastic_sparsity=False,
 ):
     """Load the model from file."""
-
     model = pickle.load(open(path, "rb"))
 
     model["model_config"] = check_model_config_and_process(model_config)

src/dcegm/simulation/random_keys.py

@@ -4,7 +4,6 @@
 
 def draw_random_keys_for_seed(n_agents, n_periods, taste_shock_scale_is_scalar, seed):
     """Draw the random keys jax uses for a given seed."""
-
     # We start by determining the number of keys per period  for the three stochastic components
     # of the model. We will draw all keys together and assign them then based on index.
     # First: The transition of the exogenous processes. This happens vectorized over all

src/dcegm/solve_single_period.py

@@ -205,7 +205,6 @@ def run_upper_envelope(
     Vectorized over all state-choice combinations.
 
     """
-
     if has_second_continuous_state:
         return vmap(
             vmap(

src/dcegm/templates/fullmodel/model_funcs/state_space_objects.py

@@ -14,7 +14,6 @@ def create_state_space_function_dict():
 
 def state_specific_choice_set(period, lagged_choice, job_offer, model_specs):
     """Return available choices depending on state."""
-
     # Retirement is absorbing
     if lagged_choice == 0:
         return [0]

tests/test_law_of_motion.py

@@ -55,7 +55,6 @@ def _calc_stochastic_income_for_experience(
     params: Dict[str, float],
 ) -> float:
     """Computes the current level of deterministic and stochastic income."""
-
     log_wage = (
         params["constant"]
         + params["exp"] * experience

tests/test_sparse_stochastic_transitions.py

@@ -17,7 +17,8 @@
 
 
 def prob_exog_health_father(health_mother):
-    """Sparse transition: can only go to 2 out of 3 states depending on mother's health."""
+    """Sparse transition: can only go to 2 out of 3 states depending on mother's
+    health."""
     # health_mother == 0: can go to states 0, 1 (not 2)
     # health_mother == 1: can go to states 1, 2 (not 0)
     # health_mother == 2: can go to states 0, 2 (not 1)
@@ -40,7 +41,8 @@ def prob_exog_health_father(health_mother):
 
 
 def prob_exog_health_mother(health_father):
-    """Sparse transition: can only go to 2 out of 3 states depending on father's health."""
+    """Sparse transition: can only go to 2 out of 3 states depending on father's
+    health."""
     # health_father == 0: can go to states 0, 1 (not 2)
     # health_father == 1: can go to states 0, 2 (not 1)
     # health_father == 2: can go to states 1, 2 (not 0)
@@ -108,7 +110,6 @@ def util_new(
 
 def test_sparse_stochastic_transitions():
     """Test that solving with sparse transitions gives same results."""
-
     params = {
         "rho": 2,
         "delta": 0.5,

tests/utils/interp2d_auxiliary.py

@@ -34,7 +34,6 @@ def custom_interp2d_quad(x_grids, y_grid, values, points):
         Output array that contains the interpolated values
 
     """
-
     n_y_grid = y_grid.shape[0]
     n_points = points.shape[0]
     z = np.empty((n_points, 4))
@@ -98,7 +97,6 @@ def custom_interp2d_quad_value_function(
         Output array that contains the interpolated values
 
     """
-
     n_y_grid = y_grid.shape[0]
     n_points = points.shape[0]
 
@@ -178,7 +176,6 @@ def compute_weights(l, m):
 def search(grid_length, sorted_grid, element):
     """This function searches the element on the specific grid that is used for the
     interpolation."""
-
     # if the element is outside the bounds
     # use the border elements for interpolation