Add ability to insist on using IDL-refined AHF catalogues (Alyson Bro…

…ok's toolset),

or to insist on ignoring those catalogues. This is in response to issue #117, where
it has become clear people have both types of catalogue on disk and this can cause
severe confusion due to the automated search for catalogues and statistics files.

To insist on using IDL-generated catalogues, use

  tangos add <name_of_sim> --handler=pynbody.ChangaUseIDLInputHandler

To insist on ignoring IDL-generated catalogues, use

  tangos add <name_of_sim> --handler=pynbody.ChangaIgnoreIDLInputHandler

tangos/input_handlers/__init__.py

@@ -35,6 +35,8 @@ class HandlerBase(object):
     Subclasses provide implementations for different formats and situations.
     """
 
+    halo_stat_file_class_name = "HaloStatFile"
+
     def __init__(self, basename):
         self.basename = self.strip_slashes(basename)
         self.quicker = False # a flag to indicate that corners may be cut in the interest of efficiency
@@ -86,7 +88,7 @@ def get_stat_file(self, ts_extension, object_typetag):
         #ts = DummyTimeStep()
         #ts.redshift = self.get_timestep_properties(ts_extension)['redshift']
         from . import caterpillar
-        statfile = halo_stat_files.HaloStatFile(self._extension_to_filename(ts_extension))
+        statfile = getattr(halo_stat_files, self.halo_stat_file_class_name)(self._extension_to_filename(ts_extension))
         return statfile
 
 

tangos/input_handlers/pynbody.py

@@ -32,6 +32,8 @@ def __repr__(self):
 
 
 class PynbodyInputHandler(finding.PatternBasedFileDiscovery, HandlerBase):
+    pynbody_halo_class_name = None
+
     def __new__(cls, *args, **kwargs):
         import pynbody as pynbody_local
 
@@ -48,14 +50,23 @@ def __init__(self, *args, **kwargs):
         pynbody_version = getattr(pynbody, "__version__","0.00")
         assert pynbody_version>="0.46", "Tangos requires pynbody 0.46 or later"
 
+    @classmethod
+    def _construct_pynbody_halos(cls, sim, *args, **kwargs):
+        if cls.pynbody_halo_class_name is None:
+            return sim.halos(*args, **kwargs)
+        else:
+            halo_class = getattr(pynbody.halo, cls.pynbody_halo_class_name)
+            return halo_class(sim, *args, **kwargs)
+
     def _is_able_to_load(self, ts_extension):
         filepath = self._extension_to_filename(ts_extension)
         try:
             f = pynbody.load(filepath)
             if self.quicker:
                 logger.warn("Pynbody was able to load %r, but because 'quicker' flag is set we won't check whether it can also load the halo files", filepath)
             else:
-                h = f.halos()
+                h = self._construct_pynbody_halos(f)
+
             return True
         except (IOError, RuntimeError):
             return False
@@ -174,9 +185,10 @@ def _construct_halo_cat(self, ts_extension, object_typetag):
         f = self.load_timestep(ts_extension)
         h = _loaded_halocats.get(id(f), lambda: None)()
         if h is None:
-            h = f.halos()
+            h = self._construct_pynbody_halos(f)
             if isinstance(h, pynbody.halo.SubfindCatalogue):
-                h = f.halos(subs=True)
+                # ugly fix - loads groups by default, wanted halos
+                h = self._construct_pynbody_halos(f, subs=True)
             _loaded_halocats[id(f)] = weakref.ref(h)
             f._db_current_halocat = h # keep alive for lifetime of simulation
         return h  # pynbody.halo.AmigaGrpCatalogue(f)
@@ -216,6 +228,7 @@ def enumerate_objects(self, ts_extension, object_typetag="halo", min_halo_partic
         else:
             logger.warn("No halo statistics file found for timestep %r",ts_extension)
 
+            snapshot_keep_alive = self.load_timestep(ts_extension)
             try:
                 h = self._construct_halo_cat(ts_extension, object_typetag)
             except:
@@ -346,7 +359,7 @@ def _construct_group_cat(self, ts_extension):
         f = self.load_timestep(ts_extension)
         h = _loaded_halocats.get(id(f)+1, lambda: None)()
         if h is None:
-            h = f.halos()
+            h = self._construct_pynbody_halos(f)
             assert isinstance(h, pynbody.halo.SubfindCatalogue)
             _loaded_halocats[id(f)+1] = weakref.ref(h)
             f._db_current_groupcat = h  # keep alive for lifetime of simulation
@@ -453,6 +466,7 @@ def _is_able_to_load(self, filepath):
 
 
 
+
 class ChangaInputHandler(PynbodyInputHandler):
     flags_include = ["dPhysDenMin", "dCStar", "dTempMax",
                      "dESN", "bLowTCool", "bSelfShield", "dExtraCoolShutoff"]
@@ -571,4 +585,13 @@ def _param_file_to_dict(param_file):
                 pass
         return out
 
+class ChangaIgnoreIDLInputHandler(ChangaInputHandler):
+    pynbody_halo_class_name = "AHFCatalogue"
+    halo_stat_file_class_name = "AHFStatFile"
+
+class ChangaUseIDLInputHandler(ChangaInputHandler):
+    pynbody_halo_class_name = "AmigaGrpCatalogue"
+    halo_stat_file_class_name = "AmigaIDLStatFile"
+    auxiliary_file_patterns = ["*.amiga.grp"]
+
 from . import caterpillar, eagle