[yt-svn] commit/yt: MatthewTurk: Merging from PR 427
Bitbucket
commits-noreply at bitbucket.org
Fri Feb 15 14:25:36 PST 2013
1 new commit in yt:
https://bitbucket.org/yt_analysis/yt/commits/6b246ba83cb3/
changeset: 6b246ba83cb3
branch: stable
user: MatthewTurk
date: 2013-02-15 23:25:27
summary: Merging from PR 427
affected #: 4 files
diff -r 9b336aa05056cb4f6aab9066dd24c3c188695f2f -r 6b246ba83cb38d6fd2467bf76e265e6ca4f78beb yt/analysis_modules/halo_finding/halo_objects.py
--- a/yt/analysis_modules/halo_finding/halo_objects.py
+++ b/yt/analysis_modules/halo_finding/halo_objects.py
@@ -1059,7 +1059,7 @@
_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
@@ -1074,6 +1074,7 @@
mylog.info("Parsing outputs")
self._parse_output()
mylog.debug("Finished. (%s)", len(self))
+ self.redshift = redshift
def __obtain_particles(self):
if self.dm_only:
@@ -1257,6 +1258,7 @@
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"
@@ -1453,18 +1455,17 @@
pass
class HOPHaloList(HaloList):
-
+ """
+ Run hop on *data_source* with a given density *threshold*. If
+ *dm_only* is set, only run it on the dark matter particles, otherwise
+ on all particles. Returns an iterable collection of *HopGroup* items.
+ """
_name = "HOP"
_halo_class = HOPHalo
_fields = ["particle_position_%s" % ax for ax in 'xyz'] + \
["ParticleMassMsun"]
def __init__(self, data_source, threshold=160.0, dm_only=True):
- """
- Run hop on *data_source* with a given density *threshold*. If
- *dm_only* is set, only run it on the dark matter particles, otherwise
- on all particles. Returns an iterable collection of *HopGroup* items.
- """
self.threshold = threshold
mylog.info("Initializing HOP")
HaloList.__init__(self, data_source, dm_only)
@@ -1502,10 +1503,10 @@
_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 = \
@@ -1653,6 +1654,11 @@
class parallelHOPHaloList(HaloList, ParallelAnalysisInterface):
+ """
+ Run hop on *data_source* with a given density *threshold*. If
+ *dm_only* is set, only run it on the dark matter particles, otherwise
+ on all particles. Returns an iterable collection of *HopGroup* items.
+ """
_name = "parallelHOP"
_halo_class = parallelHOPHalo
_fields = ["particle_position_%s" % ax for ax in 'xyz'] + \
@@ -1661,11 +1667,6 @@
def __init__(self, data_source, padding, num_neighbors, bounds, total_mass,
period, threshold=160.0, dm_only=True, rearrange=True, premerge=True,
tree='F'):
- """
- Run hop on *data_source* with a given density *threshold*. If
- *dm_only* is set, only run it on the dark matter particles, otherwise
- on all particles. Returns an iterable collection of *HopGroup* items.
- """
ParallelAnalysisInterface.__init__(self)
self.threshold = threshold
self.num_neighbors = num_neighbors
@@ -2007,6 +2008,10 @@
--------
>>> halos.write_out("HopAnalysis.out")
"""
+ # if path denoted in filename, assure path exists
+ if len(filename.split('/')) > 1:
+ mkdir_rec('/'.join(filename.split('/')[:-1]))
+
f = self.comm.write_on_root(filename)
HaloList.write_out(self, f, ellipsoid_data)
@@ -2026,6 +2031,10 @@
--------
>>> halos.write_particle_lists_txt("halo-parts")
"""
+ # if path denoted in prefix, assure path exists
+ if len(prefix.split('/')) > 1:
+ mkdir_rec('/'.join(prefix.split('/')[:-1]))
+
f = self.comm.write_on_root("%s.txt" % prefix)
HaloList.write_particle_lists_txt(self, prefix, fp=f)
@@ -2049,6 +2058,10 @@
--------
>>> halos.write_particle_lists("halo-parts")
"""
+ # if path denoted in prefix, assure path exists
+ if len(prefix.split('/')) > 1:
+ mkdir_rec('/'.join(prefix.split('/')[:-1]))
+
fn = "%s.h5" % self.comm.get_filename(prefix)
f = h5py.File(fn, "w")
for halo in self._groups:
@@ -2082,6 +2095,10 @@
--------
>>> halos.dump("MyHalos")
"""
+ # if path denoted in basename, assure path exists
+ if len(basename.split('/')) > 1:
+ mkdir_rec('/'.join(basename.split('/')[:-1]))
+
self.write_out("%s.out" % basename, ellipsoid_data)
self.write_particle_lists(basename)
self.write_particle_lists_txt(basename)
@@ -2568,6 +2585,7 @@
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)
@@ -2602,7 +2620,8 @@
#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()
diff -r 9b336aa05056cb4f6aab9066dd24c3c188695f2f -r 6b246ba83cb38d6fd2467bf76e265e6ca4f78beb yt/analysis_modules/halo_merger_tree/api.py
--- a/yt/analysis_modules/halo_merger_tree/api.py
+++ b/yt/analysis_modules/halo_merger_tree/api.py
@@ -38,5 +38,7 @@
MergerTreeTextOutput
from .enzofof_merger_tree import \
+ HaloCatalog, \
find_halo_relationships, \
- EnzoFOFMergerTree
+ EnzoFOFMergerTree, \
+ plot_halo_evolution
diff -r 9b336aa05056cb4f6aab9066dd24c3c188695f2f -r 6b246ba83cb38d6fd2467bf76e265e6ca4f78beb yt/analysis_modules/halo_merger_tree/enzofof_merger_tree.py
--- a/yt/analysis_modules/halo_merger_tree/enzofof_merger_tree.py
+++ b/yt/analysis_modules/halo_merger_tree/enzofof_merger_tree.py
@@ -1,11 +1,14 @@
"""
A very simple, purely-serial, merger tree script that knows how to parse FOF
-catalogs output by Enzo and then compare parent/child relationships.
+catalogs, either output by Enzo or output by yt's FOF halo finder, and then
+compare parent/child relationships.
Author: Matthew J. Turk <matthewturk at gmail.com>
Affiliation: Columbia University
Author: John H. Wise <jwise at astro.princeton.edu>
Affiliation: Princeton
+Author: Cameron Hummels <chummels at gmail.com>
+Affiliation: Arizona
Homepage: http://yt-project.org/
License:
Copyright (C) 2010-2011 Matthew Turk. All Rights Reserved.
@@ -25,6 +28,7 @@
You should have received a copy of the GNU General Public License
along with this program. If not, see <http://www.gnu.org/licenses/>.
"""
+# plot_halo_evolution() gives a good full example of how to use the framework
# First pass at a simplified merger tree
#
@@ -41,6 +45,7 @@
# 8. Parentage is described by a fraction of particles that pass from one to
# the other; we have both descendent fractions and ancestory fractions.
+
import numpy as np
import h5py
import time
@@ -78,39 +83,51 @@
self.order.insert(0, None)
class HaloCatalog(object):
+ r"""A catalog of halos, parsed from EnzoFOF outputs.
+
+ This class will read in catalogs output by the Enzo FOF halo finder and
+ make available their positions, radii, etc. Enzo FOF was provided
+ starting with 2.0, and can be run either inline (with the correct
+ options) or as a postprocessing step using the `-F` command line
+ option. This class is mostly useful when calculating a merger tree,
+ and when the particle IDs for members of a given halo are output as
+ well.
+
+ Parameters
+ ----------
+ output_id : int
+ This is the integer output id of the halo catalog to parse and
+ load.
+ cache : bool
+ Should we store, in between accesses, the particle IDs? If set to
+ true, the correct particle files must exist.
+ external_FOF : bool, optional
+ Are we building a tree from outputs generated by an
+ external FOF program, or an FOF internal to yt?
+ FOF_directory : str, optional
+ Directory where FOF files are located
+ """
cache = None
- def __init__(self, output_id, cache = True):
- r"""A catalog of halos, parsed from EnzoFOF outputs.
-
- This class will read in catalogs output by the Enzo FOF halo finder and
- make available their positions, radii, etc. Enzo FOF was provided
- starting with 2.0, and can be run either inline (with the correct
- options) or as a postprocessing step using the `-F` command line
- option. This class is mostly useful when calculating a merger tree,
- and when the particle IDs for members of a given halo are output as
- well.
-
- Parameters
- ----------
- output_id : int
- This is the integer output id of the halo catalog to parse and
- load.
- cache : bool
- Should we store, in between accesses, the particle IDs? If set to
- true, the correct particle files must exist.
- """
+ def __init__(self, output_id, cache = True, external_FOF=True, FOF_directory="FOF"):
self.output_id = output_id
+ self.external_FOF = external_FOF
self.redshift = 0.0
- self.particle_file = h5py.File("FOF/particles_%05i.h5" % output_id, "r")
- self.parse_halo_catalog()
+ self.FOF_directory = FOF_directory
+ self.particle_file = h5py.File("%s/particles_%05i.h5" % \
+ (FOF_directory, output_id), "r")
+ if self.external_FOF:
+ self.parse_halo_catalog_external()
+ else:
+ self.parse_halo_catalog_internal()
if cache: self.cache = dict()#MaxLengthDict()
def __del__(self):
self.particle_file.close()
- def parse_halo_catalog(self):
+ def parse_halo_catalog_external(self):
hp = []
- for line in open("FOF/groups_%05i.dat" % self.output_id):
+ for line in open("%s/groups_%05i.dat" % \
+ (self.FOF_directory, self.output_id)):
if line.strip() == "": continue # empty
if line.startswith("# Red"):
self.redshift = float(line.split("=")[1])
@@ -126,13 +143,52 @@
self.halo_kdtree = None
return hp
- def read_particle_ids(self, halo_id):
+ def parse_halo_catalog_internal(self):
+ """
+ This parser works on the files output directly out of yt's internal
+ halo_finder. The parse_halo_catalog_external works with an
+ external version of FOF.
+
+ Examples
+ --------
+ pf = load("DD0000/DD0000")
+ halo_list = FOFHaloFinder(pf)
+ halo_list.write_out("FOF/groups_00000.txt")
+ halos_COM = parse_halo_catalog_internal()
+ """
+ hp = []
+ for line in open("%s/groups_%05i.txt" % \
+ (self.FOF_directory, self.output_id)):
+ if line.startswith("# RED"):
+ self.redshift = float(line.split("=")[1])
+ continue
+ 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
+ hp.append([x,y,z])
+ if hp != []:
+ self.halo_positions = np.array(hp)
+ self.halo_kdtree = KDTree(self.halo_positions)
+ else:
+ self.halo_positions = None
+ self.halo_kdtree = None
+ return hp
+
+ def read_particle_ids(self, halo_id):
if self.cache is not None:
if halo_id not in self.cache:
- self.cache[halo_id] = self.particle_file["/Halo%08i/Particle ID" % halo_id][:]
+ if self.external_FOF:
+ self.cache[halo_id] = \
+ self.particle_file["/Halo%08i/Particle ID" % halo_id][:]
+ else:
+ self.cache[halo_id] = \
+ self.particle_file["/Halo%08i/particle_index" % halo_id][:]
ids = self.cache[halo_id]
else:
- ids = self.particle_file["/Halo%08i/Particle ID" % halo_id][:]
+ if self.external_FOF:
+ ids = self.particle_file["/Halo%08i/Particle ID" % halo_id][:]
+ else:
+ ids = self.particle_file["/Halo%08i/particle_index" % halo_id][:]
return HaloParticleList(halo_id, self.halo_positions[halo_id,:], ids)
def calculate_parentage_fractions(self, other_catalog, radius = 0.10):
@@ -197,43 +253,89 @@
class EnzoFOFMergerTree(object):
r"""Calculates the parentage relationships for halos for a series of
outputs, using the framework provided in enzofof_merger_tree.
- """
+ Parameters
+ ----------
+ zrange : tuple
+ This is the redshift range (min, max) to calculate the
+ 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,
+ ignore zrange.
+ output : bool, optional
+ If provided, both .cpkl and .txt files containing the parentage
+ relationships will be output.
+ load_saved : bool, optional
+ Flag to load previously saved parental relationships
+ save_filename : str, optional
+ Filename to save parental relationships
+ external_FOF : bool, optional
+ Are we building a tree from outputs generated by an
+ external FOF program, or an FOF internal to yt?
+ FOF_directory : str, optional
+ Directory where FOF files are located, note that the files
+ must be named according to the syntax: groups_DDDDD.txt for
+ internal yt outputs, and groups_DDDDD.dat for external FOF outputs.
+ where DDDDD are digits representing the equivalent cycle number.
+ e.g. groups_00000.txt
+
+ Examples
+ --------
+ mt = EnzoFOFMergerTree() # by default it grabs every DD in FOF dir
+ mt.build_tree(0) # Create tree for halo 0
+ mt.print_tree()
+ mt.write_dot()
+
+ See Also
+ --------
+ plot_halo_evolution()
+ """
def __init__(self, zrange=None, cycle_range=None, output=False,
- load_saved=False, save_filename="merger_tree.cpkl"):
- r"""
- Parameters
- ----------
- zrange : tuple
- This is the redshift range (min, max) to calculate the
- merger tree.
- cycle_range : tuple, optional
- This is the cycle number range (min, max) to caluclate the
- merger tree. If both zrange and cycle_number given,
- ignore zrange.
- output : bool, optional
- If provided, both .cpkl and .txt files containing the parentage
- relationships will be output.
- load_saved : bool, optional
- Flag to load previously saved parental relationships
- save_filename : str, optional
- Filename to save parental relationships
-
- Examples
- --------
- mt = EnzoFOFMergerTree((0.0, 6.0))
- mt.build_tree(0) # Create tree for halo 0
- mt.print_tree()
- mt.write_dot()
- """
+ load_saved=False, save_filename="merger_tree.cpkl",
+ external_FOF=True, FOF_directory="FOF"):
+
self.relationships = {}
self.redshifts = {}
- self.find_outputs(zrange, cycle_range, output)
+ self.external_FOF = external_FOF
+ self.FOF_directory = FOF_directory
if load_saved:
- self.load_tree(save_filename)
+ self.load_tree("%s/%s" % (self.FOF_directory, 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:
+ self.select_cycles(cycle_range)
else:
+ self.find_outputs(zrange, cycle_range, output)
self.run_merger_tree(output)
- self.save_tree(save_filename)
+ self.save_tree("%s/%s" % (self.FOF_directory, 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),
@@ -251,32 +353,44 @@
def find_outputs(self, zrange, cycle_range, output):
self.numbers = []
- files = glob.glob("FOF/groups_*.dat")
+ if self.external_FOF:
+ filenames = "%s/groups_*.dat" % (self.FOF_directory)
+ files = glob.glob(filenames)
+ else:
+ filenames = "%s/groups_*.txt" % (self.FOF_directory)
+ files = glob.glob(filenames)
# If using redshift range, load redshifts only
for f in files:
num = int(f[-9:-4])
- if cycle_range == None:
- HC = HaloCatalog(num)
+ if zrange is not None:
+ HC = HaloCatalog(num, external_FOF=self.external_FOF, \
+ FOF_directory=self.FOF_directory)
# Allow for some epsilon
diff1 = (HC.redshift - zrange[0]) / zrange[0]
diff2 = (HC.redshift - zrange[1]) / zrange[1]
if diff1 >= -1e-3 and diff2 <= 1e-3:
self.numbers.append(num)
del HC
- else:
+ elif cycle_range is not None:
if num >= cycle_range[0] and num <= cycle_range[1]:
self.numbers.append(num)
+ else:
+ self.numbers.append(num)
self.numbers.sort()
def run_merger_tree(self, output):
# Run merger tree for all outputs, starting with the last output
for i in range(len(self.numbers)-1, 0, -1):
if output:
- output = "tree-%5.5d-%5.5d" % (self.numbers[i], self.numbers[i-1])
+ output = "%s/tree-%5.5d-%5.5d" % \
+ (self.FOF_directory, self.numbers[i], self.numbers[i-1])
else:
output = None
- z0, z1, fr = find_halo_relationships(self.numbers[i], self.numbers[i-1],
- output_basename=output)
+ z0, z1, fr = find_halo_relationships(self.numbers[i], \
+ self.numbers[i-1], \
+ output_basename=output, \
+ external_FOF=self.external_FOF,
+ FOF_directory=self.FOF_directory)
self.relationships[self.numbers[i]] = fr
self.redshifts[self.numbers[i]] = z0
# Fill in last redshift
@@ -370,6 +484,7 @@
r"""Prints the merger tree to stdout.
"""
for lvl in sorted(self.levels, reverse=True):
+ if lvl not in self.redshifts: continue
print "========== Cycle %5.5d (z=%f) ==========" % \
(lvl, self.redshifts[lvl])
for br in self.levels[lvl]:
@@ -380,6 +495,79 @@
if i > self.max_children: break
print "--> Halo %8.8d :: fraction = %g" % (c[0], c[1])
+ def save_halo_evolution(self, filename):
+ """
+ Saves as an HDF5 file the relevant details about a halo
+ over the course of its evolution following the most massive
+ progenitor to have given it the bulk of its particles.
+ It stores info from the FOF_groups file: location, mass, id, etc.
+ """
+ f = h5py.File("%s/%s" % (self.FOF_directory, filename), 'a')
+ cycle_fin = self.redshifts.keys()[-1]
+ halo_id = self.levels[cycle_fin][0].halo_id
+ halo = "halo%05d" % halo_id
+ if halo in f:
+ del f["halo%05d" % halo_id]
+ g = f.create_group("halo%05d" % halo_id)
+ size = len(self.redshifts)
+ cycle = np.zeros(size)
+ redshift = np.zeros(size)
+ halo_id = np.zeros(size)
+ fraction = np.zeros(size)
+ mass = np.zeros(size)
+ densest_point = np.zeros((3,size))
+ COM = np.zeros((6,size))
+ fraction[0] = 1.
+
+ for i, lvl in enumerate(sorted(self.levels, reverse=True)):
+ if len(self.levels[lvl]) == 0: # lineage for this halo ends
+ cycle = cycle[:i] # so truncate arrays, and break
+ redshift = redshift[:i] # Not big enough.
+ halo_id = halo_id[:i]
+ fraction = fraction[:i]
+ mass = mass[:i]
+ densest_point = densest_point[:,:i]
+ COM = COM[:,:i]
+ break
+ if lvl not in self.redshifts: continue
+ mylog.info("========== Cycle %5.5d (z=%f) ==========" % \
+ (lvl, self.redshifts[lvl]))
+ cycle[i] = lvl
+ redshift[i] = self.redshifts[lvl]
+
+ br = self.levels[lvl][0]
+ mylog.info("Parent halo = %d" % br.halo_id)
+ mylog.info("-> Most massive progenitor == Halo %d" % (br.progenitor))
+ halo_id[i] = br.halo_id
+
+ if len(br.children) == 0: # lineage for this halo ends
+ cycle = cycle[:i+1] # (no children)
+ redshift = redshift[:i+1] # so truncate arrays, and break
+ halo_id = halo_id[:i+1]
+ fraction = fraction[:i+1]
+ mass = mass[:i+1]
+ densest_point = densest_point[:,:i+1]
+ COM = COM[:,:i+1]
+ break
+
+ if i < size-1:
+ fraction[i+1] = br.children[0][1]
+
+ # open up FOF file to parse for details
+ filename = "%s/groups_%05d.txt" % (self.FOF_directory, lvl)
+ mass[i], densest_point[:,i], COM[:,i] = \
+ grab_FOF_halo_info_internal(filename, br.halo_id)
+
+ # save the arrays in the hdf5 file
+ g.create_dataset("cycle", data=cycle)
+ g.create_dataset("redshift", data=redshift)
+ g.create_dataset("halo_id", data=halo_id)
+ g.create_dataset("fraction", data=fraction)
+ g.create_dataset("mass", data=mass)
+ g.create_dataset("densest_point", data=densest_point)
+ g.create_dataset("COM", data=COM)
+ f.close()
+
def write_dot(self, filename=None):
r"""Writes merger tree to a GraphViz or image file.
@@ -393,7 +581,9 @@
automatically. See GraphViz (e.g. "dot -v")
for a list of available output formats.
"""
- if filename == None: filename = "tree_halo%5.5d.gv" % self.halonum
+ if filename == None:
+ filename = "%s/tree_halo%5.5d.gv" % \
+ (self.FOF_directory, self.halonum)
# Create the pydot graph object.
self.graph = pydot.Dot('galaxy', graph_type='digraph')
self.halo_shape = "rect"
@@ -446,7 +636,8 @@
self.graph.write("%s" % filename, format=suffix)
def find_halo_relationships(output1_id, output2_id, output_basename = None,
- radius = 0.10):
+ radius = 0.10, external_FOF=True,
+ FOF_directory='FOF'):
r"""Calculate the parentage and child relationships between two EnzoFOF
halo catalogs.
@@ -456,8 +647,9 @@
particle IDs from the child halos are identified in potential parents, and
then both percent-of-parent and percent-to-child values are recorded.
- Note that this works only with catalogs constructed by Enzo's FOF halo
- finder, not with catalogs constructed by yt.
+ Note that this works with catalogs constructed by Enzo's FOF halo
+ when used in external_FOF=True mode, whereas it will work with
+ catalogs constructed by yt using external_FOF=False mode.
Parameters
----------
@@ -474,6 +666,8 @@
In absolute units, the radius to examine when guessing possible
parent/child relationships. If this value is too small, you will miss
possible relationships.
+ FOF_directory : str, optional
+ Directory where FOF files are located
Returns
-------
@@ -484,9 +678,11 @@
parent.
"""
mylog.info("Parsing Halo Catalog %04i", output1_id)
- HC1 = HaloCatalog(output1_id, False)
+ HC1 = HaloCatalog(output1_id, False, external_FOF=external_FOF, \
+ FOF_directory=FOF_directory)
mylog.info("Parsing Halo Catalog %04i", output2_id)
- HC2 = HaloCatalog(output2_id, True)
+ HC2 = HaloCatalog(output2_id, True, external_FOF=external_FOF, \
+ FOF_directory=FOF_directory)
mylog.info("Calculating fractions")
pfrac = HC1.calculate_parentage_fractions(HC2)
@@ -504,3 +700,102 @@
cPickle.dump(pfrac, open("%s.cpkl" % (output_basename), "wb"))
return HC1.redshift, HC2.redshift, pfrac
+
+def grab_FOF_halo_info_internal(filename, halo_id):
+ """
+ Finds a specific halo's information in the FOF group output information
+ and pass relevant parameters to caller.
+ """
+ # open up FOF file to parse for details
+ groups_file = open(filename, 'r')
+ for line in groups_file:
+ if line.startswith("#"): continue
+ if int(line.split()[0]) == halo_id:
+ ar = np.array(line.split()).astype('float64')
+ return ar[1], ar[4:7], ar[7:13] # mass, xyz_dens, xyzvxvyvz_COM
+
+def plot_halo_evolution(filename, halo_id, x_quantity='cycle', y_quantity='mass',
+ x_log=False, y_log=True, FOF_directory='FOF'):
+ """
+ Once you have generated a file using the
+ EnzoFOFMergerTree.save_halo_evolution function, this is a simple way of
+ plotting the evolution in the quantities of that halo over its lifetime.
+
+ Parameters
+ ----------
+ filename : str
+ The filename to which you saved the hdf5 data from save_halo_evolution
+ halo_id : int
+ The halo in 'filename' that you want to follow
+ x_quantity, y_quantity : str, optional
+ The quantity that you want to plot as the x_coord (or y_coords).
+ Valid options are:
+ cycle, mass, fraction, halo_id, redshift, dense_x, dense_y, dense_z
+ COM_x, COM_y, COM_z, COM_vx, COM_vy, COM_vz
+ x_log, y_log : bool, optional
+ Do you want the x(y)-axis to be in log or linear?
+ FOF_directory : str, optional
+ Directory where FOF files (and hdf file) are located
+
+ Examples
+ --------
+ # generates mass history plots for the 20 most massive halos at t_fin.
+
+ ts = TimeSeriesData.from_filenames("DD????/DD????")
+
+ # long step--must run FOF on each DD, but saves outputs for later use
+ for pf in ts:
+ halo_list = FOFHaloFinder(pf)
+ i = int(pf.basename[2:])
+ halo_list.write_out("FOF/groups_%05i.txt" % i)
+ halo_list.write_particle_lists("FOF/particles_%05i" % i)
+
+ mt = EnzoFOFMergerTree(external_FOF=False)
+ for i in range(20):
+ mt.build_tree(i)
+ mt.save_halo_evolution('halos.h5')
+ for i in range(20):
+ plot_halo_evolution('halos.h5', i)
+ """
+ import matplotlib.pyplot as plt
+ f = h5py.File("%s/%s" % (FOF_directory, filename), 'r')
+ basename = os.path.splitext(filename)[0]
+ halo = "halo%05d" % halo_id
+ basename = basename + "_" + halo
+ g = f[halo]
+ values = list(g)
+ index_dict = {'x' : 0, 'y' : 1, 'z' : 2, 'vx' : 3, 'vy' : 4, 'vz' : 5}
+ coords = {}
+ fields = {}
+ for i, quantity in enumerate((x_quantity, y_quantity)):
+ field = quantity
+ if quantity.startswith('COM'):
+ index = index_dict[quantity.split('_')[-1]]
+ quantity = ('COM')
+ if quantity.startswith('dense'):
+ index = index_dict[quantity.split('_')[-1]]
+ quantity = ('densest_point')
+ if quantity not in values:
+ exit('%s not in list of values in %s for halo %d' % \
+ (quantity, filename, halo_id))
+ if not field == quantity:
+ coords[i] = g[quantity][index,:]
+ else:
+ coords[i] = g[quantity]
+ if len(coords[i]) == 1:
+ # ("Only 1 value for Halo %d. Ignoring." % halo_id)
+ return
+ fields[i] = field
+
+ ax = plt.axes()
+ ax.plot(coords[0], coords[1])
+ ax.set_title(basename)
+ ax.set_xlabel(fields[0])
+ ax.set_ylabel(fields[1])
+ if x_log:
+ ax.set_xscale("log")
+ if y_log:
+ ax.set_yscale("log")
+ ofn = "%s_%s_%s.png" % (basename, fields[0], fields[1])
+ plt.savefig(ofn)
+ plt.clf()
diff -r 9b336aa05056cb4f6aab9066dd24c3c188695f2f -r 6b246ba83cb38d6fd2467bf76e265e6ca4f78beb yt/funcs.py
--- a/yt/funcs.py
+++ b/yt/funcs.py
@@ -602,3 +602,18 @@
def get_image_suffix(name):
suffix = os.path.splitext(name)[1]
return suffix if suffix in ['.png', '.eps', '.ps', '.pdf'] else ''
+
+def mkdir_rec(path):
+ """
+ Recursive mkdir, so that if you mkdir two levels deep and the first
+ one doesn't exist, it creates the first, and then any subsequent dirs.
+
+ Examples
+ --------
+ mkdir_rec("a/b/c")
+ """
+ dir_list = path.split("/")
+ basedir = "."
+ for dir in dir_list:
+ basedir = "%s/%s" % (basedir, dir)
+ if not os.path.isdir(basedir): os.mkdir(basedir)
Repository URL: https://bitbucket.org/yt_analysis/yt/
--
This is a commit notification from bitbucket.org. You are receiving
this because you have the service enabled, addressing the recipient of
this email.
More information about the yt-svn
mailing list