def test_chop_technicals_chops_correctly(self) -> None:
"""
Make sure CompartmentTransitions.chop_technical_revocations zeros technicals after the correct duration and
that table sums to the same amount (i.e. total population shifted but not removed)
"""
compartment_policies = [
SparkPolicy(
policy_fn=partial(
TransitionTable.chop_technical_revocations,
technical_outflow="prison",
release_outflow="jail",
retroactive=False,
),
sub_population={"sub_group": "test_population"},
spark_compartment="test_compartment",
policy_ts=5,
apply_retroactive=False,
)
]
transition_table = TransitionTable(
5,
compartment_policies,
self.prev_table.get_table(TransitionTableType.AFTER),
)
baseline_transitions = TransitionTable(
5,
[],
self.prev_table.get_table(TransitionTableType.AFTER),
)
transition_table.normalize_transitions()
baseline_transitions.normalize_transitions()
# check total population was preserved
assert_series_equal(
transition_table.transition_dfs[TransitionTableType.AFTER].iloc[0],
baseline_transitions.transition_dfs[
TransitionTableType.AFTER].iloc[0],
)
# check technicals chopped
transition_table.unnormalize_table(TransitionTableType.AFTER)
self.assertTrue((transition_table.transition_dfs[
TransitionTableType.AFTER].loc[3:, "prison"] == 0).all())
self.assertTrue(transition_table.transition_dfs[
TransitionTableType.AFTER].loc[1, "prison"] != 0)
def test_unnormalized_table_inverse_of_normalize_table(self) -> None:
transition_table = TransitionTable(
5,
[],
self.prev_table.get_table(TransitionTableType.AFTER),
)
original_before_table = transition_table.transition_dfs[
TransitionTableType.BEFORE].copy()
# 'normalize' table (in the classical, mathematical sense) to match scale of unnormalized table
original_before_table /= original_before_table.sum().sum()
transition_table._normalize_table( # pylint: disable=protected-access
TransitionTableType.BEFORE)
transition_table.unnormalize_table(TransitionTableType.BEFORE)
assert_frame_equal(
pd.DataFrame(original_before_table),
pd.DataFrame(
transition_table.transition_dfs[TransitionTableType.BEFORE]),
)