From 0da0dc1cd9f45be8c654c906c39aba4ce3261c63 Mon Sep 17 00:00:00 2001
From: Oskar Taubert <oskar.taubert@kit.edu>
Date: Fri, 19 Apr 2024 15:05:24 +0200
Subject: [PATCH] added a checkpointing test for unequal island sizes

---
 tests/test_island.py | 101 +++++++++++++++++++++++++++++++++++--------
 1 file changed, 84 insertions(+), 17 deletions(-)

diff --git a/tests/test_island.py b/tests/test_island.py
index 957cd6ab..a37b6174 100644
--- a/tests/test_island.py
+++ b/tests/test_island.py
@@ -36,7 +36,7 @@ def function_parameters(request):
     return request.param
 
 
-@pytest.mark.mpi(min_size=4)
+@pytest.mark.mpi(min_size=8)
 def test_island(function_parameters, mpi_tmp_path) -> None:
     """Test basic island functionality."""
     fname, expected, abs_tolerance = function_parameters
@@ -341,19 +341,86 @@ def test_checkpointing_pollination(function_parameters, mpi_tmp_path):
     )
 
 
-# @pytest.mark.mpi
-# def test_checkpointing_unequal_populations():
-#     # TODO test loading unequal islands
-#     raise
-#
-#
-# @pytest.mark.mpi
-# def test_checkpointing_setup_downgrade():
-#     # TODO write checkpoint with large setup and read it with a small one
-#     raise
-#
-#
-# @pytest.mark.mpi
-# def test_checkpointing_setup_upgrade():
-#     # TOOD write checkpoint with small setup and read it with a large one
-#     raise
+@pytest.mark.mpi
+def test_checkpointing_unequal_populations(function_parameters, mpi_tmp_path):
+    """Test island checkpointing without migration."""
+    fname, expected, abs_tolerance = function_parameters
+    rng = random.Random(
+        42 + MPI.COMM_WORLD.rank
+    )  # Separate random number generator for optimization
+    function, limits = get_function_search_space(fname)
+    set_logger_config(
+        level=logging.INFO,
+        log_file=mpi_tmp_path / "propulate.log",
+        log_to_stdout=True,
+        log_rank=False,
+        colors=True,
+    )
+
+    # Set up evolutionary operator.
+    propagator = get_default_propagator(
+        pop_size=4,
+        limits=limits,
+        crossover_prob=0.7,
+        mutation_prob=0.9,
+        random_init_prob=0.1,
+        rng=rng,
+    )
+
+    # Set up migration topology.
+    migration_topology = 1 * np.ones(  # Set up fully connected migration topology.
+        (2, 2), dtype=int
+    )
+    np.fill_diagonal(
+        migration_topology, 0
+    )  # An island does not send migrants to itself.
+
+    # Set up island model.
+    islands = Islands(
+        loss_fn=function,
+        propagator=propagator,
+        rng=rng,
+        generations=100,
+        num_islands=2,
+        island_sizes=np.array([3, 5]),
+        migration_topology=migration_topology,
+        migration_probability=0.0,
+        emigration_propagator=SelectMin,
+        immigration_propagator=SelectMax,
+        pollination=False,  # TODO fixtureize
+        checkpoint_path=mpi_tmp_path,
+    )
+
+    # Run actual optimization.
+    islands.evolve(
+        top_n=1,
+        logging_interval=10,
+        debug=2,
+    )
+
+    old_population = copy.deepcopy(islands.propulator.population)
+    del islands
+
+    islands = Islands(
+        loss_fn=function,
+        propagator=propagator,
+        rng=rng,
+        generations=100,
+        num_islands=2,
+        island_sizes=np.array([3, 5]),
+        migration_topology=migration_topology,
+        migration_probability=0.9,
+        emigration_propagator=SelectMin,
+        immigration_propagator=SelectMax,
+        pollination=True,  # TODO fixtureize
+        checkpoint_path=mpi_tmp_path,
+    )
+
+    assert (
+        len(
+            deepdiff.DeepDiff(
+                old_population, islands.propulator.population, ignore_order=True
+            )
+        )
+        == 0
+    )