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_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_preserve_normalized_outflow_behavior_preserves_normalized_outflow_behavior(
self, ) -> None:
compartment_policies = [
SparkPolicy(
policy_fn=TransitionTable.test_retroactive_policy,
sub_population={"compartment": "test_compartment"},
spark_compartment="test_compartment",
policy_ts=5,
apply_retroactive=True,
),
SparkPolicy(
policy_fn=partial(
TransitionTable.preserve_normalized_outflow_behavior,
outflows=["prison"],
state=TransitionTableType.BEFORE,
),
sub_population={"compartment": "test_compartment"},
spark_compartment="test_compartment",
policy_ts=5,
apply_retroactive=True,
),
]
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()
assert_series_equal(
baseline_transitions.transition_dfs[
TransitionTableType.BEFORE]["prison"],
transition_table.transition_dfs[TransitionTableType.BEFORE]
["prison"],
)