A Python script to run Peter Behroozi's Consistent Trees (replacing do_merger_tree.pl).

run-ctrees.py

#!/usr/bin/env python
import os
import time
from argparse import ArgumentParser

#user settings
_perl_exe = 'perl'
_default_ctrees_dir = os.getenv('HOME') + '/software/consistent-trees'


#define ctrees stages
_ctrees_stages = ('gravitational_consistency', 'find_parents_and_cleanup', 
       'resort_outputs', 'assemble_halo_trees', 'halo_trees_to_catalog.pl')
_nstage = len(_ctrees_stages)
_stage_info = 'Stages of Consistent Trees are: ' \
        + ', '.join(['({0}) {1}'.format(*t) for t in enumerate(_ctrees_stages)])


#define argument parser
parser = ArgumentParser(epilog=_stage_info)
parser.add_argument('-d', dest='ctrees_dir', metavar='CTREES_DIR', \
                    default=_default_ctrees_dir, help='path to the Consistent Trees directory')
parser.add_argument('-s', dest='start_stage', metavar='STAGE', \
                    type=int, choices=range(_nstage), \
                    help='stage to start from')
parser.add_argument('-l', dest='log', metavar='LOG_FILE', nargs='?', \
                    const='ctrees.log', help='path to the log file; file will be overwritten if exists')
parser.add_argument('--no-periodic', dest='periodic', action='store_false', \
                    help='Run without periodic boundary condition')
parser.add_argument('--stages', help='stages to run; overwrites -s')
parser.add_argument('--dry-run', action='store_true', \
                    help='print out the commands, but do not run them')
parser.add_argument('CFG_FILE', help='Consistent Trees configuration file')
args = parser.parse_args()


#process stages
s = slice(args.start_stage, None)
if args.stages is not None:
    a, __, b = args.stages.partition('-')
    try:
        a = int(a) if a else None
        if __:
            b = int(b)+1 if b else None
        else:
            b = a+1
    except ValueError:
        parser.error('Stages "{0}" not correctly specified.'.format(args.stages))
    s = slice(a, b)

stages = range(_nstage)[s]
if not stages:
    parser.error('Stage(s) not correctly specified. Nothing to run.')

exe = list(_ctrees_stages)
if not args.periodic:
    for i in [0, 1]:
        exe[i] += '_no_periodic'
stages = zip(stages, exe[s])


#check cfg file and ctrees dir
cfg_file = os.path.expanduser(args.CFG_FILE)
ctrees_dir = os.path.expanduser(args.ctrees_dir)

if not os.path.isfile(cfg_file):
    parser.error('Cannot locate cfg file "{0}".'.format(args.CFG_FILE))

if not os.path.isdir(ctrees_dir):
    parser.error('Cannot locate Consistent Trees directory "{0}".'.format(args.ctrees_dir))

if filter(lambda s: not os.path.isfile('{0}/{1}'.format(ctrees_dir, s[1])), stages):
    parser.error('Consistent Trees directory "{0}" is invalid.'.format(args.ctrees_dir))


#process dry_run and log
if args.dry_run:
    def check_call(l): print ' '.join(l)
else:
    from subprocess import check_call

if args.log is None or args.dry_run:
    def log(msg):
        print get_timestamp(), msg
else:
    log_file = os.path.expanduser(args.log)
    try:
        with open(log_file, 'w') as f:
            pass
    except OSError:
        parser.error('Cannot access log file "{0}".'.format(args.log))
    def log(msg):
        with open(log_file, 'a', 0) as f:
            f.write('{0} {1}\n'.format(get_timestamp(), msg))

def get_timestamp():
    t = time.time()
    return '[{0}.{1:03d}]'.format( \
            time.strftime("%m/%d %H:%M:%S", time.localtime(t)), \
            int(t*1000)%1000)

#start running
log('starting ctrees from <stage {0}> {1}.'.format(*stages[0]))
for stage, exe in stages:
    comm = [_perl_exe, '-I{0}/src'.format(ctrees_dir)] if exe.endswith('.pl') else []
    comm.extend(['{0}/{1}'.format(ctrees_dir, exe), cfg_file])
    check_call(comm)
    log('<stage {0}> {1} is done.'.format(stage, exe))