ANN prarameterization of horizontal momentum eddy fluxes#944
Merged
Conversation
- Computes subgrid stress using ANN in MOM_Zanna_Bolton - Uses MOM_ANN module for ANN inference Equivalent MOM_override for defaults ``` USE_ZB2020 = True ZB2020_USE_ANN = True USE_CIRCULATION_IN_HORVISC = True ZB2020_ANN_FILE_TALL = /path/to/ocean3d/subfilter/FGR3/EXP1/model/Tall.nc ```
Blank commit after squash/rebase was handled on command line
- Removed unused modules - Removed unused MOM_memory.h - Added input and output means which default to 0 and do not need to be present in the weights file - Gave defaults to means, norms, tests so that they do no need to be present in file - Added missing array notation "(:)" - Minor formatting
- Added ANN_allocate, set_layer, set_input_normalization, and set_output_normalization methods to allow reconfiguration during unit tests - Added ANN_unit_tests with some simple constructed-by-code networks with known solutions - Added config_src/drivers/unit_tests/test_MOM_ANN.F90 to drive unit tests - Added config_src/drivers/timing_tests/time_MOM_ANN.F90 as rudimentary for timing inference
- Adds inference operating on array (instead of single vector of
features)
- Implements several different versions of inference with various
loop orders
- Involves storing the transpose of A in the type
- Tested by checking inference on same inputs is identical between
variants
- Added randomizers to assist in unit testing
- Adds timing of variants to config_src/drivers/timing/time_MOM_ANN.F90
- Adds an interface (MOM_apply) to select preferred version of
inference subroutine
- Added command line args to time_MOM_ANN.F90 to allow more rapid
evaluation of performance
Variants explored, timed with gfortran (13.2) -O3 on Xeon:
- vector_v1:
- original inference from Pavel
- vector_v2:
- allocate work arrays just once, using widest layer
- loop over layers in 2's to avoid pointer calculations and copies
- speed up, x0.8 relative to v1
- vector_v3:
- transpose loops
- slow down, x1.54 relative to v1
- vector_v4:
- transpose weights with same loop order as v1
- slow down, x1.03 relative to v1
- array_v1:
- same structure as v2, working on x(space,feature) input/outputs
- speed up, x0.41 relative to v1
- array_v2:
- as for array_v1 but with transposed loop order
- apply activation function on vector of first index while in cache
- speed up, x0.35 relative to v1
- array_v3:
- same structure as v2, working on x(feature,space) input/outputs
- speed up, x0.58 relative to v1
- Added module dox - Renamed _v1, _v2 etc to labels - Added ANN_apply_array_sio to ANN_apply interface - Replaced "flops" with "MBps" in timing output
- Deleted variants of ANN that did not perform as well as the two versions that remain.
Re-factor of MOM_ANN
* Apply array_sio ANN inference for computation of momentum fluxes * remove trailing space * Initial commit
…lines address Robert Hallberg code review
Author
|
Hi @Hallberg-NOAA, Sorry for the late reply. I modified the code according to your suggestions. Answers did not change on NYU HPC. Best, |
Hallberg-NOAA
requested changes
Sep 10, 2025
Member
Hallberg-NOAA
left a comment
Hallberg-NOAA
left a comment
There was a problem hiding this comment.
Apart from restoring the default value of ZB_SCALING to its previous value so that we can follow the established MOM6 procedures for changing parameter defaults, I think that this PR is ready for acceptance.
Author
|
Thanks! I restored the default value. Now a default parameter list for ANN parameterization of mesoscale momentum fluxes is updated as follows: |
Hallberg-NOAA
approved these changes
Sep 18, 2025
Member
Hallberg-NOAA
left a comment
Hallberg-NOAA
left a comment
There was a problem hiding this comment.
This PR now looks good to me. Thank you @Pperezhogin for this contribution of a valuable new set of capabilities.
Member
|
This PR has passed pipeline testing at https://gitlab.gfdl.noaa.gov/ogrp/mom6ci/MOM6/-/pipelines/28874 with the expected warnings about a new runtime parameter. |
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
3 participants
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.
This PR contains the implementation of the neural network parameterization of mesoscale eddy effects introduced in Perezhogin, Adcroft, Zanna 2025.
Three files implement the module of neural network. In particular, its inference:
MOM_ANN.F90testing:
test_MOM_ANN.F90and timing:
time_MOM_ANN.F90A function implementing parameterization,
compute_stress_ANN_collocated, was introduced in moduleMOM_Zanna_Bolton.F90.Slight change in
MOM_hor_visc.F90implements alternative discretization of relative vorticity using contour integral (turned off by default).The netCDF file with neural network weights can be found here. In order to turn on parameterization, provide path to netCDF file and switch a couple of other parameters in
MOM_override: