Skip to content

Commit

Permalink
Merge pull request #319 from justin-richling/quick-fixes
Browse files Browse the repository at this point in the history
Quick fixes
  • Loading branch information
justin-richling authored Jul 29, 2024
2 parents b0e1c4c + e436ae4 commit 1c7364c
Show file tree
Hide file tree
Showing 2 changed files with 78 additions and 43 deletions.
29 changes: 21 additions & 8 deletions scripts/plotting/global_latlon_map.py
Original file line number Diff line number Diff line change
Expand Up @@ -125,12 +125,13 @@ def global_latlon_map(adfobj):
"MAM": [3, 4, 5],
"SON": [9, 10, 11]
}

# probably want to do this one variable at a time:
for var in var_list:
if var not in adfobj.data.ref_var_nam:
dmsg = f"No reference data found for variable `{var}`, zonal mean plotting skipped."
dmsg = f"No reference data found for variable `{var}`, global lat/lon mean plotting skipped."
adfobj.debug_log(dmsg)
print(dmsg)
continue

#Notify user of variable being plotted:
Expand All @@ -156,10 +157,19 @@ def global_latlon_map(adfobj):
vres['central_longitude'] = pf.get_central_longitude(adfobj)

# load reference data (observational or baseline)
# odata = adfobj.data.load_reference_da(var)
odata = adfobj.data.load_reference_regrid_da(adfobj.data.ref_case_label, var)
if not adfobj.compare_obs:
base_name = adfobj.data.ref_case_label
else:
base_name = adfobj.data.ref_labels[var]

# Gather reference variable data
odata = adfobj.data.load_reference_regrid_da(base_name, var)

if odata is None:
dmsg = f"No regridded test file for {base_name} for variable `{var}`, global lat/lon mean plotting skipped."
adfobj.debug_log(dmsg)
continue

o_has_dims = pf.validate_dims(odata, ["lat", "lon", "lev"]) # T iff dims are (lat,lon) -- can't plot unless we have both
if (not o_has_dims['has_lat']) or (not o_has_dims['has_lon']):
print(f"\t = skipping global map for {var} as REFERENCE does not have both lat and lon")
Expand All @@ -184,6 +194,8 @@ def global_latlon_map(adfobj):

#Skip this variable/case if the regridded climo file doesn't exist:
if mdata is None:
dmsg = f"No regridded test file for {case_name} for variable `{var}`, global lat/lon mean plotting skipped."
adfobj.debug_log(dmsg)
continue

#Determine dimensions of variable:
Expand All @@ -192,16 +204,17 @@ def global_latlon_map(adfobj):
print(f"\t = skipping global map for {var} for case {case_name} as it does not have both lat and lon")
continue
else: # i.e., has lat&lon
if pres_levs and (not has_dims['has_lev']):
print(f"\t - skipping global map for {var} as it has more than lat/lon dims, but no pressure levels were provided")
if (has_dims['has_lev']) and (not pres_levs):
print(f"\t - skipping global map for {var} as it has more than lev dimension, but no pressure levels were provided")
continue

# Check output file. If file does not exist, proceed.
# If file exists:
# if redo_plot is true: delete it now and make plot
# if redo_plot is false: add to website and move on
doplot = {}
if not pres_levs:

