def initialize(self): """Initialize the transition tables, and then the compartments, for the SubSimulation""" # TODO(#4512): allow sparse data if not self.microsim and not self.total_population_data.empty: if self.user_inputs['start_time_step'] not in self.total_population_data.time_step.values: raise ValueError(f"Start time must be included in population data input\n" f"Expected: {self.user_inputs['start_time_step']}, " f"Actual: {self.total_population_data.time_step.values}") # Initialize a default transition class for each compartment to represent the no-policy scenario transitions_per_compartment = {} for compartment in self.simulation_architecture: transition_type = self.simulation_architecture[compartment] compartment_duration_data = self.transitions_data[self.transitions_data['compartment'] == compartment] if compartment_duration_data.empty: if transition_type is not None: 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 transition_type == 'incarcerated': transition_class = IncarceratedTransitions(compartment_duration_data) elif transition_type == 'released': transition_class = ReleasedTransitions(compartment_duration_data) else: raise ValueError(f'unrecognized transition table type {transition_type}') transition_class.initialize_transition_table() transitions_per_compartment[compartment] = transition_class # Create a transition object for each compartment and year with policies applied and store shell policies shell_policies = dict() for compartment in self.simulation_architecture: # Select any policies that are applicable for this compartment compartment_policies = SparkPolicy.get_compartment_policies(self.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 # Preprocess the historical data into separate pieces per compartment historical_outflows = self._load_data() # Initialize the compartment classes self._initialize_compartments(historical_outflows, transitions_per_compartment, shell_policies)
def test_get_compartment_policies(self): expected_list = self.policy_list[1:] result_list = SparkPolicy.get_compartment_policies( self.policy_list, spark_compartment='jail') self.assertEqual(expected_list, result_list)
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: logging.warning( "Some transitions data not fed to a compartment: %s", unused_transitions_data, ) # 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_transition_tables( compartment_policies) # add shell policies to dict that gets passed to initialization else: shell_policies[compartment] = compartment_policies return transitions_per_compartment, shell_policies