Repository navigation
Conversation
LottMillerSso unpacked `nlev, ncols = state.temperature.shape`, vmapped over
axis 1 and transposed its output, so it only ran on a column-vectorized
(nlev, ncols) host and raised on a whole (nlev, nlon, nlat) grid:
ValueError: too many values to unpack (expected 2)
The drag is a per-column scheme, so it has no reason to care how the host lays
out the horizontal.
ComposablePhysics(vectorize_columns=True) does not help: __add__ inherits that
flag from its left operand, and a grid-based package cannot switch to it
because speedy_shortwave.get_zonal_average_fields computes zonal averages and
needs a real longitude axis. So composing this term with SPEEDY was impossible
in either configuration.
Flatten whatever trailing axes the host uses into one column axis, run the
existing vmap unchanged, and restore the shape on the way out. No branching on
rank; the column-vectorized path is unchanged.
Adds lott_miller_host_test.py, which builds the pressure diagnostics directly
so it exercises this term alone. It uses real T30 orography deliberately, since
flat orography gives identically zero drag and every comparison would pass
vacuously. Four of the five fail without the fix.
Contributor
|
LGTM, thanks! |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
LottMillerSsoopened withthen
vmapped over axis 1 and transposed its output, so it only ran on a column-vectorized(nlev, ncols)host and raised on a whole(nlev, nlon, nlat)grid:The drag itself is a per-column scheme, so it has no reason to care how the host lays out the horizontal.
ComposablePhysics(vectorize_columns=True)is not a workaround here.__add__inherits that flag from its left operand, and a grid-based package cannot simply switch:speedy_shortwave.get_zonal_average_fieldscomputes zonal averages and needs a real longitude axis, which flattening destroys. So composingLottMillerSsowith SPEEDY was not possible in either configuration.Fixed by flattening whatever trailing axes the host uses into one column axis, running the existing
vmapunchanged, and restoring the shape on the way out. No branching on rank, and the column-vectorized path is unchanged.Test plan
New
lott_miller_host_test.pybuilds the pressure and height diagnostics directly rather than going throughMoistAirColumnState, so it exercises this term alone:(nlev, nlon, nlat)and column(nlev, ncols)hosts both return correctly shaped tendenciesIt uses the real T30 orography deliberately. With flat orography the scheme returns identically zero and every comparison passes without testing anything, so there is an explicit assertion that the drag is non-zero.
Results:
test_column_host_returns_column_shaped_tendencies, the layout that already worked.pytest jcm/physics/gravity_waves/— 70 passed.ruff check jcm/physics/gravity_waves/sso/clean.Verified beyond unit tests by composing
speedy_physics() + MoistAirColumnState() + LottMillerSso()and running 450 simulated days at T31, finite throughout.This is a shape fix only. It makes no claim that the scheme's parameters or its
nktopg/ntoplevel indices are appropriate for any particular host resolution.Related: #587 fixes the same class of issue in
MoistAirColumnState. The two are independent and this PR does not depend on it.🤖 Generated with Claude Code