Add split-explicit AB2 time stepping capability to the ocean component

[hyungyukang]

The second-order Adams Bashforth (AB2) time stepping method is applied to the baroclinic system of MPAS-Ocean to enable faster computation of the baroclinic system. The AB2 time stepping method, one of multistep methods, computes time stepping procedure (Stage 1~3) once per time step, while the predictor-corrector which is a default scheme computes it twice per time step. In practice, the AB2 method can provide a speedup of 1.5x to 1.8x. Due to its high sensitivity to time step size, the predictor-corrector schem is only applied to the layer thickness equation.

All subroutines and the barotropic system advance (Stage 2) in the AB2 code are the same as used in the split-explicit code (mpas_ocn_time_integration_split.F).

To enable this option, users need to set config_time_integrator = 'split_explicit_ab2'.

Results of G-case (100 years) and WCYCL1850 (30 years) and more details of the method can be found here: https://acme-climate.atlassian.net/wiki/spaces/ImPACTS/pages/3826385232/AB2+time+stepping+for+the+baroclinic+system

[NML]
[non-BFB]

[github-actions]

PR Preview Action v1.4.4
🚀 Deployed preview to https://E3SM-Project.github.io/E3SM/pr-preview/pr-5989/
on branch gh-pages at 2023-11-15 20:42 UTC

[sarats]

Tangentially, we need to discuss if we wish to include this PR in our planned high res GPU simulations
https://acme-climate.atlassian.net/wiki/spaces/ImPACTS/pages/3902799937/High-res+Ocean+GPU+simulations

[hyungyukang]

Tangentially, we need to discuss if we wish to include this PR in our planned high res GPU simulations
https://acme-climate.atlassian.net/wiki/spaces/ImPACTS/pages/3902799937/High-res+Ocean+GPU+simulations

@sarats , I will run this code on GPU and report here. Thanks!

[hyungyukang]

Tangentially, we need to discuss if we wish to include this PR in our planned high res GPU simulations
https://acme-climate.atlassian.net/wiki/spaces/ImPACTS/pages/3902799937/High-res+Ocean+GPU+simulations

@sarats , After modifying some OpenACC directives, this code runs well on GPUs as tested on PM-GPU for both the standalone MPAS-O and the ocean-ice coupled (GMPAS-IAF) case.

This PR does not address the issue of running simulations on PM-GPU. I've referenced @grnydawn 's fix for PM-GPU at here: https://github.com/E3SM-Project/E3SM/tree/ykim/mpas/gcase.

[mark-petersen]

@hyungyukang can you comment here on the exact tests you did for performance and quality? Do you have scaling plots for particular resolutions and machines? If information is in slide format, you can just upload them.

[hyungyukang]

@mark-petersen , I've run the standalone MPAS-O on RRS18to6v3 mesh with a real-world configuration using atmospheric forcing to check performance and quality.

• Computational performance (Runtime was measured from 'se timestep')
  • Tested on the AF machine at OLCF, HPC11 (Miller), which has a nearly identical architecture to NERSC's Perlmutter.

• Kinetic energy field (k=1) comparison
  • AB2 uses 75% of the default del4 diffusion coefficient

[mark-petersen]

@hyungyukang, thank you for these results. They look fantastic!

I tested this branch with both optimized and debug, using gnu on perlmutter and intel on chrysalis. Using the default flags (including 'split-explicit' timestepping) everything passes and compares bfb against the master branch point, as expected. Using the new AB2 method (config_time_integrator = 'split_explicit_ab2') in debug, both gnu and intel produce the following results.

nightly suite output:

