def test_wage_nonpecs(): """Replace constants in reward functions with constants due to observables.""" point_constr = {"n_periods": 3, "n_lagged_choices": 1, "observables": [2]} params, options = generate_random_model(point_constr=point_constr) solve = get_solve_func(params, options) state_space = solve(params) # Replace constants in choices with constant utility by observables. optim_paras, _ = process_params_and_options(params, options) wage_choice = ("wage", np.random.choice(optim_paras["choices_w_wage"])) nonpec_choice = ("nonpec", np.random.choice(list(optim_paras["choices"]))) # Change specs accordingly for reward in [wage_choice, nonpec_choice]: constant = params.loc[(f"{reward[0]}_{reward[1]}", "constant"), "value"] params = params.drop(index=[(f"{reward[0]}_{reward[1]}", "constant")]) for obs in range(2): params.loc[(f"{reward[0]}_{reward[1]}", f"observable_0_{obs}"), "value"] += constant solve = get_solve_func(params, options) state_space_ = solve(params) for attribute in ["wages", "nonpecs"]: apply_to_attributes_of_two_state_spaces( getattr(state_space, attribute), getattr(state_space_, attribute), np.testing.assert_array_almost_equal, )
def get_solve_func(params, options): """Get the solve function. This function takes a model specification and returns the state space of the model along with components of the solution such as covariates, non-pecuniary rewards, wages, continuation values and expected value functions as attributes of the class. Parameters ---------- params : pandas.DataFrame DataFrame containing parameter series. options : dict Dictionary containing model attributes which are not optimized. Returns ------- solve : :func:`~respy.solve.solve` Function with partialed arguments. Examples -------- >>> import respy as rp >>> params, options = rp.get_example_model("robinson_crusoe_basic", with_data=False) >>> solve = rp.get_solve_func(params, options) >>> state_space = solve(params) """ optim_paras, options = process_params_and_options(params, options) state_space = create_state_space_class(optim_paras, options) solve_function = functools.partial(solve, options=options, state_space=state_space) return solve_function
def test_create_state_space_vs_specialized_kw97(model_or_seed): params, options = process_model_or_seed(model_or_seed) # Reduce runtime options["n_periods"] = 10 if options["n_periods"] > 10 else options[ "n_periods"] optim_paras, options = process_params_and_options(params, options) # Create old state space arguments. n_periods = options["n_periods"] n_types = optim_paras["n_types"] edu_max = optim_paras["choices"]["school"]["max"] edu_starts = np.array(list(optim_paras["choices"]["school"]["start"])) # Get states and indexer from old state space. if model_or_seed == "kw_97_basic": states_old, indexer_old = _create_state_space_kw97_base( n_periods, n_types, edu_starts, edu_max) else: states_old, indexer_old = _create_state_space_kw97_extended( n_periods, n_types, edu_starts, edu_max) states_new, indexer_new = _create_state_space(optim_paras, options) # Compare the state spaces via sets as ordering changed in some cases. states_old_set = set(map(tuple, states_old)) states_new_set = set(map(tuple, states_new.to_numpy())) assert states_old_set == states_new_set # Compare indexers via masks for valid indices. for period in range(n_periods): mask_old = indexer_old[period] != -1 mask_new = indexer_new[period] != -1 assert np.array_equal(mask_old, mask_new)
def solve(params, options, state_space): """Solve the model.""" optim_paras, options = process_params_and_options(params, options) transit_keys = None if hasattr(state_space, "dense_key_to_transit_keys"): transit_keys = state_space.dense_key_to_transit_keys wages, nonpecs = _create_param_specific_objects( state_space.dense_key_to_complex, state_space.dense_key_to_choice_set, optim_paras, options, transit_keys=transit_keys, bypass={ "dense_key_to_dense_covariates": state_space.dense_key_to_dense_covariates }, ) state_space.wages = wages state_space.nonpecs = nonpecs state_space = _solve_with_backward_induction(state_space, optim_paras, options) return state_space
def test_choice_restrictions(): """Basic first test.""" # Load model. params, options = process_model_or_seed("robinson_crusoe_extended") # Extend with observable characteristic. params.loc[("observable_health_well", "probability"), "value"] = 0.9 params.loc[("observable_health_sick", "probability"), "value"] = 0.1 # Sick people can never work. options["negative_choice_set"] = { "fishing": ["health == 'sick' & period < 2", "health == 'sick' & period >= 2"], "friday": ["period < 2", "exp_fishing == 0"], } # Create internal specification objects. optim_paras, options = process_params_and_options(params, options) state_space = create_state_space_class(optim_paras, options) for x in state_space.dense_key_to_complex.values(): if (x[0] < 2) & (x[2] == (0, )): assert x[1] == (False, False, True) elif x[2] == (0, ): assert x[1] in [(False, False, True), (False, True, True)] elif (x[0] < 2) & (x[2] == (1, )): assert x[1] == (True, False, True) elif x[2] == (1, ): assert x[1] in [(True, False, True), (True, True, True)]
def test_state_space_restrictions_by_traversing_forward(model): """Test for inadmissible states in the state space. The test is motivated by the addition of another restriction in https://github.com/OpenSourceEconomics/respy/pull/145. To ensure that similar errors do not happen again, this test takes all states of the first period and finds all their child states. Taking only the child states their children are found and so on. At last, the set of visited states is compared against the total set of states. The test can only applied to some models. Most models would need custom ``options["core_state_space_filters"]`` to remove inaccessible states from the state space. """ params, options = process_model_or_seed(model) optim_paras, options = process_params_and_options(params, options) solve = get_solve_func(params, options) state_space = solve(params) out = {} for x in state_space.child_indices.values(): array = np.concatenate(x) for state in array: if state[0] in out.keys(): if state[1] not in out[state[0]]: out[state[0]].append(state[1]) else: continue else: out[state[0]] = [state[1]] for x in out: assert len(out[x]) == len(state_space.core_key_to_core_indices[x])
def test_invariance_of_wage_calc(): """The model reproduces invariant properties of wage outcomes.""" point_constr = {"n_periods": 2, "observables": [3]} params, options = generate_random_model(point_constr=point_constr) # Add some inadmissible states optim_paras, _ = process_params_and_options(params, options) # Solve first model solve = get_solve_func(params, options) state_space = solve(params) pos = np.random.choice(range(len(state_space.dense))) dense_combination = list(state_space.dense.keys())[pos] dense_index = state_space.dense_covariates_to_dense_index[ dense_combination] idx = state_space.core_key_and_dense_index_to_dense_key[(1, dense_index)] # Solve relevant wages wages_b = state_space.wages[idx][:, 1] # Impose some restriction options["negative_choice_set"] = {"a": ["period == 1"]} solve = get_solve_func(params, options) state_space = solve(params) wages_b_alt = state_space.wages[idx][:, 0] np.testing.assert_array_equal(wages_b, wages_b_alt)
def test_child_indices(): """Testing existence of properties for calculation of child indices!""" point_constr = {"n_periods": 2, "n_lagged_choices": 1} params, options = generate_random_model(point_constr=point_constr) # Add some inadmissible states optim_paras, options = process_params_and_options(params, options) state_space = create_state_space_class(optim_paras, options) # Create all relevant columns core_columns = ["period"] + create_core_state_space_columns(optim_paras) # compose child indices of first choice initial_state = state_space.core.iloc[0][core_columns].to_numpy() # Get all the future states states = [] for i in range(len(optim_paras["choices"])): child = initial_state.copy() child[0] += 1 child[i + 1] += 1 child[-1] = i ix = state_space.indexer[tuple(child)] states.append(np.array(ix).reshape(1, 2)) manual = np.concatenate(states, axis=0) np.testing.assert_array_equal(state_space.child_indices[0][0], manual)
def test_create_state_space_vs_specialized_kw94(model_or_seed): point_constr = {"n_lagged_choices": 1, "observables": False} params, options = process_model_or_seed(model_or_seed, point_constr=point_constr) optim_paras, options = process_params_and_options(params, options) # Create old state space arguments. n_periods = options["n_periods"] n_types = optim_paras["n_types"] edu_max = optim_paras["choices"]["edu"]["max"] edu_starts = np.array(list(optim_paras["choices"]["edu"]["start"])) # Get states and indexer from old state space. states_old, indexer_old = _create_state_space_kw94(n_periods, n_types, edu_starts, edu_max) states_new, indexer_new = _create_state_space(optim_paras, options) # Compare the state spaces via sets as ordering changed in some cases. states_old_set = set(map(tuple, states_old)) states_new_set = set(map(tuple, states_new.to_numpy())) assert states_old_set == states_new_set # Compare indexers via masks for valid indices. for period in range(n_periods): mask_old = indexer_old[period] != -1 mask_new = indexer_new[period] != -1 assert np.array_equal(mask_old, mask_new)
def solve(params, options): """Solve the model. This function takes a model specification and returns the state space of the model along with components of the solution such as covariates, non-pecuniary rewards, wages, continuation values and value functions as attributes of the class. Parameters ---------- params : pandas.DataFrame DataFrame containing parameter series. options : dict Dictionary containing model attributes which are not optimized. Returns ------- state_space : :class:`~respy.state_space.StateSpace` State space of the model which is already solved via backward-induction. """ optim_paras, options = process_params_and_options(params, options) state_space = StateSpace(optim_paras, options) state_space = solve_with_backward_induction(state_space, optim_paras, options) return state_space
def test_create_state_space_vs_specialized_kw94(model): point_constr = {"n_lagged_choices": 1, "observables": False} params, options = process_model_or_seed(model, point_constr=point_constr) optim_paras, options = process_params_and_options(params, options) # Create old state space arguments. n_periods = options["n_periods"] n_types = optim_paras["n_types"] edu_max = optim_paras["choices"]["edu"]["max"] edu_starts = np.array(list(optim_paras["choices"]["edu"]["start"])) # Get states and indexer from old state space. states_old, indexer_old = _create_state_space_kw94(n_periods, n_types, edu_starts, edu_max) if n_types == 1: states_old = states_old[:, :-1] for i, idx in enumerate(indexer_old): shape = idx.shape indexer_old[i] = idx.reshape(shape[:-2] + (-1, )) states_new, indexer_new = _create_core_and_indexer(optim_paras, options) # Compare the state spaces via sets as ordering changed in some cases. states_old_set = set(map(tuple, states_old)) states_new_set = set(map(tuple, states_new.to_numpy())) assert states_old_set == states_new_set # Compare indexers via masks for valid indices. for period in range(n_periods): mask_old = indexer_old[period] != INDEXER_INVALID_INDEX mask_new = indexer_new[period] != INDEXER_INVALID_INDEX assert np.array_equal(mask_old, mask_new)
def test_invariance_of_solution(model_or_seed): """Test for the invariance of the solution. We run solve two times and check whether all attributes of the state space match. """ params, options = process_model_or_seed(model_or_seed) optim_paras, options = process_params_and_options(params, options) solve = get_solve_func(params, options) state_space = solve(params) state_space_ = solve(params) for attribute in [ "core", "wages", "nonpecs", "expected_value_functions", "base_draws_sol", ]: apply_to_attributes_of_two_state_spaces( getattr(state_space, attribute), getattr(state_space_, attribute), np.testing.assert_array_equal, )
def log_like_obs( params, choices, idx_indiv_first_obs, indices, log_wages_observed, base_draws_est, state_space, type_covariates, options, ): """Criterion function for the likelihood maximization. This function calculates the likelihood contributions of the sample. Parameters ---------- params : pandas.Series Parameter Series choices : numpy.ndarray Array with shape (n_observations * n_types) containing choices for each individual-period pair. idx_indiv_first_obs : numpy.ndarray Array with shape (n_individuals,) containing indices for the first observation of each individual in the data. This is used to aggregate probabilities of the individual over all periods. indices : numpy.ndarray Array with shape (n_observations, n_types) containing indices to map each observation to its correponding state for each type. log_wages_observed : numpy.ndarray Array with shape (n_observations, n_types) containing observed log wages. base_draws_est : numpy.ndarray Set of draws to calculate the probability of observed wages. state_space : :class:`~respy.state_space.StateSpace` State space. options : dict Contains model options. """ optim_paras, options = process_params_and_options(params, options) state_space.update_systematic_rewards(optim_paras) state_space = solve_with_backward_induction(state_space, optim_paras, options) contribs = _internal_log_like_obs( state_space, choices, idx_indiv_first_obs, indices, log_wages_observed, base_draws_est, type_covariates, optim_paras, options, ) return contribs
def test_create_state_space_vs_specialized_kw97(model): """State space reproduces invariant features of the kw97 state space.""" params, options = process_model_or_seed(model) optim_paras, options = process_params_and_options(params, options) # Create old state space arguments. n_periods = options["n_periods"] n_types = optim_paras["n_types"] edu_max = optim_paras["choices"]["school"]["max"] edu_starts = np.array(list(optim_paras["choices"]["school"]["start"])) # Get states and indexer from old state space. if "kw_97_basic" in model: states_old, indexer_old = _create_state_space_kw97_base( n_periods, n_types, edu_starts, edu_max) else: states_old, indexer_old = _create_state_space_kw97_extended( n_periods, n_types, edu_starts, edu_max) states_old = states_old[:, :-1] states_new = _create_core_state_space(optim_paras, options) core_period_choice = _create_core_period_choice(states_new, optim_paras, options) # I think here we can get more elegant! Or is this the only way? core_index_to_complex = {i: k for i, k in enumerate(core_period_choice)} core_index_to_indices = { i: core_period_choice[core_index_to_complex[i]] for i in core_index_to_complex } # Create sp indexer indexer = _create_indexer(states_new, core_index_to_indices, optim_paras) # Compare the state spaces via sets as ordering changed in some cases. states_old_set = set(map(tuple, states_old)) states_new_set = set(map(tuple, states_new.to_numpy())) assert states_old_set == states_new_set # Compare indexers via masks for valid indices. for period in range(n_periods): index_old_period = indexer_old[period] != INDEXER_INVALID_INDEX index_old_period = np.nonzero(index_old_period) indices_old = [ [period] + [index_old_period[x][i] for x in range(len(index_old_period) - 1)] for i in range(len(index_old_period[0])) ] for index in indexer.keys(): if index[0] == period: assert list(index) in indices_old for index in indices_old: assert tuple(index) in indexer.keys()
def test_check_solution(model_or_seed): params, options = process_model_or_seed(model_or_seed) state_space = rp.solve(params, options) optim_paras, options = process_params_and_options(params, options) check_model_solution(optim_paras, options, state_space)
def test_check_solution(model_or_seed): """Test solution of a random model.""" params, options = process_model_or_seed(model_or_seed) solve = get_solve_func(params, options) state_space = solve(params) optim_paras, options = process_params_and_options(params, options) check_model_solution(optim_paras, options, state_space)
def test_normalize_probabilities(): constraints = {"observables": [3]} params, options = generate_random_model(point_constr=constraints) optim_paras_1, _ = process_params_and_options(params, options) for group in ["initial_exp_edu", "observable_"]: mask = params.index.get_level_values(0).str.contains(group) params.loc[mask, "value"] = params.loc[mask, "value"].to_numpy() / 2 with pytest.warns(UserWarning, match=r"The probabilities for parameter group"): optim_paras_2, _ = process_params_and_options(params, options) for key in optim_paras_1["choices"]["edu"]["start"]: np.testing.assert_array_almost_equal( optim_paras_1["choices"]["edu"]["start"][key], optim_paras_2["choices"]["edu"]["start"][key], ) for level in optim_paras_1["observables"]["observable_0"]: np.testing.assert_array_almost_equal( optim_paras_1["observables"]["observable_0"][level], optim_paras_2["observables"]["observable_0"][level], )
def test_state_space_restrictions_by_traversing_forward(model): """Test for inadmissible states in the state space. The test is motivated by the addition of another restriction in https://github.com/OpenSourceEconomics/respy/pull/145. To ensure that similar errors do not happen again, this test takes all states of the first period and finds all their child states. Taking only the child states their children are found and so on. At last, the set of visited states is compared against the total set of states. The test can only applied to some models. Most models would need custom ``options["core_state_space_filters"]`` to remove inaccessible states from the state space. """ params, options = process_model_or_seed(model) optim_paras, options = process_params_and_options(params, options) solve = get_solve_func(params, options) state_space = solve(params) indices = np.full((state_space.core.shape[0], len(optim_paras["choices"])), INDEXER_INVALID_INDEX) core_columns = create_core_state_space_columns(optim_paras) for period in range(options["n_periods"] - 1): if period == 0: states = state_space.core.query( "period == 0")[core_columns].to_numpy(np.int) else: indices_period = state_space.indices_of_child_states[ state_space.slices_by_periods[period - 1]] indices_period = indices_period[indices_period >= 0] states = state_space.core[core_columns].to_numpy( np.int)[indices_period] indices = _insert_indices_of_child_states( indices, states, state_space.indexer[period], state_space.indexer[period + 1], state_space.is_inadmissible, len(optim_paras["choices_w_exp"]), optim_paras["n_lagged_choices"], ) # Take all valid indices and add the indices of the first period. set_valid_indices = set(indices[indices != INDEXER_INVALID_INDEX]) | set( range(state_space.core.query("period == 0").shape[0])) assert set_valid_indices == set(range(state_space.core.shape[0]))
def test_simulated_data(model_or_seed): """Test simulated data with ``check_simulated_data``. Note that, ``check_estimation_data`` is also tested in this function as these tests focus on a subset of the data. """ params, options = process_model_or_seed(model_or_seed) simulate = rp.get_simulate_func(params, options) df = simulate(params) optim_paras, _ = process_params_and_options(params, options) check_simulated_data(optim_paras, df)
def test_sorting_of_type_probability_parameters(model_or_seed): # Set configuration for random models. n_types = np.random.randint(2, 5) n_type_covariates = np.random.randint(2, 4) params, options = process_model_or_seed( model_or_seed, n_types=n_types, n_type_covariates=n_type_covariates) optim_paras, options = process_params_and_options(params, options) # Update variables if not a random model, but an example model is tested. if isinstance(model_or_seed, str): n_types = optim_paras["n_types"] n_type_covariates = (None if n_types == 1 else len( optim_paras["type_covariates"])) if optim_paras["n_types"] > 1: # Resort type probability parameters. types = [f"type_{i}" for i in range(2, optim_paras["n_types"] + 1)] params.loc[types] = params.sort_index(ascending=False).loc[types] optim_paras_, _ = process_params_and_options(params, options) assert (optim_paras["type_prob"] == optim_paras_["type_prob"]).all()
def log_like( params, df, base_draws_est, solve, type_covariates, options, return_scalar, ): """Criterion function for the likelihood maximization. This function calculates the likelihood contributions of the sample. Parameters ---------- params : pandas.Series Parameter Series df : pandas.DataFrame The DataFrame contains choices, log wages, the indices of the states for the different types. base_draws_est : numpy.ndarray Set of draws to calculate the probability of observed wages. solve : :func:`~respy.solve.solve` Function which solves the model with new parameters. options : dict Contains model options. """ optim_paras, options = process_params_and_options(params, options) state_space = solve(params) contribs, df, log_type_probabilities = _internal_log_like_obs( state_space, df, base_draws_est, type_covariates, optim_paras, options) # Return mean log likelihood or log likelihood contributions. out = contribs.mean() if not return_scalar: out = { "value": out, "contributions": contribs, "comparison_plot_data": _create_comparison_plot_data(df, log_type_probabilities, optim_paras), } return out
def solve(params, options, state_space): """Solve the model.""" optim_paras, options = process_params_and_options(params, options) states = state_space.states is_inadmissible = state_space.get_attribute("is_inadmissible") wages, nonpecs = _create_choice_rewards(states, is_inadmissible, optim_paras) state_space.set_attribute("wages", wages) state_space.set_attribute("nonpecs", nonpecs) state_space = _solve_with_backward_induction(state_space, optim_paras, options) return state_space
def test_dense_period_choice(): params, options = rp.get_example_model("kw_94_one", with_data=False) options["negative_choice_set"] = {} options["negative_choice_set"]["b"] = ["period < 5"] optim_paras, options = process_params_and_options(params, options) state_space = create_state_space_class(optim_paras, options) check = _create_dense_period_choice(state_space.core, state_space.dense, state_space.core_key_to_core_indices, state_space.core_key_to_complex, optim_paras, options) for key in check: if key[0] < 5: assert ~key[1][1]
def simulate_truncated_data(params, options, is_missings=True): """Simulate a (truncated) dataset. The data can have two more properties. First, truncated history, second, missing wages. """ optim_paras, _ = process_params_and_options(params, options) simulate = get_simulate_func(params, options) df = simulate(params) np.random.seed(options["simulation_seed"]) if is_missings: # Truncate the histories of agents. This mimics the effect of attrition. # Histories can be truncated after the first period or not at all. So, all # individuals have at least one observation. period_of_truncation = ( # noqa: F841 df.reset_index() .groupby("Identifier") .Period.transform(lambda x: np.random.choice(x.max() + 1) + 1) .to_numpy() ) data_subset = df.query("Period < @period_of_truncation").copy() # Add some missings to wage data. is_working = data_subset["Choice"].isin(optim_paras["choices_w_wage"]) num_drop_wages = int(is_working.sum() * np.random.uniform(high=0.5)) if num_drop_wages > 0: indices = data_subset["Wage"][is_working].index index_missing = np.random.choice(indices, num_drop_wages, replace=False) data_subset.loc[index_missing, "Wage"] = np.nan else: pass else: data_subset = df # We can restrict the information to observed entities only. col_dtype = generate_column_dtype_dict_for_estimation(optim_paras) data_subset = data_subset[list(col_dtype)[2:]] return data_subset
def test_simulation_with_flexible_choice_sets(): params, options = process_model_or_seed("robinson_crusoe_basic") # Extend with observable characteristic. params.loc[("observable_health_well", "probability"), "value"] = 0.9 params.loc[("observable_health_sick", "probability"), "value"] = 0.1 # Sick people can never work. options["negative_choice_set"] = { "fishing": ["health == 'sick'"], "friday": ["period < 2", "exp_fishing == 0"], } # Create internal specification objects. optim_paras, options = process_params_and_options(params, options) simulate = get_simulate_func(params, options) df = simulate(params) assert isinstance(df, pd.DataFrame)
def solve(params, options, state_space): """Solve the model.""" optim_paras, options = process_params_and_options(params, options) wages, nonpecs = _create_choice_rewards( state_space.dense_key_to_complex, state_space.dense_key_to_choice_set, optim_paras, options, ) state_space.wages = wages state_space.nonpecs = nonpecs state_space = _solve_with_backward_induction(state_space, optim_paras, options) return state_space
def create_kw_97(params, options): """Create data for Keane and Wolpin (1997). The data includes individuals labor market history and accumulated experiences in white-collar, blue-collar occupations, military and schooling. """ optim_paras, options = process_params_and_options(params, options) dtypes = { "Identifier": np.int, "Age": np.int, "Experience_School": np.uint8, "Choice": "category", "Wage": np.float, } df = pd.read_csv(TEST_RESOURCES_DIR / "kw_97_data.csv", dtype=dtypes, float_precision="high") df.Identifier = df.groupby("Identifier").ngroup().astype(np.uint16) codes_to_choices = { "3": "white_collar", "4": "blue_collar", "5": "military", "1": "school", "2": "home", } df.Choice = df.Choice.cat.set_categories( codes_to_choices).cat.rename_categories(codes_to_choices) df = _create_working_experience(df, optim_paras) df["Lagged_Choice_1"] = df.groupby("Identifier").Choice.shift(1) df["Period"] = df.Age - 16 df = df.query("Age >= 16") cd_dict = rp_shared.generate_column_dtype_dict_for_estimation(optim_paras) df = df[cd_dict].set_index(["Identifier", "Period"]) return df
def get_simulate_func(params, options): """Get the simulation function. Return :func:`simulate` where all arguments except the parameter vector are fixed with :func:`functools.partial`. Thus, the function can be directly passed into an optimizer for estimation with simulated method of moments or other techniques. Parameters ---------- params : pandas.DataFrame DataFrame containing model parameters. options : dict Dictionary containing model options. Returns ------- simulate_function : :func:`simulate` Simulation function where all arguments except the parameter vector are set. """ optim_paras, options = process_params_and_options(params, options) state_space = StateSpace(optim_paras, options) shape = ( options["n_periods"], options["simulation_agents"], len(optim_paras["choices"]), ) base_draws_sim = create_base_draws( shape, next(options["simulation_seed_startup"]), "random") base_draws_wage = create_base_draws( shape, next(options["simulation_seed_startup"]), "random") simulate_function = functools.partial( simulate, base_draws_sim=base_draws_sim, base_draws_wage=base_draws_wage, state_space=state_space, options=options, ) return simulate_function
def test_state_space_restrictions_by_traversing_forward(model_or_seed): """Test for inadmissible states in the state space. The test is motivated by the addition of another restriction in https://github.com/OpenSourceEconomics/respy/pull/145. To ensure that similar errors do not happen again, this test takes all states of the first period and finds all their child states. Taking only the child states their children are found and so on. At last, the set of visited states is compared against the total set of states. """ params, options = process_model_or_seed(model_or_seed) optim_paras, options = process_params_and_options(params, options) state_space = rp.solve(params, options) indices = np.full( (state_space.states.shape[0], len(optim_paras["choices"])), -1) for period in range(options["n_periods"] - 1): if period == 0: states = state_space.get_attribute_from_period("states", period) else: indices_period = state_space.get_attribute_from_period( "indices_of_child_states", period - 1) indices_period = indices_period[indices_period >= 0] states = state_space.states[indices_period] indices = _insert_indices_of_child_states( indices, states, state_space.indexer[period], state_space.indexer[period + 1], state_space.is_inadmissible, len(optim_paras["choices_w_exp"]), optim_paras["n_lagged_choices"], ) # Take all valid indices and add the indices of the first period. set_valid_indices = set(indices[indices >= 0]) | set( range(state_space.get_attribute_from_period("states", 0).shape[0])) assert set_valid_indices == set(range(state_space.states.shape[0]))
def test_create_state_space_vs_specialized_kw97(model): params, options = process_model_or_seed(model) # Reduce runtime options["n_periods"] = 10 if options["n_periods"] > 10 else options[ "n_periods"] optim_paras, options = process_params_and_options(params, options) # Create old state space arguments. n_periods = options["n_periods"] n_types = optim_paras["n_types"] edu_max = optim_paras["choices"]["school"]["max"] edu_starts = np.array(list(optim_paras["choices"]["school"]["start"])) # Get states and indexer from old state space. if model == "kw_97_basic": states_old, indexer_old = _create_state_space_kw97_base( n_periods, n_types, edu_starts, edu_max) else: states_old, indexer_old = _create_state_space_kw97_extended( n_periods, n_types, edu_starts, edu_max) if n_types == 1: states_old = states_old[:, :-1] for i, idx in enumerate(indexer_old): shape = idx.shape indexer_old[i] = idx.reshape(shape[:-2] + (-1, )) states_new, indexer_new = _create_core_and_indexer(optim_paras, options) states_new = pd.concat( [states_new.copy().assign(type=i) for i in range(4)]) # Compare the state spaces via sets as ordering changed in some cases. states_old_set = set(map(tuple, states_old)) states_new_set = set(map(tuple, states_new.to_numpy())) assert states_old_set == states_new_set # Compare indexers via masks for valid indices. for period in range(n_periods): mask_old = indexer_old[period] != INDEXER_INVALID_INDEX mask_new = indexer_new[period] != INDEXER_INVALID_INDEX adj_mask_new = np.repeat(mask_new, 4).reshape(mask_old.shape) assert np.array_equal(mask_old, adj_mask_new)