diff --git a/script/analysis/tracers_to_PRISM.py b/script/analysis/tracers_to_PRISM.py index d2245f5..e67ae86 100755 --- a/script/analysis/tracers_to_PRISM.py +++ b/script/analysis/tracers_to_PRISM.py @@ -1,7 +1,7 @@ #!/usr/bin/env python # Little script to convert tracer data to a form suitable for PRISM # Author: Jonah Miller (jonah.maxwell.miller@gmail.com) -# Based on the work of Matt Mumpower +# Based on the work of Matt Mumpower, mods by Kels from __future__ import print_function, division import sys; sys.dont_write_bytecode = True @@ -174,23 +174,19 @@ def frdm_comp(Ye, T9, rho, filename, frdm_path): def signal_smooth(x,window_len=11,window='hanning'): """From scipy cookbook: https://scipy-cookbook.readthedocs.io/items/SignalSmooth.html - input: - x: the input signal - window_len: the dimension of the smoothing window; - should be an odd integer - window: the type of window from - 'flat', 'hanning', 'hamming', 'bartlett', 'blackman' - flat window will produce a moving average smoothing. - + x: the input signal + window_len: the dimension of the smoothing window; + should be an odd integer + window: the type of window from + 'flat', 'hanning', 'hamming', 'bartlett', 'blackman' + flat window will produce a moving average smoothing. output: - the smoothed signal - + the smoothed signal example: - - t=linspace(-2,2,0.1) - x=sin(t)+randn(len(t))*0.1 - y=smooth(x) + t=linspace(-2,2,0.1) + x=sin(t)+randn(len(t))*0.1 + y=smooth(x) """ if x.ndim != 1: raise ValueError("smooth only accepts 1 dimension arrays.") @@ -259,7 +255,6 @@ def value_crossing(a, v): second = residual[1:] return np.logical_and((first*second)<=0,first > 0) - def cleanup_trace(trace, T9 = 10, atol = 1, p = 6, thresh = 1e-1): @@ -267,56 +262,63 @@ def cleanup_trace(trace, # output trace trsm,trace_out = {},{} Tunit = cgs['MEV']/cgs['GK'] - trsm['T'],trsm['rho'] = log_smooth(trace['T']),log_smooth(trace['rho']) - + # starting temperature - #Tgk = trace['T']*cgs['MEV']/cgs['GK'] - Tgk = trsm['T']*cgs['MEV']/cgs['GK'] + Tgk = trsm['T']*Tunit TgkT9 = value_crossing(Tgk, T9) - - try: - sidx = np.where(TgkT9)[0][-1] + 1 - # Sometimes the temperature is not actually that close to T9=10, - # so this will do a little interpolation to find a suitable - # starting point. - if T9 > Tgk[sidx:][0] > T9-atol or Tgk[0] < T9: - if Tgk[0] < T9: - print("WARNING: Tracer {} starting at temperature {} < T9 = {}".format(trace['id'][0], Tgk[0], T9)) - sidx = 0 - for k in trace.keys() - ['T','rho']: - trace_out[k] = trace[k][sidx:] - trace_out['T'],trace_out['rho'] = trsm['T'][sidx:],trsm['rho'][sidx:] - else: - interpsidx = interpolate.interp1d(trace['time'][sidx-1:sidx+1],Tgk[sidx-1:sidx+1]) - residual = lambda t: interpsidx(t) - T9 - root_results = optimize.root_scalar(residual, bracket=(trace['time'][sidx-1],trace['time'][sidx+1])) - if not root_results.converged: - raise ValueError("Root finding failed for tracer {}".format(trace['id'][0])) - new_time = root_results.root - new_val = interpsidx(new_time) - trace_out['T'] = np.insert(trsm['T'][sidx:],0,new_val/Tunit) - trace_out['time'] = np.insert(trace['time'][sidx:],0,new_time) - trace_out['time'] = np.insert(trace['time'][sidx:],0,newx[ind]) - # Still need to get corresponding other key values - for k in trace.keys() - ['time','T']: - trace_in = trsm if k in trsm.keys() else trace - interpsidx = interpolate.interp1d(trace['time'][sidx-1:sidx+1],trace_in[k][sidx-1:sidx+1],axis=0) - trace_out[k] = np.insert(trace_in[k][sidx:],0,interpsidx(new_time),axis=0) - except (ValueError, IndexError) as error: - print("WARNING: Tracer {} starting at temperature {} < T9 = {}".format(trace['id'][0], Tgk[0], T9)) + # Need to account for possibility that there is no T9 crossing point + # in which case TgkT9 will just be an array full of False values + if len(np.where(TgkT9)[0]) == 0: sidx = 0 - for k in trace.keys() - ['T','rho']: - trace_out[k] = trace[k][sidx:] - trace_out['T'],trace_out['rho'] = trsm['T'][sidx:],trsm['rho'][sidx:] - + print('No T > T9 = ',T9,' exists. Starting at T = ',Tgk[sidx]) + else: + sidx = np.where(TgkT9)[0][-1] + 1 + print('Initial is < T9 of',T9,'; Starting at crossing point T = ',Tgk[sidx]) + + # Sometimes the temperature is not actually that close to T9=10, + # so this will do a little interpolation to find a suitable + # starting point. + if T9 > Tgk[sidx:][0] > T9-atol or Tgk[0] < T9: + print(Tgk[sidx:][0]) + for k in trace.keys() - ['T','rho']: + trace_out[k] = trace[k][sidx:] + trace_out['T'],trace_out['rho'] = trsm['T'][sidx:],trsm['rho'][sidx:] + else: + interpsidx = interpolate.interp1d(trace['time'][sidx-1:sidx+1],Tgk[sidx-1:sidx+1]) + newx = np.arange(trace['time'][sidx-1],trace['time'][sidx],0.1) + tgk_interp = interpsidx(newx) + ind = np.where(np.isclose(tgk_interp,T9,atol=0.1))[0][-1] + trace_out['T'] = np.insert(trsm['T'][sidx:],0,tgk_interp[ind]/Tunit) + trace_out['time'] = np.insert(trace['time'][sidx:],0,newx[ind]) + # Still need to get corresponding other key values + for k in trace.keys() - ['time','T']: + if trace[k].size == len(trace): + if k == 'rho': + interpsidx = interpolate.interp1d(trace['time'][sidx-1:sidx+1],trsm[k][sidx-1:sidx+1]) + interp = interpsidx(newx) + trace_out[k] = np.insert(trsm[k][sidx:],0,interp[ind]) + else: + interpsidx = interpolate.interp1d(trace['time'][sidx-1:sidx+1],trace[k][sidx-1:sidx+1]) + interp = interpsidx(newx) + trace_out[k] = np.insert(trace[k][sidx:],0,interp[ind]) + else: + dim = int(trace[k].size/len(trace)) + temp,trace_out[k] = {},[] + for i in range(dim): + interpsidx = interpolate.interp1d(trace['time'][sidx-1:sidx+1],trace[k][:,i][sidx-1:sidx+1]) + interp = interpsidx(newx) + temp[i] = np.insert(trace[k][sidx:][:,i],0,interp[ind]) + trace_out[k].append(temp[i]) + trace_out[k] = np.array(trace_out[k]) + if trace_out[k].shape[0] == dim: + trace_out[k] = np.reshape(trace_out[k],(trace_out[k].shape[1],trace_out[k].shape[0])) + # Cut off at the end cutoff = get_cutoff(trace_out,p,thresh) for k in trace.keys(): trace_out[k] = trace_out[k][:cutoff] - - # smooth rho and T - trace_out['Tgk'] = trace_out['T']*cgs['MEV']/cgs['GK'] + trace_out['Tgk'] = trace_out['T'] * Tunit # heavies don't matter nu_ab = (np.abs(trace_out['rate_absorbed'][:,0]) @@ -332,9 +334,7 @@ def cleanup_trace(trace, ye_cutoff = 0 ye_avg = np.mean(trace_out['Ye'][ye_cutoff:]) ye_std = np.std(trace_out['Ye'][ye_cutoff:]) - # trace_out['Ye_avg'] = trace_out['Ye'].copy() trace_out['Ye_avg'] = ye_avg*np.ones_like(trace_out['Ye']) - # trace_out['Ye_avg'] = signal_smooth(trace_out['Ye_avg'],11) # R trace_out['R'] = np.sqrt(trace_out['Xcart'][:,0]**2 @@ -363,27 +363,17 @@ def to_file(trace, header=header, comments = '# ') -def meta_file(idx, mass, mass_unit, T_initx, Tmax, - rho_initx = None,total_mass = None): +def meta_file(idx,mass,mass_unit,T_initx,rho_initx,total_mass = None): if total_mass is not None: mass_fraction = mass / total_mass mass *= mass_unit with open('metadata.dat','w') as f: - if rho_initx is not None: if total_mass is not None: - f.write("# 1:id 2:mass[g] 3:mass fraction 4:T_initx[GK] 5:rho_initx[g/cm3] 6:T_max[GK]\n") - f.write("{} {:.6e} {:.6e} {:.6e} {:.6e} {:.6e}\n".format(idx,mass,mass_fraction,T_initx,rho_initx, Tmax)) + f.write("# 1:id 2:mass[g] 3:mass fraction 4:T_initx[GK] 5:rho_initx[g/cm3]\n") + f.write("{} {:.6e} {:.6e} {:.6e} {:.6e}\n".format(idx,mass,mass_fraction,T_initx,rho_initx)) else: - f.write("# 1:id 2:mass[g] 3:T_smooth[GK] 4:rho_initx[g/cm3] 5: T_max[GK]\n") - f.write("{} {:.6e} {:.6e} {:.6e} {:.6e}\n".format(idx,mass,T_initx,rho_initx, Tmax)) - else: - if total_mass is not None: - f.write("# 1:id 2:mass[g] 3:mass fraction 4:T_initx[GK] 5: T_max[GK]\n") - f.write("{} {:.6e} {:.6e} {:.6e} {:.6e}\n".format(idx,mass,mass_fraction,T_initx, Tmax)) - else: - f.write("# 1:id 2:mass[g] 3:T_smooth[GK] 4: T_max[GK]\n") - f.write("{} {:.6e} {:.6e} {:.6e}\n".format(idx,mass,T_initx, Tmax)) - + f.write("# 1:id 2:mass[g] 3:T_smooth[GK] 4:rho_initx[g/cm3]\n") + f.write("{} {:.6e} {:.6e} {:.6e}\n".format(idx,mass,T_initx,rho_initx)) def extrapolate_trace(trace,T9,atol,p,thresh,tfinal = 1e3): # extrapolation @@ -427,14 +417,11 @@ def convert_file(tracer_file,outdir, os.makedirs(directory,exist_ok=True) os.chdir(directory) to_file(trace) - meta_file(trace['id'],trace['mass'],mass_unit, - trace['T_smooth'],trace['Tgk'].max(), trace['rho_smooth']) + meta_file(trace['id'],trace['mass'],mass_unit,trace['T_smooth'],trace['rho_smooth']) if frdm_path is not None: Ye = trace['Ye_avg'][0] - #rho = trace['rho'][0] - rhosmooth = trace['rho_smooth'] - Tsmooth = trace['T_smooth'] - #frdm_comp(Ye,T9,rho,'initx.dat',frdm_path) + rhosmooth = trace['rho_smooth'] # This is a single value, not an array + Tsmooth = trace['T_smooth'] # This is a single value, not an array frdm_comp(Ye,Tsmooth,rhosmooth,'initx.dat',frdm_path) print('initx file written.') os.chdir(currdir) @@ -473,7 +460,6 @@ def convert_ids(tracers,directory, os.chdir(directory) for i,idx in enumerate(ids): print("\t...",i,idx) - trace = tracers.get_trace(idx) trace = extrapolate_trace(trace,T9,atol,p,thresh, tfinal) istr = "{:08d}".format(i) @@ -481,8 +467,7 @@ def convert_ids(tracers,directory, os.mkdir(istr) os.chdir(istr) to_file(trace) - meta_file(idx,trace['mass'],mass_unit, trace['T_smooth'], trace['Tgk'].max(), - trace['rho_smooth'], total_mass) + meta_file(idx,trace['mass'],mass_unit,total_mass) if frdm_path is not None: Ye = trace['Ye_avg'][0] rho = trace['rho'][0] @@ -493,8 +478,7 @@ def convert_ids(tracers,directory, def convert_path_by_time(inpath,outdir, T9 = 10, atol = 0.3, p = 6, thresh = 1e-1, - frdm_path = None, - tfinal = 1e3): + frdm_path = None): tracers = io.TracerData.frompath(inpath) convert_ids(tracers,outdir, T9=T9,atol=atol,