Fix bugs in calls to vertical mixing within timesteppers by andrewdnolan · Pull Request #488 · E3SM-Project/Omega

andrewdnolan · 2026-07-31T14:54:46Z

Addresses two issue when calling VertMix from the time stepping loop which were causing ERS tests for fail:

raw buffer indices were passed as a logical time level to VertMixImplicit
missing halo exchange after VertMixImplicit modifies state variables.

Fixes #487

Also use getters to access state instead of raw access

andrewdnolan · 2026-07-31T21:18:26Z

+      // Re-exchange halos after vertical mixing
+      Pacer::timingBarrier("ForwardBackward:vMixHaloExchBarrier", 3, Comm);
+      Pacer::start("ForwardBackward:vMixHaloExch", 3);
+      State->exchangeHalo(VelCurLevel);
+      Tracers::exchangeHalo(VelCurLevel);
+      Pacer::stop("ForwardBackward:vMixHaloExch", 3);


Alternatively we could move this to the end of VertMix::VertMixImplicit so that each time stepper did not have to manually call the halo exchange.

Seems like the more self contained approach and less error prone for when more time steppers are added in the future.

I'm just not that familiar with VertMix, so not sure if there are downsides I'm not considering.

andrewdnolan · 2026-07-31T22:11:56Z

`Frontier` CTest unit tests:

Machine: fontier
Compiler: craygnu / craygnu-mphipcc
Build type: Release
Result: All tests passed
Logs:
- craygnu: /lustre/orion/cli115/proj-shared/anolan/omega_PR488_vertmix/craygnu/build/ctest.log
- craygnu-mphipcc: /lustre/orion/cli115/proj-shared/anolan/omega_PR488_vertmix/craygnu-mphipcc/build/ctest.log

andrewdnolan · 2026-07-31T23:04:12Z

CIME Testing:

Rebasing onto #470 ERS and PEM tests now pass. Without the changes in this PR these test fail the bit-for-bit comparison step.

Testing (all-pass):

ERS_Vmct.T62_oQU240.COMEGA-IAF.frontier_craygnu
PEM_Vmct.T62_oQU240.COMEGA-IAF.frontier_craygnu

