Sharded fields dump (NanoComp#1738)

* Add support for fully local HDF5 files and shared dumping of meep::structure

* Add support for fully local HDF5 files and shared dumping of meep::structure

* Update python func docs

* Update python API documentation

* Dump/Load of 'fields'

The PR is incomplete (does not include the dft fields) and also has a
lot of debugging statements.

* Save dft chunks

* Remove debug log stmts

* Add saving of time value and also reorg tests.

* Fix dft-chunk saving for the single-parallel-file mode.

* Abort when trying to dump fields with non-null polarization state

* Clean up test_dump_load.py and add test_dump_fails_for_non_null_polarization_state test

* Add new tests/dump_load to set of files to ignore

* Clean up the test to remove unnecessary stuff from the copied over test

* Also dump/load 'f_w_prev'

* Fix typo causing build breakage

* load fields at the very end of init_sim after add_sources

* remove init_sim before loading fields since that now works.

.gitignore

@@ -90,6 +90,7 @@ tests/convergence_cyl_waveguide
 tests/cyl-ellipsoid-ll
 tests/cylindrical
 tests/dft-fields
+tests/dump_load
 tests/flux
 tests/gdsII-3d
 tests/h5test

python/Makefile.am

@@ -64,14 +64,15 @@ TESTS =                                   \
     $(TEST_DIR)/test_get_point.py              \
     $(TEST_DIR)/test_holey_wvg_bands.py        \
     $(TEST_DIR)/test_holey_wvg_cavity.py       \
-    $(KDOM_TEST)                          \
+    $(KDOM_TEST)                               \
     $(TEST_DIR)/test_ldos.py                   \
-    $(MDPYTEST)                           \
+    $(TEST_DIR)/test_dump_load.py              \
+    $(MDPYTEST)                                \
     $(TEST_DIR)/test_material_grid.py          \
     $(TEST_DIR)/test_medium_evaluations.py     \
-    $(MPBPYTEST)                          \
-    $(MODE_COEFFS_TEST)                   \
-    $(MODE_DECOMPOSITION_TEST)            \
+    $(MPBPYTEST)                               \
+    $(MODE_COEFFS_TEST)                        \
+    $(MODE_DECOMPOSITION_TEST)                 \
     $(TEST_DIR)/test_multilevel_atom.py        \
     $(TEST_DIR)/test_n2f_periodic.py           \
     $(TEST_DIR)/test_oblique_source.py         \

python/simulation.py

@@ -1250,6 +1250,7 @@ def __init__(self,
 
         self.load_single_parallel_file = True
         self.load_structure_file = None
+        self.load_fields_file = None
 
         self.special_kz = False
         if self.cell_size.z == 0 and self.k_point and self.k_point.z != 0:
@@ -1874,29 +1875,48 @@ def dump_structure(self, fname, single_parallel_file=True):
         Dumps the structure to the file `fname`.
         """
         if self.structure is None:
-            raise ValueError("Fields must be initialized before calling dump_structure")
+            raise ValueError("Structure must be initialized before calling dump_structure")
         self.structure.dump(fname, single_parallel_file)
+        print("Dumped structure to file: %s (%s)" % (fname, str(single_parallel_file)))
 
     def load_structure(self, fname, single_parallel_file=True):
         """
-        Loads a structure from the file `fname`. A file name to load can also be passed to
-        the `Simulation` constructor via the `load_structure` keyword argument.
+        Loads a structure from the file `fname`.
         """
         if self.structure is None:
-            raise ValueError("Fields must be initialized before calling load_structure")
+            raise ValueError("Structure must be initialized before loading structure from file '%s'" % fname)
         self.structure.load(fname, single_parallel_file)
+        print("Loaded structure from file: %s (%s)" % (fname, str(single_parallel_file)))
+
+    def dump_fields(self, fname, single_parallel_file=True):
+        """
+        Dumps the fields to the file `fname`.
+        """
+        if self.fields is None:
+            raise ValueError("Fields must be initialized before calling dump_fields")
+        self.fields.dump(fname, single_parallel_file)
+        print("Dumped fields to file: %s (%s)" % (fname, str(single_parallel_file)))
+
+    def load_fields(self, fname, single_parallel_file=True):
+        """
+        Loads a fields from the file `fname`.
+        """
+        if self.fields is None:
+            raise ValueError("Fields must be initialized before loading fields from file '%s'" % fname)
+        self.fields.load(fname, single_parallel_file)
+        print("Loaded fields from file: %s (%s)" % (fname, str(single_parallel_file)))
 
     def dump_chunk_layout(self, fname):
         """
         Dumps the chunk layout to file `fname`.
         """
         if self.structure is None:
-            raise ValueError("Fields must be initialized before calling load_structure")
+            raise ValueError("Structure must be initialized before calling dump_chunk_layout")
         self.structure.dump_chunk_layout(fname)
 
     def load_chunk_layout(self, br, source):
         if self.structure is None:
-            raise ValueError("Fields must be initialized before calling load_structure")
+            raise ValueError("Structure must be initialized before loading chunk layout from file '%s'" % fname)
 
         if isinstance(source, Simulation):
             vols = source.structure.get_chunk_volumes()
@@ -1918,18 +1938,22 @@ def _get_load_dump_dirname(self, dirname, single_parallel_file):
             dump_dirname = os.path.join(dirname, 'rank%02d' % mp.my_rank())
         return dump_dirname
 
-    def dump(self, dirname, structure=True, single_parallel_file=True):
+    def dump(self, dirname, dump_structure=True, dump_fields=True, single_parallel_file=True):
         """
         Dumps simulation state.
         """
         dump_dirname = self._get_load_dump_dirname(dirname, single_parallel_file)
         os.makedirs(dump_dirname, exist_ok=True)
 
-        if structure:
+        if dump_structure:
             structure_dump_filename = os.path.join(dump_dirname, 'structure.h5')
             self.dump_structure(structure_dump_filename, single_parallel_file)
 
-    def load(self, dirname, structure=True, single_parallel_file=True):
+        if dump_fields:
+            fields_dump_filename = os.path.join(dump_dirname, 'fields.h5')
+            self.dump_fields(fields_dump_filename, single_parallel_file)
+
+    def load(self, dirname, load_structure=True, load_fields=True, single_parallel_file=True):
         """
         Loads simulation state.
 
@@ -1938,8 +1962,22 @@ def load(self, dirname, structure=True, single_parallel_file=True):
         """
         dump_dirname = self._get_load_dump_dirname(dirname, single_parallel_file)
         self.load_single_parallel_file = single_parallel_file
-        if structure:
-          self.load_structure_file = os.path.join(dump_dirname, 'structure.h5')
+
+        if load_structure:
+          load_structure_file = os.path.join(dump_dirname, 'structure.h5')
+          # If structure is already initialized, load it straight away.
+          # Otherwise, do a delayed load.
+          if self.structure:
+            self.load_structure(load_structure_file, self.load_single_parallel_file)
+          else:
+            self.load_structure_file = load_structure_file
+
+        if load_fields:
+          load_fields_file = os.path.join(dump_dirname, 'fields.h5')
+          if self.fields:
+            self.load_fields(load_fields_file, self.load_single_parallel_file)
+          else:
+            self.load_fields_file = load_fields_file
 
     def init_sim(self):
         if self._is_initialized:
@@ -1988,7 +2026,11 @@ def init_sim(self):
             hook()
 
         self._is_initialized = True
-
+
+        if self.load_fields_file:
+            self.load_fields(
+                self.load_fields_file, self.load_single_parallel_file)
+
     def using_real_fields(self):
         cond1 = self.is_cylindrical and self.m != 0
         cond2 = any([s.phase.imag for s in self.symmetries])

python/tests/test_dump_load.py

@@ -0,0 +1,199 @@
+import itertools
+import os
+import re
+import sys
+import unittest
+import warnings
+import h5py
+import numpy as np
+import meep as mp
+
+try:
+    unicode
+except NameError:
+    unicode = str
+
+
+class TestLoadDump(unittest.TestCase):
+
+    fname_base = re.sub(r'\.py$', '', os.path.split(sys.argv[0])[1])
+    fname = fname_base + '-ez-000200.00.h5'
+
+    def setUp(self):
+        print("Running {}".format(self._testMethodName))
+
+    @classmethod
+    def setUpClass(cls):
+        cls.temp_dir = mp.make_output_directory()
+        print("Saving temp files to dir: {}".format(cls.temp_dir))
+
+    @classmethod
+    def tearDownClass(cls):
+        mp.delete_directory(cls.temp_dir)
+
+    # Tests various combinations of dumping/loading structure & chunk layout.
+    def _load_dump_structure(self, chunk_file=False, chunk_sim=False, single_parallel_file=True):
+        from meep.materials import Al
+        resolution = 50
+        cell = mp.Vector3(5, 5)
+        sources = mp.Source(src=mp.GaussianSource(1, fwidth=0.2), center=mp.Vector3(), component=mp.Ez)
+        one_by_one = mp.Vector3(1, 1, mp.inf)
+        geometry = [mp.Block(material=Al, center=mp.Vector3(), size=one_by_one),
+                    mp.Block(material=mp.Medium(epsilon=13), center=mp.Vector3(1), size=one_by_one)]
+        pml_layers = [mp.PML(0.5)]
+
+        symmetries = [mp.Mirror(mp.Y)]
+
+        sim1 = mp.Simulation(resolution=resolution,
+                             cell_size=cell,
+                             boundary_layers=pml_layers,
+                             geometry=geometry,
+                             symmetries=symmetries,
+                             sources=[sources])
+
+        sample_point = mp.Vector3(0.12, -0.29)
+        ref_field_points = []
+
+        def get_ref_field_point(sim):
+            p = sim.get_field_point(mp.Ez, sample_point)
+            ref_field_points.append(p.real)
+
+        sim1.run(mp.at_every(5, get_ref_field_point), until=50)
+
+        dump_dirname = os.path.join(self.temp_dir, 'test_load_dump_structure')
+        sim1.dump(dump_dirname, dump_structure=True, dump_fields=False, single_parallel_file=single_parallel_file)
+
+        dump_chunk_fname = None
+        chunk_layout = None
+        if chunk_file:
+            dump_chunk_fname = os.path.join(dump_dirname, 'chunk_layout.h5')
+            sim1.dump_chunk_layout(dump_chunk_fname)
+            chunk_layout = dump_chunk_fname
+        if chunk_sim:
+            chunk_layout = sim1
+
+        sim = mp.Simulation(resolution=resolution,
+                            cell_size=cell,
+                            boundary_layers=pml_layers,
+                            sources=[sources],
+                            symmetries=symmetries,
+                            chunk_layout=chunk_layout)
+        sim.load(dump_dirname, load_structure=True, load_fields=False, single_parallel_file=single_parallel_file)
+        field_points = []
+
+        def get_field_point(sim):
+            p = sim.get_field_point(mp.Ez, sample_point)
+            field_points.append(p.real)
+
+        sim.run(mp.at_every(5, get_field_point), until=50)
+
+        for ref_pt, pt in zip(ref_field_points, field_points):
+            self.assertAlmostEqual(ref_pt, pt)
+
+    def test_load_dump_structure(self):
+        self._load_dump_structure()
+
+    @unittest.skipIf(not mp.with_mpi(), "MPI specific test")
+    def test_load_dump_structure_sharded(self):
+        self._load_dump_structure(single_parallel_file=False)
+
+    def test_load_dump_chunk_layout_file(self):
+        self._load_dump_structure(chunk_file=True)
+
+    def test_load_dump_chunk_layout_sim(self):
+        self._load_dump_structure(chunk_sim=True)
+
+    # Tests dumping/loading of fields & structure.
+    def _load_dump_fields(self, single_parallel_file=True):
+        resolution = 50
+        cell = mp.Vector3(5, 5)
+        sources = mp.Source(src=mp.GaussianSource(1, fwidth=0.4), center=mp.Vector3(), component=mp.Ez)
+        one_by_one = mp.Vector3(1, 1, mp.inf)
+        geometry = [mp.Block(material=mp.Medium(index=3.2), center=mp.Vector3(), size=one_by_one),
+                    mp.Block(material=mp.Medium(epsilon=13), center=mp.Vector3(1), size=one_by_one)]
+        pml_layers = [mp.PML(0.5)]
+        symmetries = [mp.Mirror(mp.Y)]
+
+        sim1 = mp.Simulation(resolution=resolution,
+                             cell_size=cell,
+                             boundary_layers=pml_layers,
+                             geometry=geometry,
+                             symmetries=symmetries,
+                             sources=[sources])
+
+        sample_point = mp.Vector3(0.12, -0.29)
+
+        dump_dirname = os.path.join(self.temp_dir, 'test_load_dump_fields')
+        os.makedirs(dump_dirname, exist_ok=True)
+
+        ref_field_points = {}
+        def get_ref_field_point(sim):
+            p = sim.get_field_point(mp.Ez, sample_point)
+            ref_field_points[sim.meep_time()] = p.real
+
+        # First run until t=15 and save structure/fields
+        sim1.run(mp.at_every(1, get_ref_field_point), until=15)
+        sim1.dump(dump_dirname, dump_structure=True, dump_fields=True, single_parallel_file=single_parallel_file)
+
+        # Then continue running another 5 until t=20
+        sim1.run(mp.at_every(1, get_ref_field_point), until=5)
+
+        # Now create a new simulation and try restoring state.
+        sim = mp.Simulation(resolution=resolution,
+                            cell_size=cell,
+                            boundary_layers=pml_layers,
+                            sources=[sources],
+                            symmetries=symmetries,
+                            chunk_layout=sim1)
+        # Just restore structure first.
+        sim.load(dump_dirname, load_structure=True, load_fields=False, single_parallel_file=single_parallel_file)
+        field_points = {}
+
+        def get_field_point(sim):
+            p = sim.get_field_point(mp.Ez, sample_point)
+            field_points[sim.meep_time()] = p.real
+
+        # Now load the fields (at t=15) and then continue to t=20
+        sim.load(dump_dirname, load_structure=False, load_fields=True, single_parallel_file=single_parallel_file)
+        sim.run(mp.at_every(1, get_field_point), until=5)
+
+        for t, v in field_points.items():
+            self.assertAlmostEqual(ref_field_points[t], v)
+
+    def test_load_dump_fields(self):
+        self._load_dump_fields()
+
+    @unittest.skipIf(not mp.with_mpi(), "MPI specific test")
+    def test_load_dump_fields_sharded(self):
+        self._load_dump_fields(single_parallel_file=False)
+
+    # This assertRaisesRegex check does not play well with MPI due to the
+    # underlying call to meep::abort
+    @unittest.skipIf(mp.with_mpi(), "MPI specific test")
+    def test_dump_fails_for_non_null_polarization_state(self):
+        resolution = 50
+        cell = mp.Vector3(5, 5)
+        sources = mp.Source(src=mp.GaussianSource(1, fwidth=0.4), center=mp.Vector3(), component=mp.Ez)
+        one_by_one = mp.Vector3(1, 1, mp.inf)
+        from meep.materials import Al
+        geometry = [mp.Block(material=Al, center=mp.Vector3(), size=one_by_one),
+                    mp.Block(material=mp.Medium(epsilon=13), center=mp.Vector3(1), size=one_by_one)]
+
+        sim = mp.Simulation(resolution=resolution,
+                            cell_size=cell,
+                            boundary_layers=[],
+                            geometry=geometry,
+                            symmetries=[],
+                            sources=[sources])
+
+        dump_dirname = os.path.join(self.temp_dir, 'test_load_dump_fields')
+        os.makedirs(dump_dirname, exist_ok=True)
+
+        sim.run(until=1)
+        # NOTE: We do not yet support checkpoint/restore when there is a
+        # non-null polarization_state
+        with self.assertRaisesRegex(RuntimeError, 'meep: non-null polarization_state in fields::dump'):
+          sim.dump(dump_dirname, dump_structure=True, dump_fields=True)
+
+if __name__ == '__main__':
+    unittest.main()