More tests

Author: MaxBlesch

tests/test_partial.py tests/test_partial_and_interfaces.pytests/test_partial.py renamed to tests/test_partial_and_interfaces.py

@@ -1,5 +1,6 @@
 from pathlib import Path
 
+import numpy as np
 from numpy.testing import assert_array_almost_equal as aaae
 
 import dcegm
@@ -37,3 +38,30 @@ def test_partial_solve_func():
     aaae(model_solved.policy, partial_sol["policy"])
     aaae(model_solved.value, partial_sol["value"])
     aaae(model_solved.endog_grid, partial_sol["endog_grid"])
+
+    state_choices = model_solved.model_structure["state_choice_space"]
+    choices = state_choices[:, -1]
+    states_dict = {
+        state: state_choices[:, id]
+        for id, state in enumerate(
+            model_solved.model_structure["discrete_states_names"]
+        )
+    }
+    states_dict["assets_begin_of_period"] = model_solved.endog_grid[:, 5]
+    value_states_all_choices = model_solved.choice_values_for_states(states=states_dict)
+
+    # Take in each row the value corresponding to the choice made
+    value_choices = value_states_all_choices[
+        np.arange(value_states_all_choices.shape[0]), choices
+    ]
+
+    aaae(model_solved.value[:, 5], value_choices)
+
+    # Same for policies
+    policy_states_all_choices = model_solved.choice_policies_for_states(
+        states=states_dict
+    )
+    policy_choices = policy_states_all_choices[
+        np.arange(policy_states_all_choices.shape[0]), choices
+    ]
+    aaae(model_solved.policy[:, 5], policy_choices)