if not has_dims['has_lev']:
for s in seasons:
plot_name = plot_loc / f"{var}_{s}_LatLon_Mean.{plot_type}"
doplot[plot_name] = plot_file_op(adfobj, plot_name, var, case_name, s, web_category, redo_plot, "LatLon")
Expand Down Expand Up @@ -353,4 +366,4 @@ def plot_file_op(adfobj, plot_name, var, case_name, season, web_category, redo_p
return True

##############
#END OF SCRIPT
#END OF SCRIPT
92 changes: 57 additions & 35 deletions scripts/plotting/zonal_mean.py
Original file line number Diff line number Diff line change
Expand Up @@ -144,6 +144,7 @@ def zonal_mean(adfobj):
if var not in adfobj.data.ref_var_nam:
dmsg = f"No obs found for variable `{var}`, zonal mean plotting skipped."
adfobj.debug_log(dmsg)
print(dmsg)
continue

#Notify user of variable being plotted:
Expand All @@ -160,7 +161,21 @@ def zonal_mean(adfobj):
#End if

# load reference data (observational or baseline)
odata = adfobj.data.load_reference_regrid_da(adfobj.data.ref_case_label, var)
if not adfobj.compare_obs:
base_name = adfobj.data.ref_case_label
else:
base_name = adfobj.data.ref_labels[var]

# Gather reference variable data
odata = adfobj.data.load_reference_regrid_da(base_name, var)

#Check if regridded file exists, if not skip zonal plot for this var
if odata is None:
dmsg = f"No regridded baseline file for {adfobj.data.ref_case_label} for variable `{var}`, zonal mean plotting skipped."
adfobj.debug_log(dmsg)
continue

#Check zonal mean dimensions
has_lat_ref, has_lev_ref = pf.zm_validate_dims(odata)

#Loop over model cases:
Expand All @@ -175,6 +190,11 @@ def zonal_mean(adfobj):
# load re-gridded model files:
mdata = adfobj.data.load_regrid_da(case_name, var)

if mdata is None:
dmsg = f"No regridded test file for {case_name} for variable `{var}`, zonal mean plotting skipped."
adfobj.debug_log(dmsg)
continue

# determine whether it's 2D or 3D
# 3D triggers search for surface pressure
has_lat, has_lev = pf.zm_validate_dims(mdata) # assumes will work for both mdata & odata
Expand All @@ -200,52 +220,55 @@ def zonal_mean(adfobj):
# so we could have made one script instead of two.
# Merging would make overall timing better because looping twice will double I/O steps.
#
if not has_lev:
plot_name = plot_loc / f"{var}_{s}_Zonal_Mean.{plot_type}"

if plot_name not in zonal_skip:
# difference: each entry should be (lat, lon) or (plev, lat, lon)
# dseasons[s] = mseasons[s] - oseasons[s]
# difference will be calculated in plot_zonal_mean_and_save;
# because we can let any pressure-level interpolation happen there
# This could be re-visited for efficiency or improved code structure.

#Seasonal Averages
mseasons[s] = pf.seasonal_mean(mdata, season=s, is_climo=True)
oseasons[s] = pf.seasonal_mean(odata, season=s, is_climo=True)

#Set the file name
plot_name = plot_loc / f"{var}_{s}_Zonal_Mean.{plot_type}"
plot_name_log = None

#Seasonal Averages
mseasons[s] = pf.seasonal_mean(mdata, season=s, is_climo=True)
oseasons[s] = pf.seasonal_mean(odata, season=s, is_climo=True)
if has_lev:
#Set the file name for log-pressure plots
plot_name_log = plot_loc / f"{var}_logp_{s}_Zonal_Mean.{plot_type}"

# difference: each entry should be (lat, lon) or (plev, lat, lon)
# dseasons[s] = mseasons[s] - oseasons[s]
# difference will be calculated in plot_zonal_mean_and_save;
# because we can let any pressure-level interpolation happen there
# This could be re-visited for efficiency or improved code structure.
#Check if reference file has vertical levels
if not has_lev_ref:
print(f"Error: expecting lev for both case: {has_lev} and ref: {has_lev_ref}")
continue
#End if

#Create new plot:
pf.plot_zonal_mean_and_save(plot_name, case_nickname, adfobj.data.ref_nickname,
#Create plots
if plot_name not in zonal_skip:

#Create new plot:
pf.plot_zonal_mean_and_save(plot_name, case_nickname, adfobj.data.ref_nickname,
[syear_cases[case_idx],eyear_cases[case_idx]],
[syear_baseline,eyear_baseline],
mseasons[s], oseasons[s], has_lev, log_p=False, obs=adfobj.compare_obs, **vres)

#Add plot to website (if enabled):
adfobj.add_website_data(plot_name, var, case_name, season=s, plot_type="Zonal")
#Add plot to website (if enabled):
adfobj.add_website_data(plot_name, var, case_name, season=s, plot_type="Zonal")
#End if

#Create new plot with log-p:
# NOTE: The log-p should be an option here.
else:
if (not has_lev_ref) or (not has_lev):
print(f"Error: expecting lev for both case: {has_lev} and ref: {has_lev_ref}")
continue
if len(mdata['lev']) != len(odata['lev']):
print(f"Error: zonal mean contour expects `lev` dim to have same size, got {len(mdata['lev'])} and {len(odata['lev'])}")
continue
plot_name_log = plot_loc / f"{var}_{s}_Zonal_logp_Mean.{plot_type}"
if plot_name_log not in logp_zonal_skip:
#Seasonal Averages
mseasons[s] = pf.seasonal_mean(mdata, season=s, is_climo=True)
oseasons[s] = pf.seasonal_mean(odata, season=s, is_climo=True)
#Create log-pressure plots as well (if applicable)
if (plot_name_log) and (plot_name_log not in logp_zonal_skip):

pf.plot_zonal_mean_and_save(plot_name_log, case_nickname, adfobj.data.ref_nickname,
pf.plot_zonal_mean_and_save(plot_name_log, case_nickname, adfobj.data.ref_nickname,
[syear_cases[case_idx],eyear_cases[case_idx]],
[syear_baseline,eyear_baseline],
mseasons[s], oseasons[s], has_lev, log_p=True, obs=adfobj.compare_obs, **vres)

#Add plot to website (if enabled):
adfobj.add_website_data(plot_name_log, f"{var}_logp", case_name, season=s, plot_type="Zonal", category="Log-P")
#Add plot to website (if enabled):
adfobj.add_website_data(plot_name_log, f"{var}_logp", case_name, season=s, plot_type="Zonal", category="Log-P")
#End if

#End for (seasons loop)
#End for (case names loop)
Expand All @@ -256,5 +279,4 @@ def zonal_mean(adfobj):


##############
#END OF SCRIPT

#END OF SCRIPT

0 comments on commit 1c7364c

Please sign in to comment.