00:05 PASS ocean_baroclinic_channel_10km_default
00:03 PASS ocean_baroclinic_channel_10km_threads_test
00:03 PASS ocean_baroclinic_channel_10km_decomp_test
00:03 PASS ocean_baroclinic_channel_10km_restart_test
01:42 PASS ocean_global_ocean_QU240_mesh
00:47 PASS ocean_global_ocean_QU240_WOA23_init
00:35 PASS ocean_global_ocean_QU240_WOA23_performance_test
01:04 FAIL ocean_global_ocean_QU240_WOA23_restart_test
01:04 PASS ocean_global_ocean_QU240_WOA23_decomp_test
01:04 PASS ocean_global_ocean_QU240_WOA23_threads_test
00:37 PASS ocean_global_ocean_QU240_WOA23_analysis_test
00:35 PASS ocean_global_ocean_QU240_WOA23_RK4_performance_test
01:05 PASS ocean_global_ocean_QU240_WOA23_RK4_restart_test
01:04 PASS ocean_global_ocean_QU240_WOA23_RK4_decomp_test
01:04 PASS ocean_global_ocean_QU240_WOA23_RK4_threads_test
00:00 PASS ocean_global_ocean_QUwISC240_mesh
00:00 PASS ocean_global_ocean_QUwISC240_WOA23_init
01:11 PASS ocean_global_ocean_QUwISC240_WOA23_performance_test
00:12 PASS ocean_ice_shelf_2d_5km_z-star_restart_test
00:13 PASS ocean_ice_shelf_2d_5km_z-level_restart_test
00:06 PASS ocean_ziso_20km_default
00:05 PASS ocean_ziso_20km_with_frazil

Note that the restart test fails with AB2. Do we need to include the previous two timestep states in the restart file for AB2, rather than just one? A bit-for-bit match between fields when comparing with and without a restart is required.

[hyungyukang]

@mark-petersen , Thanks for testing! I don't think we need to include previous two steps, but I'm going to look into it today. For QU240 restart tests, AB2 passes QU240_PHC_restart_test but failed QU240_WOA23_restart_test. What's the difference between those two QU240 tests?

[xylar]

For QU240 restart tests, AB2 passes QU240_PHC_restart_test but failed QU240_WOA23_restart_test. What's the difference between two QU240 tests?

That sounds very odd because the difference between the two is just the initial condition (i.e. the T and S fields being interpolated from). Are you sure you ran the PHC restart test with config_time_integrator = 'split_explicit_ab2'?

[hyungyukang]

For QU240 restart tests, AB2 passes QU240_PHC_restart_test but failed QU240_WOA23_restart_test. What's the difference between two QU240 tests?

That sounds very odd because the difference between the two is just the initial condition (i.e. the T and S fields being interpolated from). Are you sure you ran the PHC restart test with config_time_integrator = 'split_explicit_ab2'?

@xylar , Another odd thing is that AB2 passes the baroclinic channel restart tests on both Compy and Perlmutter...
Anyway, I left my Compy RSA authenticator at home so I can't check my configuration on Compy now. But I will setup the same tests on Perlmutter and look into results and codes again. Please let me know if you have any suggestions on a test side!

[xylar]

The method obviously doesn't require storing multiple time levels in restart files or no restart tests would be passing. This is obviously a different AB2 method than the one I used when I was a grad student (involving times n and n-1).

[mark-petersen]

The restart tests work and are bfb for active ocean cells. For some reason on AB2 deeper land edges get that dummy value of -1e34, but only for restart runs and not for the very first forward run. I'm looking into it now.

cd ocean/global_ocean/QU240/WOA23/restart_test
ncdump -v normalVelocity restart_run/output.nc |tail -n3
  0, 0, 0, 0, -1e+34, -1e+34, -1e+34, -1e+34, -1e+34, -1e+34, -1e+34, -1e+34,
    -1e+34, -1e+34, -1e+34, -1e+34 ;
}

[hyungyukang]

@mark-petersen and @xylar, Thanks for reviewing and your suggestions! Based on your findings and suggestions, I added a fix to subroutine ocn_time_integrated_split_ab2_init to initialize normalVelocity=0 on inactive cells when restarting. This also fixes the G-case restart issue we discussed yesterday.

