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
import xarray as xr
import logging
def get_pv_field_name(ds):
v_names = ["pv", "PV", "pv_avg200-500"]
for v_name in v_names:
if v_name in ds.data_vars:
return v_name
return None
def pv_to_pvu(ds, config):
pv_str = get_pv_field_name(ds)
if pv_str is None:
logging.error("Could not detect PV field in dataset. Expected pv, PV")
exit(1)
if config.pv_data_unit == 'pv':
ds = ds.assign(pvu=lambda d: d[pv_str] * 1e6)
elif config.pv_data_unit == 'pvu':
ds = ds.assign(pvu=lambda d: d[pv_str] * 1)
else:
raise Exception("Unknown pv_data_unit.")
return ds
def pvu_to_pvu2(ds):
# logic: pvu>-500(def) AND ((>0N and pvu>2) OR (<0N and pvu<-2))
ds['pvu2'] = xr.ufuncs.logical_and(ds.pvu > -500.0, xr.ufuncs.logical_or(ds.pvu > 2.0, ds.pvu < -2.0))
return ds
def compute_global_fields(dataset: xr.Dataset, config):
# from PV -> PVU -> PVU2 (binary thresholding 2pvu)
dataset = pv_to_pvu(dataset, config)
dataset = pvu_to_pvu2(dataset)
return dataset
def filter_dataset(dataset: xr.Dataset, theta_range, layer_2d):
from enstools.feature.util.enstools_utils import get_latitude_dim, get_vertical_dim
l_dim = get_latitude_dim(dataset)
# also make sure top is north pole
dataset = dataset.sortby(l_dim)
# TODO only nh
if dataset.coords[l_dim].values.min() < 0:
print("Emit southern hemisphere, currently only concepted for NH.")
dataset = dataset.sel({l_dim: slice(0, 90)})
# dataset = dataset.sortby(l_dim, ascending=False)
z_dim = get_vertical_dim(dataset)
if theta_range is not None:
dataset = dataset.sel({z_dim: slice(*theta_range)})
elif layer_2d is not None:
dataset = dataset.sel({z_dim: layer_2d})
return dataset