Note: with 03e986a (from #470) the default timestepper is changed to RungeKutta4, so the PEM test passes with just these changes. In order to use ForwardBackward and get the PEM test to pass the changes in #490 are also needed.

xylar · 2026-08-01T09:04:08Z

@hyungyukang and @mwarusz, you seem like the best reviewers. Please take a look or suggest alternative reviewers if you don't have time.

hyungyukang · 2026-08-03T18:39:57Z

@andrewdnolan , thanks for opening this bug fix PR.

It looks like shear vertical mixing is the source of the non-BFB behavior. With only background and convective vertical mixing enabled, the implicit vertical mixing remains BFB across different numbers of GPUs and CPU cores.

I’m investigating it further.

hyungyukang · 2026-08-04T16:06:33Z

After digging into it further, I found that GradRichNum in the shear vertical mixing is responsible for the non-BFB behavior. The issue is that cells in the outermost halo layer are missing some of the neighboring cells that GradRichNum averages over:

Omega/components/omega/src/ocn/VertMix.h

Lines 113 to 179 in 061d0ab

    
           class GradRichardsonNum { 
        
            public: 
        
              Real RiInitValue = 100.0_Real; ///< Initial Richardson number value 
        
              /// constructor declaration 
        
              GradRichardsonNum(const HorzMesh *Mesh, const VertCoord *VCoord); 
        
              //   The functor takes the full arrays of Richardson number (inout), 
        
              //   the index ICell, and normal and tangential velocities as inputs, 
        
              //   and outputs the Richardson number. 
        
              KOKKOS_FUNCTION void 
        
              operator()(Array2DReal GradRichNum, I4 ICell, I4 KChunk, 
        
                         const Array2DReal &NormalVelocity, 
        
                         const Array2DReal &TangentialVelocity, 
        
                         const Array2DReal &BruntVaisalaFreqSq) const { 
        
                 const I4 KStart = chunkStart(KChunk, MinLayerCell(ICell) + 1); 
        
                 const I4 KLen   = chunkLength(KChunk, KStart, MaxLayerCell(ICell)); 
        
                 Real GradRichNumNorm[VecLength]; 
        
                 Real GradRichNumTmp[VecLength]; 
        
                 for (int KVec = 0; KVec < KLen; ++KVec) { 
        
                    GradRichNumNorm[KVec] = 1.0e-12_Real; 
        
                    GradRichNumTmp[KVec]  = RiInitValue; 
        
                 } 
        
                 for (int J = 0; J < NEdgesOnCell(ICell); ++J) { 
        
                    I4 JEdge = EdgesOnCell(ICell, J); 
        
                    I4 JCell = CellsOnCell(ICell, J); 
        
                    for (int KVec = 0; KVec < KLen; ++KVec) { 
        
                       const I4 K = KStart + KVec; 
        
                       I4 K1      = K - 1; 
        
                       I4 K2      = K; 
        
                       // Skip this edge contribution if it would access 
        
                       // invalid edge velocity levels. 
        
                       if (K1 > MaxLayerEdgeBot(JEdge) || K2 > MaxLayerEdgeBot(JEdge)) 
        
                          continue; 
        
                       Real DNormVel = 
        
                           NormalVelocity(JEdge, K1) - NormalVelocity(JEdge, K2); 
        
                       Real DTanVel = 
        
                           TangentialVelocity(JEdge, K1) - TangentialVelocity(JEdge, K2); 
        
                       Real DzEdge = 
        
                           0.5_Real * (GeomZMid(ICell, K1) + GeomZMid(JCell, K1) - 
        
                                       (GeomZMid(ICell, K2) + GeomZMid(JCell, K2))); 
        
                       Real ShearSquared = 
        
                           (DNormVel * DNormVel + DTanVel * DTanVel) / (DzEdge * DzEdge); 
        
                       Real RiEdge = 
        
                           Kokkos::max(0.0_Real, 
        
                                       0.5_Real * (BruntVaisalaFreqSq(ICell, K2) + 
        
                                                   BruntVaisalaFreqSq(JCell, K2))) / 
        
                           (ShearSquared + 1.0e-12_Real); 
        
                       Real Weight = 0.25_Real * DcEdge(JEdge) * DvEdge(JEdge); 
        
                       GradRichNumNorm[KVec] += Weight; 
        
                       GradRichNumTmp[KVec] += Weight * RiEdge; 
        
                    } 
        
                 } 
        
                 for (int KVec = 0; KVec < KLen; ++KVec) { 
        
                    const I4 K            = KStart + KVec; 
        
                    GradRichNum(ICell, K) = GradRichNumTmp[KVec] / GradRichNumNorm[KVec]; 
        
                 } 
        
              }

One communication-efficient solution would be to perform a halo exchange for GradRichNum after it is computed.

However, following the MPAS-Ocean approach and the suggestion in this PR, a more robust solution would be to perform halo exchanges after applying the implicit vertical mixing:

Omega/components/mpas-ocean/src/mode_forward/mpas_ocn_time_integration_split_ab2.F

Lines 2694 to 2740 in 061d0ab

    
           call ocn_vmix_implicit(dt, meshPool, statePool, forcingPool, & 
        
                                  scratchPool, err, 2) 
        
           ! Update halo on u and tracers, which were just updated for 
        
           ! implicit vertical mixing.  If not done, this leads to lack of 
        
           ! volume conservation.  It is required because halo updates in 
        
           ! stage 3 are only conducted on tendencies, not on the velocity 
        
           ! and tracer fields.  So this update is required to communicate 
        
           ! the change due to implicit vertical mixing across the boundary. 
        
           call mpas_timer_start('se vmix halos') 
        
           call mpas_timer_start('se vmix halos normalVelFld') 
        
           call mpas_dmpar_field_halo_exch(domain, 'normalVelocity', & 
        
                                           timeLevel=2) 
        
           call mpas_timer_stop('se vmix halos normalVelFld') 
        
           call mpas_pool_begin_iteration(tracersPool) 
        
           do while ( mpas_pool_get_next_member(tracersPool, groupItr) ) 
        
              if (groupItr%memberType == MPAS_POOL_FIELD) then 
        
                 ! Reset iAge to zero where mask == 0 
        
                 if (config_use_idealAgeTracers.and.trim(groupItr%memberName) == 'idealAgeTracers') then 
        
                     call mpas_pool_get_array(tracersPool, & 
        
                                              groupItr%memberName, & 
        
                                              tracersGroupNew, 2) 
        
                     call mpas_pool_get_subpool(forcingPool, & 
        
                                                'tracersIdealAgeFields', & 
        
                                                 tracersIdealAgeFieldsPool) 
        
                     call mpas_pool_get_array(tracersIdealAgeFieldsPool, & 
        
                                              'idealAgeTracersIdealAgeMask', & 
        
                                               tracerGroupIdealAgeMask) 
        
                     !$omp parallel 
        
                     !$omp do schedule(runtime) 
        
                     do iCell = 1, nCellsOwned 
        
                        tracersGroupNew(1,1,iCell) = tracerGroupIdealAgeMask(1,iCell)*tracersGroupNew(1,1,iCell) 
        
                     end do ! cells 
        
                     !$omp end do 
        
                     !$omp end parallel 
        
                  endif 
        
                 ! Halo update on all tracer groups 
        
                 call mpas_dmpar_field_halo_exch(domain, groupItr%memberName, timeLevel=2) 
        
              end if 
        
           end do

If we go with this approach, I would suggest computing the implicit vertical mixing only for owned cells/edges to avoid unnecessary computations in the halo regions. I'd be happy to help implement this.

@mwarusz , what do you think about this approach?

mwarusz

The fixes look fine to me. CTests pass on perlmutter cpu and gpu.

If we go with this approach, I would suggest computing the implicit vertical mixing only for owned cells/edges to avoid unnecessary computations in the halo regions. I'd be happy to help implement this.

@mwarusz , what do you think about this approach?

Sounds good to me.

hyungyukang

I'm approving this PR based on my testing (E3SM-Project/polaris#683 (comment)) and visual inspection.

Once this PR is merged, I will open a follow-up PR to compute implicit vertical mixing only over owned cells/edges, avoiding unnecessary computations in the halo regions.

andrewdnolan · 2026-08-06T18:59:44Z

Thanks for looking into the details @hyungyukang. Sounds like a good plan of action to me!

sbrus89 · 2026-08-06T21:42:16Z

Passes all omega_pr tests on Frontier with expected diffs on cases with vertical mixing and column test fails unrelated to this PR

Polaris `omega_pr` suite

Baseline workdir: /ccs/home/brus/run/polaris_baseline_omega_pr_omega_craygnu-mphipcc_mpich/
Baseline build: /ccs/home/brus/run/polaris_baseline_omega_pr_omega_craygnu-mphipcc_mpich/build
PR build: /ccs/home/brus/run/polaris_ers_fixes_omega_pr_omega_craygnu-mphipcc_mpich/build
PR workdir: /ccs/home/brus/run/polaris_ers_fixes_omega_pr_omega_craygnu-mphipcc_mpich
Machine: frontier
Partition: batch
Compiler: craygnu-mphipcc
Build type: Release
Log: not found
Result:
- Failures (4 of 17):
  - ocean/column/vmix_stable
  - ocean/column/vmix_unstable
  - ocean/column/ekman
  - ocean/column/inertial
- Diffs (5 of 17):
  - ocean/planar/baroclinic_channel/10km/default
  - ocean/planar/overflow/linear/zstar/smoke_test_horiz_adv_order_2_del4
  - ocean/planar/overflow/linear/zstar/smoke_test_horiz_adv_order_3
  - ocean/planar/overflow/linear/zstar/smoke_test_horiz_adv_order_4
  - ocean/planar/overflow/nonlinear/pstar/smoke_test_horiz_adv_order_4_del4

Polaris `omega_pr` suite

Baseline workdir: /ccs/home/brus/run/polaris_baseline_omega_pr_omega_craygnu_mpich/
Baseline build: /ccs/home/brus/run/polaris_baseline_omega_pr_omega_craygnu_mpich/build
PR build: /ccs/home/brus/run/polaris_ers_fixes_omega_pr_omega_craygnu_mpich/build
PR workdir: /ccs/home/brus/run/polaris_ers_fixes_omega_pr_omega_craygnu_mpich
Machine: frontier
Partition: batch
Compiler: craygnu
Build type: Release
Log: not found
Result:
- Failures (4 of 17):
  - ocean/column/vmix_stable
  - ocean/column/vmix_unstable
  - ocean/column/ekman
  - ocean/column/inertial
- Diffs (5 of 17):
  - ocean/planar/baroclinic_channel/10km/default
  - ocean/planar/overflow/linear/zstar/smoke_test_horiz_adv_order_2_del4
  - ocean/planar/overflow/linear/zstar/smoke_test_horiz_adv_order_3
  - ocean/planar/overflow/linear/zstar/smoke_test_horiz_adv_order_4
  - ocean/planar/overflow/nonlinear/pstar/smoke_test_horiz_adv_order_4_del4

Also passes with craygnu and craygnu-mphipcc:
ocean/planar/baroclinic_channel/10km/decomp
ocean/planar/baroclinic_channel/10km/restart

sbrus89 · 2026-08-07T17:54:40Z

Verified the non-BFB changes from this PR are negligible for the overflow case:
Baseline:

This PR:

sbrus89 · 2026-08-07T18:57:14Z

Passes with local merge:

SMS_Ld3.T62_oQU240.COMEGA-IAF.frontier_craygnu
ERS.T62_oQU240.COMEGA-IAF.frontier_craygnu
PEM.T62_oQU240.COMEGA-IAF.frontier_craygnu

sbrus89 · 2026-08-07T19:31:58Z

Passes CTests on pm-gpu and pm-cpu. Expected diffs and errors unrelated to this PR for omega_pr on pm-cpu:

Polaris `omega_pr` suite

Baseline workdir: /global/homes/s/sbrus/scratch/polaris_baseline_omega_pr_omega_cpu
Baseline build: /global/homes/s/sbrus/scratch/polaris_baseline_omega_pr_omega_cpu/build
PR build: /global/homes/s/sbrus/scratch/polaris_ers_fix_omega_pr_omega_cpu/build
PR workdir: /global/homes/s/sbrus/scratch/polaris_ers_fix_omega_pr_omega_cpu
Machine: pm-cpu
Compiler: gnu
Build type: Release
Log: not found
Result:
- Failures (1 of 18):
  - ocean/column/ekman
- Diffs (8 of 18):
  - ocean/column/vmix_stable
  - ocean/column/vmix_unstable
  - ocean/planar/baroclinic_channel/10km/default
  - ocean/planar/overflow/linear/zstar/smoke_test_horiz_adv_order_2_del4
  - ocean/planar/overflow/linear/zstar/smoke_test_horiz_adv_order_3
  - ocean/planar/overflow/linear/zstar/smoke_test_horiz_adv_order_4
  - ocean/planar/overflow/nonlinear/pstar/smoke_test_horiz_adv_order_4_del4
  - ocean/column/inertial

- Bring in diffs due to vertical mixing fix which addresses ERS and PEM tests (E3SM-Project/Omega#488)

This PR add missing halo exchanges between sub steps for the ForwardBackward time stepper. Without these PEM test using ForwardBackward time stepping failed, even after rebasing onto #488 . Fixes #489.

andrewdnolan added 2 commits July 31, 2026 09:51

Pass logical current level to VertMixImplicit

ddd52a5

Also use getters to access state instead of raw access

Exchange halos after vertical mixing

c05ffb8

andrewdnolan changed the title ~~Bug fixes need to support bit-for-bit restarts (through CIME ERS tests)~~ Fix bugs in calls to vertical mixing within timesteppers Jul 31, 2026

andrewdnolan mentioned this pull request Jul 31, 2026

Add halo exchanges to ForwardBackward Timestepper #490

Merged

16 tasks

andrewdnolan added the bug Something isn't working label Jul 31, 2026

andrewdnolan commented Jul 31, 2026

View reviewed changes

andrewdnolan marked this pull request as ready for review July 31, 2026 22:08

xylar requested review from hyungyukang and mwarusz August 1, 2026 09:02

xylar assigned sbrus89 Aug 1, 2026

hyungyukang mentioned this pull request Aug 5, 2026

Add bcl channel tasks to omega_pr suite E3SM-Project/polaris#683

Merged

1 task

mwarusz approved these changes Aug 6, 2026

View reviewed changes

hyungyukang approved these changes Aug 6, 2026

View reviewed changes

mwarusz mentioned this pull request Aug 6, 2026

Richardson number computation accesses inactive layers #500

Closed

sbrus89 merged commit 3cd7cdd into E3SM-Project:develop Aug 7, 2026
1 check passed

sbrus89 mentioned this pull request Aug 7, 2026

Update Omega submodule E3SM-Project/polaris#698

Merged

6 tasks

andrewdnolan deleted the omega/bfb-restart-fixes branch August 8, 2026 00:36

xylar pushed a commit to E3SM-Project/polaris that referenced this pull request Aug 8, 2026

Update Omega submodule

5d93f38

- Bring in diffs due to vertical mixing fix which addresses ERS and PEM tests (E3SM-Project/Omega#488)

Uh oh!

Conversation

andrewdnolan commented Jul 31, 2026 • edited by sbrus89 Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

andrewdnolan Jul 31, 2026

Choose a reason for hiding this comment

Uh oh!

andrewdnolan commented Jul 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Frontier CTest unit tests:

Uh oh!

andrewdnolan commented Jul 31, 2026

CIME Testing:

Uh oh!

xylar commented Aug 1, 2026

Uh oh!

hyungyukang commented Aug 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

hyungyukang commented Aug 4, 2026

Uh oh!

mwarusz left a comment

Choose a reason for hiding this comment

Uh oh!

hyungyukang left a comment

Choose a reason for hiding this comment

Uh oh!

andrewdnolan commented Aug 6, 2026

Uh oh!

sbrus89 commented Aug 6, 2026

Polaris omega_pr suite

Polaris omega_pr suite

Uh oh!

sbrus89 commented Aug 7, 2026

Uh oh!

sbrus89 commented Aug 7, 2026

Uh oh!

sbrus89 commented Aug 7, 2026

Polaris omega_pr suite

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

andrewdnolan commented Jul 31, 2026 •

edited by sbrus89

Loading

andrewdnolan commented Jul 31, 2026 •

edited

Loading

`Frontier` CTest unit tests:

hyungyukang commented Aug 3, 2026 •

edited

Loading

Polaris `omega_pr` suite

Polaris `omega_pr` suite

Polaris `omega_pr` suite