summaryrefslogtreecommitdiffstats
path: root/plot_stats.py
blob: 0ec413cf3becbd09d999408a1fa892e41880c8a1 (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
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
"""
Plot Mass Budget and Rates of Change.

@todo: Plot dM/dt as log(dM/dt+1) so we don't take log(<0)

Volker Hoffmann <volker@cheleb.net>
04 June 2013
"""

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_total
    ii += 1

# Calculate Time Derivatives
# @todo - What about negative derivatives?
# @todo - Get them in the plot somehow; color code \pm?
delta_mass_total_lost_to_star = np.diff(mass_total_lost_to_star) / \
                                np.diff(time)
delta_mass_total_lost_to_halo = np.diff(mass_total_lost_to_halo) / \
                                np.diff(time)
delta_mass_disk = np.diff(mass_disk) / np.diff(time)
delta_mass_halo = np.diff(mass_halo) / np.diff(time)
delta_mass_total = np.diff(mass_total) / np.diff(time)

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

# ------------------------
# Plot Mass Rate of Change
# ------------------------
fig2 = plt.figure()
ax3 = fig2.add_subplot(1,1,1)
ax4 = ax3.twiny()
if args.disk:
    ax3.plot(time[1::]/2./pi, np.abs(2.*pi*delta_mass_disk), 'bs-', \
             label='abs(d(Disk Mass)/dt)')
if args.halo:
    ax3.plot(time[1::]/2./pi, 2.*pi*delta_mass_halo, 'rs-', label='d(Halo Mass)/dt')
if args.box:
    ax3.plot(time[1::]/2./pi, 2.*pi*delta_mass_total, 'gs-', label='d(Box Mass)/dt')
if args.lost_halo:
    ax3.plot(time[1::]/2./pi, 2.*pi*delta_mass_total_lost_to_halo, 'ms-', \
             label='d(Total Mass Lost to Halo)/dt')
if args.lost_star:
    ax3.plot(time[1::]/2./pi, 2.*pi*delta_mass_total_lost_to_star, 'ks-', \
             label='d(Total Mass Lost to Star)/dt')

ax4.plot(time, np.ones_like(time))
ax4.cla(); ax4.set_xlabel('Time [code]')
ax3.grid(); ax3.set_yscale('log')
ax4.set_ylim(ax3.get_ylim()); ax4.set_yscale(ax3.get_yscale())
ax3.set_xlabel('Time [yr]')
ax3.set_ylabel('d(Mass)/dt [Mstar/yr]')
ax3.legend(loc='best')
fig2.suptitle('Mass Rates of Change')

if args.save:
    fig1.savefig('mass_budget.pdf')
    fig2.savefig('mass_rate_of_change.pdf')
if args.show:
    plt.show()