Enable dumping only part of the buffer for burst writer.

hoomd/GSDDequeWriter.cc

@@ -42,7 +42,7 @@ GSDDequeWriter::GSDDequeWriter(std::shared_ptr<SystemDefinition> sysdef,
         else
             {
             analyze(timestep);
-            dump();
+            dump(0, -1, true);
             }
         }
     }
@@ -59,14 +59,37 @@ void GSDDequeWriter::analyze(uint64_t timestep)
         }
     }
 
-void GSDDequeWriter::dump()
+void GSDDequeWriter::dump(long int start, long int end, bool empty_buffer)
     {
-    for (auto i {static_cast<long int>(m_frame_queue.size()) - 1}; i >= 0; --i)
+    auto buffer_length = static_cast<long int>(m_frame_queue.size());
+    if (end > buffer_length)
+        {
+        throw std::runtime_error("Burst.dump's end index is out of range.");
+        }
+    if (start < 0 || start > buffer_length)
+        {
+        throw std::runtime_error("Burst.dump's start index is out of range.");
+        }
+    if (end < 0)
+        {
+        end = buffer_length - start;
+        start = 0;
+        }
+    else
+        {
+        auto temp_end = end;
+        end = buffer_length - start;
+        start = buffer_length - temp_end;
+        }
+    for (auto i = end - 1; i >= start; --i)
         {
         write(m_frame_queue[i], m_log_queue[i]);
         }
-    m_frame_queue.clear();
-    m_log_queue.clear();
+    if (empty_buffer)
+        {
+        m_frame_queue.clear();
+        m_log_queue.clear();
+        }
     }
 
 int GSDDequeWriter::getMaxQueueSize() const

hoomd/GSDDequeWriter.h

@@ -31,7 +31,7 @@ class PYBIND11_EXPORT GSDDequeWriter : public GSDDumpWriter
 
     void analyze(uint64_t timestep) override;
 
-    void dump();
+    void dump(long int start, long int end, bool empty_buffer);
 
     int getMaxQueueSize() const;
     void setMaxQueueSize(int new_max_size);

hoomd/md/pytest/test_burst_writer.py

@@ -145,7 +145,9 @@ def test_len(sim, tmp_path):
     assert len(burst_writer) == 0
 
 
-def test_burst_dump(sim, tmp_path):
+@pytest.mark.parametrize("start, end", [(0, -1), (0, 0), (0, 1), (0, 2), (1, 1),
+                                        (2, 2), (1, 2), (1, -1), (2, -1)])
+def test_burst_dump(sim, tmp_path, start, end):
     filename = tmp_path / "temporary_test_file.gsd"
 
     burst_trigger = hoomd.trigger.Periodic(period=2, phase=1)
@@ -164,11 +166,15 @@ def test_burst_dump(sim, tmp_path):
             # First frame is always written
             assert len(traj) == 1
 
-    burst_writer.dump()
+    burst_writer.dump(start=start, end=end)
     burst_writer.flush()
+    dumped_frames = [3, 5, 7]
     if sim.device.communicator.rank == 0:
+        if end == -1:
+            end = len(dumped_frames)
         with gsd.hoomd.open(name=filename, mode='r') as traj:
-            assert [frame.configuration.step for frame in traj] == [0, 3, 5, 7]
+            assert [frame.configuration.step for frame in traj
+                    ] == [0] + dumped_frames[start:end]
 
 
 def test_burst_max_size(sim, tmp_path):
@@ -242,3 +248,36 @@ def test_write_burst_log(sim, tmp_path):
         with gsd.hoomd.open(name=filename, mode='r') as traj:
             for frame, sim_ke in zip(traj[1:], kinetic_energies):
                 assert frame.log[key] == sim_ke
+
+
+@pytest.mark.parametrize("empty_buffer", [True, False])
+def test_burst_dump_empty_buffer(sim, tmp_path, empty_buffer):
+    filename = tmp_path / "temporary_test_file.gsd"
+    burst_trigger = hoomd.trigger.Periodic(period=2, phase=1)
+    burst_writer = hoomd.write.Burst(trigger=burst_trigger,
+                                     filename=filename,
+                                     mode='wb',
+                                     dynamic=['property', 'momentum'],
+                                     max_burst_size=3,
+                                     write_at_start=True)
+    sim.operations.writers.append(burst_writer)
+    sim.run(8)
+    burst_writer.flush()
+    if sim.device.communicator.rank == 0:
+        assert Path(filename).exists()
+        with gsd.hoomd.open(filename, "r") as traj:
+            # First frame is always written
+            assert len(traj) == 1
+
+    burst_writer.dump(empty_buffer=empty_buffer)
+    burst_writer.flush()
+    if sim.device.communicator.rank == 0:
+        with gsd.hoomd.open(name=filename, mode='r') as traj:
+            assert len(traj) == 4
+
+    sim.run(4)
+    burst_writer.dump()
+    burst_writer.flush()
+    if sim.device.communicator.rank == 0:
+        with gsd.hoomd.open(name=filename, mode='r') as traj:
+            assert len(traj) == 6 if empty_buffer else 8

hoomd/write/gsd_burst.py

@@ -114,20 +114,27 @@ def _attach_hook(self):
                                               self.mode, self.write_at_start,
                                               sim.timestep)
 
-    def dump(self):
-        """Write all currently stored frames to the file and empties the buffer.
+    def dump(self, start=0, end=-1, empty_buffer=True):
+        """Write stored frames in range to the file and empties the buffer.
 
         This method alllows for custom writing of frames at user specified
         conditions.
 
+        Args:
+            start (int): The first frame to write. Defaults to 0.
+            end (int): The last frame to write, must be positive integer.
+                Defaults to -1 (last frame).
+            empty_buffer (bool): When ``True`` the buffer is emptied after
+                writing. Defaults to ``True``.
+
         .. rubric:: Example:
 
         .. code-block:: python
 
             burst.dump()
         """
         if self._attached:
-            self._cpp_obj.dump()
+            self._cpp_obj.dump(start, end, empty_buffer)
 
     def __len__(self):
         """Get the current length of the internal frame buffer.