def zipfit_rhovn_to_psin(record):
for sig_name in zipfit_sig_names:
record['zipfit_{}_rhon_basis'.format(sig_name)] = standardize_time(
record['zipfit_{}_full'.format(sig_name)]['data'],
record['zipfit_{}_full'.format(sig_name)]['times'],
record['standard_time'])
rho_to_psi = [
my_interp(record['rhovn'][time_ind], record['rhovn_full']['psi'])
for time_ind in range(len(record['standard_time']))
]
record['zipfit_{}_psi'.format(sig_name)] = []
for time_ind in range(len(record['standard_time'])):
record['zipfit_{}_psi'.format(sig_name)].append(
rho_to_psi[time_ind](
record['zipfit_{}_full'.format(sig_name)]['rhon']))
record['zipfit_{}_psi'.format(sig_name)] = np.array(
record['zipfit_{}_psi'.format(sig_name)])
zipfit_interp = fit_function_dict['linear_interp_1d']
record['zipfit_{}'.format(sig_name)] = zipfit_interp(
record['zipfit_{}_psi'.format(sig_name)], record['standard_time'],
record['zipfit_{}_rhon_basis'.format(sig_name)],
np.ones(record['zipfit_{}_rhon_basis'.format(sig_name)].shape),
standard_x)
def linear_interp_1d(in_x, in_t, value, uncertainty, out_x):
final_sig = []
for time_ind in range(len(in_t)):
excluded_inds = np.isnan(value[time_ind, :])
y = value[time_ind, ~excluded_inds]
x = in_x[time_ind, ~excluded_inds]
err = uncertainty[time_ind, ~excluded_inds]
try:
get_value = my_interp(x, y, kind='linear')
final_sig.append(get_value(out_x))
except:
final_sig.append(np.zeros(len(out_x)))
final_sig = np.array(final_sig)
return final_sig
raw_data={}
for i in range(len(records)):
record=records[i]
shot=int(record['shot'])
final_data[shot]={}
final_data[shot]['t_ip_flat']=np.array(cfg['data']['tmin'])
final_data[shot]['ip_flat_duration']=np.array(cfg['data']['tmax'])
#final_data[shot]['topology']='SNB'
for sig in record.keys():
if sig in name_map:
# for handling zipfit
if 'rhon_basis' in sig:
final_data[shot][name_map[sig]]=[]
rhon=record['zipfit_{}_full'.format(cfg['data']['zipfit_sig_names'][0])]['rhon']
for time_ind in range(len(record['standard_time'])):
rho_to_zipfit=my_interp(rhon,
record[sig][time_ind])
final_data[shot][name_map[sig]].append(rho_to_zipfit(standard_x))
final_data[shot][name_map[sig]]=np.array(final_data[shot][name_map[sig]])
else:
final_data[shot][name_map[sig]]=np.array(record[sig])
if name_map[sig]=='curr_target':
final_data[shot][name_map[sig]]=final_data[shot][name_map[sig]]*0.5e6
if cfg['data']['gather_raw'] and 'full' in sig:
final_data[shot][sig]=record[sig]
# to accomodate the old code's bug of flipping top and bottom
if False:
tmp=final_data[shot]['triangularity_top_EFIT01'].copy()
final_data[shot]['triangularity_top_EFIT01']=final_data[shot]['triangularity_bot_EFIT01'].copy()
final_data[shot]['triangularity_bot_EFIT01']=tmp
final_data[shot]['pinj_full']=record['pinj_full']
final_data[shot]['tinj_full']=record['tinj_full']
# see pencil beam theory, Rome's 1974
# "Neutral-beam injection into a tokamak,
# part I: fast- ion spatial distribution for tangential injection"
# dNb / ds = - Nb / \lambda(s)
# \lambda(s) = 1 / n\sigma
r_z_to_psi = interpolate.interp2d(full_data[shot][efit_type]['R'],
full_data[shot][efit_type]['Z'],
data[shot]['psirz'][time_ind])
if fit_psi:
psi_to_density = interpolate.interp1d(
standard_psi, data[shot][density_sig_name][time_ind])
psi_to_temp = interpolate.interp1d(
standard_psi, data[shot][temp_sig_name][time_ind])
else:
rho_to_psi = my_interp(data[shot]['rho_grid'][time_ind], efit_psi)
deposition_psi = rho_to_psi(standard_rho)
psi_to_density = my_interp(deposition_psi,
data[shot][density_sig_name][time_ind])
psi_to_temp = my_interp(deposition_psi,
data[shot][temp_sig_name][time_ind])
iBeam = np.zeros(len(points))
iBeam[0] = iBeam0[time_ind]
for point_ind in range(len(points) - 1):
r = points[point_ind][0]
z = points[point_ind][1]
psi = r_z_to_psi(r, z)
if psi < np.max(deposition_psi):
psi_index = np.searchsorted(deposition_psi, psi)
def r_z_to_psi_2d_fit(entries):
channel_times, channel_rs, channel_zs = entries
psis = np.zeros(len(channel_times))
for i in range(len(channel_times)):
time_ind = np.searchsorted(standard_times, channel_times[i])
if time_ind == len(standard_times):
time_ind -= 1
try:
psis[i] = r_z_to_psi[time_ind](channel_rs[i], channel_zs[i])
except:
psis[i] = r_z_to_psi[time_ind](channel_rs, channel_zs)
return np.array(psis)
rho_to_psi = [
my_interp(data[shot][efit_type]['rho_grid'][time_ind], efit_psi)
for time_ind in range(len(standard_times))
]
psi_to_rho = [
my_interp(efit_psi, data[shot][efit_type]['rho_grid'][time_ind])
for time_ind in range(len(standard_times))
]
psi_to_rho_2d_fit = interp_ND_rectangular(
(data[shot][efit_type]['time'], efit_psi),
data[shot][efit_type]['rho_grid'])
def psi_to_rho_2d_fit(entries):
time_arr, psi_arr = entries
rhos = np.zeros(len(time_arr))
psi_grid = np.array(data[shot]['psirz'])
rho_grid = np.array(data[shot]['rho_grid'])
basis = np.linspace(.025, .975, 20)
fit_psi = False
efit_psi = np.linspace(0, 1, 65)
num_psi = len(basis)
dr = np.diff(r)
dz = np.diff(z)
dv = np.zeros((psi_grid.shape[0], num_psi - 1))
G1 = np.zeros((psi_grid.shape[0], num_psi - 1))
for time_ind in range(psi_grid.shape[0]):
if fit_psi:
basis_psi = basis
psi_to_rho = my_interp(efit_psi, rho_grid[time_ind])
basis_rho = psi_to_rho(basis)
else:
basis_rho = basis
rho_to_psi = my_interp(rho_grid[time_ind], efit_psi)
basis_psi = rho_to_psi(basis)
dpsi = np.diff(basis_psi)
drho = np.diff(basis_rho)
dRho_dPsi = np.divide(dpsi, drho)
for psi_ind in range(num_psi - 1):
gradient = np.gradient(psi_grid[time_ind])
psi = basis_psi[psi_ind] #dpsi[psi_ind]*psi_ind
hull = np.where(
np.logical_and(psi_grid[time_ind] < psi + dpsi[psi_ind],
