Line data    Source code

       1             : module tau_neural_net_quantile
       2             : 
       3             :     use shr_kind_mod,   only: r8=>shr_kind_r8
       4             : 
       5             :     use module_neural_net, only : Dense, init_neural_net, load_quantile_scale_values
       6             :     use module_neural_net, only : quantile_transform, quantile_inv_transform, neural_net_predict
       7             : 
       8             :     implicit none
       9             :     integer, parameter, public :: i8 = selected_int_kind(18)
      10             :     integer, parameter :: num_inputs = 7
      11             :     integer, parameter :: num_outputs = 3
      12             :     integer, parameter :: batch_size = 1
      13             : 
      14             :     ! Neural networks and scale values saved within the scope of the module.
      15             :     ! Need to call initialize_tau_emulators to load weights and tables from disk.
      16             :     type(Dense), allocatable, save :: q_all(:)
      17             :     real(r8), dimension(:, :), allocatable, save :: input_scale_values
      18             :     real(r8), dimension(:, :), allocatable, save :: output_scale_values
      19             : contains
      20             : 
      21             : 
      22           0 :     subroutine initialize_tau_emulators( stochastic_emulated_filename_quantile, stochastic_emulated_filename_input_scale, &
      23           0 :                                          stochastic_emulated_filename_output_scale, iulog, errstring)
      24             : 
      25             :     ! Load neural network netCDF files and scaling values. Values are placed in to emulators,
      26             :     ! input_scale_values, and output_scale_values.
      27             :     character(len=*), intent(in) ::  stochastic_emulated_filename_quantile, stochastic_emulated_filename_input_scale, &
      28             :                                      stochastic_emulated_filename_output_scale
      29             :     integer,          intent(in)  :: iulog
      30             :     character(128),   intent(out) :: errstring  ! output status (non-blank for error return)
      31             : 
      32           0 :         errstring = ''
      33             : 
      34           0 :         write(iulog,*) "Begin loading neural nets"
      35           0 :         call init_neural_net(trim(stochastic_emulated_filename_quantile), batch_size, q_all, iulog, errstring)
      36           0 :         if (trim(errstring) /= '') return
      37           0 :         write(iulog,*) "End loading neural nets"
      38             :         ! Load the scale values from a csv file.
      39           0 :         call load_quantile_scale_values(trim(stochastic_emulated_filename_input_scale), input_scale_values, iulog, errstring)
      40           0 :         call load_quantile_scale_values(trim(stochastic_emulated_filename_output_scale), output_scale_values, iulog, errstring)
      41           0 :         write(iulog,*) "Loaded neural nets scaling values"
      42             : 
      43             :     end subroutine initialize_tau_emulators
      44             : 
      45             : 
      46           0 :     subroutine tau_emulated_cloud_rain_interactions(qc, nc, qr, nr, rho, lcldm, &
      47           0 :             precip_frac, mgncol, q_small, qc_tend, qr_tend, nc_tend, nr_tend)
      48             :         ! Calculates emulated tau microphysics tendencies from neural networks.
      49             :         !
      50             :         ! Input args:
      51             :         !   qc: cloud water mixing ratio in kg kg-1
      52             :         !   nc: cloud water number concentration in particles m-3
      53             :         !   qr: rain water mixing ratio in kg kg-1
      54             :         !   nr: rain water number concentration in particles m-3
      55             :         !   rho: density of air in kg m-3
      56             :         !   q_small: minimum cloud water mixing ratio value for running the microphysics
      57             :         !   mgncol: MG number of grid cells in vertical column
      58             :         ! Output args:
      59             :         !    qc_tend: qc tendency
      60             :         !    qr_tend: qr tendency
      61             :         !    nc_tend: nc tendency
      62             :         !    nr_tend: nr tendency
      63             :         !
      64             :         integer, intent(in) :: mgncol
      65             :         real(r8), dimension(mgncol), intent(in) :: qc, qr, nc, nr, rho, lcldm, precip_frac
      66             :         real(r8), intent(in) :: q_small
      67             :         real(r8), dimension(mgncol), intent(out) :: qc_tend, qr_tend, nc_tend, nr_tend
      68             :         integer(i8) :: i, j
      69             :         real(r8), dimension(batch_size, num_inputs) :: nn_inputs, nn_quantile_inputs
      70             :         real(r8), dimension(batch_size, num_outputs) :: nn_quantile_outputs, nn_outputs
      71             :         real(r8), parameter :: dt = 1800.0
      72           0 :         do i = 1, mgncol
      73           0 :             if (qc(i) >= q_small) then
      74           0 :                 nn_inputs(1, 1) = qc(i)
      75           0 :                 nn_inputs(1, 2) = qr(i)
      76           0 :                 nn_inputs(1, 3) = nc(i)
      77           0 :                 nn_inputs(1, 4) = nr(i)
      78           0 :                 nn_inputs(1, 5) = rho(i)
      79           0 :                 nn_inputs(1, 6) = precip_frac(i)
      80           0 :                 nn_inputs(1, 7) = lcldm(i)
      81           0 :                 call quantile_transform(nn_inputs, input_scale_values, nn_quantile_inputs)
      82           0 :                 call neural_net_predict(nn_quantile_inputs, q_all, nn_quantile_outputs)
      83           0 :                 call quantile_inv_transform(nn_quantile_outputs, output_scale_values, nn_outputs)
      84           0 :                 qr_tend(i) = (nn_outputs(1, 1) - qr(i)) / dt
      85           0 :                 qr_tend(i) = (nn_outputs(1, 1) - qr(i)) / dt
      86           0 :                 qc_tend(i) = -qr_tend(i)
      87           0 :                 nc_tend(i) = (nn_outputs(1, 2) - nc(i)) / dt
      88           0 :                 nr_tend(i) = (nn_outputs(1, 3) - nr(i)) / dt
      89             :             else
      90           0 :                 qc_tend(i) = 0._r8
      91           0 :                 qr_tend(i) = 0._r8
      92           0 :                 nc_tend(i) = 0._r8
      93           0 :                 nr_tend(i) = 0._r8
      94             :             end if
      95             :         end do
      96           0 :     end subroutine tau_emulated_cloud_rain_interactions
      97             : end module tau_neural_net_quantile