Add plotting enhancements

[justin-richling]

New feature type

New plot and/or plot enhancement

What is this new feature?

Add flexibility for plotting (contour levels, colormap, etc.) based on pressure level and/or plot type (polar vs global)

Update plots for more professional look, ie Zonal and Meridional are very bare boned, lat/lon were missing units, some variables were not getting proper scaling, fit colorbar better to axes and extend if necessary, etc.

• NOTE: There is an additional enhancement for 2d line plots for Zonal/Meridional that allows a color map for the 2d variables (See SWCF example below - variable defaults argument: colormap_2d)

Combine Lat/Lon Vector into Lat/Lon code since it was a lot of reused code.

• NOTE: the calculation for percent difference is currently off for lat/lon vectors plots, so this map is being ignored and vector plots default back to the two cases on top and difference centered below

Add option in variable defaults config for variable nickname; remove all percent diff args in defaults yaml and hard code in plotting script. The user will still have the option to supply these args in the yaml file if default is not desired.

Add option for old NCL color maps

Add units cleaning method to ensure units are always available, this was not enforced throughout the code base.

Add code to resolve levels for hemisphere and resolve colormap in prep_contour_plot

---

Core new features:

Different contours based on vertical levels

RESTOM:
  colormap: "RdBu_r"
  contour_levels_range: [-100, 101, 5]
  diff_contour_range: [-10, 10, 0.5]
  polar_map:
    colormap:
      nh: "RdBu_r"
      sh: "terrain"
    contour_levels_range:
      #nh: [-100, 101, 5]
      sh: [-100, 0, 2.5]
   scale_factor: 1
  add_offset: 0
  new_unit: "W m$^{-2}$"
  category: "TOA energy flux"
  derivable_from: ['FLNT','FSNT']

---

SWCF:
  colormap: "Reds"
  polar_map:
    colormap:
      nh: "fakemap"
      sh: "Blues"
  contour_levels_range: [-150, 51, 10]
  diff_colormap: "BrBG"
  diff_contour_range: [-20, 21, 2]
  scale_factor: 1
  add_offset: 0
  new_unit: "Wm$^{-2}$"
  obs_file: "CERES_EBAF_Ed4.1_2001-2020.nc"
  obs_name: "CERES_EBAF_Ed4.1"
  obs_var_name: "toa_cre_sw_mon"
  obs_scale_factor: 1
  obs_add_offset: 0
  category: "TOA energy flux"
  colormap_2d: True

---

U:
  global_latlon_map:
    colormap: # if not specified, default to viridis for any missing values nested (levs, ie 850) or not
      200: "ncl_default"
      #850: "Spectral"
    contour_levels_range:
      200: [-10, 51, 2.5]
      850: [-10, 31, 2.5]
    diff_contour_range:
      200: [-15, 16, 2]
      850: [-5, 5.1, 0.5]
  polar_map:
    colormap:
      nh: "terrain"
      sh: "fake_cmap"
    contour_levels_range:
      nh:
        200: [-10, 31, 2.5]
        850: [-10, 13, 1]
      sh:
        200: [-10, 51, 5]
        850: [-10, 31, 2.5]
    diff_contour_range:
      200: [-10, 10, 1]
      850: [-5, 5.1, 0.5]
      #nh:
      #  200: [-15, 16, 2]
      #  850: [-5, 5.1, 0.5]
      #sh:
      #  200: [-15, 16, 2]
      #  850: [-5, 5.1, 0.5]
  zonal_mean:
    colormap: "ncl_default"
    contour_levels_linspace: [-10, 30, 15] # These will be applied to all zonal 3d mean plots
    #diff_contour_range: # default to figuring our the contour levels from difference from max/min
  scale_factor: 1
  add_offset: 0
  new_unit: "ms$^{-1}$"
  obs_file: "U_ERA5_monthly_climo_197901-202112.nc"
  obs_name: "ERA5"
  obs_var_name: "U"
  vector_pair: "V"
  vector_name: "Wind"
  category: "State"

---

Surface_Wind_Stress:
  nickname: "Surface Wind Stress"
  new_unit: "N m$^{-2}$"
  category: "Surface variables"

TAUX:
  new_unit: "N m$^{-2}$"
  vector_pair: "TAUY"
  vector_name: "Surface_Wind_Stress"
  category: "Surface variables"
  scale_factor: -1
  add_offset: 0

