Make timestep thinning work correctly even if the interval specified …

…is greater than all existing separations
pynbody · Dec 28, 2023 · d231830 · d231830
1 parent cd816c3
commit d231830
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 9 deletions.
diff --git a/tangos/tools/timestep_thinner.py b/tangos/tools/timestep_thinner.py
@@ -59,33 +59,41 @@ def thin_simulation(self, sim):
 
         to_remove = []
 
+        dt_already_accumulated = 0.0
+
         for i, dt in enumerate(intervals, 1):
-            if dt < min_interval:
+            if dt + dt_already_accumulated < min_interval:
                 to_remove.append(timesteps[i])
-                print(f"    {timesteps[i].extension:s} (delta_t =  {dt:.3f} Gyr)")
+                if dt_already_accumulated > 0.0:
+                    print(f"    {timesteps[i].extension:s} (delta_t = {dt+dt_already_accumulated:.3f} Gyr; "
+                          f"original delta_t =  {dt:.3f} Gyr)")
+                else:
+                    print(f"    {timesteps[i].extension:s} (delta_t = {dt:.3f} Gyr)")
+                dt_already_accumulated += dt
+            else:
+                dt_already_accumulated = 0.0
 
         if len(to_remove) == 0:
             print("    None")
         else:
             print(f"  There are {len(to_remove)} timesteps to remove")
 
             if not self.options.force:
-                print("""Type "yes" to continue""")
-                ok = input(":").lower() == "yes"
+                print("""  Type "yes" to continue""")
+                ok = input("  :").lower() == "yes"
             else:
                 ok = True
 
             if ok:
                 session = core.get_default_session()
                 for ts in to_remove:
-                    print(f"    Removing {ts.extension:s}")
                     session.execute(
                        sqlalchemy.delete(core.TimeStep).filter(core.TimeStep.id == ts.id)
                     )
                 session.commit()
 
             else:
-                print("Skipping")
+                print("  Skipping")
 
         self._cleanup_orphan_objects()
         self._cleanup_orphan_links()
@@ -102,7 +110,7 @@ def _cleanup_orphan_objects(self):
                 )
             ).rowcount
             connection.commit()
-        print(f"Removed {count} orphan objects")
+        print(f"  Removed {count} orphan objects")
 
     def _cleanup_orphan_links(self):
         engine = core.get_default_engine()
@@ -117,7 +125,7 @@ def _cleanup_orphan_links(self):
                 )
             ).rowcount
             connection.commit()
-        print(f"Removed {count} orphan links")
+        print(f"  Removed {count} orphan links")
 
     def _cleanup_orphan_properties(self):
         engine = core.get_default_engine()
@@ -132,4 +140,4 @@ def _cleanup_orphan_properties(self):
             ).rowcount
             connection.commit()
 
-        print(f"Removed {count} orphan properties")
+        print(f"  Removed {count} orphan properties")
diff --git a/tests/test_timestep_thinner.py b/tests/test_timestep_thinner.py
@@ -79,6 +79,17 @@ def test_timestep_thinner_relative(fresh_database):
 
     _assert_timestep_removed(target_ts_id)
 
+def test_timestep_thinner_doesnt_over_thin(fresh_database):
+    """Check that when the threshold delta_time is more than all the delta times, we retain
+    some timesteps, just not spaced more regularly than delta_time"""
+    _assert_everything_present()
+
+    tt = timestep_thinner.TimestepThinner()
+    tt.parse_command_line(["0.1999", "-f"])
+    tt.run_calculation_loop()
+
+    assert [t.extension for t in tangos.get_simulation("sim").timesteps] == ["ts1", "ts3", "ts6"]
+
 def test_timestep_thinner_absolute(fresh_database):
     _assert_everything_present()