Skip to content

Instantly share code, notes, and snippets.

@yymao

yymao/run-ctrees.py

Last active Aug 29, 2015
Embed
What would you like to do?
A Python script to run Peter Behroozi's Consistent Trees (replacing do_merger_tree.pl).
#!/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))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.