MOM_meso_sfn_ANN.F90

1! This file is part of MOM6, the Modular Ocean Model version 6.
2! See the LICENSE file for licensing information.
3! SPDX-License-Identifier: Apache-2.0
4
5!> Implements an ANN-based mesoscale streamfunction parameterization for use
6!! with isopycnal height diffusion in MOM_thickness_diffuse.
7!!
8!! The network reads a nondimensionalized stencil of density gradients,
9!! strain rate components, and relative vorticity, and returns two density
10!! flux components at the cell center. The dimensionalization in
11!! meso_sfn_ANN_compute (multiplication by rho_grad_mag * vel_grad_mag *
12!! areaT * ann_coeff) must match the nondimensionalization used when the
13!! network was trained -- changing one without the other will produce
14!! garbage fluxes. The training procedure is the implicit contract.
15!!
16!! Density fluxes are converted to a velocity-scale streamfunction
17!! Upsilon (Ferrari et al. 2010) by dividing by the local 3-D density
18!! gradient magnitude; a configurable clamp acts on Upsilon so the cap is
19!! grid-independent. The volume-transport streamfunction passed back to
20!! thickness_diffuse is Upsilon * dy_Cu (or dx_Cv), matching MOM6's
21!! Sfn_unlim convention.
23
25use mom_diag_mediator, only : post_data, register_diag_field, diag_ctrl, time_type
26use mom_error_handler, only : mom_error, fatal
27use mom_file_parser, only : get_param, log_version, param_file_type
28use mom_grid, only : ocean_grid_type
33use mom_domains, only : pass_vector
34
35implicit none ; private
36
37#include <MOM_memory.h>
38
40
41!> Control structure for meso-scale streamfunction ANN parameterization
42type, public :: meso_sfn_ann_cs; private
43 logical :: initialized = .false. !< If true, the module has been initialized.
44 logical :: debug !< if true, write verbose checksums for debugging purposes.
45
46 real :: ann_coeff !< Coefficient to multiply the ANN output by.
47 real :: kappa_smooth !< Vertical diffusivity used to interpolate more sensible values
48 !! of T & S into thin layers [H Z T-1 ~> m2 s-1 or kg m-1 s-1]
49 integer :: ann_window !< Size of the window used in the ANN model.
50
51 type(ann_cs) :: ann_rho_flux !< ANN instance for off-diagonal and diagonal stress
52 character(len=200) :: ann_file_rho_flux !< Path to netcdf file with ANN
53 real :: min_dist_from_boundary !< Minimum distance from bottom for valid interface [Z ~> m]
54 real :: mag_grad_floor !< Floor for density gradient magnitude [R Z-1 ~> kg m-4]
55 real :: flux_clamp !< Maximum magnitude of ANN output density flux [R L T-1 ~> kg m-2 s-1]
56 real :: upsilon_clamp !< Maximum magnitude of the velocity-scale streamfunction
57 !! Upsilon (Ferrari et al. 2010) [L Z T-1 ~> m2 s-1]
58 type(diag_ctrl), pointer :: diag => null() !< structure used to regulate timing of diagnostics
59 ! Diagnostic identifiers
60 integer :: id_drdx_u !< Diagnostic id for zonal density gradient at u-points.
61 integer :: id_drdy_v !< Diagnostic id for meridional density gradient at v-points.
62 integer :: id_drdz_u !< Diagnostic id for vertical density gradient at u-points.
63 integer :: id_drdz_v !< Diagnostic id for vertical density gradient at v-points.
64 integer :: id_drdx_c !< Diagnostic id for zonal density gradient at center points.
65 integer :: id_drdy_c !< Diagnostic id for meridional density gradient at center points.
66 integer :: id_fx_c !< Diagnostic id for zonal density flux at center points.
67 integer :: id_fy_c !< Diagnostic id for meridional density flux at center points.
68 integer :: id_fx_u !< Diagnostic id for zonal density flux at u-points.
69 integer :: id_fy_v !< Diagnostic id for meridional density flux at v-points.
70 integer :: id_sfn_u !< Diagnostic id for volume streamfunction at u-points.
71 integer :: id_sfn_v !< Diagnostic id for volume streamfunction at v-points.
72end type meso_sfn_ann_cs
73
74contains
75
76!> Compute the ANN-based mesoscale streamfunction on u- and v-points.
77!!
78!! Computes density gradients and velocity gradients, feeds them through the ANN
79!! to get density fluxes at cell centers, then converts those fluxes into a
80!! streamfunction on u- and v-points for use in thickness_diffuse.
81subroutine meso_sfn_ann_compute(h, e, sfn_u, sfn_v, G, GV, US, tv, CS, dt, u, v)
82 type(ocean_grid_type), intent(in) :: g !< Ocean grid structure
83 type(verticalgrid_type), intent(in) :: gv !< Vertical grid structure
84 type(unit_scale_type), intent(in) :: us !< A dimensional unit scaling type
85 type(thermo_var_ptrs), intent(in) :: tv !< Thermodynamics structure
86 real, dimension(SZI_(G),SZJ_(G),SZK_(GV)), intent(in) :: h !< Layer thickness [Z ~> m or kg m-2]
87 real, dimension(SZI_(G),SZJ_(G),SZK_(GV)+1), intent(in) :: e !< Layer thickness [Z ~> m or kg m-2]
88 type(meso_sfn_ann_cs), intent(inout) :: cs !< Control structure for thickness_flux_ann
89 real, dimension(SZIB_(G),SZJ_(G),SZK_(GV)+1), intent(out) :: sfn_u !< Mesoscale volume streamfunction
90 !! on u-points [Z L2 T-1 ~> m3 s-1]
91 real, dimension(SZI_(G),SZJB_(G),SZK_(GV)+1), intent(out) :: sfn_v !< Mesoscale volume streamfunction
92 !! on v-points [Z L2 T-1 ~> m3 s-1]
93 real, intent(in) :: dt !< Model time step [T ~> s]
94 real, dimension(SZIB_(G),SZJ_(G),SZK_(GV)), intent(in) :: u !< Zonal velocity [L T-1 ~> m s-1].
95 real, dimension(SZI_(G),SZJB_(G),SZK_(GV)), intent(in) :: v !< Meridional velocity [L T-1 ~> m s-1].
96
97 ! Local variables
98 integer :: i, j, k, is, ie, js, je, nz, shift, stencil_points, ii, jj
99 integer :: nij, m
100
101 real, dimension(SZIB_(G),SZJ_(G),SZK_(GV)+1) :: drdx_u !< Zonal density gradient at u [R L-1 ~> kg m-4]
102 real, dimension(SZIB_(G),SZJ_(G),SZK_(GV)+1) :: drdz_u !< Vertical density gradient at u [R Z-1 ~> kg m-4]
103 real, dimension(SZI_(G),SZJB_(G),SZK_(GV)+1) :: drdy_v !< Meridional density gradient at v [R L-1 ~> kg m-4]
104 real, dimension(SZI_(G),SZJB_(G),SZK_(GV)+1) :: drdz_v !< Vertical density gradient at v [R Z-1 ~> kg m-4]
105 real, dimension(SZIB_(G),SZJ_(G),SZK_(GV)+1) :: slope_x !< Isopycnal slope in x at u [Z L-1 ~> nondim]
106 real, dimension(SZI_(G),SZJB_(G),SZK_(GV)+1) :: slope_y !< Isopycnal slope in y at v [Z L-1 ~> nondim]
107
108 real, dimension(SZIB_(G),SZJ_(G),SZK_(GV)+1) :: fx_u !< Zonal density flux at u-points [R L T-1 ~> kg m-2 s-1]
109 real, dimension(SZI_(G),SZJB_(G),SZK_(GV)+1) :: fy_v !< Meridional density flux at v-points [R L T-1 ~> kg m-2 s-1]
110
111 real, dimension(SZI_(G),SZJ_(G),SZK_(GV)+1) :: drdx_c !< Zonal density gradient at center points [R L-1 ~> kg m-4]
112 real, dimension(SZI_(G),SZJ_(G),SZK_(GV)+1) :: drdy_c !< Meridional density gradient
113 !! at center points [R L-1 ~> kg m-4]
114
115 real, dimension(SZI_(G),SZJ_(G),SZK_(GV)+1) :: fx_c !< Zonal density flux at center points [R L T-1 ~> kg m-2 s-1]
116 real, dimension(SZI_(G),SZJ_(G),SZK_(GV)+1) :: fy_c !< Meridional density flux at
117 !! center points [R L T-1 ~> kg m-2 s-1]
118
119 real, dimension(SZI_(G),SZJ_(G),SZK_(GV)) :: dudx !< du/dx at cell center [T-1 ~> s-1]
120 real, dimension(SZI_(G),SZJ_(G),SZK_(GV)) :: dvdy !< dv/dy at cell center [T-1 ~> s-1]
121 real, dimension(SZI_(G),SZJ_(G),SZK_(GV)) :: dudy !< du/dy at cell center [T-1 ~> s-1]
122 real, dimension(SZI_(G),SZJ_(G),SZK_(GV)) :: dvdx !< dv/dx at cell center [T-1 ~> s-1]
123
124 real, dimension(SZI_(G),SZJ_(G)) :: norm_y !< Scaling coefficient for ANN outputs [R L-1 T-1 ~> kg m-4 s-1]
125
126 real, allocatable :: drdx_local(:,:) !< Local stencil of drdx [R L-1 ~> kg m-4]
127 real, allocatable :: drdy_local(:,:) !< Local stencil of drdy [R L-1 ~> kg m-4]
128 real, allocatable :: dudx_local(:,:) !< Local stencil of du/dx [T-1 ~> s-1]
129 real, allocatable :: dudy_local(:,:) !< Local stencil of du/dy [T-1 ~> s-1]
130 real, allocatable :: dvdx_local(:,:) !< Local stencil of dv/dx [T-1 ~> s-1]
131 real, allocatable :: dvdy_local(:,:) !< Local stencil of dv/dy [T-1 ~> s-1]
132 real, allocatable :: sh_xx_local(:,:) !< Local stencil of normal strain [T-1 ~> s-1]
133 real, allocatable :: sh_xy_local(:,:) !< Local stencil of shear strain [T-1 ~> s-1]
134 real, allocatable :: vort_local(:,:) !< Local stencil of relative vorticity [T-1 ~> s-1]
135 real :: vel_grad_mag !< Magnitude of velocity gradient tensor over stencil [T-1 ~> s-1]
136 real :: rho_grad_mag !< Magnitude of density gradient over stencil [R L-1 ~> kg m-4]
137 real, allocatable :: x(:,:) !< Input vector to the ANN
138 real, allocatable :: y(:,:) !< Output vector from the ANN
139 real, allocatable :: yy(:) !< Local output vector from the ANN
140 real :: mag_grad !< Magnitude of 3-D density gradient [R Z-1 ~> kg m-4]
141 logical :: use_stanley
142 logical :: use_eos
143 real :: dist_from_bot_a, dist_from_bot_b ! Distance from interface to bottom [Z ~> m]
144 real :: dist_from_sfc_a, dist_from_sfc_b ! Distance from interface to surface [Z ~> m]
145 real :: upsilon_u ! Velocity-scale streamfunction at u-point (Ferrari et al. 2010) [L Z T-1 ~> m2 s-1]
146 real :: upsilon_v ! Velocity-scale streamfunction at v-point (Ferrari et al. 2010) [L Z T-1 ~> m2 s-1]
147 real :: rho_grad_neglect ! A density gradient magnitude so small it is lost in
148 ! roundoff; used to prevent division by zero [R L-1 ~> kg m-4]
149 real :: vel_grad_neglect ! A velocity gradient magnitude so small it is lost in
150 ! roundoff; used to prevent division by zero [T-1 ~> s-1]
151
152 if (.not. cs%initialized) call mom_error(fatal, &
153 "meso_sfn_ANN_compute: Module MOM_meso_sfn_ANN must be initialized before use.")
154
155 use_stanley = .false. ! Not using Stanley smoothing here.
156 use_eos = associated(tv%eqn_of_state)
157 is = g%isc ; ie = g%iec ; js = g%jsc ; je = g%jec ; nz = gv%ke
158
159 rho_grad_neglect = 1.0e-30 * us%kg_m3_to_R * us%L_to_m
160 vel_grad_neglect = 1.0e-30 * us%T_to_s
161
162 ! Allocate the local stencil variables
163 allocate(drdx_local(cs%ann_window, cs%ann_window), drdy_local(cs%ann_window, cs%ann_window), &
164 dudx_local(cs%ann_window, cs%ann_window), dudy_local(cs%ann_window, cs%ann_window), &
165 dvdx_local(cs%ann_window, cs%ann_window), dvdy_local(cs%ann_window, cs%ann_window), &
166 sh_xx_local(cs%ann_window, cs%ann_window), sh_xy_local(cs%ann_window, cs%ann_window), &
167 vort_local(cs%ann_window, cs%ann_window))
168
169 shift = (cs%ann_window-1)/2
170 stencil_points = cs%ann_window * cs%ann_window
171
172 ! Number of horizontal grid points in ANN inference loop below
173 nij = (ie - is + 3) * (je - js + 3)
174 allocate(x(nij, stencil_points*5), y(nij, 2))
175 allocate(yy(2))
176
177 slope_x(:,:,:) = 0.0
178 slope_y(:,:,:) = 0.0
179
180 sfn_u(:,:,:) = 0.0
181 sfn_v(:,:,:) = 0.0
182
183 fx_u(:,:,:) = 0.0
184 fy_v(:,:,:) = 0.0
185 fx_c(:,:,:) = 0.0
186 fy_c(:,:,:) = 0.0
187
188 drdx_u(:,:,:) = 0.0
189 drdy_v(:,:,:) = 0.0
190 drdz_u(:,:,:) = 0.0
191 drdz_v(:,:,:) = 0.0
192 drdx_c(:,:,:) = 0.0
193 drdy_c(:,:,:) = 0.0
194 ! Compute rho gradients
195 if (use_eos) then
196 call calc_isoneutral_slopes(g, gv, us, h, e, tv, dt*cs%kappa_smooth, use_stanley, slope_x, slope_y, &
197 drdx_u=drdx_u, drdy_v=drdy_v, drdz_u=drdz_u, drdz_v=drdz_v, halo=3)
198 else
199 call calc_layered_density_gradients(g, gv, us, h, e, drdx_u, drdy_v, drdz_u, drdz_v, halo=3, &
200 min_dist_from_boundary=cs%min_dist_from_boundary)
201 endif
202
203 ! Interpolate the rho gradients to the center point
204 call center_grad_rho(drdx_u, drdy_v, drdx_c, drdy_c, g, gv, cs)
205
206 ! Compute velocity gradients at center points
207 call vel_gradients(u, v, g, gv, dudx, dudy, dvdx, dvdy, cs)
208
209 ! Post diagnostics
210 if (cs%id_drdx_u > 0) call post_data(cs%id_drdx_u, drdx_u, cs%diag)
211 if (cs%id_drdy_v > 0) call post_data(cs%id_drdy_v, drdy_v, cs%diag)
212
213 if (cs%id_drdz_u > 0) call post_data(cs%id_drdz_u, drdz_u, cs%diag)
214 if (cs%id_drdz_v > 0) call post_data(cs%id_drdz_v, drdz_v, cs%diag)
215
216 if (cs%id_drdx_c > 0) call post_data(cs%id_drdx_c, drdx_c, cs%diag)
217 if (cs%id_drdy_c > 0) call post_data(cs%id_drdy_c, drdy_c, cs%diag)
218
219 ! Compute the density fluxes at center points using the ANN.
220 do k = 2, nz
221 m = 0
222 do j = js-1, je+1 ; do i = is-1, ie+1
223 m = m + 1
224 drdx_local(:,:) = drdx_c(i-shift:i+shift,j-shift:j+shift,k)
225 drdy_local(:,:) = drdy_c(i-shift:i+shift,j-shift:j+shift,k)
226 ! Take the velocity gradients below the interface K
227 dudx_local(:,:) = dudx(i-shift:i+shift,j-shift:j+shift,k)
228 dudy_local(:,:) = dudy(i-shift:i+shift,j-shift:j+shift,k)
229 dvdx_local(:,:) = dvdx(i-shift:i+shift,j-shift:j+shift,k)
230 dvdy_local(:,:) = dvdy(i-shift:i+shift,j-shift:j+shift,k)
231
232 ! Compute the strain rate tensor components and vorticity
233 sh_xx_local(:,:) = dudx_local(:,:) - dvdy_local(:,:)
234 sh_xy_local(:,:) = dudy_local(:,:) + dvdx_local(:,:)
235 vort_local(:,:) = dvdx_local(:,:) - dudy_local(:,:)
236
237 ! Compute the magnitude of the velocity gradient tensor for the local stencil
238 rho_grad_mag = 0.0
239 vel_grad_mag = 0.0
240 do jj=1, cs%ann_window
241 do ii=1, cs%ann_window
242 rho_grad_mag = (rho_grad_mag + drdx_local(ii,jj)*drdx_local(ii,jj)) + &
243 drdy_local(ii,jj)*drdy_local(ii,jj)
244 vel_grad_mag = ((vel_grad_mag + sh_xx_local(ii,jj)*sh_xx_local(ii,jj)) + &
245 sh_xy_local(ii,jj)*sh_xy_local(ii,jj)) + &
246 vort_local(ii,jj)*vort_local(ii,jj)
247 enddo
248 enddo
249 rho_grad_mag = sqrt(rho_grad_mag) + rho_grad_neglect
250 vel_grad_mag = sqrt(vel_grad_mag) + vel_grad_neglect
251 norm_y(i,j) = rho_grad_mag * vel_grad_mag
252
253 ! Normalize inputs
254 drdx_local(:,:) = drdx_local(:,:) / rho_grad_mag
255 drdy_local(:,:) = drdy_local(:,:) / rho_grad_mag
256
257 sh_xx_local(:,:) = sh_xx_local(:,:)/ vel_grad_mag
258 sh_xy_local(:,:) = sh_xy_local(:,:)/ vel_grad_mag
259 vort_local(:,:) = vort_local(:,:)/ vel_grad_mag
260
261 ! Prepare input vector for ANN
262 x(m,1:stencil_points) = reshape(drdx_local, (/stencil_points/))
263 x(m,stencil_points+1:2*stencil_points) = reshape(drdy_local, (/stencil_points/))
264 x(m,2*stencil_points+1:3*stencil_points) = reshape(sh_xx_local, (/stencil_points/))
265 x(m,3*stencil_points+1:4*stencil_points) = reshape(sh_xy_local, (/stencil_points/))
266 x(m,4*stencil_points+1:5*stencil_points) = reshape(vort_local, (/stencil_points/))
267
268 enddo ; enddo
269
270 ! Call the ANN
271 call ann_apply_array_sio(nij, x,y, cs%ann_rho_flux)
272
273 m=0
274 do j = js-1, je+1 ; do i = is-1, ie+1
275 m=m+1
276 ! Dimensionalize the output. The factors applied here must match the
277 ! nondimensionalization used when the network was trained; this is
278 ! an implicit contract with the training procedure.
279 yy(:) = ((y(m,:) * norm_y(i,j)) * g%areaT(i,j)) * cs%ann_coeff
280
281 ! Clamp ANN output to prevent extreme values
282 yy(1) = max(-cs%flux_clamp, min(cs%flux_clamp, yy(1)))
283 yy(2) = max(-cs%flux_clamp, min(cs%flux_clamp, yy(2)))
284
285 ! The sign convention is that ANN outputs -u'rho', so we negate.
286 fx_c(i,j,k) = -yy(1)
287 fy_c(i,j,k) = -yy(2)
288
289 enddo ; enddo
290 enddo
291
292 ! Interpolate the density fluxes to u and v points.
293 call center2uv(fx_c, fy_c, fx_u, fy_v, g, gv)
294
295 do k=2, nz
296 do j=js,je ; do i=is-1,ie
297 ! In layered mode, skip interfaces at the bottom or surface
298 if (.not. use_eos) then
299 dist_from_bot_a = e(i,j,k) - e(i,j,nz+1)
300 dist_from_bot_b = e(i+1,j,k) - e(i+1,j,nz+1)
301 dist_from_sfc_a = e(i,j,1) - e(i,j,k)
302 dist_from_sfc_b = e(i+1,j,1) - e(i+1,j,k)
303 if (dist_from_bot_a < cs%min_dist_from_boundary .or. &
304 dist_from_bot_b < cs%min_dist_from_boundary .or. &
305 dist_from_sfc_a < cs%min_dist_from_boundary .or. &
306 dist_from_sfc_b < cs%min_dist_from_boundary) then
307 sfn_u(i,j,k) = 0.0
308 cycle
309 endif
310 endif
311 ! Skip if density gradient is too small (prevents division by ~zero)
312 mag_grad = sqrt( us%Z_to_L**2*drdx_u(i,j,k)**2 + drdz_u(i,j,k)**2 )
313 if (mag_grad < cs%mag_grad_floor) then
314 sfn_u(i,j,k) = 0.0
315 cycle
316 endif
317 ! Velocity-scale (grid-independent) streamfunction, Upsilon in Ferrari et al. 2010.
318 upsilon_u = (fx_u(i,j,k)/mag_grad) * g%OBCmaskCu(i,j)
319 upsilon_u = max(-cs%Upsilon_clamp, min(cs%Upsilon_clamp, upsilon_u))
320 sfn_u(i,j,k) = upsilon_u * g%dy_Cu(i,j)
321
322 enddo ; enddo
323 do j=js-1,je ; do i=is,ie
324 if (.not. use_eos) then
325 dist_from_bot_a = e(i,j,k) - e(i,j,nz+1)
326 dist_from_bot_b = e(i,j+1,k) - e(i,j+1,nz+1)
327 dist_from_sfc_a = e(i,j,1) - e(i,j,k)
328 dist_from_sfc_b = e(i,j+1,1) - e(i,j+1,k)
329 if (dist_from_bot_a < cs%min_dist_from_boundary .or. &
330 dist_from_bot_b < cs%min_dist_from_boundary .or. &
331 dist_from_sfc_a < cs%min_dist_from_boundary .or. &
332 dist_from_sfc_b < cs%min_dist_from_boundary) then
333 sfn_v(i,j,k) = 0.0
334 cycle
335 endif
336 endif
337 ! Skip if density gradient is too small (prevents division by ~zero)
338 mag_grad = sqrt( us%Z_to_L**2*drdy_v(i,j,k)**2 + drdz_v(i,j,k)**2 )
339
340 if (mag_grad < cs%mag_grad_floor) then
341 sfn_v(i,j,k) = 0.0
342 cycle
343 endif
344
345 ! Velocity-scale (grid-independent) streamfunction, Upsilon in Ferrari et al. 2010.
346 upsilon_v = (fy_v(i,j,k)/mag_grad) * g%OBCmaskCv(i,j)
347 upsilon_v = max(-cs%Upsilon_clamp, min(cs%Upsilon_clamp, upsilon_v))
348 sfn_v(i,j,k) = upsilon_v * g%dx_Cv(i,j)
349
350 enddo ; enddo
351 enddo
352
353 call pass_vector(sfn_u, sfn_v, g%Domain)
354
355 if (cs%id_Fx_c > 0) call post_data(cs%id_Fx_c, fx_c, cs%diag)
356 if (cs%id_Fy_c > 0) call post_data(cs%id_Fy_c, fy_c, cs%diag)
357
358 if (cs%id_Fx_u > 0) call post_data(cs%id_Fx_u, fx_u, cs%diag)
359 if (cs%id_Fy_v > 0) call post_data(cs%id_Fy_v, fy_v, cs%diag)
360
361 if (cs%id_sfn_u > 0) call post_data(cs%id_sfn_u, sfn_u, cs%diag)
362 if (cs%id_sfn_v > 0) call post_data(cs%id_sfn_v, sfn_v, cs%diag)
363
364 deallocate(drdx_local, drdy_local, dudx_local, dudy_local, dvdx_local, dvdy_local, &
365 sh_xx_local, sh_xy_local, vort_local)
366 deallocate(x, y, yy)
367end subroutine meso_sfn_ann_compute
368
369!> Interpolate density gradients from u- and v-points to cell centers.
370subroutine center_grad_rho(drdx_u, drdy_v, drdx_c, drdy_c, G, GV, CS)
371 type(ocean_grid_type), intent(in) :: G !< Ocean grid structure
372 type(verticalgrid_type), intent(in) :: GV !< Vertical grid structure
373 type(meso_sfn_ann_cs), intent(inout) :: CS !< Control structure for thickness_flux_ann
374 real, dimension(SZIB_(G),SZJ_(G),SZK_(GV)+1), intent(in) :: drdx_u !< Zonal density gradient
375 !! at u-points [R L-1 ~> kg m-4]
376 real, dimension(SZI_(G),SZJB_(G),SZK_(GV)+1), intent(in) :: drdy_v !< Meridional density gradient
377 !! at v-points [R L-1 ~> kg m-4]
378 real, dimension(SZI_(G),SZJ_(G),SZK_(GV)+1), intent(inout) :: drdx_c !< Zonal density gradient
379 !! at center [R L-1 ~> kg m-4]
380 real, dimension(SZI_(G),SZJ_(G),SZK_(GV)+1), intent(inout) :: drdy_c !< Meridional density gradient
381 !! at center [R L-1 ~> kg m-4]
382
383 integer :: i, j, k, is, ie, js, je, nz, shift
384
385 is = g%isc ; ie = g%iec ; js = g%jsc ; je = g%jec ; nz = gv%ke
386
387 shift = (cs%ann_window-1)/2
388
389 do k=1, nz+1
390 do j=js-shift-1,je+shift+1 ; do i=is-shift-1,ie+shift+1
391 drdx_c(i,j,k) = 0.5 * (drdx_u(i-1,j,k) * g%mask2dCu(i-1,j) + drdx_u(i,j,k) * g%mask2dCu(i,j)) * g%mask2dT(i,j)
392 drdy_c(i,j,k) = 0.5 * (drdy_v(i,j-1,k) * g%mask2dCv(i,j-1) + drdy_v(i,j,k) * g%mask2dCv(i,j)) * g%mask2dT(i,j)
393 enddo ; enddo
394 enddo
395
396end subroutine center_grad_rho
397
398!> Interpolate two fields from cell centers to u- and v-points.
399subroutine center2uv(var1_c, var2_c, var1_u, var2_v, G, GV)
400 type(ocean_grid_type), intent(in) :: G !< Ocean grid structure
401 type(verticalgrid_type), intent(in) :: GV !< Vertical grid structure
402 real, dimension(SZI_(G),SZJ_(G),SZK_(GV)+1), intent(in) :: var1_c !< Variable at center points [arbitrary]
403 real, dimension(SZI_(G),SZJ_(G),SZK_(GV)+1), intent(in) :: var2_c !< Variable at center points [arbitrary]
404 real, dimension(SZIB_(G),SZJ_(G),SZK_(GV)+1), intent(inout) :: var1_u !< Variable at u points [arbitrary]
405 real, dimension(SZI_(G),SZJB_(G),SZK_(GV)+1), intent(inout) :: var2_v !< Variable at v points [arbitrary]
406
407 integer :: i, j, k, is, ie, js, je, nz
408
409 is = g%isc ; ie = g%iec ; js = g%jsc ; je = g%jec ; nz = gv%ke
410 do k=1, nz+1
411 do j=js,je ; do i=is-1,ie
412 var1_u(i,j,k) = 0.5 * (var1_c(i,j,k) * g%mask2dT(i,j) + var1_c(i+1,j,k) * g%mask2dT(i+1,j)) * g%mask2dCu(i,j)
413 enddo ; enddo
414 do j=js-1,je ; do i=is,ie
415 var2_v(i,j,k) = 0.5 * (var2_c(i,j,k) * g%mask2dT(i,j) + var2_c(i,j+1,k) * g%mask2dT(i,j+1)) * g%mask2dCv(i,j)
416 enddo ; enddo
417 enddo
418
419end subroutine center2uv
420!> Calculates the velocity gradients at the center points in 3D.
421subroutine vel_gradients(u, v, G, GV, dudx, dudy, dvdx, dvdy, CS)
422 type(ocean_grid_type), intent(in) :: G !< Ocean grid structure
423 type(verticalgrid_type), intent(in) :: GV !< The ocean's vertical grid structure.
424 real, dimension(SZIB_(G),SZJ_(G),SZK_(GV)),intent(in) :: u !< The zonal velocity [L T-1 ~> m s-1].
425 real, dimension(SZI_(G),SZJB_(G),SZK_(GV)),intent(in) :: v !< The meridional velocity [L T-1 ~> m s-1].
426 real, dimension(SZI_(G),SZJ_(G),SZK_(GV)), intent(out) :: dudx !< du/dx [T-1 ~> s-1]
427 real, dimension(SZI_(G),SZJ_(G),SZK_(GV)), intent(out) :: dvdy !< dv/dy [T-1 ~> s-1]
428 real, dimension(SZI_(G),SZJ_(G),SZK_(GV)), intent(out) :: dudy !< du/dy [T-1 ~> s-1]
429 real, dimension(SZI_(G),SZJ_(G),SZK_(GV)), intent(out) :: dvdx !< dv/dx [T-1 ~> s-1]
430 type(meso_sfn_ann_cs), intent(in) :: CS !< Control structure for thickness_flux_ann
431
432 ! Corner points
433 real, dimension(SZIB_(G), SZJB_(G),SZK_(GV)) :: dudy_q !< du/dy at corner points [T-1 ~> s-1]
434 real, dimension(SZIB_(G), SZJB_(G),SZK_(GV)) :: dvdx_q !< dv/dx at corner points [T-1 ~> s-1]
435 integer :: is, ie, js, je
436 integer :: nz
437 integer :: i, j, k
438 integer :: shift
439
440 is = g%isc ; ie = g%iec ; js = g%jsc ; je = g%jec ; nz = gv%ke
441
442 shift = (cs%ann_window-1)/2
443
444 do k=1, nz
445 ! Calculate velocity gradients at center points directly.
446 do j=js-shift-1,je+shift+1 ; do i=is-shift-1,ie+shift+1
447 dudx(i,j,k) = g%IdxT(i,j)* (u(i,j,k) * g%mask2dCu(i,j) - u(i-1,j,k) * g%mask2dCu(i-1,j)) * g%mask2dT(i,j)
448 dvdy(i,j,k) = g%IdyT(i,j)* (v(i,j,k) * g%mask2dCv(i,j) - v(i,j-1,k) * g%mask2dCv(i,j-1)) * g%mask2dT(i,j)
449 enddo ; enddo
450
451 ! Calculate velocity gradients at corner points.
452 ! Bounds extend one further on the lower side than the center-point loop above
453 ! because the 4-point corner-to-center interpolation below reads indices (I-1,J-1).
454 do j=js-shift-2,je+shift+1 ; do i=is-shift-2,ie+shift+1
455 dvdx_q(i,j,k) = g%IdxBu(i,j)*(v(i+1,j,k) - v(i,j,k) ) * g%mask2dBu(i,j)
456 dudy_q(i,j,k) = g%IdyBu(i,j)*(u(i,j+1,k) - u(i,j,k) ) * g%mask2dBu(i,j)
457 !
458 enddo ; enddo
459
460 ! interpolate corner grads to center points
461 do j = js-shift-1, je+shift+1; do i = is-shift-1, ie+shift+1
462 dvdx(i,j,k) = 0.25 * (((dvdx_q(i,j,k) + dvdx_q(i-1,j,k)) + dvdx_q(i,j-1,k)) + &
463 dvdx_q(i-1,j-1,k)) * g%mask2dT(i,j)
464 dudy(i,j,k) = 0.25 * (((dudy_q(i,j,k) + dudy_q(i-1,j,k)) + dudy_q(i,j-1,k)) + &
465 dudy_q(i-1,j-1,k)) * g%mask2dT(i,j)
466 enddo; enddo
467 enddo
468end subroutine vel_gradients
469!> Compute density gradients from fixed layer densities and interface heights
470!! This is a workaround for running the ANN parameterization in pure layered
471!! mode (USE_EOS=False) where calc_isoneutral_slopes won't work.
472subroutine calc_layered_density_gradients(G, GV, US, h, e, &
473 drdx_u, drdy_v, drdz_u, drdz_v, halo, min_dist_from_boundary)
474 type(ocean_grid_type), intent(in) :: G
475 type(verticalgrid_type), intent(in) :: GV
476 type(unit_scale_type), intent(in) :: US
477 real, dimension(SZI_(G),SZJ_(G),SZK_(GV)), intent(in) :: h ! Layer thickness [Z ~> m]
478 real, dimension(SZI_(G),SZJ_(G),SZK_(GV)+1), intent(in) :: e ! Interface heights [Z ~> m]
479 real, dimension(SZIB_(G),SZJ_(G),SZK_(GV)+1), intent(out) :: drdx_u ! [R L-1 ~> kg m-4]
480 real, dimension(SZI_(G),SZJB_(G),SZK_(GV)+1), intent(out) :: drdy_v ! [R L-1 ~> kg m-4]
481 real, dimension(SZIB_(G),SZJ_(G),SZK_(GV)+1), intent(out) :: drdz_u ! [R Z-1 ~> kg m-4]
482 real, dimension(SZI_(G),SZJB_(G),SZK_(GV)+1), intent(out) :: drdz_v ! [R Z-1 ~> kg m-4]
483 integer, intent(in) :: halo
484 real, intent(in) :: min_dist_from_boundary ! Threshold for boundaries [Z]
485
486 ! Local variables
487 real :: drho_k ! Density difference across interface K [R]
488 real :: dz_u, dz_v ! Vertical length scale at u,v points [Z]
489 real :: dedx, dedy ! Interface slope [Z L-1]
490 real :: h_neglect ! Small thickness [H]
491 real :: dist_from_bot_a, dist_from_bot_b ! Distance from interface to bottom [Z]
492 real :: dist_from_sfc_a, dist_from_sfc_b ! Distance from interface to surface [Z]
493
494 integer :: i, j, k, is, ie, js, je, nz
495
496 is = g%isc ; ie = g%iec ; js = g%jsc ; je = g%jec ; nz = gv%ke
497 h_neglect = gv%H_subroundoff
498 ! Initialize to zero
499 drdx_u(:,:,:) = 0.0
500 drdy_v(:,:,:) = 0.0
501 drdz_u(:,:,:) = 0.0
502 drdz_v(:,:,:) = 0.0
503
504 ! Loop over interfaces (K=1 is surface, K=nz+1 is bottom)
505 do k = 2, nz
506 ! Density jump across this interface (from GV%Rlay - the target layer densities)
507 drho_k = gv%Rlay(k) - gv%Rlay(k-1) ! [R ~> kg m-3]
508
509 ! --- U-points (zonal gradients) ---
510 do j = js-halo, je+halo
511 do i = is-1-halo, ie+halo
512
513 ! Check if interface is above bottom on both sides
514 ! e is negative (depth), e(nz+1) is the bottom, e(1) is the surface
515 dist_from_bot_a = e(i,j,k) - e(i,j,nz+1)
516 dist_from_bot_b = e(i+1,j,k) - e(i+1,j,nz+1)
517 dist_from_sfc_a = e(i,j,1) - e(i,j,k)
518 dist_from_sfc_b = e(i+1,j,1) - e(i+1,j,k)
519
520 if (dist_from_bot_a > min_dist_from_boundary .and. &
521 dist_from_bot_b > min_dist_from_boundary .and. &
522 dist_from_sfc_a > min_dist_from_boundary .and. &
523 dist_from_sfc_b > min_dist_from_boundary .and. &
524 g%mask2dCu(i,j) > 0.5) then
525
526 ! Average thickness of layers above and below interface at u-point
527 dz_u = 0.25 * gv%H_to_Z * ( &
528 (h(i,j,k-1) + h(i,j,k)) + (h(i+1,j,k-1) + h(i+1,j,k)) )
529 dz_u = max(dz_u, gv%H_to_Z * h_neglect)
530
531 ! Interface height gradient (slope of isopycnal)
532 dedx = (e(i+1,j,k) - e(i,j,k)) * g%IdxCu(i,j) ! [Z L-1]
533
534 ! In a layered model with tilted interfaces:
535 ! dρ/dx comes from the interface tilt: (Δρ across interface) * (∂η/∂x) / Δz
536 ! dρ/dz is simply Δρ / Δz
537 !
538 ! Physical interpretation: if interface tilts up to the east,
539 ! denser water (layer k) is lifted, creating ∂ρ/∂x < 0
540
541 drdx_u(i,j,k) = drho_k * dedx / dz_u ! [R L-1]
542 drdz_u(i,j,k) = - drho_k / dz_u ! [R Z-1]
543
544 ! Apply land mask
545 drdx_u(i,j,k) = drdx_u(i,j,k) * (g%mask2dCu(i,j) * g%mask2dT(i,j) * g%mask2dT(i+1,j))
546 drdz_u(i,j,k) = drdz_u(i,j,k) * (g%mask2dCu(i,j) * g%mask2dT(i,j) * g%mask2dT(i+1,j))
547 else
548 ! Interface is at/near bottom or surface on at least one side - set gradients to zero
549 drdx_u(i,j,k) = 0.0
550 drdz_u(i,j,k) = 0.0
551 endif
552 enddo
553 enddo
554
555 ! --- V-points (meridional gradients) ---
556 do j = js-1-halo, je+halo
557 do i = is-halo, ie+halo
558 ! Check if interface is above bottom and below surface on both sides
559 dist_from_bot_a = e(i,j,k) - e(i,j,nz+1)
560 dist_from_bot_b = e(i,j+1,k) - e(i,j+1,nz+1)
561 dist_from_sfc_a = e(i,j,1) - e(i,j,k)
562 dist_from_sfc_b = e(i,j+1,1) - e(i,j+1,k)
563
564 if (dist_from_bot_a > min_dist_from_boundary .and. &
565 dist_from_bot_b > min_dist_from_boundary .and. &
566 dist_from_sfc_a > min_dist_from_boundary .and. &
567 dist_from_sfc_b > min_dist_from_boundary .and. &
568 g%mask2dCv(i,j) > 0.5) then
569
570 ! Interface is a real isopycnal on both sides - compute gradients
571 dz_v = 0.25 * gv%H_to_Z * ( &
572 (h(i,j,k-1) + h(i,j,k)) + (h(i,j+1,k-1) + h(i,j+1,k)) )
573 dz_v = max(dz_v, gv%H_to_Z * h_neglect)
574
575 ! Interface height gradient
576 dedy = (e(i,j+1,k) - e(i,j,k)) * g%IdyCv(i,j) ! [Z L-1]
577
578 drdy_v(i,j,k) = drho_k * dedy / dz_v ! [R L-1]
579 drdz_v(i,j,k) = -drho_k / dz_v ! [R Z-1]
580 else
581 ! Interface is at/near bottom or surface on at least one side, or masked
582 drdy_v(i,j,k) = 0.0
583 drdz_v(i,j,k) = 0.0
584 endif
585 enddo
586 enddo
587 enddo
588
590
591!> Initializes the meso-scale streamfunction ANN parameterization
592!!
593subroutine meso_sfn_ann_init(Time, G, GV, US, param_file, diag, CS)
594 type(time_type), intent(in) :: time !< Current model time
595 type(ocean_grid_type), intent(in) :: g !< Ocean grid structure
596 type(verticalgrid_type), intent(in) :: gv !< Vertical grid structure
597 type(unit_scale_type), intent(in) :: us !< A dimensional unit scaling type
598 type(param_file_type), intent(in) :: param_file !< Parameter file handles
599 type(diag_ctrl), target, intent(inout) :: diag !< Diagnostics control structure
600 type(meso_sfn_ann_cs), intent(inout) :: cs !< Control structure for meso sfn ann
601
602 ! Local variables
603 character(len=40) :: mdl = "meso_sfn_ANN" ! This is module's name
604# include "version_variable.h"
605
606 cs%diag => diag
607
608 call log_version(param_file, mdl, version, &
609 "ANN-based mesoscale streamfunction parameterization.")
610
611 ! We don't need to check if use is true, because this is only called if it is.
612 call get_param(param_file, mdl, "MESO_SFN_ANN_COEFF", cs%ann_coeff, &
613 "Coefficient to multiply the mesoscale streamfunction ANN output by", default=1.0, units="nondim")
614
615 call get_param(param_file, mdl, "KD_SMOOTH", cs%kappa_smooth, &
616 "A diapycnal diffusivity that is used to interpolate "//&
617 "more sensible values of T & S into thin layers.", &
618 units="m2 s-1", default=1.0e-6, scale=gv%m2_s_to_HZ_T)
619 call get_param(param_file, mdl, "MESO_SFN_MIN_DIST_BOUNDARY", cs%min_dist_from_boundary, &
620 "Minimum distance from surface or bottom for interface to be considered valid "//&
621 "for density gradient calculations in layered mode.", &
622 units="m", default=50.0, scale=us%m_to_Z)
623 call get_param(param_file, mdl, "MESO_SFN_MAG_GRAD_FLOOR", cs%mag_grad_floor, &
624 "Minimum density gradient magnitude below which the streamfunction "//&
625 "is set to zero to avoid division by near-zero values.", &
626 units="kg m-4", default=1.0e-10, scale=us%kg_m3_to_R*us%Z_to_m)
627 call get_param(param_file, mdl, "MESO_SFN_FLUX_CLAMP", cs%flux_clamp, &
628 "Maximum magnitude of ANN output density flux before conversion "//&
629 "to streamfunction.", &
630 units="kg m-2 s-1", default=1.0e2, scale=us%kg_m3_to_R*us%m_to_L*us%T_to_s)
631 call get_param(param_file, mdl, "MESO_UPSILON_CLAMP", cs%Upsilon_clamp, &
632 "Maximum magnitude of the velocity-scale mesoscale streamfunction "//&
633 "(Upsilon in Ferrari et al. 2010).", &
634 units="m2 s-1", default=15., scale=us%m_to_L*us%m_to_Z*us%T_to_s)
635 call get_param(param_file, mdl, "MESO_SFN_ANN_WINDOW", cs%ann_window, &
636 "Number of horizontal grid points to use in the thickness flux ANN window", default=1)
637 ! The stencil reads drdx_c(i-shift:i+shift,...) with shift=(ann_window-1)/2.
638 ! halo=3 is requested in meso_sfn_ANN_compute, so shift must be <= 3.
639 if (cs%ann_window < 1 .or. cs%ann_window > 3 .or. mod(cs%ann_window, 2) == 0) &
640 call mom_error(fatal, "meso_sfn_ANN_init: MESO_SFN_ANN_WINDOW must be an odd integer in [1,3].")
641 call get_param(param_file, mdl, "MESO_SFN_ANN_FILE", cs%ann_file_rho_flux, &
642 "ANN parameters for prediction of density fluxes (netcdf)", &
643 default="INPUT/rho_flux.nc")
644 call ann_init(cs%ann_rho_flux, cs%ann_file_rho_flux)
645
646 ! Register diagnostic fields
647 cs%id_drdx_u = register_diag_field('ocean_model', 'meso_sfn_drdx_u', diag%axesCui, time, &
648 'Zonal density gradient used in meso sfn', &
649 'kg m-4', conversion=us%R_to_kg_m3*us%m_to_L)
650 cs%id_drdy_v = register_diag_field('ocean_model', 'meso_sfn_drdy_v', diag%axesCvi, time, &
651 'Meridional density gradient used in meso sfn', &
652 'kg m-4', conversion=us%R_to_kg_m3*us%m_to_L)
653 cs%id_drdz_u = register_diag_field('ocean_model', 'meso_sfn_drdz_u', diag%axesCui, time, &
654 'Vertical density gradient at u points used in meso sfn', &
655 'kg m-4', conversion=us%R_to_kg_m3*us%m_to_Z)
656 cs%id_drdz_v = register_diag_field('ocean_model', 'meso_sfn_drdz_v', diag%axesCvi, time, &
657 'Vertical density gradient at v points used in meso sfn', &
658 'kg m-4', conversion=us%R_to_kg_m3*us%m_to_Z)
659 cs%id_drdx_c = register_diag_field('ocean_model', 'meso_sfn_drdx_c', diag%axesTi, time, &
660 'Zonal density gradient at center points used in meso sfn', &
661 'kg m-4', conversion=us%R_to_kg_m3*us%m_to_L)
662 cs%id_drdy_c = register_diag_field('ocean_model', 'meso_sfn_drdy_c', diag%axesTi, time, &
663 'Meridional density gradient at center points used in meso sfn', &
664 'kg m-4', conversion=us%R_to_kg_m3*us%m_to_L)
665 cs%id_Fx_c = register_diag_field('ocean_model', 'meso_sfn_flux_x_c', diag%axesTi, time, &
666 'Zonal density flux at center points used in meso sfn', &
667 'kg m-2 s-1', conversion=us%R_to_kg_m3*us%L_to_m*us%s_to_T)
668 cs%id_Fy_c = register_diag_field('ocean_model', 'meso_sfn_flux_y_c', diag%axesTi, time, &
669 'Meridional density flux at center points used in meso sfn', &
670 'kg m-2 s-1', conversion=us%R_to_kg_m3*us%L_to_m*us%s_to_T)
671 cs%id_Fx_u = register_diag_field('ocean_model', 'meso_sfn_flux_x_u', diag%axesCui, time, &
672 'Zonal density flux at u points used in meso sfn', &
673 'kg m-2 s-1', conversion=us%R_to_kg_m3*us%L_to_m*us%s_to_T)
674 cs%id_Fy_v = register_diag_field('ocean_model', 'meso_sfn_flux_y_v', diag%axesCvi, time, &
675 'Meridional density flux at v points used in meso sfn', &
676 'kg m-2 s-1', conversion=us%R_to_kg_m3*us%L_to_m*us%s_to_T)
677 cs%id_sfn_u = register_diag_field('ocean_model', 'meso_sfn_unlim_u', diag%axesCui, time, &
678 'Meso-scale volume streamfunction at u points', &
679 'm3 s-1', conversion=us%Z_to_m*us%L_to_m**2*us%s_to_T)
680 cs%id_sfn_v = register_diag_field('ocean_model', 'meso_sfn_unlim_v', diag%axesCvi, time, &
681 'Meso-scale volume streamfunction at v points', &
682 'm3 s-1', conversion=us%Z_to_m*us%L_to_m**2*us%s_to_T)
683
684 cs%initialized = .true.
685end subroutine meso_sfn_ann_init
686!> Finalizes the meso-scale streamfunction ANN parameterization
687!!
688subroutine meso_sfn_ann_end(CS)
689 type(meso_sfn_ann_cs), intent(inout) :: cs !< Control structure
690
691 ! Deallocate anything that needs to be.
692 call ann_end(cs%ann_rho_flux)
693
694end subroutine meso_sfn_ann_end
695
696end module mom_meso_sfn_ann