diff options
Diffstat (limited to 'plot_stats.py')
| -rw-r--r-- | plot_stats.py | 145 |
1 files changed, 145 insertions, 0 deletions
diff --git a/plot_stats.py b/plot_stats.py new file mode 100644 index 0000000..0ec413c --- /dev/null +++ b/plot_stats.py @@ -0,0 +1,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() |