csiborgtools/scripts/mergertree_extract.py

# Copyright (C) 2023 Mladen Ivkovic, Richard Stiskalek
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 3 of the License, or (at your
# option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
# Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

import copy
import os
from os.path import exists, join
from os import makedirs
from sys import argv
from datetime import datetime

import numpy as np
from joblib import dump, load
from tqdm import trange

errmsg = """

------------------------------------
    mergertree-extract.py
------------------------------------


---------------
    Usage
---------------

This script extracts the masses of clumps and haloes written by the mergertree
patch.
It needs output_XXXXX/mergertree_XXXXX.txtYYYYY and
output_XXXXX/clump_XXXXX.txtYYYYY files to work.
You need to run it from the directory where the output_XXXXX directories are
in.


There are three working modes defined:

1) do for one clump only.
    You need to provide the clump ID you want it done for.
    You can provide a starting directory, but by default the script will
    search for the directory where z = 0.

    run with `python3 mergertree-extract.py <clumpid> [--options] `

    this creates the file mergertree_XXXXX_halo-<halo-ID>.txt. Its contents are
    discussed below.


2) do for one halo.
    You need to provide the halo ID you want it done for, and the flag
    -c or --children.
    The script will by itself find all the child clumps and walk through
    their main branches as well, and write them down.

    run with `python3 mergertree-extract.py <haloid> -c [--options]`
          or `python3 mergertree-extract.py <haloid> --children [--options]`

    this creates the hollowing files:

        - halo_hierarchy_XXXXX-<halo-ID>.txt
            contains the halo ID, how many children it has, and the children
            IDs

        - mergertree_XXXXX_halo-<halo-ID>.txt
            mergertree data for halo that you chose.

        - mergertree_XXXXX_subhalo-<child-ID>.txt
            mergertree data for subhalos of the halo you chose.  One file will
            be created for each subhalo.

        The contents of the mergertree_XXXXX* files are discussed below.


3) do for all haloes
    The script will just walk off all haloes in the z = 0 directory. Note:
    Haloes, not clumps!
    run with `python3 mergertree-extract.py -a [--options]`
          or `python3 mergertree-extract.py --all [--options]`

    This will create the same type of files as in mode (2), just for all
    haloes.


If only an integer is given as cmdline arg, mode (1) [one clump only] will be
run. If no cmd line argument is given, mode (3) [--all] will be run.



---------------
    Output
---------------

the mergertree_XXXXX* files have 6 columns:

snapshot            The snapshot from which this data is taken from

redshift            The redshift of that snapshot

clump_ID            The clump ID of the clump at that snapshot

mass                The mass of the clump at that snapshot, based on what's in
                    the output_XXXXX/mergertree_XXXXX.txtYYYYY files, not the
                    output_XXXXX/clump_XXXXX.txtYYYYY files.

mass_from_mergers   how much mass has been merged into this clump in this
                    snapshot, i.e. the sum of all the clump masses that have
                    been found to merge with this clump at this snapshot. This
                    does not include the mass of clumps which only seem to
                    merge with this clump, but re-emerge later.

mass_from_jumpers   The mass of all clumps that seem to merge with this clump,
                    but re-emerge at a later time.


----------------
    Options
----------------

List of all flags:

Running modes

    -a, --all:      make trees for all clumps in output where z = 0
    -c --children:  make trees for a halo and all its subhaloes. You need to
                    specify which halo via its halo ID.
    -h, --help:     print this help and exit.

Options:
    --start-at=INT      don't start at z = 0 snapshot, but with the specified
                        directory output_00INT.
    --prefix=some/path/ path where you want your output written to.
    -v, --verbose:      be more verbose about what you're doing




-----------------
  Requirements
-----------------

It needs output_XXXXX/mergertree_XXXXX.txtYYYYY and
output_XXXXX/clump_XXXXX.txtYYYYY files to work, which are created using the
mergertree patch in ramses.

Also needs numpy.
"""

###############################################################################
#                             Clump data                                      #
###############################################################################


class ClumpData:
    """
    Data from clump_XXXXX.txt

    Parameters
    ----------
    par : params object
    """
    def __init__(self, par):
        self.clumpids = np.zeros(1)     # clump ID
        self.parent = np.zeros(1)       # parent ID
        self.level = np.zeros(1)        # clump level

    def read_clumpdata(self, par):
        """Reads in the clump data for the z = 0 directory."""
        if par.verbose:
            print("Reading clump data.")

        out = p.z0

        raw_data = [None for i in range(par.ncpu)]
        dirnrstr = str(par.outputnrs[out]).zfill(5)
        dirname = 'output_' + dirnrstr

        i = 0
        for cpu in range(1):
            fname = join(par.workdir, dirname, 'clump_' + dirnrstr + '.dat')
            new_data = np.loadtxt(fname, dtype='int', skiprows=1,
                                  usecols=[0, 1, 2])
            if new_data.ndim == 2:
                raw_data[i] = new_data
                i += 1
            elif new_data.shape[0] == 3:  # if only 1 row is present in file
                raw_data[i] = np.atleast_2d(new_data)
                i += 1

        fulldata = np.concatenate(raw_data[:i], axis=0)
        self.clumpids = fulldata[:, 0]
        self.level = fulldata[:, 1]
        self.parent = fulldata[:, 2]

    def cleanup_clumpdata(self, par, mtd):
        """
        The particle unbinding can remove entire clumps from the catalogue.
        If the option isn't set in the namelist, the clumpfinder output will
        still be made not based on the clumpfinder. If that is the case, the
        clumpfinder catalogue will contain clumps which the mergertree data
        doesn't have, leading to problems. So remove those here.
        """
        for i, c in enumerate(self.clumpids):
            if c not in mtd.descendants[par.z0]:
                self.clumpids[i] = 0
                self.level[i] = 0
                self.parent[i] = -1  # don't make it the same as clumpid

    def find_children(self, clumpid):
        """Find the children for given clump ID."""
        children = []
        last_added = [clumpid]

        loopcounter = 0
        while True:
            loopcounter += 1
            this_level_parents = copy.copy(last_added)
            children += this_level_parents
            last_added = []
            for i, cid in enumerate(self.clumpids):
                if self.parent[i] in this_level_parents and cid != clumpid:
                    last_added.append(cid)

            if len(last_added) == 0:
                break

            if loopcounter == 100:
                print("Finished 100 iterations, we shouldn't be this deep")
                break

        return children[1:]  # don't return top level parent

    def write_children(self, par, clumpid, children):
        """Write the children to file."""
        hfile = join(par.outdir, f"{par.halofilename}-{str(clumpid)}.txt")

        with open(hfile, 'w') as f:
            f.write("# {0:>18} {1:>18} {2:>18}\n".format("halo", "nr_of_children", "children"))  # noqa
            nc = len(children)
            dumpstring = "  {0:18d} {1:18d}".format(clumpid, nc)
            dumpstring = "".join([dumpstring] + [" {0:18d}".format(c) for c in children] + ['\n'])  # noqa
            f.write(dumpstring)


###############################################################################
#                            Constants object                                 #
###############################################################################


class Constants:
    """
    Class holding constants.
    """
    def __init__(self):
        self.Mpc = 3.086e24                 # cm
        self.M_Sol = 1.98855e33             # g
        self.Gyr = (24 * 3600 * 365 * 1e9)  # s
        self.G = 4.492e-15                  # Mpc^3/(M_sol Gyr^2)

        self.H0 = 100                      # km/s/Mpc
        self.omega_m = 0.307000011205673
        self.omega_l = 0.693000018596649
        self.omega_k = 0.0
        self.omega_b = 0.0


###############################################################################
#                             Params object                                   #
###############################################################################


class Params:
    """
    Global parameters to be stored
    """
    def __init__(self):
        # self.workdir = f"/mnt/extraspace/hdesmond/ramses_out_{self.nsim}"
        # self.outdir = f"/mnt/extraspace/rstiskalek/CSiBORG/cleaned_mtree/ramses_out_{self.nsim}"  # noqa
        # if not exists(self.outdir):
        #     makedirs(self.outdir)
        self.lastdir = ""               # last output_XXXXX directory
        self.lastdirnr = -1             # XXXX from lastdir
        self.ncpu = 1                   # Number of CPUs used
        self.noutput = 1                # how many output_XXXXX dirs exist
        self.nout = 1                   # how many outputs we're gonna deal with. (Some might not have merger tree data)  # noqa
        self.outputnrs = None           # numpy array of output numbers
        self.output_lowest = 0          # lowest snapshot number that we're dealing with (>= 1)  # noqa
        self.z0 = 0                     # index of z=0 snapshot (or whichever you want to start with)  # noqa

        # NOTE: params.nout will be defined such that you can easily loop

        self.verbose = False            # verbosity
        self.start_at = 0               # output dir to start with, if given

        self.output_prefix = ""         # user given prefix for output files
        self.outputfilename = ""        # output filename. Stores prefix/mergertree_XXXXX part of name only  # noqa
        self.halofilename = ""          # output filename for halo hierarchy. Stores prefix/halo_hierarchy_XXXXX part of filename only  # noqa

        self.one_halo_only = False      # do the tree for one clump only
        self.halo_and_children = False  # do the tree for one halo, including subhaloes  # noqa
        self.do_all = False             # do for all clumps at z=0 output

        self.clumpid = 0                # which clump ID to work for.
        self.nsim = None

        # Dictionnary of accepted keyword command line arguments
        self.accepted_flags = {
            '-a': self.set_do_all,
            '--all': self.set_do_all,
            '-r': self.set_halo_and_children,
            '--recursive': self.set_halo_and_children,
            '-c': self.set_halo_and_children,
            '--children': self.set_halo_and_children,
            '-h': self.get_help,
            '--help': self.get_help,
            '-v': self.set_verbose,
            '--verbose': self.set_verbose,
            }

        self.accepted_flags_with_args = {
            "--nsim": self.set_nsim,
            '--start-at': self.set_startnr,
            '--prefix': self.set_prefix,
            }

    # -----------------------------
    # Setter methods
    # -----------------------------

    def set_do_all(self):
        self.do_all = True
        return

    def set_halo_and_children(self):
        self.halo_and_children = True
        return

    def get_help(self):
        print(errmsg)
        quit()
        return

    def set_verbose(self):
        self.verbose = True
        return

    def set_startnr(self, arg):
        flag, startnr = arg.split("=")
        try:
            self.start_at = int(startnr)
        except ValueError:
            print("given value for --start-at=INT isn't an integer?")

    def set_prefix(self, arg):
        flag, prefix = arg.split("=")
        #  try:
        self.output_prefix = prefix
        try:
            os.makedirs(self.output_prefix)
        except FileExistsError:
            pass
        return

    def set_nsim(self, arg):
        flag, nsim = arg.split("=")
        try:
            self.nsim = int(nsim)
        except ValueError:
            print("given value for --nsim=INT isn't an integer?")

    def read_cmdlineargs(self):
        """
        Reads in the command line arguments and store them in the
        global_params object.
        """
        nargs = len(argv)
        i = 1  # first cmdlinearg is filename of this file, so skip it

        while i < nargs:
            arg = argv[i]
            arg = arg.strip()
            if arg in self.accepted_flags.keys():
                self.accepted_flags[arg]()
            else:
                for key in self.accepted_flags_with_args.keys():
                    if arg.startswith(key):
                        self.accepted_flags_with_args[key](arg)
                        break
                else:
                    try:
                        self.clumpid = int(arg)
                    except ValueError:
                        print(f"I didn't recognize the argument '{arg}'. Use "
                              "mergertre-extract.py -h or --help to print "
                              "help message.")
                        quit()

            i += 1

        if self.nsim is None:
            raise ValueError("nsim not set. Use --nsim=INT to set it.")

    @property
    def workdir(self):
        return f"/mnt/extraspace/hdesmond/ramses_out_{self.nsim}"

    @property
    def outdir(self):
        fname = f"/mnt/extraspace/rstiskalek/CSiBORG/cleaned_mtree/ramses_out_{self.nsim}"  # noqa
        if not exists(fname):
            makedirs(fname)
        return fname

    def get_output_info(self):
        """
        Read in the output info based on the files in the current working
        directory. Reads in last directory, ncpu, noutputs. Doesn't read
        infofiles.
        """
        # self.workdir = os.getcwd()
        filelist = os.listdir(self.workdir)

        outputlist = []
        for filename in filelist:
            if filename.startswith('output_'):
                outputlist.append(filename)

        if len(outputlist) < 1:
            print("I didn't find any output_XXXXX directories in current "
                  "working directory. Are you in the correct workdir? "
                  "Use mergertree-extract.py -h or --help to print help "
                  "message.")
            quit()

        outputlist.sort()

        self.lastdir = outputlist[-1]
        self.lastdirnr = int(self.lastdir[-5:])
        self.noutput = len(outputlist)

        if (self.start_at > 0):
            # check that directory exists
            startnrstr = str(self.start_at).zfill(5)
            if 'output_' + startnrstr not in outputlist:
                print("Didn't find specified starting directory "
                      f"output_{startnrstr} use mergertree-extract.py -h or "
                      "--help to print help message.")
                quit()

        # read ncpu from infofile in last output directory
        infofile = join(self.workdir, self.lastdir,
                        f"info_{self.lastdir[-5:]}.txt")
        with open(infofile, 'r') as f:
            ncpuline = f.readline()
            line = ncpuline.split()
            self.ncpu = int(line[-1])

    def setup_and_checks(self, sd):
        """
        Do checks and additional setups once you have all the cmd line args and
        output infos

        Parameters
        ----------
        sd: snapshotdata object
        """
        # set running mode
        if not self.do_all:
            if self.clumpid <= 0:
                print("No or wrong clump id given. Setting the --all mode.")
                self.set_do_all()
            else:
                if not self.halo_and_children:
                    self.one_halo_only = True

        # generate list of outputdirnumbers
        startnr = self.lastdirnr
        self.outputnrs = np.array(range(startnr, startnr - self.noutput, -1))

        # find starting output directory
        self.z0 = np.argmin(np.absolute(sd.redshift))

        if self.start_at > 0:
            # replace z0 dir with starting dir
            self.z0 = self.lastdirnr - self.start_at

        # generate output filename
        dirnrstr = str(self.outputnrs[self.z0]).zfill(5)
        fname = "mergertree_" + dirnrstr
        self.outputfilename = join(self.output_prefix, fname)

        # generate halo output filename
        fname = "halo_hierarchy_" + dirnrstr
        self.halofilename = join(self.output_prefix, fname)

        # rename output_prefix to something if it wasn't set
        if self.output_prefix == "":
            self.output_prefix = os.path.relpath(self.workdir)

        # find self.nout; i.e. how many outputs we are actually going to have
        for out in range(self.noutput - 1, -1, -1):
            dirnrstr = str(self.outputnrs[out]).zfill(5)
            mtreefile = join(self.workdir,
                             f"output_{dirnrstr}",
                             f"mergertree_{dirnrstr}.dat")

            if os.path.exists(mtreefile):
                print("Loading mergertree data from ", mtreefile)
                # if there is a file, this is lowest snapshot number directory
                # that we'll be dealing with, and hence will have the highest
                # index number in the arrays I'm using

                # NOTE: params.nout will be defined such that you can easily
                # loop for out in range(p.z0, p.nout)
                self.nout = out + 1
                break

    def print_params(self):
        """Prints out the parameters that are set."""
        if self.do_all:
            print("Working mode:             all clumps")
        else:
            if self.halo_and_children:
                print("Working mode:             halo", self.clumpid, "and its children")  # noqa
            else:
                print("Working mode:             clump ", self.clumpid)

        print("workdir:                 ", self.workdir)
        print("snapshot of tree root:   ", self.outputnrs[self.z0])
        print("p.one_halo_only          ", p.one_halo_only)
        print("p.do_all                 ", p.do_all)
        print("p.halo_and_children      ", p.halo_and_children)
        print("p.one_halo_only          ", p.one_halo_only)


###############################################################################
#                             Merger tree data                                #
###############################################################################


class MTreeData:
    """
    Merger tree data lists

    Parameters
    ----------
    par : params object
    """
    def __init__(self, par):
        self.progenitors = [np.zeros(1) for i in range(par.noutput)]            # progenitor IDs  # noqa
        self.descendants = [np.zeros(1) for i in range(par.noutput)]            # descendant IDs  # noqa
        self.progenitor_outputnrs = [np.zeros(1) for i in range(par.noutput)]   # snapshot number of progenitor  # noqa
        self.mass = [np.zeros(1) for i in range(par.noutput)]                   # descendant mass  # noqa
        self.mass_to_remove = [np.zeros(1) for i in range(par.noutput)]         # descendant mass  # noqa

    def read_mergertree_data(self, par, sd):
        """Reads in mergertree data."""

        if par.verbose:
            print("Reading in mergertree data")

        # Preparation

        # define new datatype for mergertree output
        mtree = np.dtype([('clump', 'i4'),
                          ('prog', 'i4'),
                          ('prog_outnr', 'i4'),
                          ('mass', 'f8'),
                          ('npart', 'f8'),
                          ('x', 'f8'),
                          ('y', 'f8'),
                          ('z', 'f8'),
                          ('vx', 'f8'),
                          ('vy', 'f8'),
                          ('vz', 'f8')
                          ])

        # ---------------------------
        # Loop over directories
        # ---------------------------

        startnr = par.lastdirnr
        # READ THE ONES BEFORE z0 TOO!
        for output in trange(par.nout, desc="Reading merger"):
            dirnr = str(startnr - output).zfill(5)
            srcdir = 'output_' + dirnr

            fnames = [srcdir + '/' + "mergertree_" + dirnr + '.dat']
            fnames[0] = join(par.workdir, fnames[0])

            datalist = [np.zeros((1, 3)) for i in range(par.ncpu)]
            i = 0
            nofile = 0
            for f in fnames:
                if os.path.exists(f):
                    datalist[i] = np.atleast_1d(np.genfromtxt(f, dtype=mtree,
                                                              skip_header=1))
                    i += 1
                else:
                    nofile += 1

            if nofile == p.ncpu:
                print("Didn't find any mergertree data in", srcdir)

            # ---------------------------------
            # Sort out data
            # ---------------------------------
            if i > 0:
                fulldata = np.concatenate(datalist[:i], axis=0)

                self.descendants[output] = fulldata[:]['clump']
                self.progenitors[output] = fulldata[:]['prog']
                self.progenitor_outputnrs[output] = fulldata[:]['prog_outnr']
                self.mass[output] = fulldata[:]['mass']
                #  self.npart[output] = fulldata[:]['npart']
                #  self.x[output] = fulldata[:]['x']
                #  self.y[output] = fulldata[:]['y']
                #  self.z[output] = fulldata[:]['z']
                #  self.vx[output] = fulldata[:]['vx']
                #  self.vy[output] = fulldata[:]['vy']
                #  self.vz[output] = fulldata[:]['vz']

        # --------------------------------------
        # Transform units to physical units
        # --------------------------------------

        # transform units to physical units
        for i in range(len(self.descendants)):
            self.mass[i] *= sd.unit_m[i]
            #  self.x[i] *= sd.unit_l[i] # only transform later when needed; Need to check for periodicity first!  # noqa
            #  self.y[i] *= sd.unit_l[i]
            #  self.z[i] *= sd.unit_l[i]
            #  self.vx[i] *= sd.unit_l[i]/sd.unit_t[i]
            #  self.vy[i] *= sd.unit_l[i]/sd.unit_t[i]
            #  self.vz[i] *= sd.unit_l[i]/sd.unit_t[i]

    def clean_up_jumpers(self, par):
        """
        Remove jumpers from the merger list. Take note of how much mass should
        be removed from the descendant because the jumper is to be removed.
        """
        # First initialize mass_to_remove arrays
        self.mass_to_remove = [np.zeros(self.descendants[out].shape)
                               for out in range(par.noutput)]
        nreplaced = 0
        for out in trange(par.nout + par.z0 - 1, desc="Cleaning jumpers"):
            for i, pr in enumerate(self.progenitors[out]):
                if pr < 0:
                    # Subtract 1 here from snapind:
                    # progenitor_outputnrs gives the snapshot number where the
                    # jumper was a descendant for the last time
                    # so you need to overwrite the merging one snapshot later,
                    # where the clump is the progenitor
                    snapind = get_snap_ind(p, self.progenitor_outputnrs[out][i]) - 1  # noqa

                    # NOTE bottleneck
                    jumpind = self.progenitors[snapind] == -pr

                    # NOTE bottleneck
                    # find index of descendant into which this clump will
                    # appearingly merge into
                    mergerind = self.descendants[snapind] == - self.descendants[snapind][jumpind]  # noqa
                    # overwrite merging event so it won't count
                    self.descendants[snapind][jumpind] = 0

                    # find mass of jumper in previous snapshot
                    jumpmassind = self.descendants[snapind + 1] == -pr
                    # note how much mass might need to be removed for whatever
                    # you need it
                    self.mass_to_remove[snapind][mergerind] += self.mass[snapind + 1][jumpmassind]  # noqa

                    nreplaced += 1

        print("Cleaned out", nreplaced, "jumpers")

    def get_tree(self, par, tree, sd, clumpid):
        """Follow the main branch down."""
        if par.verbose:
            print("Computing tree for clump", clumpid)

        dind = self.descendants[par.z0] == clumpid
        desc_snap_ind = p.z0
        desc = self.descendants[p.z0][dind]
        prog = self.progenitors[p.z0][dind]

        def get_prog_indices(prog, desc_snap_ind):
            """
            Compute snapshot index at which given progenitor has been a
            descendant and its index in the array

            prog:           progenitor ID
            desc_snap_ind:  snapshot index of descendant of given prog

            returns:
            p_snap_ind:     snapshot index of the progenitor
            pind:           progenitor index (np.array mask) of progenitor in
                            array where it is descendant
            """
            if prog > 0:  # if progenitor isn't jumper
                # find progenitor's index in previous snapshot
                p_snap_ind = desc_snap_ind + 1
                pind = self.descendants[p_snap_ind] == prog

            elif prog < 0:
                p_snap_ind = get_snap_ind(
                    par, self.progenitor_outputnrs[desc_snap_ind][dind])
                pind = self.descendants[p_snap_ind] == -prog

            return p_snap_ind, pind

        while True:
            # first calculate merger mass
            mergers = self.descendants[desc_snap_ind] == -desc
            mergermass = 0.0
            if mergers.any():
                for m in self.progenitors[desc_snap_ind][mergers]:
                    # find mass of merger. That's been written down at the
                    # place where merger was descendant.
                    m_snap_ind, mergerind = get_prog_indices(m, desc_snap_ind)
                    mergermass += self.mass[m_snap_ind][mergerind]

            # add the descendant to the tree
            tree.add_snap(par.outputnrs[desc_snap_ind],
                          sd.redshift[desc_snap_ind], desc,
                          self.mass[desc_snap_ind][dind], mergermass,
                          self.mass_to_remove[desc_snap_ind][dind])

            # now descend down the main branch
            if prog != 0:
                p_snap_ind, pind = get_prog_indices(prog, desc_snap_ind)
            else:
                # stop at progenitor = 0
                break

            # prepare for next round
            desc_snap_ind = p_snap_ind
            dind = pind
            desc = abs(prog)
            prog = self.progenitors[p_snap_ind][pind]


###############################################################################
#                             Snapshot data                                   #
###############################################################################


class SnapshotData():
    """Snapshot specific data"""
    def __init__(self, par):
        # read in
        self.aexp = np.zeros(par.noutput)
        self.unit_l = np.zeros(par.noutput)
        self.unit_m = np.zeros(par.noutput)
        self.unit_t = np.zeros(par.noutput)
        self.unit_dens = np.zeros(par.noutput)
        # to be computed
        self.redshift = np.zeros(par.noutput)  # z

    def read_infofiles(self, par, const):
        """Read the info_XXXXX.txt files."""
        if par.verbose:
            print("Reading info files.")

        startnr = par.lastdirnr

        for output in range(p.noutput):
            # Start with last directory (e.g. output_00060),
            # work your way to first directory (e.g. output_00001)
            # p.z0 isn't decided yet, so just read in everything here.
            dirnr = str(startnr - output).zfill(5)
            srcdir = 'output_' + dirnr

            try:
                # ------------------------------------------------------
                # get time, redshift, and units even for output_00001
                # ------------------------------------------------------
                fileloc = srcdir + '/info_' + dirnr + '.txt'
                fileloc = join(par.workdir, fileloc)
                infofile = open(fileloc)
                for i in range(9):
                    infofile.readline()  # skip first 9 lines

                # get expansion factor
                aline = infofile.readline()
                astring, equal, aval = aline.partition("=")
                afloat = float(aval)
                sd.aexp[output] = afloat

                for i in range(5):
                    infofile.readline()  # skip 5 lines

                # get unit_l
                unitline = infofile.readline()
                unitstring, equal, unitval = unitline.partition("=")
                unitfloat = float(unitval)
                sd.unit_l[output] = unitfloat

                # get unit_dens
                unitline = infofile.readline()
                unitstring, equal, unitval = unitline.partition("=")
                unitfloat = float(unitval)
                sd.unit_dens[output] = unitfloat

                # get unit_t
                unitline = infofile.readline()
                unitstring, equal, unitval = unitline.partition("=")
                unitfloat = float(unitval)
                sd.unit_t[output] = unitfloat

                infofile.close()

            except IOError:  # If file doesn't exist
                print("Didn't find any info data in ", srcdir)
                break

        self.unit_m = self.unit_dens * self.unit_l ** 3 / const.M_Sol
        self.unit_l /= const.Mpc
        self.unit_t /= const.Gyr

        self.redshift = 1. / self.aexp - 1

###############################################################################
#                             Tree object                                     #
###############################################################################


class Tree:
    """
    Holds tree result data. It's not really a tree, it's just the values along
    the main branch, but let's call it a tree anyway. Sue me.

    Parameters
    ----------
    nelements : int
        Estimate for how many snapshots you need to allocate space for.
    """
    def __init__(self, nelements):
        self.n = 0                                              # number of elements in tree  # noqa
        self.snapshotnr = -np.ones(nelements, dtype=int)        # snapshot number of array values  # noqa
        self.redshift = -np.ones(nelements, dtype=float)        # redshift at that snapshot  # noqa
        self.clumpids = -np.ones(nelements, dtype=int)          # clump id of halo in that snapshot  # noqa
        self.mass = np.zeros(nelements, dtype=float)            # mass at that snapshot  # noqa
        self.mergermass = np.zeros(nelements, dtype=float)      # sum of mass of swallowed up clumps  # noqa
        self.mass_to_remove = np.zeros(nelements, dtype=float)  # sum of mass of swallowed up clumps  # noqa

    def add_snap(self, nr, z, ID, m, mm, mdel):
        """Add new result."""
        n = self.n
        self.snapshotnr[n] = nr
        self.redshift[n] = z
        self.clumpids[n] = ID
        self.mass[n] = m
        self.mergermass[n] = mm
        self.mass_to_remove[n] = mdel
        self.n += 1

    def write_tree(self, par, case='halo'):
        """Write the results to file."""
        resfile = join(
            par.outdir,
            f"{par.outputfilename}_{case}-{str(self.clumpids[0])}.txt")

        with open(resfile, 'w') as f:
            f.write('# {0:>12} {1:>12} {2:>16} {3:>18} {4:>18} {5:>18}\n'.format(  # noqa
                "snapshot", "redshift", "clump_ID", "mass[M_sol]",
                "mass_from_mergers", "mass_from_jumpers"))

            for i in range(self.n):
                f.write('  {0:12d} {1:12.4f} {2:16d} {3:18.6e} {4:18.6e} {5:18.6e}\n'.format(  # noqa
                    self.snapshotnr[i], self.redshift[i], self.clumpids[i],
                    self.mass[i], self.mergermass[i], self.mass_to_remove[i]))

        return


def get_snap_ind(p, snap):
    """
    Computes the snapshot index in mtreedata/halodata/snapshotdata arrays for a
    given snapshot number snap
    """
    return (p.noutput - snap).item()


if __name__ == '__main__':

    p = Params()
    c = Constants()

    # Read cmdlineargs, available output, get global parameters
    p.read_cmdlineargs()
    p.get_output_info()

    sd = SnapshotData(p)
    sd.read_infofiles(p, c)

    # finish setup
    p.setup_and_checks(sd)
    p.print_params()

    # now read in mergertree data
    fname = join(p.outdir, "mtreedata.p")
    if exists(fname):
        print(f"{datetime.now()}: loading mergertree data from `{fname}`.",
              flush=True)
        mtd = load(fname)
        print(f"{datetime.now()}: finished loading mergertree data from `{fname}`.",  # noqa
              flush=True)
    else:
        print("Generating mergertree data.", flush=True)
        mtd = MTreeData(p)
        mtd.read_mergertree_data(p, sd)
        # clean up jumpers
        mtd.clean_up_jumpers(p)

        print("Saving mergertree data.", flush=True)
        dump(mtd, fname)

    # read in clump data if required
    if p.do_all or p.halo_and_children:
        cd = ClumpData(p)
        cd.read_clumpdata(p)

        # clean up halo catalogue
        cd.cleanup_clumpdata(p, mtd)

        # find children, and write them down
        if p.verbose:
            print("Searching for child clumps.")

        if p.halo_and_children:
            children = cd.find_children(p.clumpid)
            cd.write_children(p, p.clumpid, children)

        if p.do_all:
            is_halo = cd.clumpids == cd.parent
            childlist = [None for c in cd.clumpids[is_halo]]
            for i, halo in enumerate(cd.clumpids[is_halo]):
                children = cd.find_children(halo)
                cd.write_children(p, halo, children)
                childlist[i] = children

    # finally, get the bloody tree

    if p.one_halo_only:
        newtree = Tree(p.nout)
        mtd.get_tree(p, newtree, sd, p.clumpid)
        newtree.write_tree(p, 'halo')

    if p.halo_and_children:
        newtree = Tree(p.nout)
        mtd.get_tree(p, newtree, sd, p.clumpid)
        newtree.write_tree(p, 'halo')

        for c in children:
            newtree = Tree(p.nout)
            mtd.get_tree(p, newtree, sd, c)
            newtree.write_tree(p, 'subhalo')

    if p.do_all:
        for i, halo in enumerate(cd.clumpids[is_halo]):
            newtree = Tree(p.nout)
            mtd.get_tree(p, newtree, sd, halo)
            newtree.write_tree(p, 'halo')

            for c in childlist[i]:
                newtree = Tree(p.nout)
                mtd.get_tree(p, newtree, sd, c)
                newtree.write_tree(p, 'subhalo')

    print('Finished.')