Skip to content

Commit 8135c22

Browse files
author
Liam Thompson
committed
updating cesm-se reader with some of peizhi's code
1 parent c2e5ae1 commit 8135c22

2 files changed

Lines changed: 245 additions & 23 deletions

File tree

‎monetio/models/_cesm_se_mm.py‎

Lines changed: 97 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -10,10 +10,8 @@ def open_mfdataset(
1010
earth_radius=6370000,
1111
convert_to_ppb=True,
1212
var_list=["O3", "NO", "NO2", "lat", "lon"],
13-
scrip_file="",
14-
grid_file=None,
15-
**kwargs
16-
):
13+
scrip_file= None,
14+
**kwargs):
1715
"""Method to open multiple (or single) CESM SE netcdf files.
1816
This method extends the xarray.open_mfdataset functionality
1917
It is the main method called by the driver. Other functions defined
@@ -45,22 +43,32 @@ def open_mfdataset(
4543
# check that the files are netcdf format
4644
names, netcdf = _ensure_mfdataset_filenames(fname)
4745

46+
if scrip_file is None:
47+
raise ValueError(
48+
"CESM-SE requires a scrip_file (set 'scrip_file:' in your YAML).")
49+
ux_grid_path = scrip_file
50+
4851
# open the dataset using xarray
52+
# try:
53+
# if ux_grid_path:
54+
# print(f"Opening unstructured grid with UXArray: {ux_grid_path}")
55+
# dset_load = ux.open_mfdataset(ux_grid_path, fname, **kwargs)
56+
# elif netcdf:
57+
# print("Opening Xarray...")
58+
# dset_load = xr.open_mfdataset(fname, **kwargs)
59+
# else:
60+
# raise ValueError(
61+
# "File format not recognized. Files should be in netcdf format; "
62+
# "do not mix file types."
63+
# )
64+
# except Exception as e:
65+
# print("ERROR while opening dataset:")
66+
# print(repr(e))
67+
# raise
68+
4969
try:
50-
if grid_file is not None:
51-
print("Opening unstructured grid with UXArray...")
52-
#uxds = ux.open_dataset(grid_file, fname, **kwargs)
53-
#dset_load = uxds
54-
dset_load = ux.open_mfdataset(grid_file, fname, **kwargs)
55-
# may need ux.open_dataset later
56-
elif netcdf:
57-
print("Opening Xarray...")
58-
dset_load = xr.open_mfdataset(fname, **kwargs)
59-
else:
60-
raise ValueError("File format not recognized. "
61-
"Note that files should be in netcdf format. "
62-
"Do not mix and match file types.")
63-
70+
print(f"Opening unstructured grid with UXArray: {ux_grid_path}")
71+
dset_load = ux.open_mfdataset(ux_grid_path, fname, **kwargs)
6472
except Exception as e:
6573
print("ERROR while opening dataset:")
6674
print(repr(e))
@@ -74,20 +82,89 @@ def open_mfdataset(
7482
if "lev" not in var_list:
7583
var_list.append("lev")
7684

85+
# variables for cesm-se specific derivations
86+
_cesm_se_var = []
87+
88+
# Always request the source variables needed for standardized derivations
89+
# Filtered against what's actually present, so a missing one just
90+
# skips the corresponding derivation
91+
92+
for _v in ("hyam", "hybm", "PS", "P0", "T", "PDELDRY"):
93+
if _v not in var_list:
94+
var_list.append(_v)
95+
_cesm_se_var.append(_v)
96+
97+
# filter to vars present and then warn about vars missing rather than generic rename error
98+
_requested = list(var_list)
99+
_present = [v for v in _requested if v in dset_load.variables]
100+
_missing = [v for v in _requested if v not in dset_load.variables]
101+
102+
if _missing:
103+
print(
104+
f"CESM-SE: requested variables not found in {names[0]!r}: "
105+
f"{_missing}. Continuing with what's available: {_present}."
106+
)
107+
dset = dset_load[_present]
108+
77109
# ===========================
78110
# Process the loaded data
79111
# extract variables of choice
80-
dset = dset_load.get(var_list)
112+
#dset = dset_load.get(var_list)
81113
# rename altitude variable to z for monet use
82114
dset = dset.rename({"lev": "z"})
83115

84116
# re-order so surface is associated with the first vertical index
85117
dset = dset.sortby("z", ascending=False)
86118
# ===========================
87119

120+
# Derive MM-standardized variables from CESM-SE native fields.
121+
# Source vars are optional: missing inputs then derivation skipped,
122+
123+
# pres_pa_mid: hybrid sigma-pressure midpoint (Pa)
124+
# P = hyam*P0 + hybm*PS (standard CAM hybrid coords)
125+
if "pres_pa_mid" not in dset.variables and {"hyam", "hybm", "PS"} <= set(dset.variables):
126+
_P0 = float(dset["P0"].values) if "P0" in dset.variables else 100000.0
127+
dset["pres_pa_mid"] = dset["hyam"] * _P0 + dset["hybm"] * dset["PS"]
128+
dset["pres_pa_mid"].attrs.update({
129+
"units": "Pa",
130+
"long_name": "Pressure at mid-level",
131+
"description": "hyam*P0 + hybm*PS",
132+
})
133+
134+
# temperature_k
135+
if "temperature_k" not in dset.variables and "T" in dset.variables:
136+
dset["temperature_k"] = dset["T"]
137+
dset["temperature_k"].attrs.setdefault("units", "K")
138+
139+
# dz_m: layer thickness (m) via hydrostatic + ideal gas
140+
# dz = PDELDRY * Rd * T / (P_mid * g)
141+
if (
142+
"dz_m" not in dset.variables
143+
and "PDELDRY" in dset.variables
144+
and "pres_pa_mid" in dset.variables
145+
and "temperature_k" in dset.variables
146+
):
147+
_Rd = 287.04 # J/(kg*K), dry air
148+
_g = 9.80665 # m/s^2
149+
dset["dz_m"] = (
150+
dset["PDELDRY"] * _Rd * dset["temperature_k"]
151+
/ (dset["pres_pa_mid"] * _g)
152+
)
153+
dset["dz_m"].attrs.update({
154+
"units": "m",
155+
"long_name": "Layer thickness",
156+
"description": "PDELDRY * Rd * T / (P_mid * g)",
157+
})
158+
159+
# once derivations are complete, dont keep in returned dataset
160+
for _v in _cesm_se_var:
161+
if _v in dset.variables:
162+
dset = dset.drop_vars(_v)
163+
88164
# Make sure this dataset has unstructured grid
89165
dset.attrs["mio_has_unstructured_grid"] = True
90-
dset.attrs["mio_scrip_file"] = scrip_file
166+
if scrip_file:
167+
dset.attrs["mio_scrip_file"] = scrip_file
91168

92169
# convert units
93170
if convert_to_ppb:
@@ -111,7 +188,6 @@ def open_mfdataset(
111188
# Below are internal functions to this file
112189
# -----------------------------------------
113190

114-
115191
def _ensure_mfdataset_filenames(fname):
116192
"""Checks if dataset in netcdf format
117193
Parameters

‎monetio/sat/tempo_l2.py‎

Lines changed: 148 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
from netCDF4 import Dataset
1919

2020

21-
def _open_one_dataset(fname, variable_dict):
21+
def _open_one_dataset(fname, variable_dict, i = None):
2222
"""Read locally stored INTELSAT TEMPO NO2 level 2
2323
2424
Parameters
@@ -31,7 +31,9 @@ def _open_one_dataset(fname, variable_dict):
3131
-------
3232
ds : xr.Dataset
3333
"""
34-
print("reading " + fname)
34+
# probably do not need to print every file it is reading in.
35+
if i is not None and i % 75 == 0:
36+
print("reading", fname)
3537

3638
ds = xr.Dataset()
3739

@@ -155,7 +157,151 @@ def _open_one_dataset(fname, variable_dict):
155157

156158
return ds
157159

160+
def _open_one_dataset(fname, variable_dict, i=None):
161+
"""Read locally stored INTELSAT TEMPO NO2 level 2
162+
163+
Making this function less memory intensive (hopefully)
164+
165+
"""
166+
167+
# Print only every 75 files (when index is provided)
168+
if i is not None and i % 75 == 0:
169+
print(f"reading {fname}", flush=True)
170+
171+
ds = xr.Dataset()
172+
173+
with Dataset(fname, "r") as dso:
174+
175+
lon_var = dso.groups["geolocation"]["longitude"]
176+
lat_var = dso.groups["geolocation"]["latitude"]
177+
time_var = dso.groups["geolocation"]["time"]
178+
time_units = time_var.units
158179

180+
ds["lon"] = (
181+
("x", "y"),
182+
lon_var[:].squeeze(),
183+
{"long_name": lon_var.long_name, "units": lon_var.units},
184+
)
185+
186+
ds["lat"] = (
187+
("x", "y"),
188+
lat_var[:].squeeze(),
189+
{"long_name": lat_var.long_name, "units": lat_var.units},
190+
)
191+
192+
ds["time"] = (
193+
("time",),
194+
num2pydate(time_var[:].squeeze(), time_units),
195+
{"long_name": time_var.long_name},
196+
)
197+
198+
ds = ds.set_coords(["time", "lon", "lat"])
199+
200+
ds.attrs["reference_time_string"] = dso.time_coverage_start
201+
ds.attrs["granule_number"] = dso.granule_num
202+
ds.attrs["scan_num"] = dso.scan_num
203+
204+
if ("pressure" in variable_dict) and "surface_pressure" not in variable_dict:
205+
warnings.warn(
206+
"Calculating pressure in TEMPO data requires surface_pressure. "
207+
"Adding surface_pressure to output variables"
208+
)
209+
variable_dict["surface_pressure"] = {}
210+
211+
for varname in variable_dict:
212+
213+
if varname in [
214+
"main_data_quality_flag",
215+
"vertical_column_troposphere",
216+
"vertical_column_stratosphere",
217+
"vertical_column_troposphere_uncertainty",
218+
"vertical_column",
219+
"vertical_column_uncertainty",
220+
]:
221+
values_var = dso.groups["product"][varname]
222+
223+
elif varname in [
224+
"latitude_bounds",
225+
"longitude_bounds",
226+
"solar_zenith_angle",
227+
"solar_azimuth_angle",
228+
"viewing_zenith_angle",
229+
"viewing_azimuth_angle",
230+
"relative_azimuth_angle",
231+
]:
232+
values_var = dso.groups["geolocation"][varname]
233+
234+
elif varname in [
235+
"vertical_column_total",
236+
"vertical_column_total_uncertainty",
237+
"fitted_slant_column",
238+
"fitted_slant_column_uncertainty",
239+
"snow_ice_fraction",
240+
"terrain_height",
241+
"ground_pixel_quality_flag",
242+
"surface_pressure",
243+
"tropopause_pressure",
244+
"scattering_weights",
245+
"gas_profile",
246+
"albedo",
247+
"temperature_profile",
248+
"amf",
249+
"amf_total",
250+
"amf_diagnostic_flag",
251+
"eff_cloud_fraction",
252+
"amf_cloud_fraction",
253+
"amf_cloud_pressure",
254+
"amf_troposphere",
255+
"amf_stratosphere",
256+
"background_correction",
257+
]:
258+
values_var = dso.groups["support_data"][varname]
259+
260+
elif varname in ["fit_rms_residual", "fit_convergence_flag"]:
261+
values_var = dso.groups["qa_statistics"][varname]
262+
263+
values = values_var[:].squeeze()
264+
265+
# scaling
266+
if "scale" in variable_dict[varname]:
267+
values = variable_dict[varname]["scale"] * values
268+
269+
# masking
270+
if "minimum" in variable_dict[varname]:
271+
values = np.ma.masked_less(values, variable_dict[varname]["minimum"], copy=False)
272+
273+
if "maximum" in variable_dict[varname]:
274+
values = np.ma.masked_greater(values, variable_dict[varname]["maximum"], copy=False)
275+
276+
# store
277+
if "corner" in values_var.dimensions:
278+
ds[varname] = (("x", "y", "corner"), values, values_var.__dict__)
279+
280+
elif "swt_level" in values_var.dimensions:
281+
ds[varname] = (("x", "y", "swt_level"), values, values_var.__dict__)
282+
283+
else:
284+
ds[varname] = (("x", "y"), values, values_var.__dict__)
285+
286+
if "quality_flag_max" in variable_dict[varname]:
287+
ds.attrs["quality_flag"] = varname
288+
ds.attrs["quality_thresh_max"] = variable_dict[varname]["quality_flag_max"]
289+
290+
# pressure conversion
291+
if "surface_pressure" in variable_dict:
292+
if ds["surface_pressure"].attrs["units"] == "hPa":
293+
HPA2PA = 100
294+
ds["surface_pressure"][:] = ds["surface_pressure"].values * HPA2PA
295+
ds["surface_pressure"].attrs["units"] = "Pa"
296+
ds["surface_pressure"].attrs["valid_min"] *= HPA2PA
297+
ds["surface_pressure"].attrs["valid_max"] *= HPA2PA
298+
ds["surface_pressure"].attrs["Eta_A"] *= HPA2PA
299+
300+
if "pressure" in variable_dict:
301+
ds["pressure"] = calculate_pressure(ds)
302+
303+
return ds
304+
159305
def calculate_pressure(ds):
160306
"""Calculates pressure at layer and delta_pressure of layer
161307

0 commit comments

Comments
 (0)