-
Notifications
You must be signed in to change notification settings - Fork 0
/
data.py
152 lines (116 loc) · 5.49 KB
/
data.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
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
import numpy as np
import xarray as xr
from scipy.io import netcdf
from scipy.interpolate import interp2d
from scipy.ndimage import maximum_filter
import pyproj
from pyresample.geometry import AreaDefinition
from pyresample.image import ImageContainerQuick
def load_grt(fname, name = 'altitude', dlat = None, dlon = None):
'''Load topographic data from grt file.
:param fname: Path to grt file
:param name: name of the returned DataArray
:param dlat: requested resolution in latitude (None = do not resample)
:param dlon: requested resolution in longitude (None = do not resample)
:return: xarray dataset with lon, lat as coordinates
'''
vars = netcdf.netcdf_file(fname, mmap=False).variables
z = vars['z'][:].copy().reshape(vars['dimension'][::-1])[::-1,:]
lon = np.linspace(*vars['x_range'][:].copy(), vars['dimension'][0])
lat = np.linspace(*vars['y_range'][:].copy(), vars['dimension'][1])
# TODO: check if dlon and dlat is consistent with vars['spacing']
d = xr.DataArray(z, coords={'lon': lon, 'lat': lat}, dims=['lat', 'lon'],
name=name)
d = resample(d, dlat, dlon)
# TODO: assuming the data are in webmercator projection - check
d.attrs['proj4'] = '+proj=merc +a=6378137 +b=6378137 +lat_ts=0.0 +lon_0=0.0 +x_0=0.0 +y_0=0 +k=1.0 +units=m +nadgrids=@null +wktext +no_defs'
return d
def resample(data, dx = None, dy = None):
''' Resample data using maximum filter
:param data: xr.DataArray
:param dx: resolution of the first dim
:param dy: resolution of the second dim
:return: resampled DataArray
'''
if dx is None and dy is None:
return data
dimx, dimy = data.dims
data_dx = data.coords[dimx][1] - data.coords[dimx][0]
data_dy = data.coords[dimy][1] - data.coords[dimy][0]
if dx is None or dx < 2*data_dx:
# no resampling in x
filt_size_x = 1
resample_dx = data_dx
else:
filt_size_x = int(np.ceil(dx / data_dx))
filt_size_x += 1 - (filt_size_x % 2) # make it odd
resample_dx = dx / filt_size_x
if dy is None or dy < 2*data_dy:
# no resampling in y
filt_size_y = 1
resample_dy = data_dy
else:
filt_size_y = int(np.ceil(dy / data_dy))
filt_size_y += 1 - (filt_size_y % 2) # make it odd
resample_dy = dy / filt_size_y
# resample
if resample_dx != data_dx or resample_dy != data_dy:
new_x = np.arange(data.coords[dimx][0], data.coords[dimx][-1]+resample_dx/2, resample_dx)
new_y = np.arange(data.coords[dimy][0], data.coords[dimy][-1]+resample_dy/2, resample_dy)
# TODO: replace with DataArray.interp after upgrading to newer xarrays
# TODO: use nearest neighbour
z = interp2d(data.coords[dimy], data.coords[dimx], data)
data = xr.DataArray(z(new_y, new_x), coords={dimx:new_x, dimy: new_y},
dims=[dimx, dimy], attrs = data.attrs)
# use 2d maximum filter and downsample
if filt_size_x !=1 or filt_size_y != 1:
new_data = maximum_filter(data, size=(filt_size_x,filt_size_y),
origin=(int((filt_size_x-1)/2), int((filt_size_y-1)/2)))[::filt_size_x, ::filt_size_y]
data = xr.DataArray(new_data, coords={dimx:data.coords[dimx][::filt_size_x],
dimy:data.coords[dimy][::filt_size_y]},
dims=[dimx, dimy], attrs=data.attrs)
return data
def _proj4_to_dict(proj4_str):
'''Convert proj4 string to dict
:param proj4_str: proj4 string
:return: proj4_dict
'''
# TODO: maybe such function already exists in pyresample?
res = {}
for param in proj4_str.split('+'):
k_v = param.split('=')
# ignore parameters that are not in key=value format
if len(k_v) != 2:
continue
res[k_v[0]] = k_v[1]
return res
def transform(d, area_def):
'''Transform data to new projection.
:param d: DataArray
:param area_def: destination AreaDefinition or proj4 dict
:return: d projected to area_def
'''
p1 = pyproj.Proj(proj='latlong', datum='WGS84')
p2 = pyproj.Proj(d.attrs['proj4'])
x0, y0 = pyproj.transform(p1, p2, d.lon[0], d.lat[0])
x1, y1 = pyproj.transform(p1, p2, d.lon[-1], d.lat[-1])
orig_area = AreaDefinition('orig_area', 'Original Area', 'orig_area',
proj_dict=_proj4_to_dict(d.attrs['proj4']),
x_size=d.shape[1], y_size=d.shape[0],
area_extent=[x0, y0, x1, y1])
if isinstance(area_def, dict):
pk = pyproj.Proj(area_def)
xk0, yk0 = pyproj.transform(p1, pk, d.lon[0], d.lat[0])
xk1, yk1 = pyproj.transform(p1, pk, d.lon[-1], d.lat[-1])
Nx = d.shape[1]# FIXME: this may not be a good estimate
area_def = AreaDefinition(area_def['proj'], '', proj_id='',
proj_dict=area_def, x_size=Nx,
y_size=int(np.abs(yk1-yk0)/np.abs(xk1-xk0)*Nx),
area_extent=[xk0, yk0, xk1, yk1])
# TODO: use ImageContainerNearest instead
imcont = ImageContainerQuick(d.values, orig_area)
proj_data = imcont.resample(area_def)
new_d = xr.DataArray(proj_data.image_data[::-1,:], coords={'x':area_def.proj_x_coords,
'y':area_def.proj_y_coords},
dims=['y','x'], attrs={'proj4': proj_data.geo_def.proj4_string})
return new_d