From fd6e6698771f49dba8af7bacac6ef24bd44b1afc Mon Sep 17 00:00:00 2001 From: jatkinson1000 Date: Mon, 17 Jun 2024 12:17:59 -0600 Subject: [PATCH 1/5] 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 2/5] 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 3/5] 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 4/5] 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 5/5] 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)