aboutsummaryrefslogtreecommitdiffstats
path: root/Post/Formation/extract_wmf.py
blob: 8500530707d18ba98348645660d4c1fde02db5dc (plain)
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