-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathsave_eddy_properties.py
135 lines (133 loc) · 5.23 KB
/
save_eddy_properties.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
125
126
127
128
129
130
131
132
133
134
import numpy as np
#import numpy.ma as ma
from netCDF4 import Dataset
#import matplotlib.pyplot as plt
#from scipy.interpolate import griddata
from scipy import interpolate
#import micro_inverse_utils as mutils
#import NorESM_utils as utils
#from mpl_toolkits.basemap import Basemap, addcyclic, interp, maskoceans
#from matplotlib.colors import BoundaryNorm, LogNorm, SymLogNorm, from_levels_and_colors
from joblib import Parallel, delayed
import tempfile
import shutil
import os
import dist
#
load=True
variables=['speed_average'] #['radius'] #['u','v']
scale=1
#
if load:
data=Dataset('/export/scratch/anummel1/EddyAtlas/eddy_trajectory_19930101_20160925.nc')
inds=np.where(np.diff(data.variables['track'][:]))[0]
#
inds0=np.arange(len(inds)).astype(np.int)
inds0[1:]=inds[:-1]+1
times=data['time'][inds]-data['time'][inds0]
#
times_all=data['time'][:]
lat_all=data['latitude'][:]
lon_all=data['longitude'][:]
lon_all[np.where(lon_all>180)]=lon_all[np.where(lon_all>180)]-360
#
lat=data.variables['latitude'][inds0]
lon=data.variables['longitude'][inds0]
#
lat_end=data.variables['latitude'][inds]
lon_end=data.variables['longitude'][inds]
#
#############################################
# --- CALCULATE THE SPEED OF THE CORE --- #
############################################
#
def calc_core_bin(j,lat_all,lon_all,inds0,inds,gy,gx,core_map,flag=None,vardat=None):
'''Bin eddy core properties'''
#
la = lat_all[int(inds0[j]):int(inds[j]+1)]
lo = lon_all[int(inds0[j]):int(inds[j]+1)]
#
if flag in ['u']:
var = np.sign(lo[1:]-lo[:-1])*dist.distance((la[:-1],lo[:-1]),(la[:-1],lo[1:]))*1E3/(24*3600.)
elif flag in ['v']:
var = np.sign(la[1:]-la[:-1])*dist.distance((la[:-1],lo[:-1]),(la[1:],lo[:-1]))*1E3/(24*3600.)
elif flag in ['radius', 'speed_average']:
var = vardat[int(inds0[j]):int(inds[j]+1)]
#elif flag in ['speed_average']
# var = data['speed_average'][int(inds0[j]):int(inds[j]+1)]
#
fy=interpolate.interp1d(gy,np.arange(720))
fx=interpolate.interp1d(gx,np.arange(1440))
y=fy(0.5*(la[1:]+la[:-1])).astype(int)
x=fx(0.5*(lo[1:]+lo[:-1])).astype(int)
#
for k in range(len(y)):
c=len(np.where(~np.isnan(core_map[:,y[k],x[k]]))[0])
if c>=core_map.shape[0]:
#just in case there would be a very large amount in the same bin
continue
core_map[c,y[k],x[k]]=var[k]
#
if j%1000==0:
print(j)
############################################################################################
# SORT OF NEAREST NEIGHBOUR APPROACH, WHERE VELOCITIES ARE BINNED IN A PRE-DETERMINED GRID #
#
for v,var in enumerate(variables):
print(var)
ny=int(720/scale)
nx=int(1440/scale)
ne=int(250*(scale**2))
grid_x, grid_y = np.mgrid[-180:180:complex(0,nx), -90:90:complex(0,ny)]
#
folder2 = tempfile.mkdtemp()
path1 = os.path.join(folder2, 'inds0.mmap')
path2 = os.path.join(folder2, 'inds.mmap')
path3 = os.path.join(folder2, 'lat_all.mmap')
path4 = os.path.join(folder2, 'lon_all.mmap')
path5 = os.path.join(folder2, 'core_map.mmap')
path7 = os.path.join(folder2, 'gx.mmap')
path8 = os.path.join(folder2, 'gy.mmap')
#
inds0_m = np.memmap(path1, dtype=float, shape=(len(inds)), mode='w+')
inds_m = np.memmap(path2, dtype=float, shape=(len(inds)), mode='w+')
lon_all_m = np.memmap(path3, dtype=float, shape=(len(lat_all)), mode='w+')
lat_all_m = np.memmap(path4, dtype=float, shape=(len(lat_all)), mode='w+')
core_map = np.memmap(path5, dtype=float, shape=(ne,ny,nx), mode='w+')
gx = np.memmap(path7, dtype=float, shape=(nx), mode='w+')
gy = np.memmap(path8, dtype=float, shape=(ny), mode='w+')
#
inds0_m[:] = inds0
inds_m[:] = inds
lon_all_m[:] = lon_all
lat_all_m[:] = lat_all
core_map[:] = np.ones((ne,ny,nx))*np.nan
gx[:] = grid_x[:,0]
gy[:] = grid_y[0,:]
if var in ['radius','speed_average']:
path9 = os.path.join(folder2, 'data.mmap')
vardat = np.memmap(path9, dtype=float, shape=(len(lat_all)), mode='w+')
if var in ['radius']:
vardat[:] = data['speed_radius'][:]
elif var in ['speed_average']:
vardat[:] = data['speed_average'][:]
else:
vardat=None
#
#SERIAL VERSION WORKS MUCH BETTER???
#for j in range(len(inds)):
# calc_core_bin(j,lat_all,lon_all,inds0,inds,gy,gx,core_map,flag=var,vardat=vardat)
num_cores=10
Parallel(n_jobs=num_cores)(delayed(calc_core_bin)(j,lat_all_m,lon_all_m,inds0_m,inds_m,gy,gx,core_map,flag=var,vardat=vardat) for j in range(len(inds)))
#
core_map = np.array(core_map)
core_count = core_map.copy(); core_count[np.where(~np.isnan(core_count))]=1; core_count=np.nansum(core_count,0)
mask=np.zeros(core_count.shape)
mask[np.where(core_count==0)]=1
#
np.savez('/home/anummel1/Projects/MicroInv/eddy_core_'+str(var)+'_binned_scale'+str(scale)+'.npz', grid_x=grid_x.T, grid_y=grid_y.T, var_grid=np.nanmean(core_map,0),core_count=core_count, mask=mask)
#
try:
shutil.rmtree(folder2)
except OSError:
pass