#!/usr/bin/env python
import os
import h5netcdf
from PIL import Image
import scipy.spatial
import scipy.ndimage
from veros import VerosSetup, veros_routine, veros_kernel, KernelOutput
from veros.variables import Variable
from veros.core.operators import numpy as npx, update, at
import veros.tools
BASE_PATH = os.path.dirname(os.path.realpath(__file__))
DATA_FILES = veros.tools.get_assets("north_atlantic", os.path.join(BASE_PATH, "assets.json"))
TOPO_MASK_FILE = os.path.join(BASE_PATH, "topo_mask.png")
[docs]class NorthAtlanticSetup(VerosSetup):
"""A regional model of the North Atlantic, inspired by `Smith et al., 2000`_.
Forcing and initial conditions are taken from the FLAME PyOM2 setup. Bathymetry
data from ETOPO1 (resolution of 1 arcmin).
Boundary forcings are implemented via sponge layers in the Greenland Sea, by the
Strait of Gibraltar, and in the South Atlantic. This setup runs with arbitrary resolution;
upon changing the number of grid cells, all forcing files will be interpolated to
the new grid. Default resolution corresponds roughly to :math:`0.5 \\times 0.25` degrees.
.. _Smith et al., 2000:
http://journals.ametsoc.org/doi/10.1175/1520-0485%282000%29030%3C1532%3ANSOTNA%3E2.0.CO%3B2
"""
x_boundary = 17.2
y_boundary = 70.0
max_depth = 5800.0
@veros_routine
def set_parameter(self, state):
settings = state.settings
settings.identifier = "na"
settings.nx, settings.ny, settings.nz = 250, 350, 50
settings.x_origin = -98.0
settings.y_origin = -18.0
settings.dt_mom = 3600.0 / 2.0
settings.dt_tracer = 3600.0 / 2.0
settings.runlen = 86400 * 365.0 * 10.0
settings.coord_degree = True
settings.enable_neutral_diffusion = True
settings.enable_skew_diffusion = True
settings.K_iso_0 = 1000.0
settings.K_iso_steep = 200.0
settings.iso_dslope = 1.0 / 1000.0
settings.iso_slopec = 4.0 / 1000.0
settings.enable_hor_friction = True
settings.A_h = 1e3
settings.enable_hor_friction_cos_scaling = True
settings.hor_friction_cosPower = 1
settings.enable_tempsalt_sources = True
settings.enable_implicit_vert_friction = True
settings.enable_tke = True
settings.c_k = 0.1
settings.c_eps = 0.7
settings.alpha_tke = 30.0
settings.mxl_min = 1e-8
settings.tke_mxl_choice = 2
settings.kappaM_min = 2e-4
settings.kappaH_min = 2e-5
settings.enable_kappaH_profile = True
settings.K_gm_0 = 1000.0
settings.enable_eke = False
settings.enable_idemix = False
settings.eq_of_state_type = 5
state.dimensions["nmonths"] = 12
state.var_meta.update(
{
"sss_clim": Variable("sss_clim", ("xt", "yt", "nmonths"), "g/kg", "Monthly sea surface salinity"),
"sst_clim": Variable("sst_clim", ("xt", "yt", "nmonths"), "deg C", "Monthly sea surface temperature"),
"sss_rest": Variable(
"sss_rest", ("xt", "yt", "nmonths"), "g/kg", "Monthly sea surface salinity restoring"
),
"sst_rest": Variable(
"sst_rest", ("xt", "yt", "nmonths"), "deg C", "Monthly sea surface temperature restoring"
),
"t_star": Variable(
"t_star", ("xt", "yt", "zt", "nmonths"), "deg C", "Temperature sponge layer forcing"
),
"s_star": Variable("s_star", ("xt", "yt", "zt", "nmonths"), "g/kg", "Salinity sponge layer forcing"),
"rest_tscl": Variable("rest_tscl", ("xt", "yt", "zt"), "1/s", "Forcing restoration time scale"),
"taux": Variable("taux", ("xt", "yt", "nmonths"), "N/s^2", "Monthly zonal wind stress"),
"tauy": Variable("tauy", ("xt", "yt", "nmonths"), "N/s^2", "Monthly meridional wind stress"),
}
)
@veros_routine
def set_grid(self, state):
vs = state.variables
settings = state.settings
vs.dxt = update(vs.dxt, at[2:-2], (self.x_boundary - settings.x_origin) / settings.nx)
vs.dyt = update(vs.dyt, at[2:-2], (self.y_boundary - settings.y_origin) / settings.ny)
vs.dzt = veros.tools.get_vinokur_grid_steps(settings.nz, self.max_depth, 10.0, refine_towards="lower")
@veros_routine
def set_coriolis(self, state):
vs = state.variables
settings = state.settings
vs.coriolis_t = update(
vs.coriolis_t, at[...], 2 * settings.omega * npx.sin(vs.yt[npx.newaxis, :] / 180.0 * settings.pi)
)
@veros_routine(dist_safe=False, local_variables=["kbot", "xt", "yt", "zt"])
def set_topography(self, state):
vs = state.variables
settings = state.settings
with h5netcdf.File(DATA_FILES["topography"], "r") as topo_file:
topo_x, topo_y, topo_bottom_depth = (self._get_data(topo_file, k) for k in ("x", "y", "z"))
topo_mask = npx.flipud(npx.asarray(Image.open(TOPO_MASK_FILE))).T
topo_bottom_depth = npx.where(topo_mask, 0, topo_bottom_depth)
topo_bottom_depth = scipy.ndimage.gaussian_filter(
topo_bottom_depth, sigma=(len(topo_x) / settings.nx, len(topo_y) / settings.ny)
)
interp_coords = npx.meshgrid(vs.xt[2:-2], vs.yt[2:-2], indexing="ij")
interp_coords = npx.rollaxis(npx.asarray(interp_coords), 0, 3)
z_interp = scipy.interpolate.interpn(
(topo_x, topo_y), topo_bottom_depth, interp_coords, method="nearest", bounds_error=False, fill_value=0
)
vs.kbot = update(
vs.kbot,
at[2:-2, 2:-2],
npx.where(
z_interp < 0.0,
1 + npx.argmin(npx.abs(z_interp[:, :, npx.newaxis] - vs.zt[npx.newaxis, npx.newaxis, :]), axis=2),
0,
),
)
vs.kbot = npx.where(vs.kbot < settings.nz, vs.kbot, 0)
def _get_data(self, f, var):
"""Retrieve variable from h5netcdf file"""
var_obj = f.variables[var]
return npx.array(var_obj).T
@veros_routine(
dist_safe=False,
local_variables=[
"tau",
"xt",
"yt",
"zt",
"temp",
"maskT",
"salt",
"taux",
"tauy",
"sst_clim",
"sss_clim",
"sst_rest",
"sss_rest",
"t_star",
"s_star",
"rest_tscl",
],
)
def set_initial_conditions(self, state):
vs = state.variables
with h5netcdf.File(DATA_FILES["forcing"], "r") as forcing_file:
t_hor = (vs.xt[2:-2], vs.yt[2:-2])
t_grid = (vs.xt[2:-2], vs.yt[2:-2], vs.zt)
forc_coords = [self._get_data(forcing_file, k) for k in ("xt", "yt", "zt")]
forc_coords[0] = forc_coords[0] - 360
forc_coords[2] = -0.01 * forc_coords[2][::-1]
temp_raw = self._get_data(forcing_file, "temp_ic")[..., ::-1]
temp = veros.tools.interpolate(forc_coords, temp_raw, t_grid, missing_value=-1e20)
vs.temp = update(vs.temp, at[2:-2, 2:-2, :, vs.tau], vs.maskT[2:-2, 2:-2, :] * temp)
salt_raw = self._get_data(forcing_file, "salt_ic")[..., ::-1]
salt = 35.0 + 1000 * veros.tools.interpolate(forc_coords, salt_raw, t_grid, missing_value=-1e20)
vs.salt = update(vs.salt, at[2:-2, 2:-2, :, vs.tau], vs.maskT[2:-2, 2:-2, :] * salt)
forc_u_coords_hor = [self._get_data(forcing_file, k) for k in ("xu", "yu")]
forc_u_coords_hor[0] = forc_u_coords_hor[0] - 360
taux = self._get_data(forcing_file, "taux")
tauy = self._get_data(forcing_file, "tauy")
for k in range(12):
vs.taux = update(
vs.taux,
at[2:-2, 2:-2, k],
(veros.tools.interpolate(forc_u_coords_hor, taux[..., k], t_hor, missing_value=-1e20) / 10.0),
)
vs.tauy = update(
vs.tauy,
at[2:-2, 2:-2, k],
(veros.tools.interpolate(forc_u_coords_hor, tauy[..., k], t_hor, missing_value=-1e20) / 10.0),
)
# heat flux and salinity restoring
sst_clim, sss_clim, sst_rest, sss_rest = [
forcing_file.variables[k][...].T for k in ("sst_clim", "sss_clim", "sst_rest", "sss_rest")
]
for k in range(12):
vs.sst_clim = update(
vs.sst_clim,
at[2:-2, 2:-2, k],
veros.tools.interpolate(forc_coords[:-1], sst_clim[..., k], t_hor, missing_value=-1e20),
)
vs.sss_clim = update(
vs.sss_clim,
at[2:-2, 2:-2, k],
(veros.tools.interpolate(forc_coords[:-1], sss_clim[..., k], t_hor, missing_value=-1e20) * 1000 + 35),
)
vs.sst_rest = update(
vs.sst_rest,
at[2:-2, 2:-2, k],
(veros.tools.interpolate(forc_coords[:-1], sst_rest[..., k], t_hor, missing_value=-1e20) * 41868.0),
)
vs.sss_rest = update(
vs.sss_rest,
at[2:-2, 2:-2, k],
(veros.tools.interpolate(forc_coords[:-1], sss_rest[..., k], t_hor, missing_value=-1e20) / 100.0),
)
with h5netcdf.File(DATA_FILES["restoring"], "r") as restoring_file:
rest_coords = [self._get_data(restoring_file, k) for k in ("xt", "yt", "zt")]
rest_coords[0] = rest_coords[0] - 360
# sponge layers
vs.rest_tscl = update(
vs.rest_tscl,
at[2:-2, 2:-2, :],
veros.tools.interpolate(rest_coords, self._get_data(restoring_file, "tscl")[..., 0], t_grid),
)
t_star = self._get_data(restoring_file, "t_star")
s_star = self._get_data(restoring_file, "s_star")
for k in range(12):
vs.t_star = update(
vs.t_star,
at[2:-2, 2:-2, :, k],
veros.tools.interpolate(rest_coords, t_star[..., k], t_grid, missing_value=0.0),
)
vs.s_star = update(
vs.s_star,
at[2:-2, 2:-2, :, k],
veros.tools.interpolate(rest_coords, s_star[..., k], t_grid, missing_value=0.0),
)
@veros_routine
def set_forcing(self, state):
vs = state.variables
vs.update(set_forcing_kernel(state))
@veros_routine
def set_diagnostics(self, state):
diagnostics = state.diagnostics
settings = state.settings
diagnostics["snapshot"].output_frequency = 3600.0 * 24 * 10
diagnostics["averages"].output_frequency = 3600.0 * 24 * 10
diagnostics["averages"].sampling_frequency = settings.dt_tracer
diagnostics["averages"].output_variables = [
"temp",
"salt",
"u",
"v",
"w",
"surface_taux",
"surface_tauy",
"psi",
]
diagnostics["cfl_monitor"].output_frequency = settings.dt_tracer * 10
@veros_routine
def after_timestep(self, state):
pass
@veros_kernel
def set_forcing_kernel(state):
vs = state.variables
settings = state.settings
year_in_seconds = 360 * 86400.0
(n1, f1), (n2, f2) = veros.tools.get_periodic_interval(vs.time, year_in_seconds, year_in_seconds / 12.0, 12)
vs.surface_taux = f1 * vs.taux[:, :, n1] + f2 * vs.taux[:, :, n2]
vs.surface_tauy = f1 * vs.tauy[:, :, n1] + f2 * vs.tauy[:, :, n2]
if settings.enable_tke:
vs.forc_tke_surface = update(
vs.forc_tke_surface,
at[1:-1, 1:-1],
npx.sqrt(
(0.5 * (vs.surface_taux[1:-1, 1:-1] + vs.surface_taux[:-2, 1:-1]) / settings.rho_0) ** 2
+ (0.5 * (vs.surface_tauy[1:-1, 1:-1] + vs.surface_tauy[1:-1, :-2]) / settings.rho_0) ** 2
)
** 1.5,
)
cp_0 = 3991.86795711963
vs.forc_temp_surface = (
(f1 * vs.sst_rest[:, :, n1] + f2 * vs.sst_rest[:, :, n2])
* (f1 * vs.sst_clim[:, :, n1] + f2 * vs.sst_clim[:, :, n2] - vs.temp[:, :, -1, vs.tau])
* vs.maskT[:, :, -1]
/ cp_0
/ settings.rho_0
)
vs.forc_salt_surface = (
(f1 * vs.sss_rest[:, :, n1] + f2 * vs.sss_rest[:, :, n2])
* (f1 * vs.sss_clim[:, :, n1] + f2 * vs.sss_clim[:, :, n2] - vs.salt[:, :, -1, vs.tau])
* vs.maskT[:, :, -1]
)
ice_mask = (vs.temp[:, :, -1, vs.tau] * vs.maskT[:, :, -1] <= -1.8) & (vs.forc_temp_surface <= 0.0)
vs.forc_temp_surface = npx.where(ice_mask, 0.0, vs.forc_temp_surface)
vs.forc_salt_surface = npx.where(ice_mask, 0.0, vs.forc_salt_surface)
if settings.enable_tempsalt_sources:
vs.temp_source = (
vs.maskT
* vs.rest_tscl
* (f1 * vs.t_star[:, :, :, n1] + f2 * vs.t_star[:, :, :, n2] - vs.temp[:, :, :, vs.tau])
)
vs.salt_source = (
vs.maskT
* vs.rest_tscl
* (f1 * vs.s_star[:, :, :, n1] + f2 * vs.s_star[:, :, :, n2] - vs.salt[:, :, :, vs.tau])
)
return KernelOutput(
surface_taux=vs.surface_taux,
surface_tauy=vs.surface_tauy,
temp_source=vs.temp_source,
salt_source=vs.salt_source,
forc_tke_surface=vs.forc_tke_surface,
forc_temp_surface=vs.forc_temp_surface,
forc_salt_surface=vs.forc_salt_surface,
)