Merge pull request #236 from pynbody/timestep-match

Add option to tangos write to target particular timesteps

Author: apontzen

tangos/tools/property_writer.py

@@ -44,6 +44,8 @@ def add_parser_arguments(self, parser):
                             help='Specify a simulation (or multiple simulations) to run on')
         parser.add_argument('--latest', action='store_true',
                             help='Run only on the latest timesteps')
+        parser.add_argument('--timesteps-matching', action='append', type=str,
+                            help='Run only on timesteps with extensions matching the specified string. Multiple timesteps may be specified.')
         parser.add_argument('--force', action='store_true',
                             help='Run calculations even if a value is already stored in the database')
         parser.add_argument('--debug', action='store_true',
@@ -96,10 +98,18 @@ def _build_file_list(self):
                 except IndexError:
                     pass
         else:
-            files = core.get_default_session().query(core.timestep.TimeStep).filter(
-                core.timestep.TimeStep.simulation_id.in_([q.id for q in query.all()])). \
+            timestep_filter = core.timestep.TimeStep.simulation_id.in_([q.id for q in query.all()])
+            if self.options.timesteps_matching is not None and len(self.options.timesteps_matching)>0:
+                subfilter = core.timestep.TimeStep.extension.like(self.options.timesteps_matching[0])
+                for m in self.options.timesteps_matching[1:]:
+                    subfilter |= core.timestep.TimeStep.extension.like(m)
+                timestep_filter &= subfilter
+
+            files = core.get_default_session().query(core.timestep.TimeStep).filter(timestep_filter). \
                 order_by(core.timestep.TimeStep.time_gyr).all()
 
+
+
         if self.options.backwards:
             files = files[::-1]
 

tests/test_db_writer.py

@@ -11,12 +11,12 @@
 from tangos.util import proxy_object
 
 
-def setup_func():
+def setup_func(sim="dummy_sim_1"):
     parallel_tasks.use('null')
 
     testing.init_blank_db_for_testing()
     db.config.base = os.path.join(os.path.dirname(__file__), "test_simulations")
-    manager = add_simulation.SimulationAdderUpdater(output_testing.TestInputHandler("dummy_sim_1"))
+    manager = add_simulation.SimulationAdderUpdater(output_testing.TestInputHandler(sim))
     with log.LogCapturer():
         manager.scan_simulation_and_add_all_descendants()
 
@@ -29,6 +29,12 @@ def fresh_database():
     yield
     teardown_func()
 
+@fixture
+def fresh_database_2():
+    setup_func("dummy_sim_2")
+    yield
+    teardown_func()
+
 class DummyProperty(properties.PropertyCalculation):
     names = "dummy_property",
     requires_particle_data = True
@@ -225,3 +231,10 @@ def test_writer_handles_sim_properties(fresh_database):
         ts = db.get_timestep("dummy_sim_1/step.%d"%i)
         x, = ts.calculate_all("dummy_property_accessing_simulation_property")
         npt.assert_equal(x,[1]*ts.halos.count())
+
+def test_timesteps_matching(fresh_database_2):
+    run_writer_with_args("dummy_property", "--timesteps-matching", "step.1", "--timesteps-matching", "step.2")
+    assert 'dummy_property' in  db.get_halo("dummy_sim_2/step.1/1").keys()
+    assert 'dummy_property' in db.get_halo("dummy_sim_2/step.1/2").keys()
+    assert 'dummy_property' in db.get_halo("dummy_sim_2/step.2/1").keys()
+    assert 'dummy_property' not in db.get_halo("dummy_sim_2/step.3/1").keys()