TAUY:
  new_unit: "N m$^{-2}$"
  vector_pair: "TAUX"
  vector_name: "Surface_Wind_Stress"
  category: "Surface variables"
  scale_factor: -1
  add_offset: 0

---

closes #452
closes #425
closes #373

[nusbaume]

Thanks for all of the code cleanup and plot enhancements @justin-richling! I realize it looks like I have a bazillion comments, but the vast majority should hopefully be easy or trivial to implement. Of course if you do have any questions or concerns with anything I wrote just let me know!

[nusbaume]

I personally don't think we should have an NPL environment file in the ADF (or LDF). In general it's best to only have the minimum number of packages you need to run all the ADF scripts, which I am guessing is significantly less than what is contained here.

If you need help creating a more up-to-date environment file for the ADF then please let me know, as I would be happy to help work on it. Thanks!

[nusbaume]

Total nit-pick, but I might add a space in the comment here to match the format of all the other comments in this file:

Suggested change

	#NCL RGB files
	# NCL RGB files

[nusbaume]

Another nit-pick, but I might use the same quotation marks on both sides of the equation here (in case someone who doesn't know python well looks at the code and assumes the behavior is different):

Suggested change

	vres["season"] = plot['season']
	vres["hemi"] = plot['hemi']
	vres['season'] = plot['season']
	vres['hemi'] = plot['hemi']

[nusbaume]

I realize this is purely an opinion, but in general I feel that "kwarg"-like dictionary inputs into python functions should only contain metadata (e.g. plot_type), and not required data like the fields being plotted. Given this, I think I would bring in the extra vector data fields as optional function arguments instead of keyword arguments. Of course feel free to push back here if you disagree!

[nusbaume]

If you do accept my suggestion on passing the main data fields as optional arguments, then this function call would look something like this instead:

Suggested change

	pf.plot_map_and_save(adfobj, plot_name, case_nickname, base_nickname,
	[syear_cases[case_idx],eyear_cases[case_idx]],
	[syear_baseline,eyear_baseline],
	umseasons[s], uoseasons[s],
	udseasons[s], upseasons[s],
	obs=obs, **vres)
	pf.plot_map_and_save(adfobj, plot_name, case_nickname, base_nickname,
	[syear_cases[case_idx],eyear_cases[case_idx]],
	[syear_baseline,eyear_baseline],
	umseasons[s], uoseasons[s],
	udseasons[s], upseasons[s],
	vm=vmseasons[s], vo=voseasons[s],
	vd=vdseasons[s], vp=vpseasons[s],
	obs=obs, **vres)

Then in the plot_map_and_save function itself all of the vx variables would be None in the function argument list.

[nusbaume]

Remove commented-out code?

[nusbaume]

Instead of making multiple copies, why not just make one copy that is then modified for the diff and pctdiff plots? The only reason you would need multiple copies if you needed to save diff_kwargs after creating the zonal plot for the difference field.

[nusbaume]

Meridional plot?

Suggested change

	# Generate zonal plot:
	# Generate meridional plot:

[nusbaume]

Just noting again that naming this variable to something like contour_levs might help make it obvious that this is not related to vertical levels.

[nusbaume]

Can we get away with only one copy of kwargs here as well?

[brianpm]

I think this ax2 needs to be pointed at an axs element

[brianpm]

ax3 should be in axs

[brianpm]

ax4 should be in axs

[brianpm]

Is range(0,90,10)going to work for southern hemisphere?

[brianpm]

Should ax2 be axs[...]?

[brianpm]

ax3 maybe axs[...]

[brianpm]

I think you have to use some pathlib thing here, probably: locfil = Path.cwd() / f"{cmap_pctdiff}.rgb"

[brianpm]

Is this what we actually want? Divide by zero is replaced by np.nan, but then get set to 0.0. Won't that mean that any nan value ends up looking like a valid result in the plot?

[justin-richling]

Yeah, this is very confusing as now we are misrepresenting the true values, which is not what we want.

Would something like this be more accurate:

pctdiff_mag = xr.where(obs_mag != 0,
                                  (mdl_mag - obs_mag) / np.abs(obs_mag) * 100,
                                   np.nan,
                                   )

That way we also catch anywhere obs/baseline is zero, or is this already handled with the np.isfinite?

I'm hoping catching obs_mag when zero will fix the times where the pct values are excessively large that I see occassionally.

[brianpm]

I think isfinite does it. I think it works as:

pctdiff_mag = pctdiff_mag.where(np.isfinite(pctdiff_mag))

just get rid of the fillna line.

[brianpm]

Both quiver calls are unconditional inside the if i in [0,1] block — so for both i=0 (model panel) and i=1 (baseline panel), it draws model vectors first and then baseline vectors on top. Both panels end up with overlaid arrows.

It seems like the intent is i=0 → model only, i=1 → baseline only. The fix should be:

  if i == 0:
      ax[i].quiver(..., umdlfld[skip], vmdlfld[skip], mdl_mag.values[skip], ...)
  elif i == 1:
      ax[i].quiver(..., uobsfld[skip], vobsfld[skip], obs_mag.values[skip], ...)

[brianpm]

Before this PR, pct_diff_colormap and pct_diff_contour_levels were at the top level of a variable's YAML entry and were found directly in kwargs.

Now:

• If the user has a global_latlon_map (or similar) sub-block, they must put pct_diff_colormap inside that block
  which means that the previous behavior of top-level placement is silently ignored.
• If plot_type is not passed into the function at all, plot_type_dict = {} and there is no way for the user to override hard-coded values that are given in lib/plotting_utils.py:
• Line 863: cmap_pctdiff = "PuOr_r"
• Line 882: cmap_pctdiff = 'PuOr_r'
• Line 896: levels_pctdiff = [-100,-75,-50,-40,-30,-20,-10,-8,-6,-4,-2,0,2,4,6,8,10,20,30,40,50,75,100]

So the PR description's claim that users can still supply these args is only partially true. The configurability is preserved when nested under a plot-type block but silently lost in the flat/top-level case that was standard before.

This probably doesn't matter much, but it complicates customizing plots from the YAML file. Needs documentation at minimum.

[justin-richling]

This is a great observation @brianpm, thanks for pointing it out!

[brianpm]

The hard-coded yaml default was PuOr_r right? Should that be used here?

[justin-richling]

Ah, good catch, thanks.

[brianpm]

I left a few comments where I saw some inconsistent changes. Looks good overall. My first thought is that we could probably reduce total code, but no need for now.

[nusbaume]

Review generated by Claude, following the repository's AGENTS.md (on main).

Summary

This PR reworks contour/colormap resolution so levels and colormaps can vary by plot type, hemisphere and pressure level, merges the vector map into plot_map_and_save, adds NCL colormap support, and tidies units handling. The direction is good and the unit fixes (adf_dataset.py raw_units, keep_attrs=True in zonal_mean_xr) are clean. However, there are a few correctness problems that will show up in a default ADF run, so I don't think it's ready to merge as-is.

Per AGENTS.md §3, note that CI covers none of this: no file in this PR is in the pylint testable_files set, and the pytest suite only touches adf_base/adf_config. A green checkmark here means very little — everything below was found by reading and by executing the new helpers against the shipped adf_variable_defaults.yaml in the pinned adf_v1.0.0 environment.

Also flagging up front: the PR currently reports mergeable: false against main.

---

Blocking

1. Vector-component variables get the vector's name as the plot title on every scalar plot — lib/plotting_utils.py:958-973 (add_var_to_vres)

add_var_to_vres checks vector_name first and unconditionally overrides var_name with it, even when the caller is making an ordinary scalar plot. Running it against the shipped defaults:

U                    -> figure title uses var_name = 'Wind'
V                    -> figure title uses var_name = 'Wind'
TAUX                 -> figure title uses var_name = 'Surface Wind Stress'
TAUY                 -> figure title uses var_name = 'Surface Wind Stress'
PRECT                -> figure title uses var_name = 'PRECT'

var_name feeds fig.suptitle(f"{var_name}: {season}") (plotting_functions.py:256-260, 606-610, 952, 979, 1295, 1322). U and V are in the default diag_var_list, and global_latlon_map, zonal_mean, meridional_mean and polar_map all call add_var_to_vres — so in a stock run every scalar U and V plot is titled "Wind". The vector-name branch should only be taken when the caller is actually plotting the vector (e.g. gated on the vector flag that global_latlon_vect_map.py already sets).

2. Explicit 3-entry contour_levels lists are silently discarded — lib/plotting_utils.py:632-648 (resolve_levels / process_entry)

For kind == "levels", a list of exactly three entries falls into the else branch, logs "please add more values", and returns None — so the plot silently falls back to data-range levels. On main, contour_levels was used verbatim regardless of length (plotting_utils.py:357-361 at the merge base). Any user whose variable-defaults file has e.g. contour_levels: [-1, 0, 1] gets different plots with no error and nothing on stdout — the AGENTS.md §4.3 backwards-compatibility rule. No variable in the two shipped defaults files has a 3-entry explicit list, so this only bites user-supplied files, but that's exactly where it's hardest to notice.

3. NCL colormap parser silently corrupts values and crashes on plausible inputs — lib/plotting_utils.py:399-452 (read_ncl_colormap)

The trailing-comment strip at line 441-443 uses line_str[:match.start()-1]. The -1 drops a real character whenever there's no space before the comment marker. Executed against the PR code:

"255 0 0 ; red"   -> [255.0, 0.0, 0.0]    correct
"255 0 10# red"   -> [255.0, 0.0, 1.0]    blue channel 10 silently becomes 1

Same function, other inputs:

• comment-only / header-only file → UnboundLocalError: cannot access local variable 'table'
• single data row → table is 1-D, then ncl_to_mpl raises IndexError: tuple index out of range
• except: at line 444 is bare (AGENTS.md §6.3) and leaves row bound to the previous row when a line fails to parse

The silent-truncation case is the one that matters most (AGENTS.md §5.2): wrong colors, no warning.

4. Runtime colormap download — lib/plotting_utils.py:391-412

read_ncl_colormap does filename = Path.cwd() / fil.split("/")[-1] and urlretrieve into it. Verified on a login node: calling try_load_ncl_cmap(adf, "ncl_default") writes ncl_default.rgb into the current working directory — i.e. wherever the user launched ./run_adf_diag, usually the ADF clone itself. That's what the new *.rgb line in .gitignore is papering over. Three problems:

• ADF runs in batch on Derecho/Casper compute nodes, which generally have no outbound internet. The download will fail there, and since nothing is cached on failure it retries for every variable/season that asks for the colormap.
• urlretrieve is called with no timeout, so each of those retries can block rather than fail fast. Only urllib.error.HTTPError is caught explicitly.
• Writing generated files into the source tree (or any CWD) is surprising and breaks read-only/shared installs.

Since ncl_defaults = ["ncl_default"] is a single-entry list, this is one small file — vendoring it under lib/ (or a configurable cache dir, defaulting somewhere writable) would remove the network dependency entirely.

---

Non-blocking

• Shared mutation of variable_defaults. AdfObs.variable_defaults returns copy.copy(...) — a shallow copy, so the per-variable sub-dicts are shared. vres = adfobj.variable_defaults.get(var, {}) aliases the master entry, and add_var_to_vres plus vres["plot_type"]/vres["vector"]/vres["umdlfld_nowrap"] = ... write straight into it. After global_latlon_vect_map runs, the master entry for U permanently carries vector: True and the season's umdlfld_nowrap/vobsfld_nowrap DataArrays, which then flow as **vres into later scripts. Each script does overwrite plot_type and var_name, so I did not find a visible failure from this today — but it's a live footgun and holds references to full fields. A copy.deepcopy of vres (or having add_var_to_vres return a new dict) closes it.
• Dead code, one piece of it broken. load_colormap (494) and choose_colormap_type (485) are never called. load_colormap:503 calls read_ncl_colormap(adfobj, locfil, msg) — three arguments to a two-parameter function; confirmed TypeError: read_ncl_colormap() takes 2 positional arguments but 3 were given as soon as a local .rgb exists. It also says "Defaulting to 'coolwarm'" while returning 'viridis', and unpacks cm, cmr = ncl_to_mpl(...) when ncl_to_mpl can return a bare None. Either delete these two or fix and wire them up.
• try_load_ncl_cmap (513-538). Line 536 is msg + "Something went wrong..." — an expression statement, so that message never reaches the debug log (msg += intended). If the inner urllib.error.HTTPError branch fires, data is never assigned and the following isinstance(data, ...) raises UnboundLocalError, which the outer bare except Exception then swallows — it lands on 'viridis' by accident rather than by design.
• NCL support is narrower than the PR description suggests. ncl_defaults = ["ncl_default"], and get_cmap only attempts an NCL lookup when the name is exactly "ncl_default". Any other NCL table name (BlueRed, amwg, …) falls through the final validity check and is replaced by the default colormap. Worth either generalizing or saying so in the docs.
• meridional_mean.py:181,183,210 — new warning text says "zonal mean plotting skipped" in the meridional script. Copy-paste from zonal_mean.py; user-visible.
• Stale module docstring, lib/plotting_functions.py:1-30. It still documents plot_map_vect_and_save, which this PR deletes, and still shows the old signatures for make_polar_plot / plot_map_and_save / plot_zonal_mean_and_save / plot_meridional_mean_and_save without the new leading adfobj argument. AGENTS.md §7.6.
• Docstrings on new helpers (AGENTS.md §6.2). 7 of the 11 new functions in plotting_utils.py have no docstring at all (guess_ncl_url, download_ncl_colormap, read_ncl_colormap, ncl_to_mpl, choose_colormap_type, load_colormap, add_var_to_vres); the other 4 are one-liners with no Parameters/Returns sections. None were added to the module docstring's Functions list. add_var_to_vres and resolve_levels in particular have non-obvious contracts that need documenting.
• New variable-defaults keys are undocumented. The header comment block in adf_variable_defaults.yaml isn't updated for nickname, colormap_2d, contour_levels_linspace, or the new nested per-plot-type blocks (polar_map:, global_latlon_map:, zonal_mean:) with their nh/sh/level sub-keys — the headline feature of this PR. AGENTS.md §4.3 asks for new fields to be documented there. Right now those forms exist only in the PR description; the shipped file uses nickname once and none of the rest, so nothing exercises them.
• black (AGENTS.md §6.1). Intersecting black --check --diff hunks with the lines this PR adds or changes gives roughly 1,333 new/changed lines that black would reformat, concentrated in plotting_functions.py (~640), plotting_utils.py (~408) and meridional_mean.py (~189). resolve_levels (620) is also indented as if it were a method — its whole body sits at 8 spaces.
• env/npl-2024a_environment.yaml — a new 649-line fully-pinned environment, unrelated to plotting enhancements, with no README mention. Intentional, or a leftover? It's a real maintenance surface (AGENTS.md §4.5).

Optional / follow-up

• Removing pf.plot_map_vect_and_save is right for the codebase, but it's a public helper that user-maintained scripts may call. Worth a line in the release notes.
• The PR adds 8 new "$\mathbf{...}$" string literals; these are not raw strings and produce SyntaxWarning: invalid escape sequence '\m' on import under Python 3.12+. Pre-existing pattern (24 occurrences on main), but new code could use r"...".
• cm.get_cmap at plotting_utils.py:795 is deprecated. I checked: it still works in the pinned matplotlib 3.9.4 (and 3.10), so this is not urgent — but since the block was rewritten anyway, plt.get_cmap/mpl.colormaps[...] would be cheap to adopt.
• The vector percent-difference is computed for every season and level (global_latlon_vect_map.py) and then discarded at plotting_functions.py:527 (if vector: #ignore percent diff for now). Skipping the computation would save the work and make the intent clearer.
• read_ncl_colormap:414 sets is_url = False, which is never used.

Not verified

• The plots themselves. I did not run the ADF against real CAM output, so the visual claims in the PR description (colorbar fit, extend behaviour, zonal/meridional styling, colormap_2d) are unverified — as are the new contour-level and colormap choices as science, which need a CAM/AMWG reviewer's eye rather than mine.
• Whether Derecho/Casper compute nodes can reach ncl.ucar.edu. The download succeeded from a login node; the compute-node behaviour behind blocking finding 4 is my expectation, not something I tested.
• pytest lib/test/unit_tests — pytest isn't installed in the local adf_v1.0.0 env, so I couldn't run it. No file in this PR is in the suite's coverage, so I'd expect it to pass unchanged. YAML parses cleanly for all files (check-yaml equivalent), and every changed .py file compiles.