forked from mmiller04/ml_collisions-python
-
Notifications
You must be signed in to change notification settings - Fork 0
/
get_stats.py
125 lines (95 loc) · 3.58 KB
/
get_stats.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
# -*- coding: utf-8 -*-
"""get_stats.ipynb
Automatically generated by Colaboratory.
Original file is located at
https://colab.research.google.com/drive/1-45OoiY1vB5rR_CTelK8EKOZ2cDNH0MT
"""
import h5py
import numpy as np
nsp = 2
hf_f = h5py.File('/scratch/gpfs/marcoam/ml_collisions/data/xgc1/ti272_JET_heat_load/00094/hdf_f.h5','r')
hf_df = h5py.File('/scratch/gpfs/marcoam/ml_collisions/data/xgc1/ti272_JET_heat_load/00094/hdf_df.h5','r')
nphi,nperp,nnode,npar = hf_f.get('i_f').shape
nnode = 150000 # this is the same lim variable as in the main file
mean_f = np.zeros([nsp,nperp,nperp])
mean_df = np.zeros([nsp,nperp,nperp])
mean_fdf = np.zeros([nsp,nperp,nperp])
std_f = np.zeros([nsp,nperp,nperp])
std_df = np.zeros([nsp,nperp,nperp])
std_fdf = np.zeros([nsp,nperp,nperp])
mean_each_f = np.zeros([nsp,nperp,nperp])
mean_each_df = np.zeros([nsp,nperp,nperp])
mean_each_fdf = np.zeros([nsp,nperp,nperp])
std_each_f = np.zeros([nsp,nperp,nperp])
std_each_df = np.zeros([nsp,nperp,nperp])
std_each_fdf = np.zeros([nsp,nperp,nperp])
for iphi in range(nphi):
print(iphi)
e_f = hf_f['e_f'][iphi]
i_f = hf_f['i_f'][iphi]
# replace negative values w/ 0
e_f_neg = np.where(e_f < 0)
e_f[e_f_neg] = 0
i_f_neg = np.where(i_f < 0)
i_f[i_f_neg] = 0
e_df = hf_df['e_df'][iphi]
i_df = hf_df['i_df'][iphi]
# get lists where df = 0
e_df_zero_inds = np.where(np.einsum('ijk -> j',e_df) == 0)
e_df_zero_inds = list(e_df_zero_inds[0])
i_df_zero_inds = np.where(np.einsum('ijk -> j',i_df) == 0)
i_df_zero_inds = list(i_df_zero_inds[0])
e_inds = list(np.arange(nnode))
for e in e_df_zero_inds:
if e < nnode:
e_inds.remove(e)
i_inds = list(np.arange(nnode))
for i in i_df_zero_inds:
if i < nnode:
i_inds.remove(i)
# remove them
i_f = i_f[:,i_inds,:]
i_df = i_df[:,i_inds,:]
e_f = e_f[:,e_inds,:]
e_df = e_df[:,e_inds,:]
e_fdf = e_f+e_df
i_fdf = i_f+i_df
for i in range(nperp):
for j in range(npar):
mean_each_f[0,i,j] = np.mean(e_f[i,:,j])
mean_each_f[1,i,j] = np.mean(i_f[i,:,j])
mean_each_df[0,i,j] = np.mean(e_df[i,:,j])
mean_each_df[1,i,j] = np.mean(i_df[i,:,j])
mean_each_fdf[0,i,j] = np.mean(e_fdf[i,:,j])
mean_each_fdf[1,i,j] = np.mean(i_fdf[i,:,j])
std_each_f[0,i,j] = np.std(e_f[i,:,j])
std_each_f[1,i,j] = np.std(i_f[i,:,j])
std_each_df[0,i,j] = np.std(e_df[i,:,j])
std_each_df[1,i,j] = np.std(i_df[i,:,j])
std_each_fdf[0,i,j] = np.std(e_fdf[i,:,j])
std_each_fdf[1,i,j] = np.std(i_fdf[i,:,j])
mean_each_f[:,i,nperp-1] = mean_each_f[:,i,npar-1]
mean_each_df[:,i,nperp-1] = mean_each_df[:,i,npar-1]
mean_each_fdf[:,i,nperp-1] = mean_each_fdf[:,i,npar-1]
std_each_f[:,i,nperp-1] = std_each_f[:,i,npar-1]
std_each_df[:,i,nperp-1] = std_each_df[:,i,npar-1]
std_each_fdf[:,i,nperp-1] = std_each_fdf[:,i,npar-1]
mean_f += mean_each_f
mean_df += mean_each_df
mean_fdf += mean_each_fdf
std_f += std_each_f**2
std_df += std_each_df**2
std_fdf += std_each_fdf**2
mean_f = mean_f/nphi
mean_df = mean_df/nphi
mean_fdf = mean_fdf/nphi
std_f = np.sqrt(std_f/nphi)
std_df = np.sqrt(std_df/nphi)
std_fdf = np.sqrt(std_fdf/nphi)
with h5py.File('/scratch/gpfs/marcoam/ml_collisions/data/xgc1/ti272_JET_heat_load/00094/hdf_stats.h5','w') as hf:
hf.create_dataset('mean_f',data=mean_f)
hf.create_dataset('mean_df',data=mean_df)
hf.create_dataset('mean_fdf',data=mean_fdf)
hf.create_dataset('std_f',data=std_f)
hf.create_dataset('std_df',data=std_df)
hf.create_dataset('std_fdf',data=std_fdf)