-
Notifications
You must be signed in to change notification settings - Fork 38
/
_ancillary_vars.py
210 lines (170 loc) · 6.25 KB
/
_ancillary_vars.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
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
205
206
207
208
209
210
"""Preprocessor functions for ancillary variables and cell measures."""
import logging
import dask.array as da
import iris
from esmvalcore.cmor.check import cmor_check_data, cmor_check_metadata
from esmvalcore.cmor.fix import fix_data, fix_metadata
from esmvalcore.preprocessor._io import concatenate, concatenate_callback, load
logger = logging.getLogger(__name__)
def _load_fx(var_cube, fx_info, check_level):
"""Load and CMOR-check fx variables."""
fx_cubes = iris.cube.CubeList()
project = fx_info['project']
mip = fx_info['mip']
short_name = fx_info['short_name']
freq = fx_info['frequency']
for fx_file in fx_info['filename']:
loaded_cube = load(fx_file, callback=concatenate_callback)
loaded_cube = fix_metadata(loaded_cube,
check_level=check_level,
**fx_info)
fx_cubes.append(loaded_cube[0])
fx_cube = concatenate(fx_cubes)
if not _is_fx_broadcastable(fx_cube, var_cube):
return None
fx_cube = cmor_check_metadata(fx_cube, cmor_table=project, mip=mip,
short_name=short_name, frequency=freq,
check_level=check_level)
fx_cube = fix_data(fx_cube, check_level=check_level, **fx_info)
fx_cube = cmor_check_data(fx_cube,
cmor_table=project,
mip=mip,
short_name=fx_cube.var_name,
frequency=freq,
check_level=check_level)
return fx_cube
def _is_fx_broadcastable(fx_cube, cube):
try:
da.broadcast_to(fx_cube.core_data(), cube.shape)
except ValueError as exc:
logger.debug("Dimensions of %s and %s cubes do not match. "
"Discarding use of fx_variable: %s",
cube.var_name, fx_cube.var_name, exc)
return False
return True
def add_cell_measure(cube, fx_cube, measure):
"""Broadcast fx_cube and add it as a cell_measure in the cube containing
the data.
Parameters
----------
cube: iris.cube.Cube
Iris cube with input data.
fx_cube: iris.cube.Cube
Iris cube with fx data.
measure: str
Name of the measure, can be 'area' or 'volume'.
Returns
-------
iris.cube.Cube
Cube with added ancillary variables
Raises
------
ValueError
If measure name is not 'area' or 'volume'.
"""
if measure not in ['area', 'volume']:
raise ValueError(f"measure name must be 'area' or 'volume', "
f"got {measure} instead")
try:
fx_data = da.broadcast_to(fx_cube.core_data(), cube.shape)
except ValueError:
logger.debug("Dimensions of %s and %s cubes do not match. "
"Cannot broadcast cubes.",
cube.var_name, fx_cube.var_name)
return
measure = iris.coords.CellMeasure(
fx_data,
standard_name=fx_cube.standard_name,
units=fx_cube.units,
measure=measure,
var_name=fx_cube.var_name,
attributes=fx_cube.attributes)
cube.add_cell_measure(measure, range(0, measure.ndim))
logger.debug('Added %s as cell measure in cube of %s.', fx_cube.var_name,
cube.var_name)
def add_ancillary_variable(cube, fx_cube):
"""Broadcast fx_cube and add it as an ancillary_variable in the cube
containing the data.
Parameters
----------
cube: iris.cube.Cube
Iris cube with input data.
fx_cube: iris.cube.Cube
Iris cube with fx data.
Returns
-------
iris.cube.Cube
Cube with added ancillary variables
"""
try:
fx_data = da.broadcast_to(fx_cube.core_data(), cube.shape)
except ValueError:
logger.debug("Dimensions of %s and %s cubes do not match. "
"Cannot broadcast cubes.",
cube.var_name, fx_cube.var_name)
return
ancillary_var = iris.coords.AncillaryVariable(
fx_data,
standard_name=fx_cube.standard_name,
units=fx_cube.units,
var_name=fx_cube.var_name,
attributes=fx_cube.attributes)
cube.add_ancillary_variable(ancillary_var, range(0, ancillary_var.ndim))
logger.debug('Added %s as ancillary variable in cube of %s.',
fx_cube.var_name, cube.var_name)
def add_fx_variables(cube, fx_variables, check_level):
"""Load requested fx files, check with CMOR standards and add the fx
variables as cell measures or ancillary variables in the cube containing
the data.
Parameters
----------
cube: iris.cube.Cube
Iris cube with input data.
fx_variables: dict
Dictionary with fx_variable information.
check_level: CheckLevels
Level of strictness of the checks.
Returns
-------
iris.cube.Cube
Cube with added cell measures or ancillary variables.
"""
if not fx_variables:
return cube
for fx_info in fx_variables.values():
if not fx_info:
continue
if isinstance(fx_info['filename'], str):
fx_info['filename'] = [fx_info['filename']]
fx_cube = _load_fx(cube, fx_info, check_level)
if fx_cube is None:
continue
measure_name = {
'areacella': 'area',
'areacello': 'area',
'volcello': 'volume'
}
if fx_cube.var_name in measure_name:
add_cell_measure(cube, fx_cube, measure_name[fx_cube.var_name])
else:
add_ancillary_variable(cube, fx_cube)
return cube
def remove_fx_variables(cube):
"""Remove fx variables present as cell measures or ancillary variables in
the cube containing the data.
Parameters
----------
cube: iris.cube.Cube
Iris cube with data and cell measures or ancillary variables.
Returns
-------
iris.cube.Cube
Cube without cell measures or ancillary variables.
"""
if cube.cell_measures():
for measure in cube.cell_measures():
cube.remove_cell_measure(measure.standard_name)
if cube.ancillary_variables():
for variable in cube.ancillary_variables():
cube.remove_ancillary_variable(variable.standard_name)
return cube