def diags_cp_cm(x, y): dd.read_parameters() diagdir = '/scratch/01658/drhatch/dna_out' par['diagdir'] = "\'" + diagdir + "\'" time = dd.get_time_from_gout() kx, ky, kz, hermiteNumbers = dd.get_grids() print 'kx =', np.shape(kx) print 'ky =', np.shape(ky) print 'kz =', np.shape(kz) print 'n =', np.shape(hermiteNumbers) g_in = dd.read_time_step_g(len(time) - 1) g_in = np.reshape(g_in, (par['nkx0'], par['nky0'], par['nkz0'], par['nv0']), order='F') x = x y = y dist = g_in[x, y, :, :] #distribution function: 2d slice of 4d g_in #make a new 2d array with the same dimensions as dist_ shape = np.shape(dist) dist_t = np.empty(shape) for i in range(len(kz)): for j in range(len(hermiteNumbers)): dist_t[i, j] = (1j * np.sign(kz[i]))**hermiteNumbers[j] * dist[i, j] dist_tp = np.empty(shape) dist_tm = np.empty(shape) for i in range(len(kz)): for j in range(len(hermiteNumbers) + 1): if j < (len(hermiteNumbers) - 1): dist_tp[i, j] = (dist_t[i, j] + dist_t[i, j + 1]) / 2 dist_tm[i, j] = (dist_t[i, j] - dist_t[i, j + 1]) / 2 c_p = np.empty(np.shape(dist_tp)) c_m = np.empty(np.shape(dist_tm)) for i in range(len(kz)): for j in range(len(kz)): c_p[i, j] = dist_tp[i, j] * np.conjugate(dist_tp[i, j]) c_m[i, j] = dist_tm[i, j] * np.conjugate(dist_tm[i, j]) c_ps = np.sum(c_p, axis=0) c_ms = np.sum(c_m, axis=0) xf = kx[x] yf = ky[y] return c_ps, c_ms, xf, yf
istart = np.argmin(abs(time - start_time)) iend = np.argmin(abs(time - end_time)) / 18 ntime = iend - istart + 1 #iend => 1/16 #final time=>1/18 #testing time =>1/27 entn_sum = np.zeros((par['nv0'], 96, 16), dtype='float') entnp_sum = np.zeros((par['nv0'], 96, 16), dtype='float') entnm_sum = np.zeros((par['nv0'], 96, 16), dtype='float') t_range = time[iend] - time[istart] print 'time range: ', t_range for i in range(istart, iend + 1): print 'time=', time[i], ' of ', time[iend] gt0 = dd.read_time_step_g(i) gt0 = np.reshape(gt0, (par['nkx0'], par['nky0'], par['nkz0'], par['nv0']), order='F') dist = gt0[:, :, :, :] shape = np.shape(dist) g_t = np.empty(shape) for v in range(len(kx)): for w in range(len(ky)): for b in range(len(kzgrid)): for j in range(len(herm_grid)): g_t[v, w, b, j] = (1j * np.sign( kzgrid[b]))**herm_grid[j] * dist[v, w, b, j] g_tp = np.empty(shape) g_tm = np.empty(shape) for v in range(len(kx)): for w in range(len(ky)):