def setUp(self):
self.test_outflow_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"] * 5,
})
self.historical_data = pd.DataFrame({
2015: {
"jail": 2,
"prison": 2
},
2016: {
"jail": 1,
"prison": 0
},
2017: {
"jail": 1,
"prison": 1
},
})
self.compartment_policies = []
self.test_transition_table = CompartmentTransitions(
self.test_outflow_data)
self.test_transition_table.initialize(self.compartment_policies)
def test_reallocate_outflow_preserves_total_population(self):
compartment_policies = [
SparkPolicy(
policy_fn=partial(
CompartmentTransitions.reallocate_outflow,
reallocation_df=pd.DataFrame({
"outflow": ["jail", "jail"],
"affected_fraction": [0.25, 0.25],
"new_outflow": ["prison", "treatment"],
}),
reallocation_type="+",
retroactive=True,
),
sub_population={"sub_group": "test_population"},
spark_compartment="test_compartment",
apply_retroactive=True,
)
]
compartment_transitions = CompartmentTransitions(self.test_data)
compartment_transitions.initialize(compartment_policies)
assert_series_equal(
compartment_transitions.transition_dfs["before"].sum(axis=1),
compartment_transitions.transition_dfs["after_retroactive"].sum(
axis=1),
)
def test_extend_table_extends_table(self):
"""make sure CompartmentTransitions.extend_table is actually adding empty rows"""
compartment_transitions = CompartmentTransitions(self.test_data)
compartment_transitions.extend_tables(15)
self.assertEqual(
set(compartment_transitions.transition_dfs["before"].index),
set(range(1, 16)),
)
class TestShellCompartment(unittest.TestCase):
"""Test the ShellCompartment class runs correctly"""
def setUp(self):
self.test_outflow_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"] * 5,
})
self.historical_data = pd.DataFrame({
2015: {
"jail": 2,
"prison": 2
},
2016: {
"jail": 1,
"prison": 0
},
2017: {
"jail": 1,
"prison": 1
},
})
self.compartment_policies = []
self.test_transition_table = CompartmentTransitions(
self.test_outflow_data)
self.test_transition_table.initialize(self.compartment_policies)
def test_all_edges_fed_to(self):
"""ShellCompartments require edges to the compartments defined in the outflows_data"""
starting_ts = 2015
policy_ts = 2018
test_shell_compartment = ShellCompartment(
self.test_outflow_data,
starting_ts=starting_ts,
policy_ts=policy_ts,
tag="test_shell",
constant_admissions=True,
policy_list=[],
)
test_full_compartment = FullCompartment(
self.historical_data,
self.test_transition_table,
starting_ts=starting_ts,
policy_ts=policy_ts,
tag="test_compartment",
)
with self.assertRaises(ValueError):
test_shell_compartment.initialize_edges(
[test_shell_compartment, test_full_compartment])
def test_non_retroactive_policy_cannot_affect_retroactive_table(self):
compartment_policies = [
SparkPolicy(
policy_fn=CompartmentTransitions.test_retroactive_policy,
sub_population={"compartment": "test_compartment"},
spark_compartment="test_compartment",
apply_retroactive=False,
)
]
compartment_transitions = CompartmentTransitions(self.test_data)
with self.assertRaises(ValueError):
compartment_transitions.initialize(compartment_policies)
def test_results_independent_of_data_order(self):
compartment_policies = [
SparkPolicy(
policy_fn=CompartmentTransitions.test_retroactive_policy,
sub_population={"compartment": "test_compartment"},
spark_compartment="test_compartment",
apply_retroactive=True,
),
SparkPolicy(
policy_fn=CompartmentTransitions.test_non_retroactive_policy,
sub_population={"compartment": "test_compartment"},
spark_compartment="test_compartment",
apply_retroactive=False,
),
]
compartment_transitions_default = CompartmentTransitions(
self.test_data)
compartment_transitions_shuffled = CompartmentTransitions(
self.test_data.sample(frac=1))
compartment_transitions_default.initialize(compartment_policies)
compartment_transitions_shuffled.initialize(compartment_policies)
self.assertEqual(compartment_transitions_default,
compartment_transitions_shuffled)
def test_preserve_normalized_outflow_behavior_preserves_normalized_outflow_behavior(
self, ):
compartment_policies = [
SparkPolicy(
policy_fn=CompartmentTransitions.test_retroactive_policy,
sub_population={"compartment": "test_compartment"},
spark_compartment="test_compartment",
apply_retroactive=True,
),
SparkPolicy(
policy_fn=partial(
CompartmentTransitions.
preserve_normalized_outflow_behavior,
outflows=["prison"],
state="after_retroactive",
before_state="before",
),
sub_population={"compartment": "test_compartment"},
spark_compartment="test_compartment",
apply_retroactive=True,
),
]
compartment_transitions = CompartmentTransitions(self.test_data)
compartment_transitions.initialize(compartment_policies)
baseline_transitions = CompartmentTransitions(self.test_data)
baseline_transitions.initialize([])
assert_series_equal(
baseline_transitions.transition_dfs["after_retroactive"]["prison"],
compartment_transitions.transition_dfs["after_retroactive"]
["prison"],
)
def test_alternate_transitions_data_equal_to_differently_instantiated_transition_table(
self, ):
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(
CompartmentTransitions.use_alternate_transitions_data,
alternate_historical_transitions=alternate_data,
retroactive=False,
),
spark_compartment="test_compartment",
sub_population={"sub_group": "test_population"},
apply_retroactive=False,
)
compartment_transitions = CompartmentTransitions(self.test_data)
compartment_transitions.initialize([policy_function])
alternate_data_transitions = CompartmentTransitions(alternate_data)
alternate_data_transitions.initialize([])
assert_frame_equal(
compartment_transitions.transition_dfs["after_non_retroactive"],
alternate_data_transitions.transition_dfs["after_non_retroactive"],
)
def test_rejects_remaining_as_outflow(self):
"""Tests that compartment transitions won't accept 'remaining' as an outflow"""
broken_test_data = self.test_data.copy()
broken_test_data.loc[broken_test_data["outflow_to"] == "jail",
"outflow_to"] = "remaining"
with self.assertRaises(ValueError):
CompartmentTransitions(broken_test_data)
def test_rejects_data_with_negative_populations_or_durations(self):
negative_duration_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,
})
negative_population_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,
})
with self.assertRaises(ValueError):
CompartmentTransitions(negative_duration_data)
with self.assertRaises(ValueError):
CompartmentTransitions(negative_population_data)
def test_chop_technicals_chops_correctly(self):
"""
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(
CompartmentTransitions.chop_technical_revocations,
technical_outflow="prison",
release_outflow="jail",
retroactive=False,
),
sub_population={"sub_group": "test_population"},
spark_compartment="test_compartment",
apply_retroactive=False,
)
]
compartment_transitions = CompartmentTransitions(self.test_data)
compartment_transitions.initialize(compartment_policies)
baseline_transitions = CompartmentTransitions(self.test_data)
baseline_transitions.initialize([])
# check total population was preserved
assert_series_equal(
compartment_transitions.transition_dfs["after_non_retroactive"].
iloc[0],
baseline_transitions.transition_dfs["after_non_retroactive"].
iloc[0],
)
# check technicals chopped
compartment_transitions.unnormalize_table("after_non_retroactive")
self.assertTrue(
(compartment_transitions.transition_dfs["after_non_retroactive"].
loc[3:, "prison"] == 0).all())
self.assertTrue(
compartment_transitions.transition_dfs["after_non_retroactive"].
loc[1, "prison"] != 0)
def test_normalize_transitions_requires_generated_transition_table(self):
compartment_transitions = CompartmentTransitions(self.test_data)
# manually initializing without the self._generate_transition_table()
compartment_transitions.transition_dfs = {
"before":
pd.DataFrame(
{
outflow: np.zeros(10)
for outflow in compartment_transitions.outflows
},
index=range(1, 11),
),
"transitory":
pd.DataFrame(),
"after_retroactive":
pd.DataFrame(),
"after_non_retroactive":
pd.DataFrame(),
}
with self.assertRaises(ValueError):
compartment_transitions.normalize_transitions(
state="after_retroactive")
def test_normalize_transitions_requires_non_normalized_before_table(self):
"""Tests that transitory transitions table rejects a pre-normalized 'before' table"""
transitions_table = CompartmentTransitions(self.test_data)
transitions_table.transition_dfs["after"] = deepcopy(
transitions_table.transition_dfs["before"])
transitions_table.normalize_transitions(state="before")
with self.assertRaises(ValueError):
transitions_table.normalize_transitions(
state="after",
before_table=transitions_table.transition_dfs["before"])
def test_apply_reduction_with_trivial_reductions_doesnt_change_transition_table(
self, ):
policy_mul = partial(
CompartmentTransitions.apply_reduction,
reduction_df=pd.DataFrame({
"outflow": ["prison"] * 2,
"affected_fraction": [0, 0.5],
"reduction_size": [0.5, 0],
}),
reduction_type="*",
retroactive=False,
)
policy_add = partial(
CompartmentTransitions.apply_reduction,
reduction_df=pd.DataFrame({
"outflow": ["prison"] * 2,
"affected_fraction": [0, 0.5],
"reduction_size": [0.5, 0],
}),
reduction_type="+",
retroactive=False,
)
compartment_policies = [
SparkPolicy(policy_mul, "test_compartment",
{"sub_group": "test_population"}, False),
SparkPolicy(policy_add, "test_compartment",
{"sub_group": "test_population"}, False),
]
compartment_transitions = CompartmentTransitions(self.test_data)
compartment_transitions.initialize(compartment_policies)
assert_frame_equal(
compartment_transitions.transition_dfs["before"],
compartment_transitions.transition_dfs["after_non_retroactive"],
)
def test_unnormalized_table_inverse_of_normalize_table(self):
compartment_transitions = CompartmentTransitions(self.test_data)
original_before_table = compartment_transitions.transition_dfs[
"before"].copy()
# 'normalize' table (in the classical mathematical sense) to match scale of unnormalized table
original_before_table /= original_before_table.sum().sum()
compartment_transitions.normalize_transitions("before")
compartment_transitions.unnormalize_table("before")
assert_frame_equal(
pd.DataFrame(original_before_table),
pd.DataFrame(compartment_transitions.transition_dfs["before"]),
)
def test_apply_reduction_matches_example_by_hand(self):
compartment_transitions = CompartmentTransitions(self.test_data)
compartment_policy = [
SparkPolicy(
policy_fn=partial(
CompartmentTransitions.apply_reduction,
reduction_df=pd.DataFrame({
"outflow": ["prison"],
"affected_fraction": [0.25],
"reduction_size": [0.5],
}),
reduction_type="+",
retroactive=True,
),
sub_population={"sub_group": "test_population"},
spark_compartment="test_compartment",
apply_retroactive=True,
)
]
expected_result = pd.DataFrame(
{
"jail": [4, 2, 0, 0, 0, 0, 0, 0, 0, 0],
"prison": [2, 0.5, 3.5, 0, 0, 0, 0, 0, 0.375, 2.625],
},
index=range(1, 11),
dtype=float,
)
expected_result.index.name = "compartment_duration"
expected_result.columns.name = "outflow_to"
expected_result /= expected_result.sum().sum()
compartment_transitions.initialize(compartment_policy)
compartment_transitions.unnormalize_table("after_retroactive")
assert_frame_equal(
round(compartment_transitions.transition_dfs["after_retroactive"],
8),
round(expected_result, 8),
)
class TestFullCompartment(unittest.TestCase):
"""Test the FullCompartment runs correctly"""
def setUp(self):
self.test_supervision_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"] * 5,
}
)
self.test_incarceration_data = pd.DataFrame(
{
"compartment_duration": [1, 1, 2, 2.5, 10],
"total_population": [4, 2, 2, 4, 3],
"outflow_to": [
"supervision",
"release",
"supervision",
"release",
"release",
],
"compartment": ["test"] * 5,
}
)
self.compartment_policies = []
self.incarceration_transition_table = CompartmentTransitions(
self.test_incarceration_data
)
self.incarceration_transition_table.initialize(self.compartment_policies)
self.release_transition_table = CompartmentTransitions(
self.test_supervision_data
)
self.release_transition_table.initialize(self.compartment_policies)
self.historical_data = pd.DataFrame(
{
2015: {"jail": 2, "prison": 2},
2016: {"jail": 1, "prison": 0},
2017: {"jail": 1, "prison": 1},
}
)
def test_step_forward_fails_without_initialized_edges(self):
"""Tests that step_forward() needs the initialize_edges() to have been run"""
rel_compartment = FullCompartment(
self.historical_data, self.release_transition_table, 2015, 2018, "release"
)
with self.assertRaises(ValueError):
rel_compartment.step_forward()
def test_all_edges_fed_to(self):
"""Tests that all edges in self.edges are included in self.transition_tables"""
rel_compartment = FullCompartment(
self.historical_data, self.release_transition_table, 2015, 2018, "release"
)
test_compartment = FullCompartment(
self.historical_data,
self.incarceration_transition_table,
2015,
2018,
"test_compartment",
)
compartment_list = [rel_compartment, test_compartment]
for compartment in compartment_list:
compartment.initialize_edges(compartment_list)
for compartment in compartment_list:
compartment.step_forward()
def _initialize_transition_tables(
cls,
transitions_data: pd.DataFrame,
compartments_architecture: Dict[str, str],
policy_list: List[SparkPolicy],
) -> Tuple[Dict[str, CompartmentTransitions], Dict[str,
List[SparkPolicy]]]:
"""Create and initialize all transition tables and store shell policies."""
# Initialize a default transition class for each compartment to represent the no-policy scenario
transitions_per_compartment = {}
unused_transitions_data = transitions_data
for compartment in compartments_architecture:
compartment_type = compartments_architecture[compartment]
compartment_duration_data = transitions_data[
transitions_data["compartment"] == compartment]
unused_transitions_data = unused_transitions_data.drop(
compartment_duration_data.index)
if compartment_duration_data.empty:
if compartment_type != "shell":
raise ValueError(
f"Transition data missing for compartment {compartment}. Data is required for all "
"disaggregtion axes. Even the 'release' compartment needs transition data even if "
"it's just outflow to 'release'")
else:
if compartment_type == "full":
transition_class = CompartmentTransitions(
compartment_duration_data)
elif compartment_type == "shell":
raise ValueError(
f"Cannot provide transitions data for shell compartment \n "
f"{compartment_duration_data}")
else:
raise ValueError(
f"unrecognized transition table type {compartment_type}"
)
transitions_per_compartment[compartment] = transition_class
if len(unused_transitions_data) > 0:
warn(
f"Some transitions data not fed to a compartment: {unused_transitions_data}",
Warning,
)
# Create a transition object for each compartment and year with policies applied and store shell policies
shell_policies = dict()
for compartment in compartments_architecture:
# Select any policies that are applicable for this compartment
compartment_policies = SparkPolicy.get_compartment_policies(
policy_list, compartment)
# add to the dict compartment -> transition class with policies applied
if compartment in transitions_per_compartment:
transitions_per_compartment[compartment].initialize(
compartment_policies)
# add shell policies to dict that gets passed to initialization
else:
shell_policies[compartment] = compartment_policies
return transitions_per_compartment, shell_policies
def test_normalize_transitions_requires_initialized_transition_table(self):
with self.assertRaises(ValueError):
compartment_transitions = CompartmentTransitions(self.test_data)
compartment_transitions.normalize_transitions(
state="after_retroactive")
