Add more extensive keepalive test

Author: franzpoeschel

test/python/unittest/API/APITest.py

@@ -2342,28 +2342,39 @@ def testScalarHdf5Fields(self):
         series_read_again.close()
 
     def testKeepaliveComponentExtraction(self):
-        """Test that keepalive specifications guard root objects from garbage collection."""
+        """Test that keepalive specifications
+        guard root objects from garbage collection."""
+        self.testKeepaliveMeshComponent()
+        self.testKeepaliveParticlePosition()
+        self.testKeepaliveParticlePatches()
+
+    def testKeepaliveMeshComponent(self):
+        """Test keepalive for mesh component extraction."""
         for ext in tested_file_extensions:
-            self.backend_keepalive_component_extraction(ext)
+            self.backend_keepalive_mesh_component(ext)
 
-    def backend_keepalive_component_extraction(self, file_ending):
-        """Helper function that creates an openPMD Series, extracts a component,
-        discards the parent objects, forces garbage collection, and then writes data
-        using only the returned component."""
+    def testKeepaliveParticlePosition(self):
+        """Test keepalive for particle position component extraction."""
+        for ext in tested_file_extensions:
+            self.backend_keepalive_particle_position(ext)
+
+    def testKeepaliveParticlePatches(self):
+        """Test keepalive for particle patches component extraction."""
+        for ext in tested_file_extensions:
+            self.backend_keepalive_particle_patches(ext)
+
+    def backend_keepalive_mesh_component(self, file_ending):
+        """Helper function that tests keepalive
+        for mesh component extraction."""
         import gc
 
-        filename = "../samples/unittest_py_keepalive." + file_ending
+        filename = "unittest_py_keepalive_mesh." + file_ending
         path = filename
 
         def get_component_only():
-            """
-            Create a Series, access a component, discard the Series and Iteration,
-            and return only the component. The keepalive specification should
-            guard the root objects from garbage collection.
-            """
-            series = io.Series(path, io.Access.create)
+            series = io.Series(path, io.Access.create_linear)
             backend = series.backend
-            iteration = series.iterations[0]
+            iteration = series.snapshots()[0]
             mesh = iteration.meshes["E"]
             component = mesh["x"]
 
@@ -2387,20 +2398,113 @@ def get_component_only():
 
         component.series_flush()
         if backend == "ADIOS2":
-            # need to close the step for data to become visible in ADIOS2
-            # cant close the step since we threw away the Iteration, so we need
-            # to trigger the GC here
             del component
             gc.collect()
 
         read = io.Series(path, io.Access.read_only)
-        loaded = read.iterations[0].meshes["E"]["x"][:]
+        loaded = read.snapshots()[0].meshes["E"]["x"][:]
         read.flush()
         np.testing.assert_array_equal(
             loaded,
             np.reshape(np.arange(100, dtype=np.dtype("float")), [10, 10])
         )
 
+    def backend_keepalive_particle_position(self, file_ending):
+        """Helper function that tests keepalive
+        for particle position component extraction."""
+        import gc
+
+        filename = "unittest_py_keepalive_particle." + file_ending
+        path = filename
+        num_particles = 100
+
+        def get_component_only():
+            series = io.Series(path, io.Access.create_linear)
+            backend = series.backend
+            iteration = series.snapshots()[0]
+            particles = iteration.particles["electrons"]
+            position = particles["position"]["x"]
+
+            position.reset_dataset(
+                io.Dataset(np.dtype("float"), [num_particles]))
+
+            del iteration
+            del particles
+            del series
+            gc.collect()
+
+            return position, backend
+
+        position, backend = get_component_only()
+        gc.collect()
+
+        position[:] = np.arange(num_particles, dtype=np.dtype("float"))
+
+        position.series_flush()
+        if backend == "ADIOS2":
+            del position
+            gc.collect()
+
+        read = io.Series(path, io.Access.read_only)
+        loaded = read.snapshots()[0] \
+            .particles["electrons"]["position"]["x"][:]
+        read.flush()
+        np.testing.assert_array_equal(
+            loaded,
+            np.arange(num_particles, dtype=np.dtype("float"))
+        )
+
+    def backend_keepalive_particle_patches(self, file_ending):
+        """Helper function that tests keepalive
+        for particle patches extraction."""
+        import gc
+
+        filename = "unittest_py_keepalive_patches." + file_ending
+        path = filename
+
+        def get_component_only():
+            series = io.Series(path, io.Access.create_linear)
+            backend = series.backend
+            iteration = series.snapshots()[0]
+            particles = iteration.particles["electrons"]
+
+            dset = io.Dataset(np.dtype("float"), [30])
+            position_x = particles["position"]["x"]
+            position_x.reset_dataset(dset)
+            position_x[:] = np.arange(30, dtype=np.float32)
+
+            dset = io.Dataset(np.dtype("uint64"), [2])
+            num_particles_comp = particles.particle_patches["numParticles"]
+            num_particles_comp.reset_dataset(dset)
+            num_particles_comp.store(0, np.uint64(10))
+            num_particles_comp.store(1, np.uint64(20))
+
+            del iteration
+            del particles
+            del series
+            gc.collect()
+
+            return num_particles_comp, backend
+
+        component, backend = get_component_only()
+        gc.collect()
+
+        component.store(0, np.uint64(50))
+
+        component.series_flush()
+        if backend == "ADIOS2":
+            del component
+            gc.collect()
+
+        read = io.Series(path, io.Access.read_only)
+        loaded = read.snapshots()[0] \
+            .particles["electrons"].particle_patches["numParticles"].load()
+        read.flush()
+        np.testing.assert_array_equal(
+            loaded[0],
+            np.uint64(50)
+        )
+
 
 if __name__ == '__main__':
     unittest.main()