Exemplo n.º 1
0
    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),
            )
Exemplo n.º 2
0
    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],
        )