In [ ]:
#%pdb
%matplotlib inline
import matplotlib
matplotlib.rcParams['font.family'] = 'stixgeneral'
matplotlib.rcParams['figure.dpi'] = 150
import matplotlib.pyplot as plt
import os
import yt
yt.mylog.setLevel("ERROR")
import numpy as np
from synchrotron.yt_synchrotron_emissivity import *
from mpl_toolkits.axes_grid1 import AxesGrid, ImageGrid
import pyfits

In [ ]:
def setup_cl(dirs):
    # Set up colors and label names
    colors = {}
    labels = {}
    for dirname in dirs:
        if 'M3_h1' in dirname:
            colors[dirname] = 'pink'
            labels[dirname] = 'low Mach (3)'
        elif 'h1' in dirname:
            colors[dirname] = 'r'
            labels[dirname] = 'helical'
        elif 'h0' in dirname:
            colors[dirname] = 'b'
            labels[dirname] = 'poloidal'
        elif 'hinf' in dirname:
            colors[dirname] = 'g'
            labels[dirname] = 'toroidal'
        elif 'hydro' in dirname:
            colors[dirname] = 'k'
            labels[dirname] = 'hydro'
        elif 'M24_b01' in dirname:
            colors[dirname] = 'purple'
            labels[dirname] = 'low beta (0.01)'
    return colors, labels

In [ ]:
#dirs = ['/home/ychen/data/00only_0605_hinf/',\
#        '/home/ychen/data/00only_0529_h1/',\
#        '/home/ychen/data/00only_0605_h0/',]

dirs = ['/d/d5/ychen/2015_production_runs/0204_hinf_10Myr/',\
        '/d/d5/ychen/2015_production_runs/1022_h1_10Myr/',\
        '/d/d5/ychen/2015_production_runs/0204_h0_10Myr/',]


filenumbers = [100, 200, 600, 910, 1050]

labels = ['toroidal', 'helical', 'poloidal']

ptype = 'lobe'
zoom_fac = 6
proj_axis = 'x'
nu = (1400, 'MHz')
extend_cells = 32
res = (256, 128)

stokes = StokesFieldName(ptype, nu, proj_axis, field_type='flash')

field = stokes.I[1]

data = {}
for j, dir in enumerate(dirs):
    times = np.zeros(len(filenumbers))
    I_totals = np.zeros(len(filenumbers))
    for i, filenumber in enumerate(filenumbers):

        # Load the data and create a single plot
        ds = yt.load(os.path.join(dir, 'data/MHD_Jet_10Myr_hdf5_plt_cnt_%04d' % filenumber))
        times[i] = ds.current_time.in_units('Myr')
        print(dir, filenumber, ds.current_time.in_units('Myr'))
        
        
#         ds_sync = yt.load(synchrotron_file_name(ds, extend_cells=extend_cells))
#         ds_sync.field_list
#         ds_sync.field_info[field].units = 'Jy/cm/arcsec**2'
#         ds_sync.field_info[field].output_units = 'Jy/cm/arcsec**2'
        
#         box = ds_sync.box(ds.domain_left_edge/zoom_fac, ds.domain_right_edge/zoom_fac)
#         I_total = np.sum(box[field]*box['cell_volume'])
#        print(I_total)
#        I_totals[i] = I_total

        fitsname = synchrotron_fits_filename(ds, dir, ptype, proj_axis)
        if not os.path.isfile(fitsname): continue
        ds_fits = yt.load(fitsname)

        ad = ds_fits.all_data()
        I_totals[i] = np.sum(ad[field])
    data[dir] = (times, I_totals)
    #plt.plot(times, I_total, label=dir.split('/')[-1])


#     if field not in ds_sync.field_list: continue
#     if proj_axis in ['x','y','z']:
#         p = yt.ProjectionPlot(ds_sync, proj_axis, field, center=[0,0,0], width=width, max_level=6)
#     else:
#         p = yt.OffAxisProjectionPlot(ds_sync, proj_axis, field, center=[0,0,0], width=width, north_vector=[0,1,0])
    #p.set_buff_size(res)

In [ ]:
data_x = data.copy()

In [ ]:
colors, labels = setup_cl(dirs)
for dir, (time, I_total) in data_x.items():
    plt.plot(time, I_total, 'o-', label=labels[dir]+' x', c=colors[dir])
plt.legend()
plt.semilogy()
plt.xlabel('time (Myr)')
plt.ylabel('Total Synchrotron Intensity')
plt.title('150MHz')

In [ ]:
colors, labels = setup_cl(dirs)
for dir, (time, I_total) in data_x.items():
    plt.plot(time, I_total, 'o-', label=labels[dir]+' x', c=colors[dir])
for dir, (time, I_total) in data_30.items():
    plt.plot(time, I_total, '^:', label=labels[dir]+' 30 deg', c=colors[dir])
plt.plot(time, 2E6/time)
plt.legend()
plt.semilogy()
plt.semilogx()
plt.xlabel('time (Myr)')
plt.ylabel('Total Synchrotron Intensity')
plt.title('150MHz')

In [ ]:
colors, labels = setup_cl(dirs)
for dir, (time, I_total) in data_x.items():
    plt.plot(time, I_total, 'o-', label=labels[dir]+' x', c=colors[dir])
for dir, (time, I_total) in data_30.items():
    plt.plot(time, I_total, '^:', label=labels[dir]+' 30 deg', c=colors[dir])
plt.legend()
plt.semilogy()
plt.xlabel('time (Myr)')
plt.ylabel('Total Synchrotron Intensity')
plt.title('1400MHz')

In [ ]: