Skip to content

Commit e410c91

Browse files
lizziellizziel
committed
Get stretched grid params from ref and dev datasets in plot functions
This update removes having to pass a separate argument for file containing the stretch parameters. This is no longer necessary because all stretch files, both history output and checkpoints, contain stretch parameters in recent model versions. Signed-off-by: Lizzie Lundgren <elundgren@seas.harvard.edu>
1 parent 6300ed8 commit e410c91

File tree

2 files changed

+66
-88
lines changed

2 files changed

+66
-88
lines changed

‎gcpy/plot/compare_single_level.py‎

Lines changed: 33 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -64,8 +64,6 @@ def compare_single_level(
6464
second_ref=None,
6565
second_dev=None,
6666
spcdb_dir=os.path.dirname(__file__),
67-
sg_ref_path='',
68-
sg_dev_path='',
6967
ll_plot_func='imshow',
7068
**extra_plot_args
7169
):
@@ -176,14 +174,6 @@ def compare_single_level(
176174
spcdb_dir: str
177175
Directory containing species_database.yml file.
178176
Default value: Path of GCPy code repository
179-
sg_ref_path: str
180-
Path to NetCDF file containing stretched-grid info
181-
(in attributes) for the ref dataset
182-
Default value: '' (will not be read in)
183-
sg_dev_path: str
184-
Path to NetCDF file containing stretched-grid info
185-
(in attributes) for the dev dataset
186-
Default value: '' (will not be read in)
187177
ll_plot_func: str
188178
Function to use for lat/lon single level plotting with
189179
possible values 'imshow' and 'pcolormesh'. imshow is much
@@ -248,40 +238,39 @@ def compare_single_level(
248238
quiet=True
249239
)
250240

251-
sg_ref_params = [1, 170, -90]
252-
sg_dev_params = [1, 170, -90]
253-
# Get stretched-grid info if passed
254-
if sg_ref_path != '':
255-
sg_ref_attrs = xr.open_dataset(sg_ref_path).attrs
256-
if 'stretch_factor' in sg_ref_attrs:
257-
sg_ref_params = [
258-
sg_ref_attrs['stretch_factor'],
259-
sg_ref_attrs['target_lon'],
260-
sg_ref_attrs['target_lat']]
261-
elif 'STRETCH_FACTOR' in sg_ref_attrs:
262-
sg_ref_params = [
263-
sg_ref_attrs['STRETCH_FACTOR'],
264-
sg_ref_attrs['TARGET_LON'],
265-
sg_ref_attrs['TARGET_LAT']]
266-
else:
267-
msg = "Stretched grid global parameters missing from ref file"
268-
raise ValueError(msg)
269-
270-
if sg_dev_path != '':
271-
sg_dev_attrs = xr.open_dataset(sg_dev_path).attrs
272-
if 'stretch_factor' in sg_dev_attrs:
273-
sg_dev_params = [
274-
sg_dev_attrs['stretch_factor'],
275-
sg_dev_attrs['target_lon'],
276-
sg_dev_attrs['target_lat']]
277-
elif 'STRETCH_FACTOR' in sg_dev_attrs:
278-
sg_dev_params = [
279-
sg_dev_attrs['STRETCH_FACTOR'],
280-
sg_dev_attrs['TARGET_LON'],
281-
sg_dev_attrs['TARGET_LAT']]
282-
else:
283-
msg = "Stretched grid global parameters missing from dev file"
284-
raise ValueError(msg)
241+
# Get stretched grid info, if any.
242+
# Parameter order is stretch factor, target longitude, target latitude.
243+
# Stretch factor 1 corresponds with no stretch.
244+
245+
# Ref stretch attributes
246+
print(refdata.attrs)
247+
if 'stretch_factor' in refdata.attrs:
248+
sg_ref_params = [
249+
refdata.attrs['stretch_factor'],
250+
refdata.attrs['target_lon'],
251+
refdata.attrs['target_lat']]
252+
elif 'STRETCH_FACTOR' in refdata.attrs:
253+
sg_ref_params = [
254+
refdata.attrs['STRETCH_FACTOR'],
255+
refdata.attrs['TARGET_LON'],
256+
refdata.attrs['TARGET_LAT']]
257+
else:
258+
sg_ref_params = [1, 170, -90]
259+
260+
# Dev stretch attributes
261+
print(devdata.attrs)
262+
if 'stretch_factor' in devdata.attrs:
263+
sg_dev_params = [
264+
devdata.attrs['stretch_factor'],
265+
devdata.attrs['target_lon'],
266+
devdata.attrs['target_lat']]
267+
elif 'STRETCH_FACTOR' in devdata.attrs:
268+
sg_dev_params = [
269+
devdata.attrs['STRETCH_FACTOR'],
270+
devdata.attrs['TARGET_LON'],
271+
devdata.attrs['TARGET_LAT']]
272+
else:
273+
sg_dev_params = [1, 170, -90]
285274

286275
# Get grid info and regrid if necessary
287276
[refres, refgridtype, devres, devgridtype, cmpres, cmpgridtype, regridref,

‎gcpy/plot/compare_zonal_mean.py‎

Lines changed: 33 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -64,8 +64,6 @@ def compare_zonal_mean(
6464
second_ref=None,
6565
second_dev=None,
6666
spcdb_dir=os.path.dirname(__file__),
67-
sg_ref_path='',
68-
sg_dev_path='',
6967
ref_vert_params=None,
7068
dev_vert_params=None,
7169
**extra_plot_args
@@ -178,14 +176,6 @@ def compare_zonal_mean(
178176
spcdb_dir: str
179177
Directory containing species_database.yml file.
180178
Default value: Path of GCPy code repository
181-
sg_ref_path: str
182-
Path to NetCDF file containing stretched-grid info
183-
(in attributes) for the ref dataset
184-
Default value: '' (will not be read in)
185-
sg_dev_path: str
186-
Path to NetCDF file containing stretched-grid info
187-
(in attributes) for the dev dataset
188-
Default value: '' (will not be read in)
189179
ref_vert_params: list(AP, BP) of list-like types
190180
Hybrid grid parameter A in hPa and B (unitless).
191181
Needed if ref grid is not 47 or 72 levels.
@@ -311,40 +301,39 @@ def compare_zonal_mean(
311301
fracrefdata = fracrefdata.isel(lev=ref_pmid_ind)
312302
fracdevdata = fracdevdata.isel(lev=dev_pmid_ind)
313303

314-
sg_ref_params = [1, 170, -90]
315-
sg_dev_params = [1, 170, -90]
316-
# Get stretched-grid info if passed
317-
if sg_ref_path != '':
318-
sg_ref_attrs = xr.open_dataset(sg_ref_path).attrs
319-
if 'stretch_factor' in sg_ref_attrs:
320-
sg_ref_params = [
321-
sg_ref_attrs['stretch_factor'],
322-
sg_ref_attrs['target_lon'],
323-
sg_ref_attrs['target_lat']]
324-
elif 'STRETCH_FACTOR' in sg_ref_attrs:
325-
sg_ref_params = [
326-
sg_ref_attrs['STRETCH_FACTOR'],
327-
sg_ref_attrs['TARGET_LON'],
328-
sg_ref_attrs['TARGET_LAT']]
329-
else:
330-
msg = "Stretched grid global parameters missing from ref file"
331-
raise ValueError(msg)
332-
333-
if sg_dev_path != '':
334-
sg_dev_attrs = xr.open_dataset(sg_dev_path).attrs
335-
if 'stretch_factor' in sg_dev_attrs:
336-
sg_dev_params = [
337-
sg_dev_attrs['stretch_factor'],
338-
sg_dev_attrs['target_lon'],
339-
sg_dev_attrs['target_lat']]
340-
elif 'STRETCH_FACTOR' in sg_dev_attrs:
341-
sg_dev_params = [
342-
sg_dev_attrs['STRETCH_FACTOR'],
343-
sg_dev_attrs['TARGET_LON'],
344-
sg_dev_attrs['TARGET_LAT']]
345-
else:
346-
msg = "Stretched grid global parameters missing from dev file"
347-
raise ValueError(msg)
304+
# Get stretched grid info, if any.
305+
# Parameter order is stretch factor, target longitude, target latitude.
306+
# Stretch factor 1 corresponds with no stretch.
307+
308+
# Ref stretch attributes
309+
print(refdata.attrs)
310+
if 'stretch_factor' in refdata.attrs:
311+
sg_ref_params = [
312+
refdata.attrs['stretch_factor'],
313+
refdata.attrs['target_lon'],
314+
refdata.attrs['target_lat']]
315+
elif 'STRETCH_FACTOR' in refdata.attrs:
316+
sg_ref_params = [
317+
refdata.attrs['STRETCH_FACTOR'],
318+
refdata.attrs['TARGET_LON'],
319+
refdata.attrs['TARGET_LAT']]
320+
else:
321+
sg_ref_params = [1, 170, -90]
322+
323+
# Dev stretch attributes
324+
print(devdata.attrs)
325+
if 'stretch_factor' in devdata.attrs:
326+
sg_dev_params = [
327+
devdata.attrs['stretch_factor'],
328+
devdata.attrs['target_lon'],
329+
devdata.attrs['target_lat']]
330+
elif 'STRETCH_FACTOR' in devdata.attrs:
331+
sg_dev_params = [
332+
devdata.attrs['STRETCH_FACTOR'],
333+
devdata.attrs['TARGET_LON'],
334+
devdata.attrs['TARGET_LAT']]
335+
else:
336+
sg_dev_params = [1, 170, -90]
348337

349338
[refres, refgridtype, devres, devgridtype, cmpres, cmpgridtype,
350339
regridref, regriddev, regridany, refgrid, devgrid, cmpgrid,

0 commit comments

Comments
 (0)