code change

      if ( config_do_restart ) then
         block => domain % blocklist
         call mpas_pool_get_subpool(block%structs, 'state', statePool)
         call mpas_pool_get_subpool(statePool, 'tracers', tracersPool)
         call mpas_pool_get_subpool(block%structs, 'mesh', meshPool)

         call mpas_pool_get_dimension(block % dimensions, 'nVertLevels', nVertLevels)
         call mpas_pool_get_dimension(block % dimensions, 'nEdges', nEdges)

         call mpas_pool_get_array(statePool, 'normalVelocity', &
                                              normalVelocity, 1)

         call mpas_pool_get_array(meshPool, 'minLevelEdgeBot', minLevelEdgeBot)
         call mpas_pool_get_array(meshPool, 'maxLevelEdgeTop', maxLevelEdgeTop)

         ! normalVelocity=0 on inactive cells
         do iEdge = 1, nEdges
            ! at top
            do k = 1, minLevelEdgeBot(iEdge)-1
               normalVelocity(k,iEdge) = 0.0_RKIND
            enddo

            ! at bottom
            do k = maxLevelEdgeTop(iEdge)+1, nVertLevels
               normalVelocity(k,iEdge) = 0.0_RKIND
            enddo
         enddo ! edge loop

      end if

nightly suite output

00:15 PASS ocean_baroclinic_channel_10km_default
00:15 PASS ocean_baroclinic_channel_10km_threads_test
00:09 PASS ocean_baroclinic_channel_10km_decomp_test
00:09 PASS ocean_baroclinic_channel_10km_restart_test
04:24 PASS ocean_global_ocean_QU240_mesh
03:48 PASS ocean_global_ocean_QU240_WOA23_init
03:46 PASS ocean_global_ocean_QU240_WOA23_performance_test
07:41 PASS ocean_global_ocean_QU240_WOA23_restart_test
06:09 PASS ocean_global_ocean_QU240_WOA23_decomp_test
05:44 PASS ocean_global_ocean_QU240_WOA23_threads_test
06:30 PASS ocean_global_ocean_QU240_WOA23_analysis_test
02:38 PASS ocean_global_ocean_QU240_WOA23_RK4_performance_test
05:50 PASS ocean_global_ocean_QU240_WOA23_RK4_restart_test
05:27 PASS ocean_global_ocean_QU240_WOA23_RK4_decomp_test
05:28 PASS ocean_global_ocean_QU240_WOA23_RK4_threads_test
00:00 PASS ocean_global_ocean_QUwISC240_mesh
00:00 PASS ocean_global_ocean_QUwISC240_WOA23_init
05:31 PASS ocean_global_ocean_QUwISC240_WOA23_performance_test
00:48 PASS ocean_ice_shelf_2d_5km_z-star_restart_test
00:45 PASS ocean_ice_shelf_2d_5km_z-level_restart_test
00:18 PASS ocean_ziso_20km_default
00:12 PASS ocean_ziso_20km_with_frazil
Total runtime 66:37
PASS: All passed successfully!

And I double checked I'm using config_time_integrator = 'split_explicit_ab2' in namelist files for the nightly suite.

[hyungyukang]

The attached is slides about this AB2 implementation.
MPASO_AB2_timestep.pdf

[mark-petersen]

Thanks! I tested with the nightly suite with gnu and intel, and everything passes.

[jonbob]

Testing of this PR in conjunction with PR #6035 is being undertaken by the coupled model group. The run with these two PRs is 20231106.v3b01-AB2.piControl.chrysalis and the corresponding case without AB2 for an apples-to-apples comparison is 20231105.v3b01.piControl.chrysalis. Both runs have gone more than 200 years

[mark-petersen]

Thanks @jonbob for the B-case comparisons. I spent an hour carefully comparing the two simulations. Here are my conclusions. In all screenshots we have MPAS-Analysis from years 151-200 with:

• control: left or top 20231105.v3b01.piControl.chrysalis
• AB2: right or bottom 20231106.v3b01-AB2.piControl.chrysalis

