miniWeather_openacc.f90 26 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574
  1. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  2. !! miniWeather
  3. !! Author: Matt Norman <normanmr@ornl.gov> , Oak Ridge National Laboratory
  4. !! This code simulates dry, stratified, compressible, non-hydrostatic fluid flows
  5. !! For documentation, please see the attached documentation in the "documentation" folder
  6. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  7. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  8. !! Copyright (c) 2018, National Center for Computational Sciences, Oak Ridge National Laboratory. All rights reserved.
  9. !!
  10. !! Portions Copyright (c) 2020, NVIDIA Corporation. All rights reserved.
  11. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  12. program miniweather
  13. use nvtx
  14. !implicit none
  15. !Declare the precision for the real constants (at least 15 digits of accuracy = double precision)
  16. integer , parameter :: rp = selected_real_kind(15)
  17. !Define some physical constants to use throughout the simulation
  18. real(rp), parameter :: pi = 3.14159265358979323846264338327_rp !Pi
  19. real(rp), parameter :: grav = 9.8_rp !Gravitational acceleration (m / s^2)
  20. real(rp), parameter :: cp = 1004._rp !Specific heat of dry air at constant pressure
  21. real(rp), parameter :: rd = 287._rp !Dry air constant for equation of state (P=rho*rd*T)
  22. real(rp), parameter :: p0 = 1.e5_rp !Standard pressure at the surface in Pascals
  23. real(rp), parameter :: C0 = 27.5629410929725921310572974482_rp !Constant to translate potential temperature into pressure (P=C0*(rho*theta)**gamma)
  24. real(rp), parameter :: gamma = 1.40027894002789400278940027894_rp !gamma=cp/Rd
  25. !Define domain and stability-related constants
  26. real(rp), parameter :: xlen = 2.e4_rp !Length of the domain in the x-direction (meters)
  27. real(rp), parameter :: zlen = 1.e4_rp !Length of the domain in the z-direction (meters)
  28. real(rp), parameter :: hv_beta = 0.25_rp !How strong to diffuse the solution: hv_beta \in [0:1]
  29. real(rp), parameter :: cfl = 1.50_rp !"Courant, Friedrichs, Lewy" number (for numerical stability)
  30. real(rp), parameter :: max_speed = 450 !Assumed maximum wave speed during the simulation (speed of sound + speed of wind) (meter / sec)
  31. integer , parameter :: hs = 2 !"Halo" size: number of cells needed for a full "stencil" of information for reconstruction
  32. integer , parameter :: sten_size = 4 !Size of the stencil used for interpolation
  33. !Parameters for indexing and flags
  34. integer , parameter :: NUM_VARS = 4 !Number of fluid state variables
  35. integer , parameter :: ID_DENS = 1 !index for density ("rho")
  36. integer , parameter :: ID_UMOM = 2 !index for momentum in the x-direction ("rho * u")
  37. integer , parameter :: ID_WMOM = 3 !index for momentum in the z-direction ("rho * w")
  38. integer , parameter :: ID_RHOT = 4 !index for density * potential temperature ("rho * theta")
  39. integer , parameter :: DIR_X = 1 !Integer constant to express that this operation is in the x-direction
  40. integer , parameter :: DIR_Z = 2 !Integer constant to express that this operation is in the z-direction
  41. !Gauss-Legendre quadrature points and weights on the domain [0:1]
  42. integer , parameter :: nqpoints = 3
  43. real(rp) :: qpoints (nqpoints) = (/ 0.112701665379258311482073460022E0_rp , 0.500000000000000000000000000000E0_rp , 0.887298334620741688517926539980E0_rp /)
  44. real(rp) :: qweights(nqpoints) = (/ 0.277777777777777777777777777779E0_rp , 0.444444444444444444444444444444E0_rp , 0.277777777777777777777777777779E0_rp /)
  45. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  46. !! Variables that are initialized but remain static over the course of the simulation
  47. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  48. real(rp) :: sim_time !total simulation time in seconds
  49. real(rp) :: output_freq !frequency to perform output in seconds
  50. real(rp) :: dt !Model time step (seconds)
  51. integer :: nx, nz !Number of local grid cells in the x- and z- dimensions
  52. real(rp) :: dx, dz !Grid space length in x- and z-dimension (meters)
  53. integer :: nx_glob, nz_glob !Number of total grid cells in the x- and z- dimensions
  54. integer :: i_beg, k_beg !beginning index in the x- and z-directions
  55. integer :: nranks, myrank !my rank id
  56. integer :: left_rank, right_rank
  57. real(rp), allocatable :: hy_dens_cell (:) !hydrostatic density (vert cell avgs). Dimensions: (1-hs:nz+hs)
  58. real(rp), allocatable :: hy_dens_theta_cell(:) !hydrostatic rho*t (vert cell avgs). Dimensions: (1-hs:nz+hs)
  59. real(rp), allocatable :: hy_dens_int (:) !hydrostatic density (vert cell interf). Dimensions: (1:nz+1)
  60. real(rp), allocatable :: hy_dens_theta_int (:) !hydrostatic rho*t (vert cell interf). Dimensions: (1:nz+1)
  61. real(rp), allocatable :: hy_pressure_int (:) !hydrostatic press (vert cell interf). Dimensions: (1:nz+1)
  62. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  63. !! Variables that are dynamics over the course of the simulation
  64. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  65. real(rp) :: etime !Elapsed model time
  66. real(rp) :: output_counter !Helps determine when it's time to do output
  67. !Runtime variable arrays
  68. real(rp), allocatable :: state (:,:,:) !Fluid state. Dimensions: (1-hs:nx+hs,1-hs:nz+hs,NUM_VARS)
  69. real(rp), allocatable :: state_tmp (:,:,:) !Fluid state. Dimensions: (1-hs:nx+hs,1-hs:nz+hs,NUM_VARS)
  70. real(rp), allocatable :: flux (:,:,:) !Cell interface fluxes. Dimensions: (nx+1,nz+1,NUM_VARS)
  71. real(rp), allocatable :: tend (:,:,:) !Fluid state tendencies. Dimensions: (nx,nz,NUM_VARS)
  72. integer(8) :: t1, t2, rate !For CPU Timings
  73. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  74. !! THE MAIN PROGRAM STARTS HERE
  75. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  76. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  77. !! BEGIN USER-CONFIGURABLE PARAMETERS
  78. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  79. !The x-direction length is twice as long as the z-direction length
  80. !So, you'll want to have nx_glob be twice as large as nz_glob
  81. nx_glob = 40 !Number of total cells in the x-dirction
  82. nz_glob = 20 !Number of total cells in the z-dirction
  83. sim_time = 1000 !How many seconds to run the simulation
  84. output_freq = 100 !How frequently to output data to file (in seconds)
  85. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  86. !! END USER-CONFIGURABLE PARAMETERS
  87. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  88. !Allocate arrays, initialize the grid and the data
  89. call nvtxStartRange("Total")
  90. call init()
  91. !Output the initial state
  92. !call output(state,etime)
  93. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  94. !! MAIN TIME STEP LOOP
  95. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  96. call system_clock(t1)
  97. call nvtxStartRange("while")
  98. do while (etime < sim_time)
  99. !If the time step leads to exceeding the simulation time, shorten it for the last step
  100. if (etime + dt > sim_time) dt = sim_time - etime
  101. !Perform a single time step
  102. call nvtxStartRange("perform_timestep")
  103. call perform_timestep(state,state_tmp,flux,tend,dt)
  104. call nvtxEndRange
  105. !Inform the user
  106. write(*,*) 'Elapsed Time: ', etime , ' / ' , sim_time
  107. !Update the elapsed time and output counter
  108. etime = etime + dt
  109. output_counter = output_counter + dt
  110. !If it's time for output, reset the counter, and do output
  111. if (output_counter >= output_freq) then
  112. output_counter = output_counter - output_freq
  113. !call output(state,etime)
  114. endif
  115. enddo
  116. call nvtxEndRange
  117. call system_clock(t2,rate)
  118. write(*,*) "CPU Time: ",dble(t2-t1)/dble(rate)
  119. !Deallocate
  120. call finalize()
  121. call nvtxEndRange
  122. contains
  123. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  124. !! HELPER SUBROUTINES AND FUNCTIONS GO HERE
  125. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  126. !Performs a single dimensionally split time step using a simple low-storate three-stage Runge-Kutta time integrator
  127. !The dimensional splitting is a second-order-accurate alternating Strang splitting in which the
  128. !order of directions is alternated each time step.
  129. !The Runge-Kutta method used here is defined as follows:
  130. ! q* = q[n] + dt/3 * rhs(q[n])
  131. ! q** = q[n] + dt/2 * rhs(q* )
  132. ! q[n+1] = q[n] + dt/1 * rhs(q** )
  133. subroutine perform_timestep(state,state_tmp,flux,tend,dt)
  134. implicit none
  135. real(rp), intent(inout) :: state (1-hs:nx+hs,1-hs:nz+hs,NUM_VARS)
  136. real(rp), intent( out) :: state_tmp(1-hs:nx+hs,1-hs:nz+hs,NUM_VARS)
  137. real(rp), intent( out) :: flux (nx+1,nz+1,NUM_VARS)
  138. real(rp), intent( out) :: tend (nx,nz,NUM_VARS)
  139. real(rp), intent(in ) :: dt
  140. logical, save :: direction_switch = .true.
  141. if (direction_switch) then
  142. !x-direction first
  143. call semi_discrete_step( state , state , state_tmp , dt / 3 , DIR_X , flux , tend )
  144. call semi_discrete_step( state , state_tmp , state_tmp , dt / 2 , DIR_X , flux , tend )
  145. call semi_discrete_step( state , state_tmp , state , dt / 1 , DIR_X , flux , tend )
  146. !z-direction second
  147. call semi_discrete_step( state , state , state_tmp , dt / 3 , DIR_Z , flux , tend )
  148. call semi_discrete_step( state , state_tmp , state_tmp , dt / 2 , DIR_Z , flux , tend )
  149. call semi_discrete_step( state , state_tmp , state , dt / 1 , DIR_Z , flux , tend )
  150. else
  151. !z-direction second
  152. call semi_discrete_step( state , state , state_tmp , dt / 3 , DIR_Z , flux , tend )
  153. call semi_discrete_step( state , state_tmp , state_tmp , dt / 2 , DIR_Z , flux , tend )
  154. call semi_discrete_step( state , state_tmp , state , dt / 1 , DIR_Z , flux , tend )
  155. !x-direction first
  156. call semi_discrete_step( state , state , state_tmp , dt / 3 , DIR_X , flux , tend )
  157. call semi_discrete_step( state , state_tmp , state_tmp , dt / 2 , DIR_X , flux , tend )
  158. call semi_discrete_step( state , state_tmp , state , dt / 1 , DIR_X , flux , tend )
  159. endif
  160. direction_switch = .not. direction_switch
  161. end subroutine perform_timestep
  162. !Perform a single semi-discretized step in time with the form:
  163. !state_out = state_init + dt * rhs(state_forcing)
  164. !Meaning the step starts from state_init, computes the rhs using state_forcing, and stores the result in state_out
  165. subroutine semi_discrete_step( state_init , state_forcing , state_out , dt , dir , flux , tend )
  166. implicit none
  167. real(rp), intent(in ) :: state_init (1-hs:nx+hs,1-hs:nz+hs,NUM_VARS)
  168. real(rp), intent(inout) :: state_forcing(1-hs:nx+hs,1-hs:nz+hs,NUM_VARS)
  169. real(rp), intent( out) :: state_out (1-hs:nx+hs,1-hs:nz+hs,NUM_VARS)
  170. real(rp), intent( out) :: flux (nx+1,nz+1,NUM_VARS)
  171. real(rp), intent( out) :: tend (nx,nz,NUM_VARS)
  172. real(rp), intent(in ) :: dt
  173. integer , intent(in ) :: dir
  174. integer :: i,k,ll
  175. if (dir == DIR_X) then
  176. !Set the halo values in the x-direction
  177. call set_halo_values_x(state_forcing)
  178. !Compute the time tendencies for the fluid state in the x-direction
  179. call compute_tendencies_x(state_forcing,flux,tend)
  180. elseif (dir == DIR_Z) then
  181. !Set the halo values in the z-direction
  182. call set_halo_values_z(state_forcing)
  183. !Compute the time tendencies for the fluid state in the z-direction
  184. call compute_tendencies_z(state_forcing,flux,tend)
  185. endif
  186. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  187. !! TODO: THREAD ME
  188. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  189. !Apply the tendencies to the fluid state
  190. !$acc parallel loop
  191. do ll = 1 , NUM_VARS
  192. do k = 1 , nz
  193. do i = 1 , nx
  194. state_out(i,k,ll) = state_init(i,k,ll) + dt * tend(i,k,ll)
  195. enddo
  196. enddo
  197. enddo
  198. end subroutine semi_discrete_step
  199. !Compute the time tendencies of the fluid state using forcing in the x-direction
  200. !First, compute the flux vector at each cell interface in the x-direction (including hyperviscosity)
  201. !Then, compute the tendencies using those fluxes
  202. subroutine compute_tendencies_x(state,flux,tend)
  203. implicit none
  204. real(rp), intent(in ) :: state(1-hs:nx+hs,1-hs:nz+hs,NUM_VARS)
  205. real(rp), intent( out) :: flux (nx+1,nz+1,NUM_VARS)
  206. real(rp), intent( out) :: tend (nx,nz,NUM_VARS)
  207. integer :: i,k,ll,s
  208. real(rp) :: r,u,w,t,p, stencil(4), d3_vals(NUM_VARS), vals(NUM_VARS), hv_coef
  209. !Compute the hyperviscosity coeficient
  210. hv_coef = -hv_beta * dx / (16*dt)
  211. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  212. !! TODO: THREAD ME
  213. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  214. !Compute fluxes in the x-direction for each cell
  215. !$acc parallel loop
  216. do k = 1 , nz
  217. !$acc loop private(stencil,vals,d3_vals)
  218. do i = 1 , nx+1
  219. !Use fourth-order interpolation from four cell averages to compute the value at the interface in question
  220. do ll = 1 , NUM_VARS
  221. do s = 1 , sten_size
  222. stencil(s) = state(i-hs-1+s,k,ll)
  223. enddo
  224. !Fourth-order-accurate interpolation of the state
  225. vals(ll) = -stencil(1)/12 + 7*stencil(2)/12 + 7*stencil(3)/12 - stencil(4)/12
  226. !First-order-accurate interpolation of the third spatial derivative of the state (for artificial viscosity)
  227. d3_vals(ll) = -stencil(1) + 3*stencil(2) - 3*stencil(3) + stencil(4)
  228. enddo
  229. !Compute density, u-wind, w-wind, potential temperature, and pressure (r,u,w,t,p respectively)
  230. r = vals(ID_DENS) + hy_dens_cell(k)
  231. u = vals(ID_UMOM) / r
  232. w = vals(ID_WMOM) / r
  233. t = ( vals(ID_RHOT) + hy_dens_theta_cell(k) ) / r
  234. p = C0*(r*t)**gamma
  235. !Compute the flux vector
  236. flux(i,k,ID_DENS) = r*u - hv_coef*d3_vals(ID_DENS)
  237. flux(i,k,ID_UMOM) = r*u*u+p - hv_coef*d3_vals(ID_UMOM)
  238. flux(i,k,ID_WMOM) = r*u*w - hv_coef*d3_vals(ID_WMOM)
  239. flux(i,k,ID_RHOT) = r*u*t - hv_coef*d3_vals(ID_RHOT)
  240. enddo
  241. enddo
  242. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  243. !! TODO: THREAD ME
  244. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  245. !Use the fluxes to compute tendencies for each cell
  246. !$acc parallel loop
  247. do ll = 1 , NUM_VARS
  248. do k = 1 , nz
  249. do i = 1 , nx
  250. tend(i,k,ll) = -( flux(i+1,k,ll) - flux(i,k,ll) ) / dx
  251. enddo
  252. enddo
  253. enddo
  254. end subroutine compute_tendencies_x
  255. !Compute the time tendencies of the fluid state using forcing in the z-direction
  256. !First, compute the flux vector at each cell interface in the z-direction (including hyperviscosity)
  257. !Then, compute the tendencies using those fluxes
  258. subroutine compute_tendencies_z(state,flux,tend)
  259. implicit none
  260. real(rp), intent(in ) :: state(1-hs:nx+hs,1-hs:nz+hs,NUM_VARS)
  261. real(rp), intent( out) :: flux (nx+1,nz+1,NUM_VARS)
  262. real(rp), intent( out) :: tend (nx,nz,NUM_VARS)
  263. integer :: i,k,ll,s
  264. real(rp) :: r,u,w,t,p, stencil(4), d3_vals(NUM_VARS), vals(NUM_VARS), hv_coef
  265. !Compute the hyperviscosity coeficient
  266. hv_coef = -hv_beta * dz / (16*dt)
  267. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  268. !! TODO: THREAD ME
  269. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  270. !Compute fluxes in the x-direction for each cell
  271. !$acc parallel loop
  272. do k = 1 , nz+1
  273. !$acc loop private(stencil,vals,d3_vals)
  274. do i = 1 , nx
  275. !Use fourth-order interpolation from four cell averages to compute the value at the interface in question
  276. do ll = 1 , NUM_VARS
  277. do s = 1 , sten_size
  278. stencil(s) = state(i,k-hs-1+s,ll)
  279. enddo
  280. !Fourth-order-accurate interpolation of the state
  281. vals(ll) = -stencil(1)/12 + 7*stencil(2)/12 + 7*stencil(3)/12 - stencil(4)/12
  282. !First-order-accurate interpolation of the third spatial derivative of the state
  283. d3_vals(ll) = -stencil(1) + 3*stencil(2) - 3*stencil(3) + stencil(4)
  284. enddo
  285. !Compute density, u-wind, w-wind, potential temperature, and pressure (r,u,w,t,p respectively)
  286. r = vals(ID_DENS) + hy_dens_int(k)
  287. u = vals(ID_UMOM) / r
  288. w = vals(ID_WMOM) / r
  289. t = ( vals(ID_RHOT) + hy_dens_theta_int(k) ) / r
  290. p = C0*(r*t)**gamma - hy_pressure_int(k)
  291. !Enforce vertical boundary condition and exact mass conservation
  292. if (k == 1 .or. k == nz+1) then
  293. w = 0
  294. d3_vals(ID_DENS) = 0
  295. endif
  296. !Compute the flux vector with hyperviscosity
  297. flux(i,k,ID_DENS) = r*w - hv_coef*d3_vals(ID_DENS)
  298. flux(i,k,ID_UMOM) = r*w*u - hv_coef*d3_vals(ID_UMOM)
  299. flux(i,k,ID_WMOM) = r*w*w+p - hv_coef*d3_vals(ID_WMOM)
  300. flux(i,k,ID_RHOT) = r*w*t - hv_coef*d3_vals(ID_RHOT)
  301. enddo
  302. enddo
  303. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  304. !! TODO: THREAD ME
  305. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  306. !Use the fluxes to compute tendencies for each cell
  307. !$acc parallel loop
  308. do ll = 1 , NUM_VARS
  309. do k = 1 , nz
  310. do i = 1 , nx
  311. tend(i,k,ll) = -( flux(i,k+1,ll) - flux(i,k,ll) ) / dz
  312. if (ll == ID_WMOM) then
  313. tend(i,k,ID_WMOM) = tend(i,k,ID_WMOM) - state(i,k,ID_DENS)*grav
  314. endif
  315. enddo
  316. enddo
  317. enddo
  318. end subroutine compute_tendencies_z
  319. !Set halo values in the x-direction.
  320. subroutine set_halo_values_x(state)
  321. implicit none
  322. real(rp), intent(inout) :: state(1-hs:nx+hs,1-hs:nz+hs,NUM_VARS)
  323. integer :: k, ll
  324. real(rp) :: z
  325. !$acc parallel loop
  326. do ll = 1 , NUM_VARS
  327. do k = 1 , nz
  328. state(-1 ,k,ll) = state(nx-1,k,ll)
  329. state(0 ,k,ll) = state(nx ,k,ll)
  330. state(nx+1,k,ll) = state(1 ,k,ll)
  331. state(nx+2,k,ll) = state(2 ,k,ll)
  332. enddo
  333. enddo
  334. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  335. if (myrank == 0) then
  336. !$acc parallel loop
  337. do k = 1 , nz
  338. z = (k_beg-1 + k-0.5_rp)*dz
  339. if (abs(z-3*zlen/4) <= zlen/16) then
  340. state(-1:0,k,ID_UMOM) = (state(-1:0,k,ID_DENS)+hy_dens_cell(k)) * 50._rp
  341. state(-1:0,k,ID_RHOT) = (state(-1:0,k,ID_DENS)+hy_dens_cell(k)) * 298._rp - hy_dens_theta_cell(k)
  342. endif
  343. enddo
  344. endif
  345. end subroutine set_halo_values_x
  346. !Set halo values in the z-direction.
  347. !decomposition in the vertical direction
  348. subroutine set_halo_values_z(state)
  349. implicit none
  350. real(rp), intent(inout) :: state(1-hs:nx+hs,1-hs:nz+hs,NUM_VARS)
  351. integer :: i, ll
  352. real(rp), parameter :: mnt_width = xlen/8
  353. real(rp) :: x, xloc, mnt_deriv
  354. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  355. !! TODO: THREAD ME
  356. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  357. !$acc parallel loop private(xloc,mnt_deriv,x)
  358. do ll = 1 , NUM_VARS
  359. do i = 1-hs,nx+hs
  360. if (ll == ID_WMOM) then
  361. state(i,-1 ,ll) = 0
  362. state(i,0 ,ll) = 0
  363. state(i,nz+1,ll) = 0
  364. state(i,nz+2,ll) = 0
  365. else
  366. state(i,-1 ,ll) = state(i,1 ,ll)
  367. state(i,0 ,ll) = state(i,1 ,ll)
  368. state(i,nz+1,ll) = state(i,nz,ll)
  369. state(i,nz+2,ll) = state(i,nz,ll)
  370. endif
  371. enddo
  372. enddo
  373. end subroutine set_halo_values_z
  374. !Initialize some grid settings, allocate data, and initialize the full grid and data
  375. subroutine init()
  376. implicit none
  377. integer :: i, k, ii, kk, ll, ierr
  378. real(rp) :: x, z, r, u, w, t, hr, ht
  379. !Set the cell grid size
  380. dx = xlen / nx_glob
  381. dz = zlen / nz_glob
  382. nranks = 1
  383. myrank = 0
  384. i_beg = 1
  385. nx = nx_glob
  386. left_rank = 0
  387. right_rank = 0
  388. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  389. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  390. !! YOU DON'T NEED TO ALTER ANYTHING BELOW THIS POINT IN THE CODE
  391. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  392. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  393. k_beg = 1
  394. nz = nz_glob
  395. !Allocate the model data
  396. allocate(state (1-hs:nx+hs,1-hs:nz+hs,NUM_VARS))
  397. allocate(state_tmp (1-hs:nx+hs,1-hs:nz+hs,NUM_VARS))
  398. allocate(flux (nx+1,nz+1,NUM_VARS))
  399. allocate(tend (nx,nz,NUM_VARS))
  400. allocate(hy_dens_cell (1-hs:nz+hs))
  401. allocate(hy_dens_theta_cell(1-hs:nz+hs))
  402. allocate(hy_dens_int (nz+1))
  403. allocate(hy_dens_theta_int (nz+1))
  404. allocate(hy_pressure_int (nz+1))
  405. !Define the maximum stable time step based on an assumed maximum wind speed
  406. dt = min(dx,dz) / max_speed * cfl
  407. !Set initial elapsed model time and output_counter to zero
  408. etime = 0
  409. output_counter = 0
  410. !Display some grid information
  411. write(*,*) 'nx_glob, nz_glob:', nx_glob, nz_glob
  412. write(*,*) 'dx,dz: ',dx,dz
  413. write(*,*) 'dt: ',dt
  414. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  415. !! Initialize the cell-averaged fluid state via Gauss-Legendre quadrature
  416. !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  417. state = 0
  418. do k = 1-hs , nz+hs
  419. do i = 1-hs , nx+hs
  420. !Initialize the state to zero
  421. !Use Gauss-Legendre quadrature to initialize a hydrostatic balance + temperature perturbation
  422. do kk = 1 , nqpoints
  423. do ii = 1 , nqpoints
  424. !Compute the x,z location within the global domain based on cell and quadrature index
  425. x = (i_beg-1 + i-0.5_rp) * dx + (qpoints(ii)-0.5_rp)*dx
  426. z = (k_beg-1 + k-0.5_rp) * dz + (qpoints(kk)-0.5_rp)*dz
  427. !Set the fluid state based on the user's specification
  428. call injection (x,z,r,u,w,t,hr,ht)
  429. !Store into the fluid state array
  430. state(i,k,ID_DENS) = state(i,k,ID_DENS) + r * qweights(ii)*qweights(kk)
  431. state(i,k,ID_UMOM) = state(i,k,ID_UMOM) + (r+hr)*u * qweights(ii)*qweights(kk)
  432. state(i,k,ID_WMOM) = state(i,k,ID_WMOM) + (r+hr)*w * qweights(ii)*qweights(kk)
  433. state(i,k,ID_RHOT) = state(i,k,ID_RHOT) + ( (r+hr)*(t+ht) - hr*ht ) * qweights(ii)*qweights(kk)
  434. enddo
  435. enddo
  436. do ll = 1 , NUM_VARS
  437. state_tmp(i,k,ll) = state(i,k,ll)
  438. enddo
  439. enddo
  440. enddo
  441. !Compute the hydrostatic background state over vertical cell averages
  442. hy_dens_cell = 0
  443. hy_dens_theta_cell = 0
  444. do k = 1-hs , nz+hs
  445. do kk = 1 , nqpoints
  446. z = (k_beg-1 + k-0.5_rp) * dz + (qpoints(kk)-0.5_rp)*dz
  447. !Set the fluid state based on the user's specification
  448. call injection (0._rp,z,r,u,w,t,hr,ht)
  449. hy_dens_cell(k) = hy_dens_cell(k) + hr * qweights(kk)
  450. hy_dens_theta_cell(k) = hy_dens_theta_cell(k) + hr*ht * qweights(kk)
  451. enddo
  452. enddo
  453. !Compute the hydrostatic background state at vertical cell interfaces
  454. do k = 1 , nz+1
  455. z = (k_beg-1 + k-1) * dz
  456. call injection (0._rp,z,r,u,w,t,hr,ht)
  457. hy_dens_int (k) = hr
  458. hy_dens_theta_int(k) = hr*ht
  459. hy_pressure_int (k) = C0*(hr*ht)**gamma
  460. enddo
  461. end subroutine init
  462. !This test case is initially balanced but injects fast, cold air from the left boundary near the model top
  463. subroutine injection(x,z,r,u,w,t,hr,ht)
  464. implicit none
  465. real(rp), intent(in ) :: x, z !x- and z- location of the point being sampled
  466. real(rp), intent( out) :: r, u, w, t !Density, uwind, wwind, and potential temperature
  467. real(rp), intent( out) :: hr, ht !Hydrostatic density and potential temperature
  468. call hydro_const_theta(z,hr,ht)
  469. r = 0
  470. t = 0
  471. u = 0
  472. w = 0
  473. end subroutine injection
  474. !Establish hydrstatic balance using constant potential temperature (thermally neutral atmosphere)
  475. subroutine hydro_const_theta(z,r,t)
  476. implicit none
  477. real(rp), intent(in ) :: z !x- and z- location of the point being sampled
  478. real(rp), intent( out) :: r, t !Density and potential temperature at this point
  479. real(rp), parameter :: theta0 = 300._rp !Background potential temperature
  480. real(rp), parameter :: exner0 = 1._rp !Surface-level Exner pressure
  481. real(rp) :: p,exner,rt
  482. !Establish hydrostatic balance first using Exner pressure
  483. t = theta0 !Potential Temperature at z
  484. exner = exner0 - grav * z / (cp * theta0) !Exner pressure at z
  485. p = p0 * exner**(cp/rd) !Pressure at z
  486. rt = (p / c0)**(1._rp / gamma) !rho*theta at z
  487. r = rt / t !Density at z
  488. end subroutine hydro_const_theta
  489. !Deallocate and call Finalize
  490. subroutine finalize()
  491. implicit none
  492. integer :: ierr
  493. deallocate(state )
  494. deallocate(state_tmp )
  495. deallocate(flux )
  496. deallocate(tend )
  497. deallocate(hy_dens_cell )
  498. deallocate(hy_dens_theta_cell)
  499. deallocate(hy_dens_int )
  500. deallocate(hy_dens_theta_int )
  501. deallocate(hy_pressure_int )
  502. end subroutine finalize
  503. end program miniweather