1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
|
"""
Extract Water Mass Fraction, Source Distribution Statistics.
"""
import io_helpers as ioh
import formation_helpers as fh
import multiprocessing as mp
import pandas as pd
import argparse
import sys
import glob
import numpy as np
###############################################################################
# MAIN PROGRAM STARTS HERE
###############################################################################
# Parse Arguments
parser = argparse.ArgumentParser()
parser.add_argument('-np', type=int, default=1, \
help='Number of Processes')
parser.add_argument('-fout', '--output_file', default='Coordinates_WMF.hdf5', \
help='Name of Output File.')
args = parser.parse_args()
print "// Using %i Subprocesses" % args.np
# List of Directories
if sys.stdin.isatty():
print "!! No Directory List (Use Stdin)."
sys.exit()
else:
lines = sys.stdin.read().rstrip("\n").split("\n")
dirs = []
for line in lines:
dirs.append(line)
print "// Reading %i Directories" % len(dirs)
# Should Only Be One Directory
if len(dirs) == 1:
cdir = dirs[0]
else:
print "!! Script can only process a single output directory."
print "!! Terminateing."
sys.exit()
# Define Steps
nsteps = np.array([ 0, 6e6, 6e7, 6e8, 3e9, 6e9, 9e9], dtype=np.int64)
# The Globbit
globs = glob.glob("%s/Out_*.dat" % cdir)
# Extract run names
# In: Out_run_03_000057000000
# Out: run_03
run_name = globs[0].strip().split("/")[-1][:-4][4:-13]
# Load Coordinate Outputs
dfo_all = []
for nstep in nsteps:
fname = "%s/Out_%s_%012d.dat" % (cdir, run_name, nstep)
dfo = ioh.read_output_and_stack([fname], frame='heliocentric')
dfo = dfo[dfo.mass < 12.0]
dfo_all.append(dfo)
# Load Collisions
fname_c = "%s/Collisions_%s.dat" % (cdir, run_name)
dfc = ioh.read_collisions_and_stack([fname_c])
# Compute WMF, Source Regions
targets = []
for istep, nstep in enumerate(nsteps):
target = [ dfo_all[istep], dfo_all[0], dfc ]
targets.append(target)
# Loop Steps
print "// Computing WMF & Sources for %i Steps" % len(targets)
pool = mp.Pool(processes=args.np)
result = pool.map(fh.compute_wmf_wrapper, targets, 1)
pool.close()
pool.join()
# Join Data Frame
df = pd.concat(result)
df.reset_index(drop=True, inplace=True)
# Save
print "// Saving to %s" % args.output_file
with pd.HDFStore("%s" % args.output_file, 'w') as store:
store['df'] = df
|