1. Dynamic variables are identical by eye, within variability: transports, MHT, SST, EKE.

plots:

2. AMOC is low in both, but slightly higher in AB2 (~9 Sv) versus control (~8 Sv). This difference might lie within the variability.

plots:

Right plot is squashed for axis comparisons:

3. Sea ice fields look identical by eye.

plots:

4. THE ONLY VISIBLE DIFFERENCE that I found is in the heat content trends. AB2 OHC drift is smaller, which is better for an 1850 control run.

plots:

Right plot is squashed for axis comparisons:

5. Despite the difference in OHC trends, the global comparisons at these depths still look identical by eye, so there are no obvious problems with a particular region or process.

plots:

For example, Antarctic bottom temperature is very warm in both, but identical by eye between the control and AB2. Also for salinity.

plots:

In summary, the only visible differences with AB2 are the temperature trends. The temperature drifts are smaller than the control, so are preferred. In addition, the spatial distribution of the temperatures look nearly identical by eye, so there are no problems that stand out.

[mark-petersen]

I am approving again based on the B-case simulations and results above. I feel comfortable with AB2 as the default baroclinic time stepping scheme for V3.

[xylar]

I focused on the differences in temperature (and heat content) in the shallower ocean. I'm seeing small differences but they're much less than the biases compared with Argo:

AB2:

SE:

So I'm comfortable with the AB2 results. The performance gain is remarkable! Excellent work!

[cbegeman]

@hyungyukang and @mark-petersen I'm glad to hear the 200 year comparison turned out so well (similar to control). Great work!

[vanroekel]

@hyungyukang can we modify this comment to clarify that the AB2 stepping is for momentum only? The tracer is euler forward.

[hyungyukang]

@vanroekel , Sure I will change it. Thanks a lot!

[vanroekel]

@hyungyukang this is excellent, great work! I have one tiny request, when I looked through the code it appears AB2 is applied to momentum only and I believe you confirmed that. Could we clarify that in the header of the ab2 file?

[vanroekel]

approving based on visual inspection and testing from @hyungyukang @mark-petersen and @jonbob

[jonbob]

passes sanity testing with expected DIFFs -- merged to next

[hyungyukang]

@mark-petersen , @sarats , @xylar , @cbegeman , @vanroekel , and @jonbob , thank you so much for your review, testing, suggestions, comments, and approval of this PR!

[jonbob]

merged to master -- expected DIFFs will be blessed with PR #6035

[amametjanov]

It appears that after the merge nightly OpenACC test

./cime/CIME/scripts/create_test SMS_Ld1.T62_oEC60to30v3.CMPASO-NYF

stopped working with either --machine ascent --compiler pgigpu or --machine pm-gpu --compiler nvidiagpu.
I took a stab at fixing initial build errors, but there's still quite a few variables that need to be copied CPU-to-GPU: e.g.

FATAL ERROR: data in PRESENT clause was not found on device 1: name=layerthickness

[hyungyukang]

It appears that after the merge nightly OpenACC test

./cime/CIME/scripts/create_test SMS_Ld1.T62_oEC60to30v3.CMPASO-NYF

stopped working with either --machine ascent --compiler pgigpu or --machine pm-gpu --compiler nvidiagpu. I took a stab at fixing initial build errors, but there's still quite a few variables that need to be copied CPU-to-GPU: e.g.

FATAL ERROR: data in PRESENT clause was not found on device 1: name=layerthickness

@amametjanov , Thanks for reporting it. I was able to run the AB2 code on both pm-gpu and Frontier-gpu, but I had to do some modifications for a macro -DMPAS_OPENACC which will not be activated just for using pm-gpu. I've referenced @grnydawn 's fix for PM-GPU at here: https://github.com/E3SM-Project/E3SM/tree/ykim/mpas/gcase, but I believe this modification is not merged yet.

I also found that I made a mistake in mpas_ocn_diagnostics_variables.F when I commit this code for GPU runs. I will open a new PR with some modifications to run the code on GPUs properly.