Cameron Hummels avatar Cameron Hummels committed aba16d4

Adding in feature for passing redshift information through halo finding process and outputting as comment to text file when using .write_out() function. Also adding functions for limiting redshift/cycle range in halo_merger_tree.

Comments (0)

Files changed (2)

yt/analysis_modules/halo_finding/halo_objects.py

 
     _fields = ["particle_position_%s" % ax for ax in 'xyz']
 
-    def __init__(self, data_source, dm_only=True):
+    def __init__(self, data_source, dm_only=True, redshift=-1):
         """
         Run hop on *data_source* with a given density *threshold*.  If
         *dm_only* is set, only run it on the dark matter particles, otherwise
         mylog.info("Parsing outputs")
         self._parse_output()
         mylog.debug("Finished. (%s)", len(self))
+        self.redshift = redshift
 
     def __obtain_particles(self):
         if self.dm_only:
         else:
             f = open(filename, "w")
         f.write("# HALOS FOUND WITH %s\n" % (self._name))
+        f.write("# REDSHIFT OF OUTPUT = %f\n" % (self.redshift))
 
         if not ellipsoid_data:
             f.write("\t".join(["# Group","Mass","# part","max dens"
     _name = "FOF"
     _halo_class = FOFHalo
 
-    def __init__(self, data_source, link=0.2, dm_only=True):
+    def __init__(self, data_source, link=0.2, dm_only=True, redshift=-1):
         self.link = link
         mylog.info("Initializing FOF")
-        HaloList.__init__(self, data_source, dm_only)
+        HaloList.__init__(self, data_source, dm_only, redshift=redshift)
 
     def _run_finder(self):
         self.tags = \
         self.period = pf.domain_right_edge - pf.domain_left_edge
         self.pf = pf
         self.hierarchy = pf.h
+        self.redshift = pf.current_redshift
         self._data_source = pf.h.all_data()
         GenericHaloFinder.__init__(self, pf, self._data_source, dm_only,
             padding)
         #self._reposition_particles((LE, RE))
         # here is where the FOF halo finder is run
         mylog.info("Using a linking length of %0.3e", linking_length)
-        FOFHaloList.__init__(self, self._data_source, linking_length, dm_only)
+        FOFHaloList.__init__(self, self._data_source, linking_length, dm_only,
+                             redshift=self.redshift)
         self._parse_halolist(1.)
         self._join_halolists()
 

yt/analysis_modules/halo_merger_tree/enzofof_merger_tree.py

         """
         hp = []
         for line in open("%s_%05i.txt" % (filename_prefix, self.output_id)):
+            if line.startswith("# RED"):
+                self.redshift = float(line.split("=")[1])
             if line.strip() == "": continue # empty
             if line[0] == "#": continue # comment
             x,y,z = [float(f) for f in line.split()[7:10]] # COM x,y,z
         ----------
         zrange : tuple
             This is the redshift range (min, max) to calculate the
-            merger tree.
+            merger tree. E.g. (0, 2) for z=2 to z=0
         cycle_range : tuple, optional
             This is the cycle number range (min, max) to calculate the
             merger tree.  If both zrange and cycle_number given,
         self.relationships = {}
         self.redshifts = {}
         self.external_FOF = external_FOF
-        self.find_outputs(zrange, cycle_range, output)
         if load_saved:
             self.load_tree(save_filename)
+            # make merger tree work within specified cycle/z limits
+            # on preloaded halos
+            if zrange is not None:
+                self.select_redshifts(zrange)
             if cycle_range is not None:
-                # actually make merger tree work within specified cycle limits
-                self.redshifts = {}
-                for i in range(cycle_range[0],cycle_range[1]+1):
-                    self.redshifts[i] = 0.0 # don't have redshift info
+                self.select_cycles(cycle_range)
         else:
+            self.find_outputs(zrange, cycle_range, output)
             self.run_merger_tree(output)
             self.save_tree(save_filename)
         
+    def select_cycles(self, cycle_range):
+        """
+        Takes an existing tree and pares it to only include a subset of
+        cycles.  Useful in paring a loaded tree. 
+        """
+        # N.B. Does not delete info from self.relationships to save space
+        # just removes it from redshift dict for indexing
+        for cycle in self.redshifts.keys():
+            if cycle <= cycle_range[0] and cycle >= cycle_range[1]:
+                del self.redshifts[cycle]
+
+    def select_redshifts(self, zrange):
+        """
+        Takes an existing tree and pares it to only include a subset of
+        redshifts.  Useful in paring a loaded tree. 
+        """
+        # N.B. Does not delete info from self.relationships to save space
+        # just removes it from redshift dict for indexing
+        for redshift in self.redshifts.values():
+            if redshift <= zrange[0] and redshift >= zrange[1]:
+                # some reverse lookup magic--assumes unique cycle/z pairs
+                cycle = [key for key,value in mt.redshifts.items() \
+                         if value == redshift][0]
+                del self.redshifts[cycle]
 
     def save_tree(self, filename):
         cPickle.dump((self.redshifts, self.relationships),
     Examples
     --------
     ts = TimeSeriesData.from_filenames("DD????/DD????")
-    if not os.path.isdir('FOF'): os.mkdir('FOF')
     for pf in ts:
         halo_list = FOFHaloFinder(pf)
         i = int(pf.basename[2:])
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.