Source code for veros.setups.north_atlantic.north_atlantic

#!/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, )