import athenaReader3d as reader3d
import athenaTools as tools
import athenaReaderPhst as readerPhst
import planOutputReader as readerPlan
from matplotlib.backends.backend_pdf import PdfPages
import multiprocessing as mp
################################################################################
myNpc = int(sys.argv[1])
G = float(sys.argv[2])
pathBase = str(sys.argv[3])
dt3d = 1.0
################################################################################
path3d = pathBase + '3d/'
pathSave = pathBase + 'plots/masterAnim/'
if not os.path.exists(pathSave): os.makedirs(pathSave)
do3d = reader3d.Data3d(path3d, dt=dt3d)
pathPlan = pathBase + 'planOutput/'
if do3d.nt > 15: nAvgEnd = 10*int(1.0/dt3d)
else: nAvgEnd = 2
################################################################################
# 1 2 3 4 5 6
# npar mass r_hill xcom ycom zcom
planPlots = True
doPlan = readerPlan.DataPlan(pathPlan, G=G, dt=0.1)
#try:
#doPlan = readerPlan.DataPlan(pathPlan, G=G, dt=0.1)
#except:
#planPlots = False
if np.sum(do3d.get3d('dpar', 0)) > 1.e-20:
dparPlots = True
else:
import matplotlib.colors as colors
import os
import math
import sys
sys.path.append('../python')
import athenaReader3d as reader3d
import athenaTools as tools
from matplotlib.backends.backend_pdf import PdfPages
import multiprocessing as mp
################################################################################
pathBase = str(sys.argv[1])
################################################################################
path3d = pathBase + '3d/'
pathSave = pathBase + 'plots/plots3d/'
if not os.path.exists(pathSave): os.makedirs(pathSave)
do3d = reader3d.Data3d(path3d)
################################################################################
# simple averages of quantities
def realAvgProfile(do3d, key, pathSave):
reader3d.profile(do3d, key, figNum=0, absAvg=0, absPlot=0)
tools.saveAndClear(pathSave + 'profileRealAvg_' + key + '.png', figNum=0)
def realAvgTimeEvo(do3d, key, pathSave):
reader3d.timeEvo(do3d, key, figNum=0, absAvg=0, absPlot=0)
tools.saveAndClear(pathSave + 'timeEvoRealAvg_' + key + '.png', figNum=0)
def absAvgProfile(do3d, key, pathSave):
#colorList = ['r', 'g', 'm', 'b', 'k']
#pathSave = pathBase + 'plots/pspec_160/'
################################################################################
pathBase = '../../data/prodRuns/'
runNameList = ['run301_noPar', 'run302_noPar', 'run303_noPar']
labelList = ['lr_strong_np', 'lr_moderate_np', 'lr_weak_np']
colorList = ['r', 'g', 'b', 'm', 'k']
pathSave = pathBase + 'plots/pspec_300/'
################################################################################
if not os.path.exists(pathSave): os.makedirs(pathSave)
plt.figure(0)
################################################################################
do3dList = []
for n in range(len(runNameList)):
path3d = pathBase + runNameList[n] + '/3d/'
do3dList.append(reader3d.Data3d(path3d))
################################################################################
vExpo = -1.833
eExpo = 2.0*(vExpo)+2.0
nStart = 4
for n in range(len(do3dList)):
do3d = do3dList[n]
color = colorList[n]
psk_vx, freqs = reader3d.psProfileMean(do3d, 'rootRhoDvx', nStart=nStart)
psk_vy, freqs = reader3d.psProfileMean(do3d, 'rootRhoDvy', nStart=nStart)
psk_vz, freqs = reader3d.psProfileMean(do3d, 'rootRhoDvz', nStart=nStart)
psk = psk_vx + psk_vy + psk_vz
psk*=np.power(freqs, -eExpo)
psk/=np.mean(psk)
plt.loglog(freqs, psk, color=color, label=labelList[n])
