class TestTransitionTable(unittest.TestCase):
"""A base class for other transition test classes"""
def setUp(self) -> None:
self.test_data = pd.DataFrame({
"compartment_duration": [1, 1, 2, 2.5, 10],
"total_population": [4, 2, 2, 4, 3],
"outflow_to": ["jail", "prison", "jail", "prison", "prison"],
"compartment": ["test_compartment"] * 5,
})
self.prev_table = TransitionTable(-9999, [])
self.prev_table.generate_transition_table(TransitionTableType.AFTER,
self.test_data)
def test_normalize_transitions_requires_non_normalized_before_table(
self) -> None:
"""Tests that transitory transitions table rejects a pre-normalized 'previous' table"""
# uses its own prev_table because we don't want to normalize the general-use one
prev_table = TransitionTable(-9999, [])
prev_table.generate_transition_table(TransitionTableType.AFTER,
self.test_data)
prev_table.normalize_transitions()
with self.assertRaises(ValueError):
TransitionTable(
0,
[],
prev_table.get_table(TransitionTableType.AFTER),
)
def test_alternate_transitions_data_equal_to_differently_instantiated_transition_table(
self, ) -> None:
alternate_data = self.test_data.copy()
alternate_data.compartment_duration *= 2
alternate_data.total_population = 10 - alternate_data.total_population
policy_function = SparkPolicy(
policy_fn=partial(
TransitionTable.use_alternate_transitions_data,
alternate_historical_transitions=alternate_data,
retroactive=False,
),
spark_compartment="test_compartment",
sub_population={"sub_group": "test_population"},
policy_ts=5,
apply_retroactive=False,
)
transition_table = TransitionTable(
5,
[policy_function],
self.prev_table.get_table(TransitionTableType.AFTER),
)
alternate_prev_table = TransitionTable(-9999, [])
alternate_prev_table.generate_transition_table(
TransitionTableType.AFTER, alternate_data)
alternate_data_table = TransitionTable(
5,
[],
alternate_prev_table.get_table(TransitionTableType.AFTER),
)
assert_frame_equal(
transition_table.transition_dfs[TransitionTableType.AFTER],
alternate_data_table.transition_dfs[TransitionTableType.AFTER],
)