From fd6e6698771f49dba8af7bacac6ef24bd44b1afc Mon Sep 17 00:00:00 2001 From: jatkinson1000 Date: Mon, 17 Jun 2024 12:17:59 -0600 Subject: [PATCH 01/22] Add variable for the convect_dp_gw net filepath to be read from namelist. --- src/physics/cam/gw_drag.F90 | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/physics/cam/gw_drag.F90 b/src/physics/cam/gw_drag.F90 index a9cf446b05..cd3ebf4cc9 100644 --- a/src/physics/cam/gw_drag.F90 +++ b/src/physics/cam/gw_drag.F90 @@ -195,6 +195,7 @@ module gw_drag ! Switch for using ML GW parameterisation for deep convection source logical :: gw_convect_dp_ml = .false. logical :: gw_convect_dp_ml_compare = .false. + character(len=132) :: gw_convect_dp_ml_net_path !========================================================================== contains @@ -237,7 +238,7 @@ subroutine gw_drag_readnl(nlfile) gw_oro_south_fac, gw_limit_tau_without_eff, & gw_lndscl_sgh, gw_prndl, gw_apply_tndmax, gw_qbo_hdepth_scaling, & gw_top_taper, front_gaussian_width, & - gw_convect_dp_ml, gw_convect_dp_ml_compare + gw_convect_dp_ml, gw_convect_dp_ml_compare, gw_convect_dp_ml_net_path !---------------------------------------------------------------------- if (use_simple_phys) return @@ -347,6 +348,9 @@ subroutine gw_drag_readnl(nlfile) call mpi_bcast(gw_convect_dp_ml_compare, 1, mpi_logical, mstrid, mpicom, ierr) if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: gw_convect_dp_ml_compare") + call mpi_bcast(gw_convect_dp_ml_net_path, len(gw_convect_dp_ml_net_path), mpi_character, mstrid, mpicom, ierr) + if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: gw_convect_dp_ml_net_path") + ! Check if fcrit2 was set. call shr_assert(fcrit2 /= unset_r8, & "gw_drag_readnl: fcrit2 must be set via the namelist."// & From bed301e262d6c44c080ceaf4da7fd17c3e4c6de0 Mon Sep 17 00:00:00 2001 From: jatkinson1000 Date: Mon, 17 Jun 2024 15:25:46 -0600 Subject: [PATCH 02/22] WIP Adapt gw_drag.F90 to import ftorch and read in a net from file. TODO: Move reading net to init routine. --- src/physics/cam/gw_drag.F90 | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/src/physics/cam/gw_drag.F90 b/src/physics/cam/gw_drag.F90 index cd3ebf4cc9..efd849d3e0 100644 --- a/src/physics/cam/gw_drag.F90 +++ b/src/physics/cam/gw_drag.F90 @@ -29,6 +29,8 @@ module gw_drag use cam_logfile, only: iulog use cam_abortutils, only: endrun + use ftorch + use ref_pres, only: do_molec_diff, nbot_molec, press_lim_idx use physconst, only: cpair @@ -1414,6 +1416,17 @@ subroutine gw_tend(state, pbuf, dt, ptend, cam_in, flx_heat) real(r8) :: zi(state%ncol,pver+1) !------------------------------------------------------------------------ + type(torch_tensor), dimension(1) :: in_tensors + type(torch_module) :: gw_convect_dp_nn + + gw_convect_dp_nn = torch_module_load(gw_convect_dp_ml_net) + + if (masterproc) then + write(iulog,*) 'Made a torch tensor and loaded a net!' + endif + call torch_tensor_delete(in_tensors(1)) + call torch_module_delete(gw_convect_dp_nn) + ! Make local copy of input state. call physics_state_copy(state, state1) From 7968fa95c2dcc1919b3adf7dab1af5d5c8a61cb4 Mon Sep 17 00:00:00 2001 From: jatkinson1000 Date: Mon, 24 Jun 2024 03:55:54 -0600 Subject: [PATCH 03/22] Move reading of NN to initialisation routine for gw_drag. --- src/physics/cam/gw_drag.F90 | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/physics/cam/gw_drag.F90 b/src/physics/cam/gw_drag.F90 index efd849d3e0..1b16318162 100644 --- a/src/physics/cam/gw_drag.F90 +++ b/src/physics/cam/gw_drag.F90 @@ -198,6 +198,7 @@ module gw_drag logical :: gw_convect_dp_ml = .false. logical :: gw_convect_dp_ml_compare = .false. character(len=132) :: gw_convect_dp_ml_net_path + type(torch_module) :: gw_convect_dp_nn !========================================================================== contains @@ -980,6 +981,12 @@ subroutine gw_init() end if + ! Set up neccessary attributes if using ML scheme for convective drag + if ((gw_convect_dp_ml == 'on') .or. (gw_convect_dp_ml == 'bothon')) then + ! Load the convective drag net from TorchScript file + gw_convect_dp_nn = torch_module_load(gw_convect_dp_ml_net) + endif + if (use_gw_convect_sh) then ttend_sh_idx = pbuf_get_index('TTEND_SH') @@ -1417,9 +1424,6 @@ subroutine gw_tend(state, pbuf, dt, ptend, cam_in, flx_heat) !------------------------------------------------------------------------ type(torch_tensor), dimension(1) :: in_tensors - type(torch_module) :: gw_convect_dp_nn - - gw_convect_dp_nn = torch_module_load(gw_convect_dp_ml_net) if (masterproc) then write(iulog,*) 'Made a torch tensor and loaded a net!' From b3ac3e9026c1a6c7b76aec649a8f971f98eccc7a Mon Sep 17 00:00:00 2001 From: jatkinson1000 Date: Mon, 24 Jun 2024 05:00:53 -0600 Subject: [PATCH 04/22] Add a gw_final() subroutine to destroy NN instance. --- src/physics/cam/gw_drag.F90 | 10 ++++++++++ src/physics/cam/physpkg.F90 | 3 +++ 2 files changed, 13 insertions(+) diff --git a/src/physics/cam/gw_drag.F90 b/src/physics/cam/gw_drag.F90 index 1b16318162..cd4aee8596 100644 --- a/src/physics/cam/gw_drag.F90 +++ b/src/physics/cam/gw_drag.F90 @@ -54,6 +54,7 @@ module gw_drag public :: gw_drag_readnl ! Read namelist public :: gw_init ! Initialization public :: gw_tend ! interface to actual parameterization + public :: gw_final ! Finalization ! ! PRIVATE: Rest of the data and interfaces are private to this module @@ -1249,6 +1250,15 @@ end subroutine handle_pio_error !========================================================================== +subroutine gw_final() + ! Destroy neccessary attributes if using ML scheme for convective drag + if ((gw_convect_dp_ml == 'on') .or. (gw_convect_dp_ml == 'bothon')) then + call torch_module_delete(gw_convect_dp_nn) + endif +end subroutine gw_final + +!========================================================================== + subroutine gw_tend(state, pbuf, dt, ptend, cam_in, flx_heat) !----------------------------------------------------------------------- ! Interface for multiple gravity wave drag parameterization. diff --git a/src/physics/cam/physpkg.F90 b/src/physics/cam/physpkg.F90 index c7d5b2524b..66a9d8433a 100644 --- a/src/physics/cam/physpkg.F90 +++ b/src/physics/cam/physpkg.F90 @@ -1337,6 +1337,7 @@ subroutine phys_final( phys_state, phys_tend, pbuf2d ) use microp_aero, only : microp_aero_final use phys_grid_ctem, only : phys_grid_ctem_final use nudging, only: Nudge_Model, nudging_final + use gw_drag, only: gw_final !----------------------------------------------------------------------- ! @@ -1367,6 +1368,8 @@ subroutine phys_final( phys_state, phys_tend, pbuf2d ) call HCOI_Chunk_Final endif + call gw_final() + end subroutine phys_final From 16b53ac5304c6debcf13e0ee79e10fed9bc1f7b9 Mon Sep 17 00:00:00 2001 From: jatkinson1000 Date: Mon, 5 Aug 2024 08:20:58 -0600 Subject: [PATCH 05/22] Remove spurious test code. --- src/physics/cam/gw_drag.F90 | 8 -------- 1 file changed, 8 deletions(-) diff --git a/src/physics/cam/gw_drag.F90 b/src/physics/cam/gw_drag.F90 index cd4aee8596..49e1c36edb 100644 --- a/src/physics/cam/gw_drag.F90 +++ b/src/physics/cam/gw_drag.F90 @@ -1433,14 +1433,6 @@ subroutine gw_tend(state, pbuf, dt, ptend, cam_in, flx_heat) real(r8) :: zi(state%ncol,pver+1) !------------------------------------------------------------------------ - type(torch_tensor), dimension(1) :: in_tensors - - if (masterproc) then - write(iulog,*) 'Made a torch tensor and loaded a net!' - endif - call torch_tensor_delete(in_tensors(1)) - call torch_module_delete(gw_convect_dp_nn) - ! Make local copy of input state. call physics_state_copy(state, state1) From dfd451dd54544270c8ab262e1b7c0cc06c38731a Mon Sep 17 00:00:00 2001 From: jatkinson1000 Date: Tue, 30 Jul 2024 06:57:01 -0600 Subject: [PATCH 06/22] Update source to use latest version of FTorch following API Changes. --- src/physics/cam/gw_drag.F90 | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/physics/cam/gw_drag.F90 b/src/physics/cam/gw_drag.F90 index 49e1c36edb..bb8e5f0b0e 100644 --- a/src/physics/cam/gw_drag.F90 +++ b/src/physics/cam/gw_drag.F90 @@ -199,7 +199,7 @@ module gw_drag logical :: gw_convect_dp_ml = .false. logical :: gw_convect_dp_ml_compare = .false. character(len=132) :: gw_convect_dp_ml_net_path - type(torch_module) :: gw_convect_dp_nn + type(torch_model) :: gw_convect_dp_nn !========================================================================== contains @@ -985,7 +985,7 @@ subroutine gw_init() ! Set up neccessary attributes if using ML scheme for convective drag if ((gw_convect_dp_ml == 'on') .or. (gw_convect_dp_ml == 'bothon')) then ! Load the convective drag net from TorchScript file - gw_convect_dp_nn = torch_module_load(gw_convect_dp_ml_net) + call torch_model_load(gw_convect_dp_nn, gw_convect_dp_ml_net) endif if (use_gw_convect_sh) then @@ -1253,7 +1253,7 @@ end subroutine handle_pio_error subroutine gw_final() ! Destroy neccessary attributes if using ML scheme for convective drag if ((gw_convect_dp_ml == 'on') .or. (gw_convect_dp_ml == 'bothon')) then - call torch_module_delete(gw_convect_dp_nn) + call torch_delete(gw_convect_dp_nn) endif end subroutine gw_final From 0d334db4c236773cab6edcc895f9bfda82701efd Mon Sep 17 00:00:00 2001 From: jatkinson1000 Date: Mon, 5 Aug 2024 07:35:54 -0600 Subject: [PATCH 07/22] WIP add gw_ml file with starting point for a routine to run the ML parameterisation via FTorch. * Implements reshaping of the data to the form expected by the net as a starting point. * Consider using torch layout argument to avoid transposition at a later date. * Does not implement any normalisation of features as we need to discuss how to do this on a chunked grid when data is provided for the global grid. * Not currently called from anywhere. --- src/physics/cam/gw_ml.F90 | 113 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 113 insertions(+) create mode 100644 src/physics/cam/gw_ml.F90 diff --git a/src/physics/cam/gw_ml.F90 b/src/physics/cam/gw_ml.F90 new file mode 100644 index 0000000000..bb2204116a --- /dev/null +++ b/src/physics/cam/gw_ml.F90 @@ -0,0 +1,113 @@ +module gw_ml + +! +! This module handles machine learnt gravity wave schemes developed by +! the DataWave Project +! + +use gw_utils, only: r8 +use ppgrid, only: pver + +use ftorch + +implicit none +private +save + +public :: gw_drag_convect_dp_ml + +contains + +!========================================================================== + +subroutine gw_drag_convect_dp_ml(convect_net, & + ncol, dt, & + u, v, t, dse, nm, netdt, zm, rhoi, ps, lat, lon, & + utgw, vtgw) + + ! Take data from CAM, normalise and concatenate before passing it to the Torch neural + ! net to calculate u and v tendencies. + + ! Neural Net as read in by FTorch + type(torch_model) :: convect_net + + ! Column dimension. + integer, intent(in) :: ncol + + ! Time step. + real(r8), intent(in) :: dt + + ! Midpoint zonal/meridional winds. + real(r8), intent(in) :: u(ncol,pver), v(ncol,pver) + ! Midpoint and interface temperatures. + real(r8), intent(in) :: t(ncol,pver) + ! Dry static energy. + real(r8), intent(in) :: dse(ncol,pver) + ! Midpoint and interface Brunt-Vaisalla frequencies. + real(r8), intent(in) :: nm(ncol,pver) + ! Heating rate due to convection. + real(r8), intent(in) :: netdt(:,:) + ! Midpoint geopotential altitudes. + real(r8), intent(in) :: zm(ncol,pver) + ! Interface densities. + real(r8), intent(in) :: rhoi(ncol,pver+1) + ! Surface Pressure + real(r8), intent(in) :: ps(ncol) + ! Latitude in radians. + real(r8), intent(in) :: lat(ncol) + ! Longitude in radians. + real(r8), intent(in) :: lon(ncol) + + ! Zonal/meridional wind tendencies. + real(r8), intent(out) :: utgw(ncol,pver), vtgw(ncol,pver) + + + !---------------------------Local storage------------------------------- + + ! Level, wavenumber, constituent and column loop indices. + integer :: k, l, m, i + + real(r8), dimension(:,:), target :: net_inputs(8*pver+4, ncol) + real(r8), dimension(:,:), target :: net_outputs(2*pver, ncol) + type(torch_tensor) :: net_input_tensors(1), net_output_tensors(1) + integer :: ninputs = 1 + integer :: noutputs = 1 + integer, dimension(1) :: layout = [1] + + ! Concatenate data into input array and map to a torch_tensor + net_inputs(:pver, :) = transpose(u(:, :)) + net_inputs(pver+1:2*pver, :) = transpose(v(:, :)) + net_inputs(2*pver+1:3*pver, :) = transpose(t(:, :)) + net_inputs(3*pver+1:4*pver, :) = transpose(dse(:, :)) + net_inputs(4*pver+1:5*pver, :) = transpose(nm(:, :)) + net_inputs(5*pver+1:6*pver, :) = transpose(netdt(:, :)) + net_inputs(6*pver+1:7*pver, :) = transpose(zm(:, :)) + net_inputs(7*pver+1:8*pver+1, :) = transpose(rhoi(:, :)) + net_inputs(8*pver+2, :) = ps(:) + net_inputs(8*pver+3, :) = lat(:) + net_inputs(8*pver+4, :) = lon(:) + + ! Loop over columns, create input and infer + do i = 1, ncol + + call torch_tensor_from_array(net_input_tensors(1), net_inputs(:, i), layout, torch_kCPU) + call torch_tensor_from_array(net_output_tensors(1), net_outputs(:, i), layout, torch_kCPU) + + ! Run net forward on data + call torch_model_forward(convect_net, net_input_tensors, net_output_tensors) + + end do + + ! Clean up the tensors + call torch_delete(net_input_tensors) + call torch_delete(net_output_tensors) + + ! Extract data and return + do i = 1, ncol + utgw(i, :) = net_outputs(1:pver, i) + vtgw(i, :) = net_outputs(pver+1:2*pver, i) + end do + +end subroutine gw_drag_convect_dp_ml + +end module gw_ml From 01950bca2f11ea298a1f203b32f8a74b980964d2 Mon Sep 17 00:00:00 2001 From: jatkinson1000 Date: Fri, 16 Aug 2024 03:17:18 -0600 Subject: [PATCH 08/22] Add normalisation file to namelist variables, read, and distribute. --- bld/namelist_files/namelist_defaults_cam.xml | 1 + bld/namelist_files/namelist_definition.xml | 6 ++++++ src/physics/cam/gw_drag.F90 | 9 +++++++-- 3 files changed, 14 insertions(+), 2 deletions(-) diff --git a/bld/namelist_files/namelist_defaults_cam.xml b/bld/namelist_files/namelist_defaults_cam.xml index 0f7dff9325..af39315dce 100644 --- a/bld/namelist_files/namelist_defaults_cam.xml +++ b/bld/namelist_files/namelist_defaults_cam.xml @@ -758,6 +758,7 @@ .false. .false. NONE +NONE 0.03D0 diff --git a/bld/namelist_files/namelist_definition.xml b/bld/namelist_files/namelist_definition.xml index a7ee1cea13..90df443a30 100644 --- a/bld/namelist_files/namelist_definition.xml +++ b/bld/namelist_files/namelist_definition.xml @@ -1405,6 +1405,12 @@ Absolute filepath to the deep convection gravity wave neural net used when `gw_convect_dp_ml` is set to `.true.`. + + Absolute filepath to the deep convection gravity wave normalisation weights (NetCDF) + used when `gw_convect_dp_ml` is set to `.true.`. + + Whether or not to run a piggybacking comparison of the ML deep convection gravity diff --git a/src/physics/cam/gw_drag.F90 b/src/physics/cam/gw_drag.F90 index bb8e5f0b0e..619c1d2f69 100644 --- a/src/physics/cam/gw_drag.F90 +++ b/src/physics/cam/gw_drag.F90 @@ -199,6 +199,7 @@ module gw_drag logical :: gw_convect_dp_ml = .false. logical :: gw_convect_dp_ml_compare = .false. character(len=132) :: gw_convect_dp_ml_net_path + character(len=132) :: gw_convect_dp_ml_norms type(torch_model) :: gw_convect_dp_nn !========================================================================== @@ -242,7 +243,8 @@ subroutine gw_drag_readnl(nlfile) gw_oro_south_fac, gw_limit_tau_without_eff, & gw_lndscl_sgh, gw_prndl, gw_apply_tndmax, gw_qbo_hdepth_scaling, & gw_top_taper, front_gaussian_width, & - gw_convect_dp_ml, gw_convect_dp_ml_compare, gw_convect_dp_ml_net_path + gw_convect_dp_ml, gw_convect_dp_ml_compare, & + gw_convect_dp_ml_net_path, gw_convect_dp_ml_norms !---------------------------------------------------------------------- if (use_simple_phys) return @@ -355,6 +357,9 @@ subroutine gw_drag_readnl(nlfile) call mpi_bcast(gw_convect_dp_ml_net_path, len(gw_convect_dp_ml_net_path), mpi_character, mstrid, mpicom, ierr) if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: gw_convect_dp_ml_net_path") + call mpi_bcast(gw_convect_dp_ml_norms, len(gw_convect_dp_ml_norms), mpi_character, mstrid, mpicom, ierr) + if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: gw_convect_dp_ml_norms") + ! Check if fcrit2 was set. call shr_assert(fcrit2 /= unset_r8, & "gw_drag_readnl: fcrit2 must be set via the namelist."// & @@ -985,7 +990,7 @@ subroutine gw_init() ! Set up neccessary attributes if using ML scheme for convective drag if ((gw_convect_dp_ml == 'on') .or. (gw_convect_dp_ml == 'bothon')) then ! Load the convective drag net from TorchScript file - call torch_model_load(gw_convect_dp_nn, gw_convect_dp_ml_net) + call torch_model_load(gw_convect_dp_nn, gw_convect_dp_ml_net_path) endif if (use_gw_convect_sh) then From 603c78a13c1e188603dcd56ed6eee37b10a80e60 Mon Sep 17 00:00:00 2001 From: jatkinson1000 Date: Mon, 19 Aug 2024 02:26:05 -0600 Subject: [PATCH 09/22] Create gw_ml initialisation routine to read in net and normalisation variables from file. --- src/physics/cam/gw_drag.F90 | 3 +- src/physics/cam/gw_ml.F90 | 261 +++++++++++++++++++++++++++++++++++- 2 files changed, 257 insertions(+), 7 deletions(-) diff --git a/src/physics/cam/gw_drag.F90 b/src/physics/cam/gw_drag.F90 index 619c1d2f69..ecff52bf55 100644 --- a/src/physics/cam/gw_drag.F90 +++ b/src/physics/cam/gw_drag.F90 @@ -42,6 +42,7 @@ module gw_drag use gw_common, only: GWBand use gw_convect, only: BeresSourceDesc use gw_front, only: CMSourceDesc + use gw_ml, only: gw_drag_convect_dp_ml_init, gw_drag_convect_dp_ml ! Typical module header implicit none @@ -990,7 +991,7 @@ subroutine gw_init() ! Set up neccessary attributes if using ML scheme for convective drag if ((gw_convect_dp_ml == 'on') .or. (gw_convect_dp_ml == 'bothon')) then ! Load the convective drag net from TorchScript file - call torch_model_load(gw_convect_dp_nn, gw_convect_dp_ml_net_path) + call gw_drag_convect_dp_ml_init(gw_convect_dp_ml_net_path, gw_convect_dp_ml_norms) endif if (use_gw_convect_sh) then diff --git a/src/physics/cam/gw_ml.F90 b/src/physics/cam/gw_ml.F90 index bb2204116a..1adca43648 100644 --- a/src/physics/cam/gw_ml.F90 +++ b/src/physics/cam/gw_ml.F90 @@ -7,6 +7,8 @@ module gw_ml use gw_utils, only: r8 use ppgrid, only: pver +use spmd_utils, only: mpicom, mstrid=>masterprocid, masterproc, mpi_real8 +use cam_abortutils, only: endrun use ftorch @@ -14,22 +16,49 @@ module gw_ml private save -public :: gw_drag_convect_dp_ml +public :: gw_drag_convect_dp_ml, gw_drag_convect_dp_ml_init + +! Neural Net as read in by FTorch +type(torch_model) :: convect_net + +! Means for normalisation +real(r8) :: utgw_mean(pver), vtgw_mean(pver) +real(r8) :: u_mean(pver), v_mean(pver) +real(r8) :: t_mean(pver) +real(r8) :: dse_mean(pver) +real(r8) :: nm_mean(pver) +real(r8) :: netdt_mean(pver) +real(r8) :: zm_mean(pver) +real(r8) :: rhoi_mean(pver+1) +real(r8) :: ps_mean +real(r8) :: lat_mean +real(r8) :: lon_mean +! Standard deviations for normalisation +real(r8) :: utgw_std(pver), vtgw_std(pver) +real(r8) :: u_std(pver), v_std(pver) +real(r8) :: t_std(pver) +real(r8) :: dse_std(pver) +real(r8) :: nm_std(pver) +real(r8) :: netdt_std(pver) +real(r8) :: zm_std(pver) +real(r8) :: rhoi_std(pver+1) +real(r8) :: ps_std +real(r8) :: lat_std +real(r8) :: lon_std contains !========================================================================== -subroutine gw_drag_convect_dp_ml(convect_net, & - ncol, dt, & +subroutine gw_drag_convect_dp_ml(ncol, dt, & u, v, t, dse, nm, netdt, zm, rhoi, ps, lat, lon, & utgw, vtgw) ! Take data from CAM, normalise and concatenate before passing it to the Torch neural ! net to calculate u and v tendencies. - ! Neural Net as read in by FTorch - type(torch_model) :: convect_net + + ! Column dimension. integer, intent(in) :: ncol @@ -46,7 +75,7 @@ subroutine gw_drag_convect_dp_ml(convect_net, & ! Midpoint and interface Brunt-Vaisalla frequencies. real(r8), intent(in) :: nm(ncol,pver) ! Heating rate due to convection. - real(r8), intent(in) :: netdt(:,:) + real(r8), intent(in) :: netdt(ncol,pver) ! Midpoint geopotential altitudes. real(r8), intent(in) :: zm(ncol,pver) ! Interface densities. @@ -110,4 +139,224 @@ subroutine gw_drag_convect_dp_ml(convect_net, & end subroutine gw_drag_convect_dp_ml + +subroutine gw_drag_convect_dp_ml_init(neural_net_path, norms_path) + + character(len=132), intent(in) :: neural_net_path ! Filepath to PyTorch Torchscript net + character(len=132), intent(in) :: norms_path ! Filepath to NetCDF normalisation weights + + ! Load the convective drag net from TorchScript file + call torch_model_load(convect_net, neural_net_path) + ! read in normalisation weights + call read_norms(norms_path) + +end subroutine gw_drag_convect_dp_ml_init + + +subroutine read_norms(norms_path) + + use netcdf + use error_messages, only: handle_ncerr + + character(len=132), intent(in) :: norms_path ! Filepath to NetCDF normalisation weights + + integer :: ncid, varid, retva, ierr + character(len=*), parameter :: sub = 'gw_ml/F90 read_norms: ' + + ! Load weights from file in master process then broadcast + if (masterproc) then + ! Open the NetCDF file + call handle_ncerr( nf90_open(trim(norms_path), NF90_NOWRITE, ncid), & + "Error opening NetCDF norms file in gw_ml.F90") + + ! We do not need to read in dimensions here as we assume inputs match the grid. + + ! Read in variables (means and deviations). + call handle_ncerr( nf90_inq_varid(ncid, 'U_mean', varid), & + "Error getting U_mean varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_get_var(ncid, varid, u_mean), & + "Error getting U_mean varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_inq_varid(ncid, 'U_std', varid), & + "Error getting U_std varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_get_var(ncid, varid, u_std), & + "Error getting U_std varid from NetCDF Norms file in gw_ml.F90") + + call handle_ncerr( nf90_inq_varid(ncid, 'V_mean', varid), & + "Error getting V_mean varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_get_var(ncid, varid, v_mean), & + "Error getting V_mean varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_inq_varid(ncid, 'V_std', varid), & + "Error getting V_std varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_get_var(ncid, varid, v_std), & + "Error getting V_std varid from NetCDF Norms file in gw_ml.F90") + + call handle_ncerr( nf90_inq_varid(ncid, 'T_mean', varid), & + "Error getting T_mean varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_get_var(ncid, varid, t_mean), & + "Error getting t_mean varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_inq_varid(ncid, 'T_std', varid), & + "Error getting T_std varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_get_var(ncid, varid, t_std), & + "Error getting T_std varid from NetCDF Norms file in gw_ml.F90") + + call handle_ncerr( nf90_inq_varid(ncid, 'DSE_mean', varid), & + "Error getting DSE_mean varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_get_var(ncid, varid, dse_mean), & + "Error getting U_mean varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_inq_varid(ncid, 'DSE_std', varid), & + "Error getting DSE_std varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_get_var(ncid, varid, dse_std), & + "Error getting DSE_std varid from NetCDF Norms file in gw_ml.F90") + + call handle_ncerr( nf90_inq_varid(ncid, 'NMBV_mean', varid), & + "Error getting NMBV_mean varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_get_var(ncid, varid, nm_mean), & + "Error getting NMBV_mean varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_inq_varid(ncid, 'NMBV_std', varid), & + "Error getting NMBV_std varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_get_var(ncid, varid, nm_std), & + "Error getting NMBV_std varid from NetCDF Norms file in gw_ml.F90") + + call handle_ncerr( nf90_inq_varid(ncid, 'NETDT_mean', varid), & + "Error getting NETDT_mean varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_get_var(ncid, varid, netdt_mean), & + "Error getting NETDT_mean varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_inq_varid(ncid, 'NETDT_std', varid), & + "Error getting NETDT_std varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_get_var(ncid, varid, netdt_std), & + "Error getting NETDT_std varid from NetCDF Norms file in gw_ml.F90") + + call handle_ncerr( nf90_inq_varid(ncid, 'Z3_mean', varid), & + "Error getting Z3_mean varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_get_var(ncid, varid, zm_mean), & + "Error getting Z3_mean varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_inq_varid(ncid, 'Z3_std', varid), & + "Error getting Z3_std varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_get_var(ncid, varid, zm_std), & + "Error getting Z3_std varid from NetCDF Norms file in gw_ml.F90") + + call handle_ncerr( nf90_inq_varid(ncid, 'RHOI_mean', varid), & + "Error getting RHOI_mean varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_get_var(ncid, varid, rhoi_mean), & + "Error getting RHOI_mean varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_inq_varid(ncid, 'RHOI_std', varid), & + "Error getting RHOI_std varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_get_var(ncid, varid, rhoi_std), & + "Error getting RHOI_std varid from NetCDF Norms file in gw_ml.F90") + + call handle_ncerr( nf90_inq_varid(ncid, 'PS_mean', varid), & + "Error getting PS_mean varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_get_var(ncid, varid, ps_mean), & + "Error getting PS_mean varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_inq_varid(ncid, 'PS_std', varid), & + "Error getting PS_std varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_get_var(ncid, varid, ps_std), & + "Error getting PS_std varid from NetCDF Norms file in gw_ml.F90") + + call handle_ncerr( nf90_inq_varid(ncid, 'lat_mean', varid), & + "Error getting lat_mean varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_get_var(ncid, varid, lat_mean), & + "Error getting lat_mean varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_inq_varid(ncid, 'lat_std', varid), & + "Error getting lat_std varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_get_var(ncid, varid, lat_std), & + "Error getting lat_std varid from NetCDF Norms file in gw_ml.F90") + + call handle_ncerr( nf90_inq_varid(ncid, 'lon_mean', varid), & + "Error getting lon_mean varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_get_var(ncid, varid, lon_mean), & + "Error getting lon_mean varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_inq_varid(ncid, 'lon_std', varid), & + "Error getting lon_std varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_get_var(ncid, varid, lon_std), & + "Error getting lon_std varid from NetCDF Norms file in gw_ml.F90") + + call handle_ncerr( nf90_inq_varid(ncid, 'UTGWSPEC_mean', varid), & + "Error getting UTGWSPEC_mean varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_get_var(ncid, varid, utgw_mean), & + "Error getting UTGWSPEC_mean varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_inq_varid(ncid, 'UTGWSPEC_std', varid), & + "Error getting UTGWSPEC_std varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_get_var(ncid, varid, utgw_std), & + "Error getting UTGWSPEC_std varid from NetCDF Norms file in gw_ml.F90") + + call handle_ncerr( nf90_inq_varid(ncid, 'VTGWSPEC_mean', varid), & + "Error getting VTGWSPEC_mean varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_get_var(ncid, varid, vtgw_mean), & + "Error getting VTGWSPEC_mean varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_inq_varid(ncid, 'VTGWSPEC_std', varid), & + "Error getting VTGWSPEC_std varid from NetCDF Norms file in gw_ml.F90") + call handle_ncerr( nf90_get_var(ncid, varid, vtgw_std), & + "Error getting VTGWSPEC_std varid from NetCDF Norms file in gw_ml.F90") + + endif + + ! Broadcast normalisation variables to other processes + call mpi_bcast(utgw_mean, pver, mpi_real8, mstrid, mpicom, ierr) + if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: utgw_mean from gw_ml.F90") + call mpi_bcast(utgw_std, pver, mpi_real8, mstrid, mpicom, ierr) + if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: utgw_std from gw_ml.F90") + + call mpi_bcast(vtgw_mean, pver, mpi_real8, mstrid, mpicom, ierr) + if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: vtgw_mean from gw_ml.F90") + call mpi_bcast(vtgw_std, pver, mpi_real8, mstrid, mpicom, ierr) + if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: vtgw_std from gw_ml.F90") + + call mpi_bcast(u_mean, pver, mpi_real8, mstrid, mpicom, ierr) + if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: u_mean from gw_ml.F90") + call mpi_bcast(u_std, pver, mpi_real8, mstrid, mpicom, ierr) + if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: u_std from gw_ml.F90") + + call mpi_bcast(v_mean, pver, mpi_real8, mstrid, mpicom, ierr) + if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: v_mean from gw_ml.F90") + call mpi_bcast(v_std, pver, mpi_real8, mstrid, mpicom, ierr) + if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: v_std from gw_ml.F90") + + call mpi_bcast(t_mean, pver, mpi_real8, mstrid, mpicom, ierr) + if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: t_mean from gw_ml.F90") + call mpi_bcast(t_std, pver, mpi_real8, mstrid, mpicom, ierr) + if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: t_std from gw_ml.F90") + + call mpi_bcast(dse_mean, pver, mpi_real8, mstrid, mpicom, ierr) + if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: dse_mean from gw_ml.F90") + call mpi_bcast(dse_std, pver, mpi_real8, mstrid, mpicom, ierr) + if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: dse_std from gw_ml.F90") + + call mpi_bcast(nm_mean, pver, mpi_real8, mstrid, mpicom, ierr) + if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: nm_mean from gw_ml.F90") + call mpi_bcast(nm_std, pver, mpi_real8, mstrid, mpicom, ierr) + if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: nm_std from gw_ml.F90") + + call mpi_bcast(netdt_mean, pver, mpi_real8, mstrid, mpicom, ierr) + if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: netdt_mean from gw_ml.F90") + call mpi_bcast(netdt_std, pver, mpi_real8, mstrid, mpicom, ierr) + if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: netdt_std from gw_ml.F90") + + call mpi_bcast(zm_mean, pver, mpi_real8, mstrid, mpicom, ierr) + if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: zm_mean from gw_ml.F90") + call mpi_bcast(zm_std, pver, mpi_real8, mstrid, mpicom, ierr) + if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: zm_std from gw_ml.F90") + + call mpi_bcast(rhoi_mean, pver+1, mpi_real8, mstrid, mpicom, ierr) + if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: rhoi_mean from gw_ml.F90") + call mpi_bcast(rhoi_std, pver+1, mpi_real8, mstrid, mpicom, ierr) + if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: rhoi_std from gw_ml.F90") + + call mpi_bcast(ps_mean, pver, mpi_real8, mstrid, mpicom, ierr) + if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: ps_mean from gw_ml.F90") + call mpi_bcast(ps_std, pver, mpi_real8, mstrid, mpicom, ierr) + if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: ps_std from gw_ml.F90") + + call mpi_bcast(lat_mean, 1, mpi_real8, mstrid, mpicom, ierr) + if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: lat_mean from gw_ml.F90") + call mpi_bcast(lat_std, 1, mpi_real8, mstrid, mpicom, ierr) + if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: lat_std from gw_ml.F90") + + call mpi_bcast(lon_mean, 1, mpi_real8, mstrid, mpicom, ierr) + if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: lon_mean from gw_ml.F90") + call mpi_bcast(lon_std, 1, mpi_real8, mstrid, mpicom, ierr) + if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: lon_std from gw_ml.F90") + +end subroutine read_norms + end module gw_ml From 72a50a9eed16da2ddfe5b551c55694f657486b05 Mon Sep 17 00:00:00 2001 From: jatkinson1000 Date: Mon, 19 Aug 2024 02:57:52 -0600 Subject: [PATCH 10/22] Add finalisation routine to gw_ml to destroy the net. --- src/physics/cam/gw_drag.F90 | 5 +++-- src/physics/cam/gw_ml.F90 | 10 +++++++++- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/src/physics/cam/gw_drag.F90 b/src/physics/cam/gw_drag.F90 index ecff52bf55..03f7414d35 100644 --- a/src/physics/cam/gw_drag.F90 +++ b/src/physics/cam/gw_drag.F90 @@ -42,7 +42,8 @@ module gw_drag use gw_common, only: GWBand use gw_convect, only: BeresSourceDesc use gw_front, only: CMSourceDesc - use gw_ml, only: gw_drag_convect_dp_ml_init, gw_drag_convect_dp_ml + use gw_ml, only: gw_drag_convect_dp_ml_init, gw_drag_convect_dp_ml_final, & + gw_drag_convect_dp_ml ! Typical module header implicit none @@ -1259,7 +1260,7 @@ end subroutine handle_pio_error subroutine gw_final() ! Destroy neccessary attributes if using ML scheme for convective drag if ((gw_convect_dp_ml == 'on') .or. (gw_convect_dp_ml == 'bothon')) then - call torch_delete(gw_convect_dp_nn) + call gw_drag_convect_dp_ml_final() endif end subroutine gw_final diff --git a/src/physics/cam/gw_ml.F90 b/src/physics/cam/gw_ml.F90 index 1adca43648..609bf6af97 100644 --- a/src/physics/cam/gw_ml.F90 +++ b/src/physics/cam/gw_ml.F90 @@ -16,7 +16,7 @@ module gw_ml private save -public :: gw_drag_convect_dp_ml, gw_drag_convect_dp_ml_init +public :: gw_drag_convect_dp_ml, gw_drag_convect_dp_ml_init, gw_drag_convect_dp_ml_final ! Neural Net as read in by FTorch type(torch_model) :: convect_net @@ -153,6 +153,14 @@ subroutine gw_drag_convect_dp_ml_init(neural_net_path, norms_path) end subroutine gw_drag_convect_dp_ml_init +subroutine gw_drag_convect_dp_ml_final() + + ! Destroy the convective drag net + call torch_delete(convect_net) + +end subroutine gw_drag_convect_dp_ml_final + + subroutine read_norms(norms_path) use netcdf From a9f95ee9df8fbee6f94aad03710036e70a136e45 Mon Sep 17 00:00:00 2001 From: jatkinson1000 Date: Mon, 19 Aug 2024 05:16:35 -0600 Subject: [PATCH 11/22] Add normalisation routine for inputs to net. --- src/physics/cam/gw_ml.F90 | 56 ++++++++++++++++++++++++++++++--------- 1 file changed, 44 insertions(+), 12 deletions(-) diff --git a/src/physics/cam/gw_ml.F90 b/src/physics/cam/gw_ml.F90 index 609bf6af97..3f80974ae0 100644 --- a/src/physics/cam/gw_ml.F90 +++ b/src/physics/cam/gw_ml.F90 @@ -103,18 +103,9 @@ subroutine gw_drag_convect_dp_ml(ncol, dt, & integer :: noutputs = 1 integer, dimension(1) :: layout = [1] - ! Concatenate data into input array and map to a torch_tensor - net_inputs(:pver, :) = transpose(u(:, :)) - net_inputs(pver+1:2*pver, :) = transpose(v(:, :)) - net_inputs(2*pver+1:3*pver, :) = transpose(t(:, :)) - net_inputs(3*pver+1:4*pver, :) = transpose(dse(:, :)) - net_inputs(4*pver+1:5*pver, :) = transpose(nm(:, :)) - net_inputs(5*pver+1:6*pver, :) = transpose(netdt(:, :)) - net_inputs(6*pver+1:7*pver, :) = transpose(zm(:, :)) - net_inputs(7*pver+1:8*pver+1, :) = transpose(rhoi(:, :)) - net_inputs(8*pver+2, :) = ps(:) - net_inputs(8*pver+3, :) = lat(:) - net_inputs(8*pver+4, :) = lon(:) + ! Normalise and concatenate the data + call normalise_data(ncol, u, v, t, dse, nm, netdt, zm, rhoi, ps, lat, lon, & + net_inputs) ! Loop over columns, create input and infer do i = 1, ncol @@ -367,4 +358,45 @@ subroutine read_norms(norms_path) end subroutine read_norms +subroutine normalise_data(ncol, u, v, t, dse, nm, netdt, zm, rhoi, ps, lat, lon, & + nn_input) + + integer, intent(in) :: ncol + real(r8), intent(in) :: u(ncol,pver), v(ncol,pver) + real(r8), intent(in) :: t(ncol,pver) + real(r8), intent(in) :: dse(ncol,pver) + real(r8), intent(in) :: nm(ncol,pver) + real(r8), intent(in) :: netdt(:,:) + real(r8), intent(in) :: zm(ncol,pver) + real(r8), intent(in) :: rhoi(ncol,pver+1) + real(r8), intent(in) :: ps(ncol) + real(r8), intent(in) :: lat(ncol) + real(r8), intent(in) :: lon(ncol) + + real(r8), intent(out) :: nn_input(8*pver+4, ncol) + + integer :: i + + ! Loop over each column. + ! Normalise data (subtract mean, divide by deviation), transpose into format + ! expected by the NN, and concatenate into a single input tensor as expected by the NN. + do i = 1,ncol + + nn_input(:pver, i) = (u(i, :) - u_mean(:))/u_std(:) + nn_input(pver+1:2*pver, i) = (v(i, :) - v_mean(:))/v_std(:) + nn_input(2*pver+1:3*pver, i) = (t(i, :) - t_mean(:))/t_std(:) + nn_input(3*pver+1:4*pver, i) = (dse(i, :) - dse_mean(:))/dse_std(:) + nn_input(4*pver+1:5*pver, i) = (nm(i, :) - nm_mean(:))/nm_std(:) + nn_input(5*pver+1:6*pver, i) = (netdt(i, :) - netdt_mean(:))/netdt_std(:) + nn_input(6*pver+1:7*pver, i) = (zm(i, :) - zm_mean(:))/zm_std(:) + nn_input(7*pver+1:8*pver+1, i) = (rhoi(i, :) - rhoi_mean(:))/rhoi_std(:) + nn_input(8*pver+2, i) = (ps(i) - ps_mean)/ps_std + nn_input(8*pver+3, i) = (lat(i) - lat_mean)/lat_std + nn_input(8*pver+4, i) = (lon(i) - lon_mean)/lon_std + + end do + +end subroutine normalise_data + + end module gw_ml From 19fa91fff6a7c6d62e9ee92f70322475824bf942 Mon Sep 17 00:00:00 2001 From: jatkinson1000 Date: Mon, 19 Aug 2024 05:25:56 -0600 Subject: [PATCH 12/22] Add denormalisation routine for the gw_convect_dp net. --- src/physics/cam/gw_ml.F90 | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/src/physics/cam/gw_ml.F90 b/src/physics/cam/gw_ml.F90 index 3f80974ae0..e9085779a5 100644 --- a/src/physics/cam/gw_ml.F90 +++ b/src/physics/cam/gw_ml.F90 @@ -122,11 +122,8 @@ subroutine gw_drag_convect_dp_ml(ncol, dt, & call torch_delete(net_input_tensors) call torch_delete(net_output_tensors) - ! Extract data and return - do i = 1, ncol - utgw(i, :) = net_outputs(1:pver, i) - vtgw(i, :) = net_outputs(pver+1:2*pver, i) - end do + ! Denormalise outputs and extract the data + call denormalise_data(ncol, utgw, vtgw, net_outputs) end subroutine gw_drag_convect_dp_ml @@ -398,5 +395,20 @@ subroutine normalise_data(ncol, u, v, t, dse, nm, netdt, zm, rhoi, ps, lat, lon, end subroutine normalise_data +subroutine denormalise_data(ncol, utgw, vtgw, nn_output) + + integer, intent(in) :: ncol + real(r8), intent(out) :: utgw(ncol,pver), vtgw(ncol,pver) + real(r8), intent(in) :: nn_output(2*pver, ncol) + + integer :: i + + ! Extract data, denormalise, and deconcatenate from NN output tensor + do i = 1, ncol + utgw(i, :) = (nn_output(1:pver, i) * utgw_std(:)) + utgw_mean(:) + vtgw(i, :) = (nn_output(pver+1:2*pver, i) * vtgw_std(:)) + vtgw_mean(:) + end do + +end subroutine denormalise_data end module gw_ml From 7bf43e178e696d027eabffff2ef4958e6f704f06 Mon Sep 17 00:00:00 2001 From: jatkinson1000 Date: Mon, 19 Aug 2024 07:48:16 -0600 Subject: [PATCH 13/22] BugFix: Remove redundant net variable --- src/physics/cam/gw_drag.F90 | 1 - 1 file changed, 1 deletion(-) diff --git a/src/physics/cam/gw_drag.F90 b/src/physics/cam/gw_drag.F90 index 03f7414d35..e7ffa96955 100644 --- a/src/physics/cam/gw_drag.F90 +++ b/src/physics/cam/gw_drag.F90 @@ -202,7 +202,6 @@ module gw_drag logical :: gw_convect_dp_ml_compare = .false. character(len=132) :: gw_convect_dp_ml_net_path character(len=132) :: gw_convect_dp_ml_norms - type(torch_model) :: gw_convect_dp_nn !========================================================================== contains From c01db7e1c25e903717d855feeed5f4b733f1f276 Mon Sep 17 00:00:00 2001 From: jatkinson1000 Date: Mon, 19 Aug 2024 08:21:56 -0600 Subject: [PATCH 14/22] Couple gw_drag to the ML scheme --- src/physics/cam/gw_drag.F90 | 25 +++++++++++++++++-------- src/physics/cam/gw_ml.F90 | 3 +-- 2 files changed, 18 insertions(+), 10 deletions(-) diff --git a/src/physics/cam/gw_drag.F90 b/src/physics/cam/gw_drag.F90 index e7ffa96955..e17e91ff94 100644 --- a/src/physics/cam/gw_drag.F90 +++ b/src/physics/cam/gw_drag.F90 @@ -1284,6 +1284,7 @@ subroutine gw_tend(state, pbuf, dt, ptend, cam_in, flx_heat) use gw_oro, only: gw_oro_src use gw_front, only: gw_cm_src use gw_convect, only: gw_beres_src + use gw_ml, only: gw_drag_convect_dp_ml !------------------------------Arguments-------------------------------- type(physics_state), intent(in) :: state ! physics state structure @@ -1437,6 +1438,9 @@ subroutine gw_tend(state, pbuf, dt, ptend, cam_in, flx_heat) real(r8) :: piln(state%ncol,pver+1) real(r8) :: zm(state%ncol,pver) real(r8) :: zi(state%ncol,pver+1) + real(r8) :: ps(state%ncol) + real(r8) :: lat(state%ncol) + real(r8) :: lon(state%ncol) !------------------------------------------------------------------------ ! Make local copy of input state. @@ -1458,6 +1462,9 @@ subroutine gw_tend(state, pbuf, dt, ptend, cam_in, flx_heat) piln = state1%lnpint(:ncol,:) zm = state1%zm(:ncol,:) zi = state1%zi(:ncol,:) + ps = state1%ps(:ncol) + lat = state1%lat(:ncol) + lon = state1%lon(:ncol) lq = .true. call physics_ptend_init(ptend, state1%psetcols, "Gravity wave drag", & @@ -1550,8 +1557,8 @@ subroutine gw_tend(state, pbuf, dt, ptend, cam_in, flx_heat) write(iulog,*) "Using the ML scheme for convective gravity waves." end if - ! Solve for the drag profile with Beres source spectrum. - ! Placeholder to be replaced with the ML scheme + ! Solve for the drag profile with Beres source spectrum as per original CAM. + ! This is required to obtain values for qtgw, ttgw, and egwdffi. call gw_drag_prof(ncol, band_mid, p, src_level, tend_level, dt, & t, vramp, & piln, rhoi, nm, ni, ubm, ubi, xv, yv, & @@ -1559,14 +1566,18 @@ subroutine gw_tend(state, pbuf, dt, ptend, cam_in, flx_heat) ttgw_temp, qtgw_temp, egwdffi, gwut, dttdf, dttke, & lapply_effgw_in=gw_apply_tndmax) + call gw_drag_convect_dp_ml(ncol, dt, & + u, v, t, dse, nm, ttend_dp, zm, rhoi, ps, & + lat, lon, & + utgw_temp, vtgw_temp) + if (gw_convect_dp_ml) then ! Save the results to apply to ptend for simulation updates - ! TODO: Check how to handle tendencies not output by ML scheme - qtgw = qtgw_temp ! in the ml scheme there is no qtgw so use qtgw = 0.0 - ttgw = ttgw_temp ! in the ml scheme there is no ttgw so use ttgw = 0.0 + qtgw = qtgw_temp ! not output by NN so use qtgw from original scheme + ttgw = ttgw_temp ! not output by NN so use ttgw from original scheme utgw = utgw_temp vtgw = vtgw_temp - ! in the ml scheme there is not egwdffi set, so use egwdffi = 0.0 + ! egwdffi is not output by the NN so use egwdffi from original scheme end if end if @@ -1574,13 +1585,11 @@ subroutine gw_tend(state, pbuf, dt, ptend, cam_in, flx_heat) taucd = calc_taucd(ncol, band_mid%ngwv, tend_level, tau, c, xv, yv, ubi) ! add the diffusion coefficients - ! TODO: Check how to handle egwdffi not output by ML scheme do k = 1, pver+1 egwdffi_tot(:,k) = egwdffi_tot(:,k) + egwdffi(:,k) end do ! Store constituents tendencies - ! TODO: Check how to handle qtgw not output by ML scheme do m=1, pcnst do k = 1, pver ptend%q(:ncol,k,m) = ptend%q(:ncol,k,m) + qtgw(:,k,m) diff --git a/src/physics/cam/gw_ml.F90 b/src/physics/cam/gw_ml.F90 index e9085779a5..278882e812 100644 --- a/src/physics/cam/gw_ml.F90 +++ b/src/physics/cam/gw_ml.F90 @@ -93,8 +93,7 @@ subroutine gw_drag_convect_dp_ml(ncol, dt, & !---------------------------Local storage------------------------------- - ! Level, wavenumber, constituent and column loop indices. - integer :: k, l, m, i + integer :: i real(r8), dimension(:,:), target :: net_inputs(8*pver+4, ncol) real(r8), dimension(:,:), target :: net_outputs(2*pver, ncol) From 921a189703cca9796f0d707e9a1e459f5e1b46b9 Mon Sep 17 00:00:00 2001 From: jatkinson1000 Date: Mon, 19 Aug 2024 10:50:13 -0600 Subject: [PATCH 15/22] Bugfix: Update ML switches to be logicals when calling gw_drag init and finalize routines. --- src/physics/cam/gw_drag.F90 | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/physics/cam/gw_drag.F90 b/src/physics/cam/gw_drag.F90 index e17e91ff94..0a293eab76 100644 --- a/src/physics/cam/gw_drag.F90 +++ b/src/physics/cam/gw_drag.F90 @@ -989,7 +989,7 @@ subroutine gw_init() end if ! Set up neccessary attributes if using ML scheme for convective drag - if ((gw_convect_dp_ml == 'on') .or. (gw_convect_dp_ml == 'bothon')) then + if ((gw_convect_dp_ml) .or. (gw_convect_dp_ml_compare)) then ! Load the convective drag net from TorchScript file call gw_drag_convect_dp_ml_init(gw_convect_dp_ml_net_path, gw_convect_dp_ml_norms) endif @@ -1258,7 +1258,7 @@ end subroutine handle_pio_error subroutine gw_final() ! Destroy neccessary attributes if using ML scheme for convective drag - if ((gw_convect_dp_ml == 'on') .or. (gw_convect_dp_ml == 'bothon')) then + if ((gw_convect_dp_ml) .or. (gw_convect_dp_ml_compare)) then call gw_drag_convect_dp_ml_final() endif end subroutine gw_final From aaecc4a0eb8868999fac320809198aba01bde78f Mon Sep 17 00:00:00 2001 From: jatkinson1000 Date: Mon, 9 Sep 2024 04:26:34 -0600 Subject: [PATCH 16/22] [DEBUG: Run net with ones as the input to compare to same operation using pure python (Same result obtained).] --- src/physics/cam/gw_ml.F90 | 74 ++++++++++++++++++++++++++++++++++++++- 1 file changed, 73 insertions(+), 1 deletion(-) diff --git a/src/physics/cam/gw_ml.F90 b/src/physics/cam/gw_ml.F90 index 278882e812..5bc0a8e4f1 100644 --- a/src/physics/cam/gw_ml.F90 +++ b/src/physics/cam/gw_ml.F90 @@ -8,6 +8,7 @@ module gw_ml use gw_utils, only: r8 use ppgrid, only: pver use spmd_utils, only: mpicom, mstrid=>masterprocid, masterproc, mpi_real8 +use cam_logfile, only: iulog use cam_abortutils, only: endrun use ftorch @@ -59,7 +60,6 @@ subroutine gw_drag_convect_dp_ml(ncol, dt, & - ! Column dimension. integer, intent(in) :: ncol @@ -103,13 +103,36 @@ subroutine gw_drag_convect_dp_ml(ncol, dt, & integer, dimension(1) :: layout = [1] ! Normalise and concatenate the data + + if (masterproc) then + write(iulog,*) "===================================================================" + write(iulog,*) "Start of Call to gw_drag_convect_dp_ml" + write(iulog,*) "Using input of Ones." + write(iulog,*) "===================================================================" +! write(iulog,*) "lat_mean is: ", lat_mean +! write(iulog,*) "lat_std is: ", lat_std +! write(iulog,*) "lon_mean is: ", lon_mean +! write(iulog,*) "lon_std is: ", lon_std + end if + call normalise_data(ncol, u, v, t, dse, nm, netdt, zm, rhoi, ps, lat, lon, & net_inputs) + net_inputs(:, :) = 1.0 + + if (masterproc) then + write(iulog,*) "Data has been normalised." + write(iulog,*) "First line of input tensor has shape: ", shape(net_inputs(:, 1)) + write(iulog,*) "First line of input tensor is:" + write(iulog,*) net_inputs(:, 1) + end if + ! Loop over columns, create input and infer do i = 1, ncol call torch_tensor_from_array(net_input_tensors(1), net_inputs(:, i), layout, torch_kCPU) + !write(iulog,*) "torch_tensor_from_array is:" + !torch_tensor_print(net_input_tensors(1) call torch_tensor_from_array(net_output_tensors(1), net_outputs(:, i), layout, torch_kCPU) ! Run net forward on data @@ -117,6 +140,15 @@ subroutine gw_drag_convect_dp_ml(ncol, dt, & end do + if (masterproc) then + write(iulog,*) "Net has been run over all columns." +! write(iulog,*) "First line of output tensor has shape: ", shape(net_outputs(:, 1)) + write(iulog,*) "First line of output tensor is:" + write(iulog,*) net_outputs(1:5, 1) + write(iulog,*) "First line of ugtw before denormalisation is:" + write(iulog,*) utgw(1, 1:5) + end if + ! Clean up the tensors call torch_delete(net_input_tensors) call torch_delete(net_output_tensors) @@ -124,6 +156,16 @@ subroutine gw_drag_convect_dp_ml(ncol, dt, & ! Denormalise outputs and extract the data call denormalise_data(ncol, utgw, vtgw, net_outputs) + if (masterproc) then + write(iulog,*) "First line of utgw tensor back in main routine is:" + write(iulog,*) utgw(1, 1:5) + + write(iulog,*) "===================================================================" + write(iulog,*) "End of Call to gw_drag_convect_dp_ml" + write(iulog,*) "===================================================================" + end if + + end subroutine gw_drag_convect_dp_ml @@ -158,6 +200,14 @@ subroutine read_norms(norms_path) integer :: ncid, varid, retva, ierr character(len=*), parameter :: sub = 'gw_ml/F90 read_norms: ' + if (masterproc) then + write(iulog,*) "before reading norms:" + write(iulog,*) "lat_mean is: ", lat_mean + write(iulog,*) "lat_std is: ", lat_std + write(iulog,*) "lon_mean is: ", lon_mean + write(iulog,*) "lon_std is: ", lon_std + end if + ! Load weights from file in master process then broadcast if (masterproc) then ! Open the NetCDF file @@ -352,6 +402,14 @@ subroutine read_norms(norms_path) call mpi_bcast(lon_std, 1, mpi_real8, mstrid, mpicom, ierr) if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: lon_std from gw_ml.F90") + if (masterproc) then + write(iulog,*) "after reading norms:" + write(iulog,*) "lat_mean is: ", lat_mean + write(iulog,*) "lat_std is: ", lat_std + write(iulog,*) "lon_mean is: ", lon_mean + write(iulog,*) "lon_std is: ", lon_std + end if + end subroutine read_norms subroutine normalise_data(ncol, u, v, t, dse, nm, netdt, zm, rhoi, ps, lat, lon, & @@ -402,12 +460,26 @@ subroutine denormalise_data(ncol, utgw, vtgw, nn_output) integer :: i + if (masterproc) then + write(iulog,*) "Denormalising data:" + write(iulog,*) "First line of output tensor is:" + write(iulog,*) nn_output(1:5, 1) + write(iulog,*) "Expected first line of output denormalised is:" + write(iulog,*) (nn_output(1:5, 1) * utgw_std(1:5)) + utgw_mean(1:5) + end if + ! Extract data, denormalise, and deconcatenate from NN output tensor do i = 1, ncol utgw(i, :) = (nn_output(1:pver, i) * utgw_std(:)) + utgw_mean(:) vtgw(i, :) = (nn_output(pver+1:2*pver, i) * vtgw_std(:)) + vtgw_mean(:) end do + if (masterproc) then + write(iulog,*) " Actual First line of output tensor is:" + write(iulog,*) utgw(1, 1:5) + write(iulog,*) "Denormalisation complete" + end if + end subroutine denormalise_data end module gw_ml From 5990616001de400399da99de706d7ee87c1f8536 Mon Sep 17 00:00:00 2001 From: jatkinson1000 Date: Mon, 9 Sep 2024 05:15:01 -0600 Subject: [PATCH 17/22] Revert changes to debug commit using ones as net input. --- src/physics/cam/gw_ml.F90 | 76 ++------------------------------------- 1 file changed, 2 insertions(+), 74 deletions(-) diff --git a/src/physics/cam/gw_ml.F90 b/src/physics/cam/gw_ml.F90 index 5bc0a8e4f1..83575e74e5 100644 --- a/src/physics/cam/gw_ml.F90 +++ b/src/physics/cam/gw_ml.F90 @@ -8,7 +8,6 @@ module gw_ml use gw_utils, only: r8 use ppgrid, only: pver use spmd_utils, only: mpicom, mstrid=>masterprocid, masterproc, mpi_real8 -use cam_logfile, only: iulog use cam_abortutils, only: endrun use ftorch @@ -60,6 +59,7 @@ subroutine gw_drag_convect_dp_ml(ncol, dt, & + ! Column dimension. integer, intent(in) :: ncol @@ -103,36 +103,13 @@ subroutine gw_drag_convect_dp_ml(ncol, dt, & integer, dimension(1) :: layout = [1] ! Normalise and concatenate the data - - if (masterproc) then - write(iulog,*) "===================================================================" - write(iulog,*) "Start of Call to gw_drag_convect_dp_ml" - write(iulog,*) "Using input of Ones." - write(iulog,*) "===================================================================" -! write(iulog,*) "lat_mean is: ", lat_mean -! write(iulog,*) "lat_std is: ", lat_std -! write(iulog,*) "lon_mean is: ", lon_mean -! write(iulog,*) "lon_std is: ", lon_std - end if - call normalise_data(ncol, u, v, t, dse, nm, netdt, zm, rhoi, ps, lat, lon, & net_inputs) - net_inputs(:, :) = 1.0 - - if (masterproc) then - write(iulog,*) "Data has been normalised." - write(iulog,*) "First line of input tensor has shape: ", shape(net_inputs(:, 1)) - write(iulog,*) "First line of input tensor is:" - write(iulog,*) net_inputs(:, 1) - end if - ! Loop over columns, create input and infer do i = 1, ncol call torch_tensor_from_array(net_input_tensors(1), net_inputs(:, i), layout, torch_kCPU) - !write(iulog,*) "torch_tensor_from_array is:" - !torch_tensor_print(net_input_tensors(1) call torch_tensor_from_array(net_output_tensors(1), net_outputs(:, i), layout, torch_kCPU) ! Run net forward on data @@ -140,15 +117,6 @@ subroutine gw_drag_convect_dp_ml(ncol, dt, & end do - if (masterproc) then - write(iulog,*) "Net has been run over all columns." -! write(iulog,*) "First line of output tensor has shape: ", shape(net_outputs(:, 1)) - write(iulog,*) "First line of output tensor is:" - write(iulog,*) net_outputs(1:5, 1) - write(iulog,*) "First line of ugtw before denormalisation is:" - write(iulog,*) utgw(1, 1:5) - end if - ! Clean up the tensors call torch_delete(net_input_tensors) call torch_delete(net_output_tensors) @@ -156,16 +124,6 @@ subroutine gw_drag_convect_dp_ml(ncol, dt, & ! Denormalise outputs and extract the data call denormalise_data(ncol, utgw, vtgw, net_outputs) - if (masterproc) then - write(iulog,*) "First line of utgw tensor back in main routine is:" - write(iulog,*) utgw(1, 1:5) - - write(iulog,*) "===================================================================" - write(iulog,*) "End of Call to gw_drag_convect_dp_ml" - write(iulog,*) "===================================================================" - end if - - end subroutine gw_drag_convect_dp_ml @@ -200,15 +158,7 @@ subroutine read_norms(norms_path) integer :: ncid, varid, retva, ierr character(len=*), parameter :: sub = 'gw_ml/F90 read_norms: ' - if (masterproc) then - write(iulog,*) "before reading norms:" - write(iulog,*) "lat_mean is: ", lat_mean - write(iulog,*) "lat_std is: ", lat_std - write(iulog,*) "lon_mean is: ", lon_mean - write(iulog,*) "lon_std is: ", lon_std - end if - - ! Load weights from file in master process then broadcast + ! Load normalisation weights from file in master process then broadcast if (masterproc) then ! Open the NetCDF file call handle_ncerr( nf90_open(trim(norms_path), NF90_NOWRITE, ncid), & @@ -402,14 +352,6 @@ subroutine read_norms(norms_path) call mpi_bcast(lon_std, 1, mpi_real8, mstrid, mpicom, ierr) if (ierr /= 0) call endrun(sub//": FATAL: mpi_bcast: lon_std from gw_ml.F90") - if (masterproc) then - write(iulog,*) "after reading norms:" - write(iulog,*) "lat_mean is: ", lat_mean - write(iulog,*) "lat_std is: ", lat_std - write(iulog,*) "lon_mean is: ", lon_mean - write(iulog,*) "lon_std is: ", lon_std - end if - end subroutine read_norms subroutine normalise_data(ncol, u, v, t, dse, nm, netdt, zm, rhoi, ps, lat, lon, & @@ -460,26 +402,12 @@ subroutine denormalise_data(ncol, utgw, vtgw, nn_output) integer :: i - if (masterproc) then - write(iulog,*) "Denormalising data:" - write(iulog,*) "First line of output tensor is:" - write(iulog,*) nn_output(1:5, 1) - write(iulog,*) "Expected first line of output denormalised is:" - write(iulog,*) (nn_output(1:5, 1) * utgw_std(1:5)) + utgw_mean(1:5) - end if - ! Extract data, denormalise, and deconcatenate from NN output tensor do i = 1, ncol utgw(i, :) = (nn_output(1:pver, i) * utgw_std(:)) + utgw_mean(:) vtgw(i, :) = (nn_output(pver+1:2*pver, i) * vtgw_std(:)) + vtgw_mean(:) end do - if (masterproc) then - write(iulog,*) " Actual First line of output tensor is:" - write(iulog,*) utgw(1, 1:5) - write(iulog,*) "Denormalisation complete" - end if - end subroutine denormalise_data end module gw_ml From 83cb8c9fc62f9c57ce4859ae9d146244ba057a73 Mon Sep 17 00:00:00 2001 From: jatkinson1000 Date: Mon, 9 Sep 2024 07:43:35 -0600 Subject: [PATCH 18/22] Add output for U and V tendencies from GW from Net and Beres schemes. --- src/physics/cam/gw_drag.F90 | 29 +++++++++++++++++++++++++++-- src/physics/cam/gw_ml.F90 | 5 +++++ 2 files changed, 32 insertions(+), 2 deletions(-) diff --git a/src/physics/cam/gw_drag.F90 b/src/physics/cam/gw_drag.F90 index 0a293eab76..08a1b0b236 100644 --- a/src/physics/cam/gw_drag.F90 +++ b/src/physics/cam/gw_drag.F90 @@ -990,8 +990,23 @@ subroutine gw_init() ! Set up neccessary attributes if using ML scheme for convective drag if ((gw_convect_dp_ml) .or. (gw_convect_dp_ml_compare)) then - ! Load the convective drag net from TorchScript file - call gw_drag_convect_dp_ml_init(gw_convect_dp_ml_net_path, gw_convect_dp_ml_norms) + ! Load the convective drag net from TorchScript file + call gw_drag_convect_dp_ml_init(gw_convect_dp_ml_net_path, gw_convect_dp_ml_norms) + + ! Register fields with the output buffer + call addfld ('UTGW_NN ', (/ 'lev' /), 'A', 'm/s2', & + 'U tendency due to convective gravity wave drag from NN.') + call addfld ('VTGW_NN ', (/ 'lev' /), 'A', 'm/s2', & + 'V tendency due to convective gravity wave drag from NN.') + call register_vector_field('UTGW_NN', 'VTGW_NN') + + ! Register fields with the output buffer + call addfld ('UTGW_BERES ', (/ 'lev' /), 'A', 'm/s2', & + 'U tendency due to convective gravity wave drag from BERES.') + call addfld ('VTGW_BERES ', (/ 'lev' /), 'A', 'm/s2', & + 'V tendency due to convective gravity wave drag from NN.') + call register_vector_field('UTGW_BERES', 'VTGW_BERES') + endif if (use_gw_convect_sh) then @@ -1543,6 +1558,12 @@ subroutine gw_tend(state, pbuf, dt, ptend, cam_in, flx_heat) ttgw_temp, qtgw_temp, egwdffi, gwut, dttdf, dttke, & lapply_effgw_in=gw_apply_tndmax) + if (gw_convect_dp_ml_compare) then + ! write fields out for comparison + call outfld('UTGW_BERES', utgw_temp, ncol, lchnk) + call outfld('VTGW_BERES', vtgw_temp, ncol, lchnk) + end if + if (.not. gw_convect_dp_ml) then ! Save the results to apply to ptend for simulation updates qtgw = qtgw_temp @@ -1571,6 +1592,10 @@ subroutine gw_tend(state, pbuf, dt, ptend, cam_in, flx_heat) lat, lon, & utgw_temp, vtgw_temp) + ! write fields out for comparison + call outfld('UTGW_NN', utgw_temp, ncol, lchnk) + call outfld('VTGW_NN', vtgw_temp, ncol, lchnk) + if (gw_convect_dp_ml) then ! Save the results to apply to ptend for simulation updates qtgw = qtgw_temp ! not output by NN so use qtgw from original scheme diff --git a/src/physics/cam/gw_ml.F90 b/src/physics/cam/gw_ml.F90 index 83575e74e5..c0ba4c2bd9 100644 --- a/src/physics/cam/gw_ml.F90 +++ b/src/physics/cam/gw_ml.F90 @@ -137,6 +137,11 @@ subroutine gw_drag_convect_dp_ml_init(neural_net_path, norms_path) ! read in normalisation weights call read_norms(norms_path) + if (masterproc) then + write(iulog,*)'gw_convect_net loaded from: ', neural_net_path + write(iulog,*)'Normalisation weights loaded from: ', norms_path + endif + end subroutine gw_drag_convect_dp_ml_init From 03bd7f53fc735f2d5cda456d3598ed3da1bf34cd Mon Sep 17 00:00:00 2001 From: jatkinson1000 Date: Tue, 15 Oct 2024 05:35:29 -0600 Subject: [PATCH 19/22] Set tendencies not predicted by the net to 0.0 for the ml scheme rather than running the physics scheme beforehand. --- src/physics/cam/gw_drag.F90 | 26 ++++++++++++-------------- 1 file changed, 12 insertions(+), 14 deletions(-) diff --git a/src/physics/cam/gw_drag.F90 b/src/physics/cam/gw_drag.F90 index 08a1b0b236..ae7bbfca6f 100644 --- a/src/physics/cam/gw_drag.F90 +++ b/src/physics/cam/gw_drag.F90 @@ -1549,6 +1549,9 @@ subroutine gw_tend(state, pbuf, dt, ptend, cam_in, flx_heat) u, v, ttend_dp(:ncol,:), zm, src_level, tend_level, tau, & ubm, ubi, xv, yv, c, hdepth, maxq0) + ! TODO: If we are running with the ML scheme save tau to a temp variable so we + ! Can reset to it instead of having it updated in the physics scheme. + if ((.not. gw_convect_dp_ml) .or. (gw_convect_dp_ml_compare)) then ! Solve for the drag profile with Beres source spectrum. call gw_drag_prof(ncol, band_mid, p, src_level, tend_level, dt, & @@ -1578,17 +1581,8 @@ subroutine gw_tend(state, pbuf, dt, ptend, cam_in, flx_heat) write(iulog,*) "Using the ML scheme for convective gravity waves." end if - ! Solve for the drag profile with Beres source spectrum as per original CAM. - ! This is required to obtain values for qtgw, ttgw, and egwdffi. - call gw_drag_prof(ncol, band_mid, p, src_level, tend_level, dt, & - t, vramp, & - piln, rhoi, nm, ni, ubm, ubi, xv, yv, & - effgw, c, kvtt, q, dse, tau, utgw_temp, vtgw_temp, & - ttgw_temp, qtgw_temp, egwdffi, gwut, dttdf, dttke, & - lapply_effgw_in=gw_apply_tndmax) - call gw_drag_convect_dp_ml(ncol, dt, & - u, v, t, dse, nm, ttend_dp, zm, rhoi, ps, & + u, v, t, dse, nm, ttend_dp(:ncol,:), zm, rhoi, ps, & lat, lon, & utgw_temp, vtgw_temp) @@ -1597,12 +1591,16 @@ subroutine gw_tend(state, pbuf, dt, ptend, cam_in, flx_heat) call outfld('VTGW_NN', vtgw_temp, ncol, lchnk) if (gw_convect_dp_ml) then - ! Save the results to apply to ptend for simulation updates - qtgw = qtgw_temp ! not output by NN so use qtgw from original scheme - ttgw = ttgw_temp ! not output by NN so use ttgw from original scheme + ! Save the results to apply to ptend for simulation updates. + ! Some tendencies are not supplied by the NN, set these to 0.0 as minor. + ttgw = 0.0 + qtgw = 0.0 + egwdffi = 0.0 + gwut = 0.0 + dttdf = 0.0 + dttke = 0.0 utgw = utgw_temp vtgw = vtgw_temp - ! egwdffi is not output by the NN so use egwdffi from original scheme end if end if From bf2a8bb14eb460e7e1f555460fbd26023ad725b4 Mon Sep 17 00:00:00 2001 From: jatkinson1000 Date: Tue, 15 Oct 2024 05:43:03 -0600 Subject: [PATCH 20/22] Add iulog to gw_ml to write outputs to diagnostics in init scheme. --- src/physics/cam/gw_ml.F90 | 1 + 1 file changed, 1 insertion(+) diff --git a/src/physics/cam/gw_ml.F90 b/src/physics/cam/gw_ml.F90 index c0ba4c2bd9..d9b7623775 100644 --- a/src/physics/cam/gw_ml.F90 +++ b/src/physics/cam/gw_ml.F90 @@ -9,6 +9,7 @@ module gw_ml use ppgrid, only: pver use spmd_utils, only: mpicom, mstrid=>masterprocid, masterproc, mpi_real8 use cam_abortutils, only: endrun +use cam_logfile, only: iulog use ftorch From 218242f736ff15e5d5722f3767b5713c2b8e3284 Mon Sep 17 00:00:00 2001 From: jatkinson1000 Date: Mon, 30 Sep 2024 10:05:18 -0600 Subject: [PATCH 21/22] Add code to produce output required by the testcase. --- src/physics/cam/gw_drag.F90 | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/src/physics/cam/gw_drag.F90 b/src/physics/cam/gw_drag.F90 index ae7bbfca6f..cc39d3bad6 100644 --- a/src/physics/cam/gw_drag.F90 +++ b/src/physics/cam/gw_drag.F90 @@ -1063,6 +1063,15 @@ subroutine gw_init() call add_default('EKGW', 1, ' ') end if + call addfld ('RHOI', (/ 'ilev' /), 'A', 'kg/m3', & + 'density at interfaces') + + call addfld ('DSE', (/ 'lev' /), 'I', 'J/kg', & + 'dry static energy') + + call addfld ('NMBV', (/ 'lev' /), 'I', 'J/kg', & + 'Brunt Vaisala Frequency') + call addfld ('UTGW_TOTAL', (/ 'lev' /), 'A','m/s2', & 'Total U tendency due to gravity wave drag') call addfld ('VTGW_TOTAL', (/ 'lev' /), 'A','m/s2', & @@ -2096,6 +2105,10 @@ subroutine gw_tend(state, pbuf, dt, ptend, cam_in, flx_heat) call outfld('CLDLIQTGW', ptend%q(:,:,ixcldliq), pcols, lchnk) call outfld('CLDICETGW', ptend%q(:,:,ixcldice), pcols, lchnk) + call outfld('RHOI', rhoi, pcols, lchnk) + call outfld('DSE', dse, ncol, lchnk) + call outfld('NMBV', nm, ncol, lchnk) + ! Destroy objects. call p%finalize() From a2d881c5af4ccb4003e65fa4344b3cede7d01595 Mon Sep 17 00:00:00 2001 From: jatkinson1000 Date: Tue, 15 Oct 2024 06:02:11 -0600 Subject: [PATCH 22/22] Update rhoi output to be ncol rather than pcol. --- src/physics/cam/gw_drag.F90 | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/physics/cam/gw_drag.F90 b/src/physics/cam/gw_drag.F90 index cc39d3bad6..d4e0601014 100644 --- a/src/physics/cam/gw_drag.F90 +++ b/src/physics/cam/gw_drag.F90 @@ -2105,7 +2105,7 @@ subroutine gw_tend(state, pbuf, dt, ptend, cam_in, flx_heat) call outfld('CLDLIQTGW', ptend%q(:,:,ixcldliq), pcols, lchnk) call outfld('CLDICETGW', ptend%q(:,:,ixcldice), pcols, lchnk) - call outfld('RHOI', rhoi, pcols, lchnk) + call outfld('RHOI', rhoi, ncol, lchnk) call outfld('DSE', dse, ncol, lchnk) call outfld('NMBV', nm, ncol, lchnk)