def set_initial_conditions(self, vs):
# initial conditions
vs.temp[:, :, :, 0:2] = ((1 - vs.zt[None, None, :] / vs.zw[0]) * 15 * vs.maskT)[..., None]
vs.salt[:, :, :, 0:2] = 35.0 * vs.maskT[..., None]
# wind stress forcing
yt_min = global_min(vs, vs.yt.min())
yu_min = global_min(vs, vs.yu.min())
yt_max = global_max(vs, vs.yt.max())
yu_max = global_max(vs, vs.yu.max())
taux = allocate(vs, ('yt',))
taux[vs.yt < -20] = 0.1 * np.sin(vs.pi * (vs.yu[vs.yt < -20] - yu_min) / (-20.0 - yt_min))
taux[vs.yt > 10] = 0.1 * (1 - np.cos(2 * vs.pi * (vs.yu[vs.yt > 10] - 10.0) / (yu_max - 10.0)))
vs.surface_taux[:, :] = taux * vs.maskU[:, :, -1]
# surface heatflux forcing
vs._t_star = allocate(vs, ('yt',), fill=15)
vs._t_star[vs.yt < -20] = 15 * (vs.yt[vs.yt < -20] - yt_min) / (-20 - yt_min)
vs._t_star[vs.yt > 20] = 15 * (1 - (vs.yt[vs.yt > 20] - 20) / (yt_max - 20))
vs._t_rest = vs.dzt[None, -1] / (30. * 86400.) * vs.maskT[:, :, -1]
if vs.enable_tke:
vs.forc_tke_surface[2:-2, 2:-2] = np.sqrt((0.5 * (vs.surface_taux[2:-2, 2:-2] + vs.surface_taux[1:-3, 2:-2]) / vs.rho_0)**2
+ (0.5 * (vs.surface_tauy[2:-2, 2:-2] + vs.surface_tauy[2:-2, 1:-3]) / vs.rho_0)**2)**(1.5)
if vs.enable_idemix:
vs.forc_iw_bottom[...] = 1e-6 * vs.maskW[:, :, -1]
vs.forc_iw_surface[...] = 1e-7 * vs.maskW[:, :, -1]
def set_initial_conditions(self, vs):
# initial conditions
vs.temp[:, :, :, 0:2] = ((1 - vs.zt[None, None, :] / vs.zw[0]) * 15 *
vs.maskT)[..., None]
vs.salt[:, :, :, 0:2] = 35.0 * vs.maskT[..., None]
# wind stress forcing
yt_min = global_min(vs, vs.yt.min())
yu_min = global_min(vs, vs.yu.min())
yt_max = global_max(vs, vs.yt.max())
yu_max = global_max(vs, vs.yu.max())
taux = allocate(vs, ('yt', ))
north = vs.yt > 30
subequatorial_north_n = (vs.yt >= 15) & (vs.yt < 30)
subequatorial_north_s = (vs.yt > 0) & (vs.yt < 15)
equator = (vs.yt > -5) & (vs.yt < 5)
subequatorial_south_n = (vs.yt > -15) & (vs.yt < 0)
subequatorial_south_s = (vs.yt <= -15) & (vs.yt > -30)
south = vs.yt < -30
taux[north] = -5e-2 * np.sin(np.pi * (vs.yu[north] - yu_max) /
(yt_max - 30.))
taux[subequatorial_north_s] = 5e-2 * np.sin(
np.pi * (vs.yu[subequatorial_north_s] - 30.) / 30.)
taux[subequatorial_north_n] = 5e-2 * np.sin(
np.pi * (vs.yt[subequatorial_north_n] - 30.) / 30.)
taux[subequatorial_south_n] = -5e-2 * np.sin(
np.pi * (vs.yu[subequatorial_south_n] - 30.) / 30.)
taux[subequatorial_south_s] = -5e-2 * np.sin(
np.pi * (vs.yt[subequatorial_south_s] - 30.) / 30.)
taux[equator] = -1.5e-2 * np.cos(np.pi *
(vs.yu[equator] - 10.) / 10.) - 2.5e-2
taux[south] = 15e-2 * np.sin(np.pi * (vs.yu[south] - yu_min) /
(-30. - yt_min))
vs.surface_taux[:, :] = taux * vs.maskU[:, :, -1]
# surface heatflux forcing
DELTA_T, TS, TN = 25., 0., 5.
vs._t_star = allocate(vs, ('yt', ), fill=DELTA_T)
vs._t_star[vs.yt < 0] = TS + DELTA_T * np.sin(
np.pi * (vs.yt[vs.yt < 0] + 60.) / np.abs(2 * vs.y_origin))
vs._t_star[vs.yt > 0] = TN + (DELTA_T + TS - TN) * np.sin(
np.pi * (vs.yt[vs.yt > 0] + 60.) / np.abs(2 * vs.y_origin))
vs._t_rest = vs.dzt[None, -1] / (10. * 86400.) * vs.maskT[:, :, -1]
if vs.enable_tke:
vs.forc_tke_surface[2:-2, 2:-2] = np.sqrt(
(0.5 *
(vs.surface_taux[2:-2, 2:-2] + vs.surface_taux[1:-3, 2:-2]) /
vs.rho_0)**2 +
(0.5 *
(vs.surface_tauy[2:-2, 2:-2] + vs.surface_tauy[2:-2, 1:-3]) /
vs.rho_0)**2)**(1.5)
if vs.enable_idemix:
vs.forc_iw_bottom[...] = 1e-6 * vs.maskW[:, :, -1]
vs.forc_iw_surface[...] = 1e-7 * vs.maskW[:, :, -1]
def set_initial_conditions(self, vs):
vs.u[:, :, :, vs.tau] = np.random.rand(54, 104, 10)
# wind stress forcing
yt_min = global_min(vs, vs.yt.min())
yu_min = global_min(vs, vs.yu.min())
yt_max = global_max(vs, vs.yt.max())
yu_max = global_max(vs, vs.yu.max())
# surface heatflux forcing
vs._t_star = allocate(vs, ('yt',), fill=15)
vs._t_star[vs.yt < -20] = 15 * (vs.yt[vs.yt < -20] - yt_min) / (-20 - yt_min)
vs._t_star[vs.yt > 20] = 15 * (1 - (vs.yt[vs.yt > 20] - 20) / (yt_max - 20))
vs._t_rest = vs.dzt[None, -1] / (30. * 86400.) * vs.maskT[:, :, -1]
pass
def set_topography(self, vs):
with h5netcdf.File(DATA_FILES['topography'], 'r') as topography_file:
topo_x, topo_y, topo_z = (np.array(topography_file.variables[k],
dtype='float').T
for k in ('x', 'y', 'z'))
topo_z[topo_z > 0] = 0.
# smooth topography to match grid resolution
gaussian_sigma = (0.5 * len(topo_x) / vs.nx, 0.5 * len(topo_y) / vs.ny)
topo_z_smoothed = scipy.ndimage.gaussian_filter(topo_z,
sigma=gaussian_sigma)
topo_z_smoothed[topo_z >= -1] = 0
topo_x_shifted, topo_z_shifted = self._shift_longitude_array(
vs, topo_x, topo_z_smoothed)
coords = (vs.xt[2:-2], vs.yt[2:-2])
z_interp = allocate(vs, ('xt', 'yt'), local=False)
z_interp[2:-2, 2:-2] = veros.tools.interpolate(
(topo_x_shifted, topo_y),
topo_z_shifted,
coords,
kind='nearest',
fill=False)
depth_levels = 1 + np.argmin(np.abs(z_interp[:, :, np.newaxis] -
vs.zt[np.newaxis, np.newaxis, :]),
axis=2)
vs.kbot[2:-2, 2:-2] = np.where(z_interp < 0., depth_levels, 0)[2:-2,
2:-2]
vs.kbot *= vs.kbot < vs.nz
enforce_boundaries(vs, vs.kbot)
# remove marginal seas
# (dilate to close 1-cell passages, fill holes, undo dilation)
marginal = (scipy.ndimage.binary_erosion(
scipy.ndimage.binary_fill_holes(
scipy.ndimage.binary_dilation(vs.kbot == 0))))
vs.kbot[marginal] = 0
def npzd(vs):
r"""
Main driving function for NPZD functionality
Computes transport terms and biological activity separately
\begin{equation}
\dfrac{\partial C_i}{\partial t} = T + S
\end{equation}
"""
if not vs.enable_npzd:
return
# TODO: Refactor transportation code to be defined only once and also used by thermodynamics
# TODO: Dissipation on W-grid if necessary
npzd_changes = biogeochemistry(vs)
"""
For vertical mixing
"""
a_tri = allocate(vs, ('xt', 'yt', 'zt'), include_ghosts=False)
b_tri = allocate(vs, ('xt', 'yt', 'zt'), include_ghosts=False)
c_tri = allocate(vs, ('xt', 'yt', 'zt'), include_ghosts=False)
d_tri = allocate(vs, ('xt', 'yt', 'zt'), include_ghosts=False)
delta = allocate(vs, ('xt', 'yt', 'zt'), include_ghosts=False)
ks = vs.kbot[2:-2, 2:-2] - 1
delta[:, :, :-1] = vs.dt_tracer / vs.dzw[np.newaxis, np.newaxis, :-1]\
* vs.kappaH[2:-2, 2:-2, :-1]
delta[:, :, -1] = 0
a_tri[:, :, 1:] = -delta[:, :, :-1] / vs.dzt[np.newaxis, np.newaxis, 1:]
b_tri[:, :,
1:] = 1 + (delta[:, :, 1:] +
delta[:, :, :-1]) / vs.dzt[np.newaxis, np.newaxis, 1:]
b_tri_edge = 1 + delta / vs.dzt[np.newaxis, np.newaxis, :]
c_tri[:, :, :-1] = -delta[:, :, :-1] / vs.dzt[np.newaxis, np.newaxis, :-1]
for tracer in vs.npzd_transported_tracers:
tracer_data = vs.npzd_tracers[tracer]
"""
Advection of tracers
"""
thermodynamics.advect_tracer(
vs, tracer_data[:, :, :, vs.tau],
vs.npzd_advection_derivatives[tracer][:, :, :, vs.tau])
# Adam-Bashforth timestepping
tracer_data[:, :, :, vs.taup1] = tracer_data[:, :, :, vs.tau] + vs.dt_tracer \
* ((1.5 + vs.AB_eps) * vs.npzd_advection_derivatives[tracer][:, :, :, vs.tau]
- (0.5 + vs.AB_eps) * vs.npzd_advection_derivatives[tracer][:, :, :, vs.taum1])\
* vs.maskT
"""
Diffusion of tracers
"""
if vs.enable_hor_diffusion:
horizontal_diffusion_change = np.zeros_like(tracer_data[:, :, :,
0])
diffusion.horizontal_diffusion(vs, tracer_data[:, :, :, vs.tau],
horizontal_diffusion_change)
tracer_data[:, :, :,
vs.taup1] += vs.dt_tracer * horizontal_diffusion_change
if vs.enable_biharmonic_mixing:
biharmonic_diffusion_change = np.empty_like(tracer_data[:, :, :,
0])
diffusion.biharmonic(vs, tracer_data[:, :, :, vs.tau],
np.sqrt(abs(vs.K_hbi)),
biharmonic_diffusion_change)
tracer_data[:, :, :,
vs.taup1] += vs.dt_tracer * biharmonic_diffusion_change
"""
Restoring zones
"""
# TODO add restoring zones to general tracers
"""
Isopycnal diffusion
"""
if vs.enable_neutral_diffusion:
dtracer_iso = np.zeros_like(tracer_data[..., 0])
isoneutral.isoneutral_diffusion_tracer(vs,
tracer_data,
dtracer_iso,
iso=True,
skew=False)
if vs.enable_skew_diffusion:
dtracer_skew = np.zeros_like(tracer_data[..., 0])
isoneutral.isoneutral_diffusion_tracer(vs,
tracer_data,
dtracer_skew,
iso=False,
skew=True)
"""
Vertical mixing of tracers
"""
d_tri[:, :, :] = tracer_data[2:-2, 2:-2, :, vs.taup1]
# TODO: surface flux?
# d_tri[:, :, -1] += surface_forcing
sol, mask = utilities.solve_implicit(vs,
ks,
a_tri,
b_tri,
c_tri,
d_tri,
b_edge=b_tri_edge)
tracer_data[2:-2, 2:-2, :, vs.taup1] = utilities.where(
vs, mask, sol, tracer_data[2:-2, 2:-2, :, vs.taup1])
# update by biogeochemical changes
for tracer, change in npzd_changes.items():
vs.npzd_tracers[tracer][:, :, :, vs.taup1] += change
# prepare next timestep with minimum tracer values
for tracer in vs.npzd_tracers.values():
tracer[:, :, :, vs.taup1] = np.maximum(tracer[:, :, :, vs.taup1],
vs.trcmin * vs.maskT)
for tracer in vs.npzd_tracers.values():
utilities.enforce_boundaries(vs, tracer)
def set_initial_conditions(self, vs):
rpart_shortwave = 0.58
efold1_shortwave = 0.35
efold2_shortwave = 23.0
# initial conditions
temp_data = self._read_forcing(vs, 'temperature')
vs.temp[2:-2, 2:-2, :,
0] = temp_data[..., ::-1] * vs.maskT[2:-2, 2:-2, :]
vs.temp[2:-2, 2:-2, :,
1] = temp_data[..., ::-1] * vs.maskT[2:-2, 2:-2, :]
salt_data = self._read_forcing(vs, 'salinity')
vs.salt[2:-2, 2:-2, :,
0] = salt_data[..., ::-1] * vs.maskT[2:-2, 2:-2, :]
vs.salt[2:-2, 2:-2, :,
1] = salt_data[..., ::-1] * vs.maskT[2:-2, 2:-2, :]
# wind stress on MIT grid
vs.taux[2:-2, 2:-2, :] = self._read_forcing(vs, 'tau_x')
vs.tauy[2:-2, 2:-2, :] = self._read_forcing(vs, 'tau_y')
qnec_data = self._read_forcing(vs, 'dqdt')
vs.qnec[2:-2,
2:-2, :] = qnec_data * vs.maskT[2:-2, 2:-2, -1, np.newaxis]
qsol_data = self._read_forcing(vs, 'swf')
vs.qsol[2:-2,
2:-2, :] = -qsol_data * vs.maskT[2:-2, 2:-2, -1, np.newaxis]
# SST and SSS
sst_data = self._read_forcing(vs, 'sst')
vs.t_star[2:-2,
2:-2, :] = sst_data * vs.maskT[2:-2, 2:-2, -1, np.newaxis]
sss_data = self._read_forcing(vs, 'sss')
vs.s_star[2:-2,
2:-2, :] = sss_data * vs.maskT[2:-2, 2:-2, -1, np.newaxis]
if vs.enable_idemix:
tidal_energy_data = self._read_forcing(vs, 'tidal_energy')
mask = np.maximum(0, vs.kbot[2:-2, 2:-2] -
1)[:, :,
np.newaxis] == np.arange(vs.nz)[np.newaxis,
np.newaxis, :]
tidal_energy_data[:, :] *= vs.maskW[2:-2, 2:-2, :][mask].reshape(
vs.nx, vs.ny) / vs.rho_0
vs.forc_iw_bottom[2:-2, 2:-2] = tidal_energy_data
wind_energy_data = self._read_forcing(vs, 'wind_energy')
wind_energy_data[:, :] *= vs.maskW[2:-2, 2:-2, -1] / vs.rho_0 * 0.2
vs.forc_iw_surface[2:-2, 2:-2] = wind_energy_data
"""
Initialize penetration profile for solar radiation and store divergence in divpen
note that pen is set to 0.0 at the surface instead of 1.0 to compensate for the
shortwave part of the total surface flux
"""
swarg1 = vs.zw / efold1_shortwave
swarg2 = vs.zw / efold2_shortwave
pen = rpart_shortwave * np.exp(swarg1) + (
1.0 - rpart_shortwave) * np.exp(swarg2)
pen[-1] = 0.
vs.divpen_shortwave = allocate(vs, ('zt', ))
vs.divpen_shortwave[1:] = (pen[1:] - pen[:-1]) / vs.dzt[1:]
vs.divpen_shortwave[0] = pen[0] / vs.dzt[0]
