Class ChunkBalancerTest

python/tests/test_chunk_balancer.py:200–334

class ChunkBalancerTest(unittest.TestCase):
    @parameterized.parameterized.expand(
        [
            (TEST_SIM_1, False),
            (TEST_SIM_DUPLICATE_PROC_ID, True),
        ]
    )
    def test_validate_sim(self, test_sim_constructor, should_raise_exception):
        test_sim = test_sim_constructor()
        test_sim.init_sim()

        chunk_balancer = ChunkBalancer()

        if should_raise_exception:
            with self.assertRaises(ValueError):
                chunk_balancer._validate_sim(test_sim)
        else:
            chunk_balancer._validate_sim(test_sim)

    @parameterized.parameterized.expand(
        [
            (TEST_SIM_1,),
            (TEST_SIM_2,),
            (TEST_SIM_3,),
            (TEST_SIM_4,),
        ]
    )
    def test_chunk_layout_improvement(self, test_sim_constructor):
        """Tests that chunk_balancer improves balance after 1 iteration."""
        test_sim = test_sim_constructor()
        test_sim.init_sim()

        old_timing_measurements = MeepTimingMeasurements.new_from_simulation(
            test_sim, -1
        )

        chunk_balancer = ChunkBalancer()

        chunk_balancer.adjust_chunk_layout(test_sim, sensitivity=1.0)

        new_timing_measurements = MeepTimingMeasurements.new_from_simulation(
            test_sim, -1
        )

        old_step_times = np.array(old_timing_measurements.measurements["time_stepping"])
        new_step_times = np.array(new_timing_measurements.measurements["time_stepping"])

        old_max_time = np.max(old_step_times)
        new_max_time = np.max(new_step_times)

        old_min_time = np.min(old_step_times)
        new_min_time = np.min(new_step_times)

        self.assertLess(new_max_time, old_max_time)
        self.assertGreater(new_min_time, old_min_time)

    @parameterized.parameterized.expand(
        [