Zonal mpsi function by jukent · Pull Request #773 · NCAR/geocat-comp

jukent · 2025-10-09T18:05:51Z

PR Summary

Zonal_MPSI function for unstructured grid datasets.

Related Tickets & Documents

Closes #774

PR Checklist

General

PR includes a summary of changes
Link relevant issues, make one if none exist
Add a brief summary of changes to docs/release-notes.rst in a relevant section for the upcoming release.
Add appropriate labels to this PR
PR follows the Contributor's Guide

Functionality

New function(s) intended for public API added to geocat/comp/__init__.py file

Testing

Update or create tests in appropriate test file

Documentation

Docstrings have been created and/or updated in accordance with Documentation Standards.
Internal functions have a preceding underscore (_) and have been added to docs/internal_api/index.rst
User facing functions have been added to docs/user_api/index.rst under their module

[pull] main from NCAR:main

…_mpsi

jukent · 2026-02-12T17:34:06Z

build_envs/environment.yml

  - pandas
  - xarray
  - xskillscore>=0.0.17
+  - uxarray>=2025.12.0


I'd say yes. There were updates to UXarray's where and some wrappers that were previously roadblocks to this function. Perhaps conda will naturally resolve the environment to the latest version, but I'd like to let more time pass before trusting that.

erogluorhan

Great to see this in progress, thanks @jukent ! I've made a few comments below.

erogluorhan · 2026-02-12T17:41:22Z

geocat/comp/meteorology.py

 _generate_wrapper_docstring(dpres_plev, delta_pressure)
+
+
+def zonal_mpsi(uxds, lat=list(range(-90, 91, 10))):


A couple questions:

Is the name "zonal_mpsi" what we want to go forward with? Any alternatives discussed? I recognize finsing something fancy to replace the "mpsi" abbreviation wouldn't be easy though.

Would it make sense to have lat argument default to a tuple of (start, step, end) and handle the list creation in the code rather than the func signature? Also, it would take tuple, array-like, or just a scalar, i.e. something similar to zonal_mean?

We haven't discussed alternatives yet. We aren't bound by character limit so zonal_meridional_stream() could work (is including the word function redundant?)

Absolutely, I'll make that change right away.

build_envs/environment.yml

erogluorhan · 2026-02-12T17:51:40Z

geocat/comp/meteorology.py

+    g = 9.80665  # gravity (m/s^2)
+
+    # Basic input validation
+    if not hasattr(uxds, "V") or not hasattr(uxds, "PS"):


I am wondering if this is too limiting, i.e. if there was a meridional wind variable in a dataset without the name (or any attribute) "V", but having the standard_name being "northward_wind", w'd be missing it. Would it be a possible case? Similar applies to "PS".

Fair. I had some checks in but it became very convoluted so I removed them. They are necessary, so I added them back in using 2 new helper functions in geocat.util that I imagine could be useful throughout the package. Unfortunately our dev and test data doesn't seem to be cf compliant for meridional wind or surface pressure. I tried using cf_xarray first.

Would cf-xarray be of help on this?

geocat/comp/meteorology.py

erogluorhan · 2026-02-12T17:55:03Z

geocat/comp/meteorology.py

+        )
+
+    # Check if interpolation needs to be done
+    if "plev" in uxds["V"].dims:


Similar to above comment on "V" and "PS", maybe we need to look for some CF compliant checks that'd cover any case, e.g. which would use standard name, units, etc?

geocat/comp/meteorology.py

erogluorhan · 2026-02-12T20:44:53Z

geocat/comp/meteorology.py

+    ]
+    hybridA_names = ['hyam', 'hybrid_A_midpoints']
+    hybridB_names = ['hybm', 'hybrid_B_midpoints']
+    plev_names = ['plev', 'pressure_lev', 'pressure_levels']


This possible names-based scheme would probably cover majority of use cases, but I guess I still have some concerns about some easy misses it would have. Also some of these possible name values might not help as expected.

For instance, 'northward_wind' can mostly show up in dataarray.attrs.standard_name rather than being the name of the variable. If the wind variable is simply named "wind" (don't know if data creators would do that though) or anything other than those in this list and has the "standard_name" attribute being "northward_wind", I think _find_coord() would miss it.

So, there needs to be some kind of CF-compliant check mechanism here that'd look into not only the var names but also attrs like "standard_name", "units", etc.

anissa111

Okay, here's my initial review. Also, I'm aware that this is not exactly the place to bring this up, but I feel the need to lodge a comment that I think that geocat-comp ingesting uxarray objects is risky and possibly not in the best interest of geocat-comp.

That being said here's my comments:

I still feel like we should make uxarray an optional dependency
pyproject.toml and other packaging should be updated to include the optional/required dependency
uxarray will need to be added to the benchmark environment
uxarray should also probably be added to upstream testing (as in the upstream test should install uxarray from source in ci/install-upstream.sh)
the helper functions should be added to the internal api documentation
the helper functions should probably have their own tests
uxarray specific test(s) for the functions in interpolation that will have uxarray objects passing through them from zonal_meridional_psi
somewhere in the docstring (and also maybe in the code) there should be a citation for the calculation we’re preforming
at least a handful more tests using inputs with different/missing names/descriptions/etc

anissa111 · 2026-02-13T22:41:46Z

geocat/comp/gc_util.py

+                    return var_name
+        error_parts.append(f"Tried units: {units}. ")
+
+    raise KeyError(f"Could not find {description} in dataset. {' '.join(error_parts)}")


Suggested change

raise KeyError(f"Could not find {description} in dataset. {' '.join(error_parts)}")

raise KeyError(f"Could not find {description} in dataset. {''.join(error_parts)}")

anissa111 · 2026-02-13T22:45:03Z

geocat/comp/gc_util.py

+                return name
+        error_parts.append(f"Tried names: {possible_names}. ")
+
+    # Finally try units match (less reliable)


Using units here still concerns me. I'm not convinced that units are a good way of determining a match.

anissa111 · 2026-02-13T22:50:50Z

geocat/comp/gc_util.py

+                        f"This is unreliable - multiple variables may share the same units. "
+                        f"Please verify this is correct and add CF standard_name. {error_parts}",
+                        UserWarning,
+                        stacklevel=3,


Why stacklevel=3 here? If my understanding is correct, this would make sense for when _find_var is called by _find_optional_var, but not when _find_var is called directly

anissa111 · 2026-02-13T22:59:48Z

test/test_meteorology.py

+        assert out.sizes["latitudes"] == len(self.lat)
+        assert out.sizes["plev"] == len(uxds.V.plev)
+
+        # ---- numerical sanity ----


Suggested change

# ---- numerical sanity ----

# ---- numerical coherence ----

Two bits:

small language best practice note

I'm not necessarily objecting to having the isinfinite or allclose to 0 asserts, but I don't think I understand why are we doing them? Is there a risk that something will go wrong and we'll get inf/nan/0 values erroneously?

anissa111 · 2026-02-13T23:03:13Z

test/test_meteorology.py

+    lat = np.arange(36, 45, 1)
+
+    def test_zonal_meridional_psi_pressure_levels(self) -> None:
+        uxds = ux.open_dataset("test/grid_subset.nc", "test/plev_subset.nc")


This could be a fixture depending on the size/potential for reuse. Also, reading in the files should probably use the same format as the other file readings in our codebase with something like

try: uxds = ux.open_dataset("grid_subset.nc", "plev_subset.nc") except FileNotFoundError: uxds = ux.open_dataset("test/grid_subset.nc", "test/plev_subset.nc")

anissa111 · 2026-02-13T23:38:28Z

geocat/comp/meteorology.py

+    return da_mpsi
+
+
+def zonal_mpsi(


We have an entire utility function for exactly this that generates wrappers and copies the docstrings for ncl-like function names

anissa111 · 2026-02-13T23:39:17Z

geocat/comp/meteorology.py

+    Returns
+    -------
+    da_mpsi : xarray.DataArray
+        Zonal mean meridional streamf unction, scaled by Earth's geometry and gravity.


Suggested change

Zonal mean meridional streamf unction, scaled by Earth's geometry and gravity.

Zonal mean meridional stream function, scaled by Earth's geometry and gravity.

anissa111 · 2026-02-13T23:45:16Z

geocat/comp/meteorology.py

+    # apply scaling factor
+    da_scaling_factor = 2 * np.pi * a * np.cos(lats) / g
+    da_mpsi = da_mpsi * da_scaling_factor


~~This could be my misunderstanding, but isn't this scaling only valid for certain grid types? Like, specifically only structured ones with equally spaced latitudes?~~

@kafitzgerald, I think I'm remembering something that you flagged on my nmse PR, though completely possible I'm just not familiar enough with the calculation to know what I'm talking about here.

edit: sorry, coming back to say I get that the weighting is for the lat bands, not the unstructured grid itself, but wouldn't this still assume that the given latitudes are evenly spaced?

anissa111 · 2026-02-13T23:47:30Z

geocat/comp/meteorology.py

+    except Exception:
+        pass


(nitpick) bare except makes me a bit nervous here, is there a reason?

Also, don't we create da_mpsi? Would updating the attributes ever raise an exception even if we didn't?

anissa111 · 2026-02-13T23:52:21Z

geocat/comp/meteorology.py

+    if not hyam_coordname:
+        hyam_coordname = _find_optional_var(
+            uxds,
+            long_name='hybrid A coefficient at layer midpoints',
+            possible_names=['hyam', 'hya', 'hybrid_A_midpoints'],
+            description='coordinate,',
+        )
+    if not hybm_coordname:
+        hybm_coordname = _find_optional_var(
+            uxds,
+            long_name='hybrid B coefficient at layer midpoints',
+            possible_names=['hybm', 'hyb', 'hybrid_B_midpoints'],
+            description='coordinate,',
+        )
+    if not plev_coordname:
+        plev_coordname = _find_optional_var(
+            uxds,
+            standard_name='air_pressure',
+            possible_names=['plev', 'pressure_lev', 'pressure_levels'],
+            description='coordinate,',
+        )


Suggested change

if not hyam_coordname:

hyam_coordname = _find_optional_var(

uxds,

long_name='hybrid A coefficient at layer midpoints',

possible_names=['hyam', 'hya', 'hybrid_A_midpoints'],

description='coordinate,',

)

if not hybm_coordname:

hybm_coordname = _find_optional_var(

uxds,

long_name='hybrid B coefficient at layer midpoints',

possible_names=['hybm', 'hyb', 'hybrid_B_midpoints'],

description='coordinate,',

)

if not plev_coordname:

plev_coordname = _find_optional_var(

uxds,

standard_name='air_pressure',

possible_names=['plev', 'pressure_lev', 'pressure_levels'],

description='coordinate,',

)

if not hyam_coordname:

hyam_coordname = _find_optional_var(

uxds,

long_name='hybrid A coefficient at layer midpoints',

possible_names=['hyam', 'hya', 'hybrid_A_midpoints'],

description='coordinate',

)

if not hybm_coordname:

hybm_coordname = _find_optional_var(

uxds,

long_name='hybrid B coefficient at layer midpoints',

possible_names=['hybm', 'hyb', 'hybrid_B_midpoints'],

description='coordinate',

)

if not plev_coordname:

plev_coordname = _find_optional_var(

uxds,

standard_name='air_pressure',

possible_names=['plev', 'pressure_lev', 'pressure_levels'],

description='coordinate',

)

Given the way _find_var is written, say hya and hyb both have a description of coordinate and nothing else, wouldn't this mean that whichever one was found first would become the coord name for both?

Also, I don't have a better suggestion, but is just 'coordinate' the best default description to look for for all of these?

anissa111 · 2026-02-14T03:25:57Z

geocat/comp/meteorology.py

+            - hyam : Hybrid A coefficients (if meridional_wind is on hybrid sigma-pressure levels)
+            - hybm : Hybrid B coefficients (if meridional_wind is on hybrid sigma-pressure levels)
+            - uxgrid : Grid information for uxarray
+    meridonal_wind_varname : str, optional


Suggested change

meridonal_wind_varname : str, optional

meridional_wind_varname : str, optional

anissa111 · 2026-02-14T03:27:13Z

geocat/comp/meteorology.py

 # NCL NAME WRAPPER FUNCTIONS BELOW
 def dpres_plev(pressure_lev, surface_pressure, pressure_top=None):
    return delta_pressure(pressure_lev, surface_pressure, pressure_top)


 _generate_wrapper_docstring(dpres_plev, delta_pressure)


This should stay at the bottom of the file (and be used to create the wrapper function made in this PR)

anissa111 · 2026-02-14T04:00:09Z

geocat/comp/meteorology.py

+    ----------
+    uxds : uxarray.UxDataset
+        Input dataset containing the following required fields:
+            - meridional_wind : CF Compliant Meridional wind component (on pressure or hybrid sigma-pressure levels)


What is the CF complaint name for meridional wind (is there one?)? Is it just northward_wind? That seems slightly strange to me.

I'm not finding something explicitly called meridional wind in the CF standard names, though, so maybe

jukent added 4 commits May 8, 2024 14:31

Merge pull request #56 from NCAR/main

4c01f3f

[pull] main from NCAR:main

Merge pull request #57 from NCAR/main

45256b4

[pull] main from NCAR:main

Merge branch 'main' of https://github.com/jukent/geocat-comp

c70cf22

zonal_mpsi fx

92d121c

jukent added dependencies Pull requests that update a dependency file feature a new feature that is going to be developed labels Oct 9, 2025

jukent added 17 commits October 9, 2025 12:09

init and release-notes

383443d

index.rst

ac52e9f

Merge branch 'main' of https://github.com/NCAR/geocat-comp into zonal…

b68952a

…_mpsi

update env

66fa9a7

hybrid_pressure kwarg and input validation

f28b382

add tests

8003b7a

precommit rm

3db32b7

pre-commit release-notes

8747e81

add uxarray to min-deps

84e191d

add uxarrary to docs

bb53840

test

f73f9ad

add test files

8226fa3

use custom datafiles

ac66118

update test

19c6827

Merge branch 'main' of https://github.com/NCAR/geocat-comp into zonal…

5862a3f

…_mpsi

cleaned up function, using new testing data, fleshed out tests

f6ce027

pre-commit and release notes

0d44645

jukent marked this pull request as ready for review February 12, 2026 17:03

jukent requested review from anissa111, cyschneck and kafitzgerald and removed request for anissa111 and cyschneck February 12, 2026 17:03

jukent requested a review from cyschneck February 12, 2026 17:03

jukent commented Feb 12, 2026

View reviewed changes

erogluorhan reviewed Feb 12, 2026

View reviewed changes

jukent added 5 commits February 12, 2026 12:55

find coords

54906c9

rm .values call

18dd80c

some style

728337f

text

51b15ee

rename

7408af9

erogluorhan reviewed Feb 12, 2026

View reviewed changes

jukent added 6 commits February 12, 2026 15:24

CF error handling

d582684

List what was tried in units warning

ccd117b

update an error message

886412b

change fx name, let users add input variables, add wrapper

1bd1b0d

VSCode thought it was being helpful

275b3ad

docstring update

0e13cf2

anissa111 added run-benchmark Add tag to a PR to run ASV comparison on new commits and removed run-benchmark Add tag to a PR to run ASV comparison on new commits labels Feb 13, 2026

anissa111 reviewed Feb 14, 2026

View reviewed changes

anissa111 requested changes Feb 14, 2026

View reviewed changes

anissa111 reviewed Feb 14, 2026

View reviewed changes

		_generate_wrapper_docstring(dpres_plev, delta_pressure)


		def zonal_mpsi(uxds, lat=list(range(-90, 91, 10))):

	raise KeyError(f"Could not find {description} in dataset. {' '.join(error_parts)}")
	raise KeyError(f"Could not find {description} in dataset. {''.join(error_parts)}")

	Zonal mean meridional streamf unction, scaled by Earth's geometry and gravity.
	Zonal mean meridional stream function, scaled by Earth's geometry and gravity.

	meridonal_wind_varname : str, optional
	meridional_wind_varname : str, optional

Conversation

jukent commented Oct 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR Summary

Related Tickets & Documents

PR Checklist

Uh oh!

Choose a reason for hiding this comment

Uh oh!

erogluorhan left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

anissa111 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

anissa111 Feb 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

jukent commented Oct 9, 2025 •

edited

Loading

anissa111 Feb 13, 2026 •

edited

Loading