Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Load Neural Net #12

Open
wants to merge 5 commits into
base: datawave_ml
Choose a base branch
from
Open

Load Neural Net #12

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
fd6e669
Add variable for the convect_dp_gw net filepath to be read from namel…
jatkinson1000 Jun 17, 2024
bed301e
WIP Adapt gw_drag.F90 to import ftorch and read in a net from file. T…
jatkinson1000 Jun 17, 2024
7968fa9
Move reading of NN to initialisation routine for gw_drag.
jatkinson1000 Jun 24, 2024
b3ac3e9
Add a gw_final() subroutine to destroy NN instance.
jatkinson1000 Jun 24, 2024
16b53ac
Remove spurious test code.
jatkinson1000 Aug 5, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 24 additions & 1 deletion src/physics/cam/gw_drag.F90
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -52,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
Expand Down Expand Up @@ -195,6 +198,8 @@ 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
type(torch_module) :: gw_convect_dp_nn

!==========================================================================
contains
Expand Down Expand Up @@ -237,7 +242,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
Expand Down Expand Up @@ -347,6 +352,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."// &
Expand Down Expand Up @@ -974,6 +982,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')
Expand Down Expand Up @@ -1236,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.
Expand Down
3 changes: 3 additions & 0 deletions src/physics/cam/physpkg.F90
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

!-----------------------------------------------------------------------
!
Expand Down Expand Up @@ -1367,6 +1368,8 @@ subroutine phys_final( phys_state, phys_tend, pbuf2d )
call HCOI_Chunk_Final
endif

call gw_final()

end subroutine phys_final


Expand Down