\documentclass[useAMS,usenatbib,usegraphicx,onecolumn]{mn2e}
\usepackage{amsmath,amsfonts,amssymb}
\usepackage[latin1]{inputenc}
\usepackage{times} 
\usepackage{color}     
\usepackage{graphicx}

\usepackage[normalem]{ulem} 
\usepackage{epstopdf}
\usepackage{epsfig}

\newcommand{\hdf}{HDF5}

\title[Merger Tree Data Format]
{Sussing Merger Trees: A proposed Merger Tree data format}
\author[Sussing Merger Trees Workshop participants]
{Peter~A.~Thomas et al.} 

\setcounter{page}{6}
\setcounter{figure}{3}
\begin{document}

\appendix

\section{Example {\sc Python} codes}
\label{sec:python}

In this appendix we present some example {\sc Python} codes to search merger
trees and to read/write them in \hdf\ format.  These routines can all be
downloaded from {\tt https://bitbucket.org/ProfPAThomas/mergertree/}.

The examples use the tree shown in Figure~\ref{fig:testtree} as recorded in HDF5 format
in test.hdf5.

\begin{figure}
  \centering
  \includegraphics[angle=0,width=1\linewidth]{TestTree.png}
  \caption{A test merger tree that is used for the examples in this appendix.}
\label{fig:testtree}
\end{figure}

\subsection{To read a tree written in the default format}
\label{sec:readtree}
We first present an example of how to read in data from an HDF5 file and to dump
some properties to the screen.  To simply the script, this example makes use of
mt.py that defines a MergerTree class and some helper functions, but this is not
really necessary and could easily be omitted.

\begin{verbatim}
"""readtree.py
   Example of how to read in a merger tree in HDF5 format
   Usage: python readtree.py <tree.hdf>
"""

#-------------------------------------------------------------------------------
# Imports

# The following imports the python3 print function into python2;
# It seems to be quietly ignored in python3
from __future__ import print_function

import sys

# The merger tree module contains data structures, methods and global variables
import mt

#-------------------------------------------------------------------------------
# Open file
if len(sys.argv) != 2: sys.exit("Usage: " + sys.argv[0] + " filename.hd5")
mtdata=mt.MergerTree(sys.argv[1], "r")
for (name,value) in mtdata.attrs('/'): print(name,value)

# Read in snap table
for (name,value) in mtdata.attrs("Snapshots"): print(name,value)
snap=mtdata.read_snaptable()
print('Snap =')
print('{0:>8} {1:>6} {2:>6} {3:>6}'.format(*snap.dtype.names))
for row in range(len(snap)): print('{0:>8} {1:>6.3f} {2:>6.3f} {3:>6.3f}'.format(*snap[row]))
del snap

# Read merger tree from HDF file
for (name,value) in mtdata.attrs("MergerTree"): print(name,value)
tree=mtdata.read_mergertree()
print('Tree =')
for dset in tree: 
    print(dset,end=' ')
print()
for ntree in range(min(10,len(tree['HaloID']))): 
    for prop in tree: 
        print(tree[prop][ntree],end=' ')
    print()
del tree

#Close file
mtdata.close()

#-------------------------------------------------------------------------------

\end{verbatim}

\subsection{Merger tree class and associated methods}
\label{sec:mt}
Here is file mt.py containing the MergerTree class used in the above example.

\begin{verbatim}
"""mt.py
   Merger tree class to simplfy some i/o functions for merger trees.
   It's not entirely obvious that a separate class is needed 
   - could just put these calls directly into the main program.
"""

#-------------------------------------------------------------------------------
import h5py

#-------------------------------------------------------------------------------
class MergerTree:
    
    def __init__(self, file=None, mode='r'):
        """Initialise the class, opening the file with the mode given"""
        if file: 
            self.open(file, mode)
        else:
            print('mt usage: mt(file, mode=mode)')

    # Return set of attributes associated with object
    def attrs(self, object):
        return self.fid[object].attrs.items()

    # Close the HDF5 file
    def close(self):
        self.fid.close()

    # Open the HDF5 file
    def open(self, file, mode):
        self.fid = h5py.File(file, mode)

    # Return a handle to the MergerTree group
    def read_mergertree(self):
        return self.fid["MergerTree"]

    # Return a handle to the Snapshot table
    def read_snaptable(self):
        return self.fid["Snapshots/Snap"]

    # Return a handle to the SnapProp table
    def read_snapprop(self):
        return self.fid["Snapshots/SnapProp"]

#-------------------------------------------------------------------------------
\end{verbatim}

\subsection{Tree-searching}
\label{sec:tracetree}
The following example tracetree.py reads in an HDF5 file then shows how to loop over
(or 'trace') a tree using a spatial and/or temporal search.  It makes use of a
Tree class defined in tree.py

\begin{verbatim}
"""tracetree.py
   Example of how to read in a merger tree in HDF5 format and to use a
   specially-defined Tree class to allow spatial and/or temporal searching
   of the tree.
   Usage: python tracetree.py
"""

#-------------------------------------------------------------------------------
# Imports

# The following imports the python3 print function into python2;
# It seems to be quietly ignored in python3
from __future__ import print_function

import sys
import numpy
import h5py

# The tree module contains a class definition for a combined spatial and temporal tree
import tree

#-------------------------------------------------------------------------------
# Parameters

# The following data set contains a simple example with both a spatial and merger trees.
infile='data/test.hdf5'
# We are going to open read only
mode='r'

#-------------------------------------------------------------------------------
# Open Merger tree file for reading
fid = h5py.File(infile, mode)

# This shows how easy it is to extract attributes from the file
for (name,value) in fid['/'].attrs.items(): print(name,value)

# Next we are going to construct a tree from the information contained in the HDF5 file

# First let's extract the number of halos
gid = fid['/MergerTree']
nHalo=gid.attrs.get('NHalo')
print('nHalo =',nHalo)

# and now the links that we want.  Without the [:] we get a pointer to the array.
firstSubhalo=gid.get('FirstSubhaloIndex')[:]
neighbour=gid.get('NeighbourIndex')[:]
firstProgenitor=gid.get('FirstProgenitorIndex')[:]
nextSibling=gid.get('NextSiblingIndex')[:]
hostHalo=gid.get('HostHaloIndex')[:]
descendant=gid.get('DescendantIndex')[:]
    
# Initialise tree nodes...
# Trick: create an extra non-existent node to receive references to index -1 (= last index)
nodes=[tree.Tree(halo) for halo in range(nHalo+1)]
nodes[nHalo]=None

# Now we run through each of the halos extracting those pointers that we need for
# efficient tree traversal
for halo in range(nHalo): 
    nodes[halo].load(nodes[firstSubhalo[halo]],
                     nodes[neighbour[halo]],
                     nodes[firstProgenitor[halo]],
                     nodes[nextSibling[halo]],
                     nodes[hostHalo[halo]],
                     nodes[descendant[halo]])

# The end main halos are those that have no Descendants and no HostHalo
# In this simple example there is only one such halo
endMainHalos=numpy.where((descendant==-1) & (hostHalo==-1))[0]
print(endMainHalos)

# Let's trace all the halos in the combined spatial and temporal tree
nodes[0].printheader()
for halo in nodes[0].next():
    print(nodes[halo])
    
# Here's a trace of just the spatial tree
nodes[0].printheader()
for halo in nodes[0].next('spatial'):
    print(nodes[halo])

# and here are the temporal trees for each end halo
endHalos=numpy.where(descendant==-1)[0]
for endHalo in endHalos:
    nodes[endHalo].printheader()
    for halo in nodes[endHalo].next('temporal'):
        print(nodes[halo])

#-------------------------------------------------------------------------------
\end{verbatim}

\subsection{Tree-search algorithm}
\label{sec:tree}
We present below an example tree class with methods for generating iterators
that can perform a combined (default behaviour), spatial, or temporal tree
search.
\begin{verbatim}
"""tree.py
   Defines a Tree class for dealing with combined temporal and merger trees
"""

#-------------------------------------------------------------------------------
# The following imports the python3 print function into python2;
# It seesm to be quietly ignored in python3
from __future__ import print_function

# In Python 3 the loops of the form
#    for __ in self.something():
#       yield __
# can all be replaced by
#    yield from self.something()
# which makes for much cleaner code.

#-------------------------------------------------------------------------------
class Tree:
    """Tree class for dealing with combined temporal and merger trees"""

    # Constructor: creates an empty tree node.
    def __init__(self,node):
        self.node=node

    # Load pointer arrays.
    def load(self,firstSubhalo,neighbour,firstProgenitor,nextSibling,hostHalo=-1,descendant=-1):
        self.firstSubhalo=firstSubhalo
        self.neighbour=neighbour
        self.firstProgenitor=firstProgenitor
        self.nextSibling=nextSibling
        self.hostHalo=hostHalo
        self.descendant=descendant

    # Default iterator that takes no arguments:
    # Note the asymmetry between spatial and temporal: because there is more
    # than one way to reach a node, can only iterate over one of them.
    # Because of the nature of the trees, space has to be the outer loop.
    def __next__(self):
        yield self.node
        for __ in self.nextSpatial():
            yield __
        for __ in self.nextTemporal('temporal'):
            yield __
    # More general iterator method with optional argument specifying tree type:
    def next(self,treeType=None):
        yield self.node
        if treeType is None:
            for __ in self.nextSpatial():
                yield __
            for __ in self.nextTemporal('temporal'):
                yield __
        elif treeType=='spatial':
            for __ in self.nextSpatial('spatial'):
                yield __
        elif treeType=='temporal':
            for __ in self.nextTemporal('temporal'):
                yield __
        else:
            raise valueError('Invalid treeType for tree')
    # Spatial tree search
    def nextSpatial(self,treeType=None):
        if self.firstSubhalo is not None:
            for __ in self.firstSubhalo.next(treeType):
                yield __
        if self.neighbour is not None:
            for __ in self.neighbour.next(treeType):
                yield __
        #raise StopIteration()
    # Temporal tree search
    def nextTemporal(self,treeType=None):
        if self.firstProgenitor is not None:
            for __ in self.firstProgenitor.next(treeType):
                yield __
        if self.nextSibling is not None:
            for __ in self.nextSibling.next(treeType):
                yield __
        #raise StopIteration()

    # Printing
    def printheader(self):
        print ('node,firstSubhalo,neighbour,firstProgenitor,nextSibling,hostHalo,descendant=')
    def __str__(self):
        s1='-1'
        if self.firstSubhalo: s1=str(self.firstSubhalo.node)
        s2='-1'
        if self.neighbour: s2=str(self.neighbour.node)
        s3='-1'
        if self.firstProgenitor: s3=str(self.firstProgenitor.node)
        s4='-1'
        if self.nextSibling: s4=str(self.nextSibling.node) 
        s5='-1'
        if self.hostHalo: s5=str(self.hostHalo.node)
        s6='-1'
        if self.descendant: s6=str(self.descendant.node)
        return  str(self.node)+', '+s1+', '+s2+', '+s3+', '+s4+', '+s5+', '+s6

#-------------------------------------------------------------------------------
\end{verbatim}

\subsection{To extract a tree from an HDF5 file}
\label{sec:extracttree}
This example, extracttree.py, reads in an HDF5 file, then extracts some data for
a particular tree into a structured numpy array.
\begin{verbatim}
"""extracttree.py
   Example of how to read in a merger tree file in HDF5 format
   and to create a table containing all the data for a particular tree.
"""

# Note: only if we pick up whole trees will the indices be guaranteed
# to be consecutive.  As we only have a single tree in our test data set,
# this example takes halo 7 to be the root halo as that does have
# consecutive indices for its subtree.

# If the indices are consecutive then there is no need to use a special
# Tree class and methods to allow tree searching.  However, we do first 
# do this as a check.

#-------------------------------------------------------------------------------
# Imports

# The following imports the python3 print function into python2;
# It seesm to be quietly ignored in python3
from __future__ import print_function

import sys
import numpy
import h5py

# The tree module contains a class definition for a combined spatial and temporal tree
import tree

#-------------------------------------------------------------------------------
# Parameters

# The following data set contains a simple example with both a spatial and merger trees.
infile='data/test.hdf5'
# We are going to open read only
mode='r'

#-------------------------------------------------------------------------------
# Open Merger tree file for reading
fid = h5py.File(infile, mode)

# This shows how easy it is to extract attributes from the file
for (name,value) in fid['/'].attrs.items(): print(name,value)

# Let's extract the data that we want from the HDF5 file.

# Define numpy data type for the table rows
# Maybe there is a clever way to do this, but this example will just do it by hand
treedata_dtype=numpy.dtype([
        ('haloID','int64'),
        ('descendant','int32'),
        ('firstProgenitor','int32'),
        ('firstSubhalo','int32'),
        ('hostHalo','int32'),
        ('neighbour','int32'),
        ('nextSibling','int32'),
        ('snapNum','int32'),
        ('density','f8'),
        ('mass','f8')
        ])

# First let's extract the number of halos
gid = fid['/MergerTree']
nHalo=gid.attrs.get('NHalo')
print('nHalo =',nHalo)

# and now the link that we want.  Without the [:] we get a pointer to the array.
haloID=gid.get('HaloID')[:]
descendant=gid.get('DescendantIndex')[:]
firstProgenitor=gid.get('FirstProgenitorIndex')[:]
firstSubhalo=gid.get('FirstSubhaloIndex')[:]
hostHalo=gid.get('HostHaloIndex')[:]
neighbour=gid.get('NeighbourIndex')[:]
nextSibling=gid.get('NextSiblingIndex')[:]
snapNum=gid.get('Snapshot')[:]
density=gid.get('Density')[:]
mass=gid.get('Mass')[:]

# Next we are going to construct a tree from the information contained in the HDF5 file
# This is only needed for trees that are not depth-first ordered, but it's a nice example
# of how to use iterators in python anyway.  It uses the Tree class in the separate tree.py file.

# Initialise tree nodes...
# Trick: create an extra non-existent node to receive references to index -1 (= last index)
nodes=[tree.Tree(halo) for halo in range(nHalo+1)]
nodes[nHalo]=None

# Now we run through each of the halos extracting those pointers that we need for
# efficient tree traversal
for halo in range(nHalo): 
    nodes[halo].load(nodes[firstSubhalo[halo]],
                     nodes[neighbour[halo]],
                     nodes[firstProgenitor[halo]],
                     nodes[nextSibling[halo]],
                     nodes[hostHalo[halo]],
                     nodes[descendant[halo]])

# Let's pick a halo to act as the root of our tree.
# Ideally, we should use a full tree, but we only have one tree in this example.
# However halo 7 has consecutive indices in its subtree, so we'll use that.
root=7
nodes[root].printheader()
nTreeHalo=0
nodeMin=nodes[root].node
nodeMax=-1
for halo in nodes[root].next():
    nTreeHalo+=1
    nodeMin=min(halo,nodeMin)
    nodeMax=max(halo,nodeMax)
    print(nodes[halo])
print('nTreeHalo =',nTreeHalo)
print('nodeMin, nodeMax =',nodeMin,nodeMax)

# Now let's loop over the tree again, recording the data that we want in our table
# Create a numpty structured array to hold the table
treedata=numpy.empty(nTreeHalo,dtype=treedata_dtype)
ihalo=-1
consecutive=True
firstHalo=nodes[root].node
for halo in nodes[root].next():
    ihalo+=1
    if halo != firstHalo+ihalo: consecutive=False
    treedata[ihalo]=(haloID[halo],descendant[halo],firstProgenitor[halo],firstSubhalo[halo],
                     hostHalo[halo],neighbour[halo],nextSibling[halo],
                     snapNum[halo],density[halo],mass[halo])
print('Consecutive halos = ',consecutive)

# If we have a consecutive list of halos, adjust indices to point to the relevant location
# in our new table.
# Setting the maximum to -1 is not really necessary but helps to make things look tidier.
if consecutive:
    numpy.maximum(treedata['descendant']-nodeMin,-1,treedata['descendant']); 
    numpy.maximum(treedata['firstProgenitor']-nodeMin,-1,treedata['firstProgenitor']); 
    numpy.maximum(treedata['firstSubhalo']-nodeMin,-1,treedata['firstSubhalo']); 
    numpy.maximum(treedata['hostHalo']-nodeMin,-1,treedata['hostHalo']); 
    numpy.maximum(treedata['neighbour']-nodeMin,-1,treedata['neighbour']); 
    numpy.maximum(treedata['nextSibling']-nodeMin,-1,treedata['nextSibling']); 

# We have established that we have a consecutive (depth-first ordered) tree,
# so no need to bother with the Tree class again.  Let's just dump the new
# properties in consecutive order as a check.
for ihalo in range(nTreeHalo):
    print(ihalo,treedata[['firstSubhalo','neighbour','firstProgenitor',
                         'nextSibling','hostHalo','descendant']][ihalo])

# Or rather more simply, but in a different order
print(treedata)

#-------------------------------------------------------------------------------
\end{verbatim}

\subsection{To save a tree in the default format}
Finally, we have an example of how to create an HDF5 file in the correct
format.  Because we need some data to write out, this example also reads in data
from an existing HDF5 file: in practice the data may have come from a mergertree
creation algorithm.

\begin{verbatim}
"""dumptree.py
   Example of how to write a merger tree in HDF5 format.
   This example needs some data which it reads in from an existing HDF file,
   which makes this whole exercise a bit pointless, but pretend that the data
   was created some other way and now needs to be written out.
   This is not complete, but it gives the general idea.
"""

#-------------------------------------------------------------------------------
# Imports

# The following imports the python3 print function into python2;
# It seesm to be quietly ignored in python3
from __future__ import print_function

import sys
import numpy as np

# The merger tree module contains data structures, methods and global variables
import mt

#-------------------------------------------------------------------------------
# Parameters

infile='data/test.hdf5'
outfile='data/dumptest_Python.hdf5'

# For the purposes of this example, we first have to read in some data.
# Actually, the following creates pointers to the data, to be read below.

mt_in=mt.MergerTree(infile, "r")
attributes = mt_in.attrs('/')

# Read in snap tables
snap_in=mt_in.read_snaptable()
snapprop_in=mt_in.read_snapprop()
snap_attributes=mt_in.attrs("Snapshots")

# Read merger tree.
tree_in=mt_in.read_mergertree()
tree_attributes=mt_in.attrs("MergerTree")

# Now we are going to create a new HDF5 file to write out the data
print('Creating output file')
mt_out=mt.MergerTree(outfile,"w")
print('File =',mt_out.fid)

# Write out the top-level attributes
print('In group',mt_out.fid.name,'attributes are:')
for (name,value) in attributes: 
    print('  ',name,value)
    mt_out.fid.attrs.modify(name,value)

# Create MergerTree directory and set attributes
tree=mt_out.fid.create_group("MergerTree")
print('In group',tree.name,'attributes are:')
for (name,value) in tree_attributes: 
    print('name, value =',name,value)
    tree.attrs.modify(name,value)

# Write out merger tree datasets
print('MergerTree datasets:')
for dset in tree_in:
    # extract name, dtype and data from input dataset
    name=tree_in[dset].name
    dtype=tree_in[dset].dtype
    data=tree_in[dset][()]
    attrs=tree_in[dset].attrs.items()
    # and now recreate in the output file
    dset=tree.create_dataset(name=name,dtype=dtype,data=data)
    print('  ',name)
    print('    ',dset[0:9])
    # and set attributes
    for (name,value) in attrs: 
        print('    ',name,value)
        dset.attrs.modify(name,value)

# Create Snapshot directory and set attributes
snap=mt_out.fid.create_group("Snapshots")
print('In group',snap.name,'attributes are:')
for (name,value) in snap_attributes: 
    print('  ',name,value)
    snap.attrs.modify(name,value)

# Write out Snap table
# Tables require use of a complex dtype.  Fortuntately, that's already set up
# us here in the input file.
# Extract data and dtype from input file
name=snap_in.name
dtype=snap_in.dtype
nsnap=len(snap_in)
data=np.empty(nsnap,dtype=dtype)
for irow in range(nsnap): data[irow]=snap_in[irow]
# and recreate in output file
dset=tree.create_dataset(name=name,dtype=dtype,data=data)
print('Snap table:')
print('  name =',dset.name)
print('  dtype =',dset.dtype)
print('  data =',dset[:])

# Write out SnapProp table
# us here in the input file.
# Extract data and dtype from input file
name=snapprop_in.name
dtype=snapprop_in.dtype
nsnapprop=len(snapprop_in)
data=np.empty(nsnapprop,dtype=dtype)
for irow in range(nsnapprop): data[irow]=snapprop_in[irow]
# Rename the Units property to lower case (for consistency).
# Cannot rename in place: have to make a copy of both dtype and data.
temp=dtype.__repr__()[6:-1]
temp=temp.replace("Units","units")
dtypenew=np.dtype(eval(temp))
print('dtypenew=',dtypenew)
datanew=np.empty(nsnapprop,dtype=dtypenew)
for irow in range(nsnapprop): datanew[irow]=data[irow]
# and recreate in output file
dset=tree.create_dataset(name=name,dtype=dtypenew,data=datanew)
print('SnapProp table:')
print('  name =',dset.name)
print('  dtype =',dset.dtype)
print('  data =',dset[:])

# The output file is not yet a complete examle of the merger tree data
# format, but hopefuly that is enough to illustrate how do do things.

# Close the files
mt_in.close()
mt_out.close()

#-------------------------------------------------------------------------------
\end{verbatim}


\end{document}
