summaryrefslogtreecommitdiffstats
path: root/plot_stats.py
blob: 1802d8db422de280b8c20c81064d7d596473a206 (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
90
91
92
93
94
95
import matplotlib.pyplot as plt
from DiskXYZ import Disk
import numpy as np
import argparse
import sys
from math import pi

# Parse Arguments
parser = argparse.ArgumentParser()
parser.add_argument("imin", type=int, help='First Output')
parser.add_argument("imax", type=int, help='Last Output')
parser.add_argument("--show", action="store_true", \
                    help='Show Figure')
parser.add_argument("--save", action="store_true", \
                    help='Save Figure')
parser.add_argument("--disk", action="store_true", \
                    help="Show Total Disk Mass")
parser.add_argument("--halo", action="store_true", \
                    help="Show Total Halo Mass")
parser.add_argument("--box", action="store_true", \
                    help="Show Total Box Mass")
parser.add_argument("--lost_halo", action="store_true", \
                    help="Show (Accumulated) Mass Lost to Halo")
parser.add_argument("--lost_star", action="store_true", \
                    help="Show (Accumulated) Mass Lost to Star")
args = parser.parse_args()

# Sanity Checks
if args.imin > args.imax:
    print("Cannot Work With Imin > Imax. Use -h for help.")
    sys.exit(-1)
if not (args.show or args.save):
    print("Specify --save or --show. Use -h for help.")
    sys.exit(-1)
if not (args.disk or args.halo or args.box or args.lost_halo or args.lost_star):
    print("Let me show you something, please? Use -h for help.")
    sys.exit(-1)

# Build Output Range
iouts = range(args.imin, args.imax+1)

# Allocate Arrays
time = np.zeros_like(iouts, dtype='float64')
mass_total_lost_to_star = np.zeros_like(iouts, dtype='float64')
mass_total_lost_to_halo = np.zeros_like(iouts, dtype='float64')
mass_disk = np.zeros_like(iouts, dtype='float64')
mass_halo = np.zeros_like(iouts, dtype='float64')
mass_total = np.zeros_like(iouts, dtype='float64')

# Load Data
ii = 0
for iout in iouts:
    print "Loading Data for Output %i." % iout
    disk = Disk(iout)
    disk.load_npz_stats()
    # print disk.info["time"]
    time[ii] = disk.info["time"]
    mass_total_lost_to_star[ii] = disk.mass_total_lost_to_star
    mass_total_lost_to_halo[ii] = disk.mass_total_lost_to_halo
    mass_disk[ii] = disk.mass_disk
    mass_halo[ii] = disk.mass_halo
    mass_total[ii] = disk.mass_halo
    ii += 1

# Make Plots
fig = plt.figure()
ax = fig.add_subplot(1,1,1)
ax2 = ax.twiny()
if args.disk:
    ax.plot(time/2./pi, mass_disk, 'bs-', label='Disk Mass')
if args.halo:
    ax.plot(time/2./pi, mass_halo, 'rs-', label='Halo Mass')
if args.box:
    ax.plot(time/2./pi, mass_total, 'gs-', label='Box Mass')
if args.lost_halo:
    ax.plot(time/2./pi, mass_total_lost_to_halo, 'ms-', \
        label='Total Mass Lost to Halo')
if args.lost_star:
    ax.plot(time/2./pi, mass_total_lost_to_star, 'ks-', \
        label='Total Mass Lost to Star')
fig.suptitle('Mass Budget')
ax2.plot(time, np.ones_like(time))
ax2.cla()
ax2.set_xlabel('Time [code]')
ax.grid()
ax.set_yscale('log')
ax2.set_ylim(ax.get_ylim())
ax2.set_yscale(ax.get_yscale())
ax.set_xlabel('Time [yr]')
ax.set_ylabel('Mass [Mstar]')
ax.legend(loc='best')
if args.save:
    plt.savefig('Stats.pdf')
if args.show:
    plt.show()