Newer
Older
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
import logging
import six
from enstools.misc import is_additional_coordinate_variable
import numpy as np
# this one is already in enstools.misc, but not in the version on the w2w cluster. simple copy.
def distance(lat1, lat2, lon1, lon2, radius=6371229.0):
"""
distance between two points on a globe.
Parameters
----------
lat1:
latitude of first point in radian.
lat2:
latitude of second point in radian.
lon1:
longitude of first point in radian.
lon2:
logitude of second point in radian.
radius:
radius of the globe in m.
Returns
-------
distance in m
"""
dlon = lon2 - lon1
dlat = lat2 - lat1
a = np.sin(dlat / 2)**2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon / 2)**2
c = 2 * np.arctan2(np.sqrt(a), np.sqrt(1 - a))
return radius * c
def get_latitude_dim(ds):
"""
get the name of the latitude dimension from a dataset or array
Parameters
----------
ds : xarray.Dataset or xarray.DataArray
Returns
-------
str or None :
if no latitude dimension was found, None is returned.
"""
lat_names = ["lat", "latitude"]
for lat_name in lat_names:
if lat_name in ds.dims:
logging.debug("get_latitude_dim: found name '%s'" % lat_name)
return lat_name
return None
def get_longitude_dim(ds):
"""
get the name of the longitude dimension from a dataset or array
Parameters
----------
ds : xarray.Dataset or xarray.DataArray
Returns
-------
str or None :
if no longitude dimension was found, None is returned.
"""
lon_names = ["lon", "longitude"]
for lon_name in lon_names:
if lon_name in ds.dims:
logging.debug("get_longitude_dim: found name '%s'" % lon_name)
return lon_name
return None
# TODO get_X_dim functions can be written as one single abstract getter
def get_vertical_dim(ds):
"""
get the name of the vertical dimension from a dataset or array.
Needs to be expanded.
Parameters
----------
ds : xarray.Dataset or xarray.DataArray
Returns
-------
str or None :
if no vertical dimension was found, None is returned.
"""
v_names = ["pres", "p", "lev", "level", "isobaric", "layer", "hybrid", "height", "height_2", "th_level", "plev"]
for v_name in v_names:
if v_name in ds.dims:
logging.debug("get_vertical_dim: found name '%s'" % v_name)
return v_name
return None
def get_init_time_dim(ds):
"""
get the name of the init_time dimension from a dataset or array.
Parameters
----------
ds : xarray.Dataset or xarray.DataArray
Returns
-------
str or None :
if no init_time dimension was found, None is returned.
"""
v_names = ["init_time"]
for v_name in v_names:
if v_name in ds.dims:
logging.debug("get_longitude_dim: found name '%s'" % v_name)
return v_name
return None
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
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
def add_time_dim(ds, time, inplace=True):
"""
create an ensemble dimension with in the dataset
Parameters
----------
ds : xarray.Dataset
time : datetime
time of the dataset
inplace : bool
modify the dataset directly?
Returns
-------
xarray.Dataset:
A copy of the dataset expanded by the time dimension or the expanded dataset
if the expansion was done inplace.
"""
# create a lazy copy of the dataset
if inplace:
new_ds = ds
else:
new_ds = ds.copy()
if time is None:
import datetime
time = datetime.datetime.now()
# loop over all data variables.
# those with time dimension are extended behind the time dimension, others at the front
for one_name, one_var in six.iteritems(ds.data_vars):
if is_additional_coordinate_variable(one_var):
continue
new_ds[one_name] = one_var.expand_dims("time")
new_ds.coords["time"] = [time]
return new_ds
def add_vertical_dim(ds, vertical_value=0, vertical_name='level', inplace=True):
"""
create a vertical dimension with in the dataset
Parameters
----------
ds : xarray.Dataset
vertical_value : float
value for the added vertical dimension of the dataset, default 0
vertical_name : str
name for the added vertical dimension of the dataset, default 'level'
inplace : bool
modify the dataset directly?
Returns
-------
xarray.Dataset:
A copy of the dataset expanded by the vertical dimension or the expanded dataset
if the expansion was done inplace.
"""
# create a lazy copy of the dataset
if inplace:
new_ds = ds
else:
new_ds = ds.copy()
# loop over all data variables.
# those with time dimension are extended behind the time dimension, others at the front
for one_name, one_var in six.iteritems(ds.data_vars):
if is_additional_coordinate_variable(one_var):
continue
if len(one_var.dims) > 0 and one_var.dims[0] == "time":
new_ds[one_name] = one_var.expand_dims(vertical_name, 1)
else:
new_ds[one_name] = one_var.expand_dims(vertical_name)
new_ds.coords[vertical_name] = [vertical_value]
return new_ds