accelerate_crm.F90 14.6 KB
Newer Older
Youngsung Kim's avatar
Youngsung Kim committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
!KGEN-generated Fortran source file 
  
!Generated at : 2020-05-29 19:47:49 
!KGEN version : 0.9.0 
  
! -----------------------------------------------------------------------------
! MODULE  accelerate_crm_mod
!   This module provides functionality to apply mean-state acceleration (MSA)
!   (Jones et al., 2015, doi:10.1002/2015MS000488) to the CRMs when using
!   superparameterization.
! PUBLIC SUBROUTINES:
!   crm_accel_init: initialize quantities needed to apply MSA
!   crm_accel_nstop: adjusts 'nstop' in crm_module based on crm_accel_factor
!   accelerate_crm: calculates and applies MSA tendency to CRM
! PUBLIC MODULE VARIABLES:
!   logical  :: use_crm_accel - apply MSA if true (cam namelist variable)
!   real(r8) :: crm_accel_factor - MSA factor to use (cam namelist variable)
! REVISION HISTORY:
!   2018-Nov-01: Initial implementation
!   2019-Jan-30: Initial subroutine port to GPU using openacc directives
! CONTACT: Christopher Jones (christopher.jones@pnnl.gov)
! -----------------------------------------------------------------------------


!
!
!
!
module accelerate_crm_mod
    USE grid, ONLY: nx, ny 
    USE shr_kind_mod, ONLY: r8=>shr_kind_r8 
    USE params, ONLY: rc=>crm_rknd 
    USE kgen_utils_mod, ONLY: kgen_dp, kgen_array_sumcheck 
    USE tprof_mod, ONLY: tstart, tstop, tnull, tprnt 

    IMPLICIT NONE 
    ! private module variables

    real(r8), parameter :: coef = 1._r8 / dble(nx * ny)  ! coefficient for horizontal averaging
    logical :: crm_accel_uv  ! (false) apply MSA only to scalar fields (T and QT)
                             ! (true) apply MSA to winds (U/V) and scalar fields
    ! public module variables

    logical :: use_crm_accel  ! use MSA if true
    real(r8) :: crm_accel_factor  ! 1 + crm_accel_factor = 'a' in Jones etal (2015)

    PRIVATE coef, crm_accel_uv 
    PUBLIC use_crm_accel, crm_accel_factor 
    PUBLIC accelerate_crm 
    PUBLIC kr_externs_in_accelerate_crm_mod 
    PUBLIC kr_externs_out_accelerate_crm_mod 

  contains


    subroutine accelerate_crm(ncrms, nstep, nstop, ceaseflag)
      ! Applies mean-state acceleration (MSA) to CRM
      ! Applies MSA to the following crm fields:
      !   t, qv, qcl, qci, micro_field(:,:,:,index_water_vapor,:),
      !   u (optional), v (optional)
      ! Raises ceaseflag and aborts MSA if the magnitude of 
      ! the change in "t" exceeds 5K at any point.
      ! Arguments:
      !   ncrms (in) - number of crm columns in this group
      !   nstep (in) - current crm iteration, needed only if 
      !                ceaseflag is triggered
      !   nstop (inout) - number of crm iterations, adjusted only
      !                   if ceaseflag is triggered
      !   ceaseflag (inout) - returns true if accelerate_crm aborted
      !                       before MSA applied; otherwise false
      ! Notes:
      !   micro_field(:,:,:,index_water_vapor,:) is the non-precipitating
      !     _total_ water mixing ratio for None microphysics.
      !   Intended to be called from crm subroutine in crm_module
      ! -----------------------------------------------------------------------
      !
      ! 
        USE grid, ONLY: nzm 
        USE vars, ONLY: u, v, u0, v0, t0, q0, t, qcl, qci, qv 
        USE microphysics, ONLY: micro_field, idx_qt=>index_water_vapor 
      implicit none
      integer, intent(in   ) :: ncrms
      integer, intent(in   ) :: nstep
      integer, intent(inout) :: nstop
      logical, intent(inout) :: ceaseflag
      real(rc), allocatable :: ubaccel  (:,:)   ! u before applying MSA tendency
      real(rc), allocatable :: vbaccel  (:,:)   ! v before applying MSA tendency
      real(rc), allocatable :: tbaccel  (:,:)   ! t before applying MSA tendency
      real(rc), allocatable :: qtbaccel (:,:)  ! Non-precipitating qt before applying MSA tendency
      real(rc), allocatable :: ttend_acc(:,:) ! MSA adjustment of t
      real(rc), allocatable :: qtend_acc(:,:) ! MSA adjustment of qt
      real(rc), allocatable :: utend_acc(:,:) ! MSA adjustment of u
      real(rc), allocatable :: vtend_acc(:,:) ! MSA adjustment of v
      real(r8), allocatable :: qpoz     (:,:) ! total positive micro_field(:,:,k,idx_qt,:) in level k
      real(r8), allocatable :: qneg     (:,:) ! total negative micro_field(:,:,k,idx_qt,:) in level k
      real(rc) :: tmp  ! temporary variable for atomic updates
      integer i, j, k, icrm  ! iteration variables
      real(r8) :: factor, qt_res ! local variables for redistributing moisture
      real(rc) :: ttend_threshold ! threshold for ttend_acc at which MSA aborts
      real(rc) :: tmin  ! mininum value of t allowed (sanity factor)

      ttend_threshold = 5.  ! 5K, following UP-None implementation
      tmin = 50.  ! should never get below 50K in crm, following UP-None implementation

      allocate( ubaccel  (ncrms,nzm) )
      allocate( vbaccel  (ncrms,nzm) )
      allocate( tbaccel  (ncrms,nzm) )
      allocate( qtbaccel (ncrms,nzm) )
      allocate( ttend_acc(ncrms,nzm) )
      allocate( qtend_acc(ncrms,nzm) )
      allocate( utend_acc(ncrms,nzm) )
      allocate( vtend_acc(ncrms,nzm) )
      allocate( qpoz     (ncrms,nzm) )
      allocate( qneg     (ncrms,nzm) )
      !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
      !! Compute the average among horizontal columns for each variable
      !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
      !$omp target enter data map(alloc: ubaccel)
      !$omp target enter data map(alloc: vbaccel)
      !$omp target enter data map(alloc: tbaccel)
      !$omp target enter data map(alloc: qtbaccel)
      !$omp target enter data map(alloc: ttend_acc)
      !$omp target enter data map(alloc: qtend_acc)
      !$omp target enter data map(alloc: utend_acc)
      !$omp target enter data map(alloc: vtend_acc)
      !$omp target enter data map(alloc: qpoz)
      !$omp target enter data map(alloc: qneg)
      !$omp target teams distribute parallel do collapse(2)
      do k = 1, nzm
        do icrm = 1, ncrms
          tbaccel(icrm,k) = 0
          qtbaccel(icrm,k) = 0
          if (crm_accel_uv) then
            ubaccel(icrm,k) = 0
            vbaccel(icrm,k) = 0
          endif
        enddo
      enddo
      !$omp target teams distribute parallel do collapse(4)
      do k = 1, nzm
        do j = 1 , ny
          do i = 1 , nx
            do icrm = 1, ncrms
              ! calculate tendency * dtn
              tmp = t(icrm,i,j,k) * coef
              !$omp atomic update
              tbaccel(icrm,k) = tbaccel(icrm,k) + tmp
              tmp = (qcl(icrm,i,j, k) + qci(icrm,i,j, k) + qv(icrm,i,j, k)) * coef
              !$omp atomic update
              qtbaccel(icrm,k) = qtbaccel(icrm,k) + tmp
              if (crm_accel_uv) then
                tmp = u(icrm,i,j,k) * coef
                !$omp atomic update
                ubaccel(icrm,k) = ubaccel(icrm,k) + tmp
                tmp = v(icrm,i,j,k) * coef
                !$omp atomic update
                vbaccel(icrm,k) = vbaccel(icrm,k) + tmp
              endif
            enddo
          enddo
        enddo
      enddo
      !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
      !! Compute the accelerated tendencies
      !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

      !$omp target teams distribute parallel do collapse(2)
      do k = 1, nzm
        do icrm = 1, ncrms
          ttend_acc(icrm,k) = tbaccel(icrm,k) - t0(icrm,k)
          qtend_acc(icrm,k) = qtbaccel(icrm,k) - q0(icrm,k)
          if (crm_accel_uv) then
            utend_acc(icrm,k) = ubaccel(icrm,k) - u0(icrm,k)
            vtend_acc(icrm,k) = vbaccel(icrm,k) - v0(icrm,k)
          endif
          if (abs(ttend_acc(icrm,k)) > ttend_threshold) then
            ceaseflag = .true.
          endif
        enddo
      enddo
      !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
      !! Make sure it isn't insane
      !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

      !$omp taskwait
      if (ceaseflag) then ! special case for dT/dt too large
        ! MSA will not be applied here or for the remainder of the CRM integration.
        ! nstop must be updated to ensure the CRM integration duration is unchanged.
        ! The effective MSA timestep is dt_a = crm_dt * (1 + crm_accel_factor). When
        ! ceaseflag is triggered at nstep, we've taken (nstep - 1) previous steps of
        ! size crm_dt * (1 + crm_accel_factor). The current step, and all future
        ! steps, will revert to size crm_dt. Therefore, the total crm integration
        ! time remaining after this step is
        !     time_remaining = crm_run_time - (nstep - 1)* dt_a + crm_dt
        !     nsteps_remaining = time_remaining / crm_dt
        !     updated nstop = nstep + nsteps_remaining
        ! Because we set nstop = crm_run_time / dt_a in crm_accel_nstop, subbing
        ! crm_run_time = nstop * dt_a and working through algebra yields 
        !     updated nstop = nstop + (nstop - nstep + 1) * crm_accel_factor.
        ! 
        nstop = nstop + (nstop - nstep + 1)*crm_accel_factor ! only can happen once
        return
      endif
      !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
      !! Apply the accelerated tendencies
      !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

      !$omp target teams distribute parallel do collapse(4)
      do k = 1, nzm
        do j = 1, ny
          do i = 1, nx
            do icrm = 1, ncrms
              ! don't let T go negative!
              t(icrm,i,j,k) = max(tmin, t(icrm,i,j,k) + crm_accel_factor * ttend_acc(icrm,k))
              if (crm_accel_uv) then
                !$omp atomic update
                u(icrm,i,j,k) = u(icrm,i,j,k) + crm_accel_factor * utend_acc(icrm,k) 
                !$omp atomic update
                v(icrm,i,j,k) = v(icrm,i,j,k) + crm_accel_factor * vtend_acc(icrm,k) 
              endif
              !$omp atomic update
              micro_field(icrm,i,j,k,idx_qt) = micro_field(icrm,i,j,k,idx_qt) + crm_accel_factor * qtend_acc(icrm,k)
            enddo
          enddo
        enddo
      enddo
      !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
      !! Fix negative micro and readjust among separate water species
      !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

      !$omp target teams distribute parallel do collapse(2)
      do k = 1, nzm
        do icrm = 1, ncrms
          qpoz(icrm,k) = 0.
          qneg(icrm,k) = 0.
        enddo
      enddo
      ! separately accumulate positive and negative qt values in each layer k
      !$omp target teams distribute parallel do collapse(4)
      do k = 1, nzm
        do j = 1, ny
          do i = 1, nx
            do icrm = 1, ncrms
              if (micro_field(icrm,i,j,k,idx_qt) < 0.) then
                !$omp atomic update
                qneg(icrm,k) = qneg(icrm,k) + micro_field(icrm,i,j,k,idx_qt)
              else
                !$omp atomic update
                qpoz(icrm,k) = qpoz(icrm,k) + micro_field(icrm,i,j,k,idx_qt)
              endif
            enddo
          enddo
        enddo
      enddo
      !$omp target teams distribute parallel do collapse(4)
      do k = 1, nzm
        do j = 1 , ny
          do i = 1 , nx
            do icrm = 1, ncrms
              if (qpoz(icrm,k) + qneg(icrm,k) <= 0.) then
                ! all moisture depleted in layer
                micro_field(icrm,i,j,k,idx_qt) = 0.
                qv(icrm,i,j,k    ) = 0.
                qcl(icrm,i,j,k    ) = 0.
                qci(icrm,i,j,k    ) = 0.
              else
                ! Clip qt values at 0 and remove the negative excess in each layer
                ! proportionally from the positive qt fields in the layer
                factor = 1._r8 + qneg(icrm,k) / qpoz(icrm,k)
                micro_field(icrm,i,j,k,idx_qt) = max(0._rc, micro_field(icrm,i,j,k,idx_qt) * factor)
                ! Partition micro_field == qv + qcl + qci following these rules:
                !    (1) attempt to satisfy purely by adjusting qv
                !    (2) adjust qcl and qci only if needed to ensure positivity
                if (micro_field(icrm,i,j,k,idx_qt) <= 0._rc) then
                  qv(icrm,i,j,k) = 0.
                  qcl(icrm,i,j,k) = 0.
                  qci(icrm,i,j,k) = 0.
                else
                  ! deduce qv as residual between qt - qcl - qci
                  qt_res = micro_field(icrm,i,j,k,idx_qt) - qcl(icrm,i,j,k) - qci(icrm,i,j,k)
                  qv(icrm,i,j,k) = max(0._rc, qt_res)
                  if (qt_res < 0._r8) then
                    ! qv was clipped; need to reduce qcl and qci accordingly
                    factor = 1._r8 + qt_res / (qcl(icrm,i,j,k) + qci(icrm,i,j,k))
                    !$omp atomic update
                    qcl(icrm,i,j,k) = qcl(icrm,i,j,k) * factor
                   !$omp atomic update
                    qci(icrm,i,j,k) = qci(icrm,i,j,k) * factor
                  endif
                endif
              endif ! qpoz + qneg < 0.
            enddo ! i = 1, nx
          enddo ! j = 1, ny
        enddo ! k = 1, nzm
      enddo ! icrm = 1, ncrms
      !$omp target exit data map(delete: ubaccel)
      !$omp target exit data map(delete: vbaccel)
      !$omp target exit data map(delete: tbaccel)
      !$omp target exit data map(delete: qtbaccel)
      !$omp target exit data map(delete: ttend_acc)
      !$omp target exit data map(delete: qtend_acc)
      !$omp target exit data map(delete: utend_acc)
      !$omp target exit data map(delete: vtend_acc)
      !$omp target exit data map(delete: qpoz)
      !$omp target exit data map(delete: qneg)
      deallocate( ubaccel   )
      deallocate( vbaccel   )
      deallocate( tbaccel   )
      deallocate( qtbaccel  )
      deallocate( ttend_acc )
      deallocate( qtend_acc )
      deallocate( utend_acc )
      deallocate( vtend_acc )
      deallocate( qpoz      )
      deallocate( qneg      )

    end subroutine accelerate_crm
    
    !read state subroutine for kr_externs_in_accelerate_crm_mod 
    SUBROUTINE kr_externs_in_accelerate_crm_mod(kgen_unit) 
        INTEGER, INTENT(IN) :: kgen_unit 
        LOGICAL :: kgen_istrue 
        REAL(KIND=8) :: kgen_array_sum 
          
        READ (UNIT = kgen_unit) crm_accel_uv 
        READ (UNIT = kgen_unit) use_crm_accel 
        READ (UNIT = kgen_unit) crm_accel_factor 
    END SUBROUTINE kr_externs_in_accelerate_crm_mod 
      
    !read state subroutine for kr_externs_out_accelerate_crm_mod 
    SUBROUTINE kr_externs_out_accelerate_crm_mod(kgen_unit) 
        INTEGER, INTENT(IN) :: kgen_unit 
          
        LOGICAL :: kgen_istrue 
        REAL(KIND=8) :: kgen_array_sum 
    END SUBROUTINE kr_externs_out_accelerate_crm_mod 
      
end module accelerate_crm_mod