refactored class names, added documentation for adjusting chunk layou…

…ts between runs of a program

doc/docs/Parallel_Meep.md

@@ -189,7 +189,7 @@ The chunk balancer interface provides four main methods:
 - `adjust_chunk_layout()` is syntactic sugar which will compute a new chunk layout, apply it to the simulation object, reset the simulation, and re-initialize the simulation.
 
 ```py
-class ChunkyMonkey(abc.ABC):
+class AbstractChunkBalancer(abc.ABC):
   """Chunk balancer for dynamically load-balanced chunk layouts."""
 
   def make_initial_chunk_layout(self, sim) -> mp.BinaryPartition:
@@ -216,9 +216,9 @@ class ChunkyMonkey(abc.ABC):
 Using the chunk balancer is very straightforward, and it can typically be integrated into existing Meep simulations with only a few lines of code. Here is a simple example:
 
 ```py
-from meep.chunk_balancer import DefaultChunkBalancer
+from meep.chunk_balancer import ChunkBalancer
 
-chunk_balancer = DefaultChunkBalancer()
+chunk_balancer = ChunkBalancer()
 
 # Compute an initial chunk layout
 initial_chunk_layout = chunk_balancer.make_initial_chunk_layout()
@@ -228,11 +228,42 @@ sim.init_sim()
 
 for iteration in range(epochs):
   sim.run(...)
-
-  # Adjust the chunk layout if needed
+  # Adjust the chunk layout for the next iteration if needed
   chunk_balancer.adjust_chunk_layout(sim, sensitivity=0.4)
 ```
 
+Chunks can also be rebalanced between runs of a program by dumping and loading the chunk layout from a pickled object. For example:
+
+```py
+import meep as mp
+from meep.chunk_balancer import ChunkBalancer
+from meep.timing_measurements import MeepTimingMeasurements
+import pickle
+
+# Fetch chunk layout from a previous run
+with open("path/to/chunk_layout.pkl", "rb") as f:
+    chunk_layout = pickle.load(f)
+
+sim = mp.Simulation(..., chunk_layout=chunk_layout)
+sim.init_sim()
+sim.run(...)
+
+# Compute and save chunk layout for next run
+timings = MeepTimingMeasurements.new_from_simulation(sim)
+chunk_volumes = sim.structure.get_chunk_volumes()
+chunk_owners = sim.structure.get_chunk_owners()
+next_chunk_layout = ChunkBalancer().compute_new_chunk_layout(
+    timings,
+    chunk_layout,
+    chunk_volumes,
+    chunk_owners,
+    sensitivity=0.4)
+
+# Save chunk layout for next run
+with open("path/to/chunk_layout.pkl", "wb") as f:
+    pickle.dump(next_chunk_layout, f)
+```
+
 ### Chunk adjustment algorithm
 
 The default chunk adjustment algorithm recursively traverses the `BinaryPartition` object and resizes the chunk volumes of the left and right children of each node to have an equal per-node simulation time. The new chunk layout has split positions which are a weighted average of the previous iteration's chunk layout and the newly computed layout, and the `sensitivity` parameter adjusts how fast the chunk sizes are adjusted. (`sensitivity=0.0` means no adjustment, `sensitivity=1.0` means an immediate snap to the predicted split positions, and `sensitivity=0.5` is the average of the old and new layouts.) The adjustment algorithm is summarized in pseudocode below:

python/chunk_balancer.py

@@ -9,7 +9,7 @@
 import numpy as np
 
 
-class ChunkyMonkey(abc.ABC):
+class AbstractChunkBalancer(abc.ABC):
     """Abstract chunk balancer for adaptive chunk layouts in Meep simulations.
 
     This class defines interfaces for a chunk balancer, which adjusts chunk
@@ -135,7 +135,7 @@ def _validate_sim(self, sim: mp.Simulation):
                     set(proc_ids) - set(chunk_owners)))
 
 
-class DefaultChunkBalancer(ChunkyMonkey):
+class ChunkBalancer(AbstractChunkBalancer):
     """A chunk balancer for adaptive chunk layouts in Meep simulations.
 
     This class generates initial chunk layouts using Meep's built-in scheme, and

python/tests/test_chunk_balancer.py

@@ -7,7 +7,7 @@
 import meep.binary_partition_utils as bpu
 
 from meep.timing_measurements import MeepTimingMeasurements, TIMING_MEASUREMENT_IDS
-from meep.chunk_balancer import DefaultChunkBalancer
+from meep.chunk_balancer import ChunkBalancer
 
 
 class MockSimulation(mp.Simulation):
@@ -200,7 +200,7 @@ def test_structure_get_chunk_owners(self, test_sim_constructor,
         self.assertListEqual(list(chunk_owners), expected_chunk_owners)
 
 
-class DefaultChunkBalancerTest(unittest.TestCase):
+class ChunkBalancerTest(unittest.TestCase):
 
     @parameterized.parameterized.expand([
         (TEST_SIM_1, False),
@@ -210,7 +210,7 @@ def test_validate_sim(self, test_sim_constructor, should_raise_exception):
         test_sim = test_sim_constructor()
         test_sim.init_sim()
 
-        chunk_balancer = DefaultChunkBalancer()
+        chunk_balancer = ChunkBalancer()
 
         if should_raise_exception:
             with self.assertRaises(ValueError):
@@ -232,7 +232,7 @@ def test_chunk_layout_improvement(self, test_sim_constructor):
         old_timing_measurements = MeepTimingMeasurements.new_from_simulation(
             test_sim, -1)
 
-        chunk_balancer = DefaultChunkBalancer()
+        chunk_balancer = ChunkBalancer()
 
         chunk_balancer.adjust_chunk_layout(test_sim, sensitivity=1.0)
 
@@ -264,7 +264,7 @@ def test_chunk_layout_convergence(self, test_sim_constructor):
         test_sim = test_sim_constructor()
         test_sim.init_sim()
 
-        chunk_balancer = DefaultChunkBalancer()
+        chunk_balancer = ChunkBalancer()
 
         num_iterations = 25
 
@@ -301,7 +301,7 @@ def test_split_pos_adjustment(self, test_chunk_data):
         chunk_volumes = sim.structure.get_chunk_volumes()
         chunk_owners = sim.structure.get_chunk_owners()
 
-        chunk_balancer = DefaultChunkBalancer()
+        chunk_balancer = ChunkBalancer()
 
         measurements = {}
         num_procs = len(test_chunk_data["time_stepping"])

python/tests/test_timing_measurements.py

@@ -15,11 +15,7 @@ def test_timing_measurements(self):
         )
         time_start = time.time()
         sim.run(until=5)
-        timing_measurements = timing.MeepTimingMeasurements.new_from_simulation(
-            sim=sim,
-            elapsed_time=time.time() - time_start,
-            time_per_step=[0.1],
-        )
+        timing_measurements = timing.MeepTimingMeasurements.new_from_simulation(sim)
 
         # Check for expected names after updating
         self.assertSetEqual(

python/timing_measurements.py

@@ -93,7 +93,7 @@ def __init__(self,
     def new_from_simulation(
             cls,
             sim: mp.Simulation,
-            elapsed_time: float,
+            elapsed_time: Optional[float] = -1,
             time_per_step: Optional[List[float]] = None,
             dft_relative_change: Optional[List[float]] = None,
             overlap_relative_change: Optional[List[float]] = None,