included momentum restoring in momentum module

side_tools.py

+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Tue Mar 24 11:14:11 2020
+
+@author: veitlueschow
+"""
+
+# clean director for new run
+
+import os
+import sys
+
+def clean():
+    directory = "./"
+    test = os.listdir( directory )
+    print(test)
+
+    for item in test:
+        if item.endswith(".nc") or item.endswith('.h5'):
+            os.remove( os.path.join( directory, item ) )
+            print('deleted ', item)
\ No newline at end of file

veros/__init__.py

@@ -6,6 +6,7 @@ import types
 # black magic: ensure lazy imports for public API by overriding module.__class__
 
 class _PublicAPI(types.ModuleType):
+    
     @property
     def __version__(self):
         from veros._version import get_versions

veros/core/momentum.py

@@ -69,6 +69,12 @@ def momentum(vs):
     momentum_advection(vs)
     vs.du[:, :, :, vs.tau] += vs.du_adv
     vs.dv[:, :, :, vs.tau] += vs.dv_adv
+    
+    """
+    add momentum restoring that is computed in set_forcing
+    """
+    vs.du[:, :, :, vs.tau] += vs.du_rest
+    vs.dv[:, :, :, vs.tau] += vs.dv_rest
 
     with vs.timers['friction']:
         """

veros/setup/acc/acc.py

@@ -134,7 +134,7 @@ class ACCSetup(VerosSetup):
     def set_diagnostics(self, vs):
         vs.diagnostics['snapshot'].output_frequency = 86400 * 10
         vs.diagnostics['averages'].output_variables = (
-            'salt', 'temp', 'u', 'v', 'w', 'psi', 'surface_taux', 'surface_tauy'
+            'salt', 'temp', 'u', 'v', 'w', 'psi', 'surface_taux', 'surface_tauy','rho'
         )
         vs.diagnostics['averages'].output_frequency = 365 * 86400.
         vs.diagnostics['averages'].sampling_frequency = vs.dt_tracer * 10

veros/setup/zonjet/__init__.py

+from veros.setup.zonjet.zonjet import zonjet

veros/setup/zonjet/notes

+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Tue Mar 24 15:47:32 2020
+
+@author: veitlueschow
+"""
+NOTES
+# Add the restoring of flow at the boundary here. example: 
+#function restoring!(m::ModelSetup, bc)
+  # restore b to bi
+  #bc[:] = m.bic  .+ (bc .-m.bic) .*exp.(-m.cc .*m.dt)
+#end
\ No newline at end of file

veros/setup/zonjet/plot.py

+import xarray as xr
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+
+#xr.open_dataset('zonjet.snapshot.nc').u.isel(x,Time=20).plot()
+
+#plt.interactive("False")
+#plt.show()
+
+class section:
+    def __init__(self, filename, varname):
+        self.filename = filename
+        self.varname = varname
+        try:
+            self.var = self.load()[varname]
+        except KeyError:
+            print('Variable does not exist, choose new from:')
+            print(self.load().data_vars)
+            
+        self.data = 'something'
+        
+    def load(self):
+        return xr.open_dataset(self.filename)
+    
+    
+    def select(self, timestep, axis, position):
+        try: 
+            self.var[axis]
+            if axis == 'xu':
+                self.data = self.var.isel(xu=position).isel(Time=timestep)
+            elif axis == 'yt':
+                self.data = self.var.isel(yt=position).isel(Time=timestep)
+            elif axis == 'zt':
+                self.data = self.var.isel(zt=position).isel(Time=timestep)
+            elif axis == 'xt':
+                self.data = self.var.isel(xt=position).isel(Time=timestep)
+        except KeyError:
+            print('Axis does not exist, choose new from: \n', self.var.coords)
+            
+
+            
+        
+        return self.data
+        
+    def figure(self):
+        pass
+    
+    def fileinfo(self):
+        print('Info: ', self.load().data_vars)
+
+
+try:
+    del p
+except NameError:
+    pass
+
+p = section('zonjet.snapshot.nc','u')
+
+p.fileinfo()
+
+p.select(20,'zt',9).plot()
+#a.plot(vmax=0.5, vmin=-0.5)

veros/setup/zonjet/test.py

+import abc
+
+#from loguru import logger
+import veros
+import matplotlib.pyplot as plt
+
+print('done')
\ No newline at end of file

veros/setup/zonjet/zonjet.py

+#!/usr/bin/env python
+
+from veros import VerosSetup, veros_method
+from veros.tools import cli
+from veros.variables import allocate
+from veros.distributed import global_min, global_max
+import numpy as np
+
+from veros.core import density
+
+class ZONJETSetup(VerosSetup):
+    """A model using spherical coordinates with a partially closed domain representing the Atlantic and ACC.
+
+    Wind forcing over the channel part and buoyancy relaxation drive a large-scale meridional overturning circulation.
+
+    This setup demonstrates:
+     - setting up an idealized geometry
+     - updating surface forcings
+     - basic usage of diagnostics
+
+    `Adapted from pyOM2 <https://wiki.cen.uni-hamburg.de/ifm/TO/pyOM2/ACC%202>`_.
+    """
+    @veros_method
+    def set_parameter(self, vs):
+        vs.identifier = 'zonjet'
+
+        vs.nx, vs.ny, vs.nz = 15, 30, 10
+        vs.dt_mom = 4800
+        vs.dt_tracer = 86400 / 2.
+        vs.runlen = 86400 * 365
+        vs.runlen = 86400 * 51
+        #vs.runlen = 100
+
+        vs.coord_degree = True
+        vs.enable_cyclic_x = True
+
+        vs.congr_epsilon = 1e-12
+        vs.congr_max_iterations = 5000
+
+        vs.enable_neutral_diffusion = True
+        vs.K_iso_0 = 1000.0
+        vs.K_iso_steep = 500.0
+        vs.iso_dslope = 0.005
+        vs.iso_slopec = 0.01
+        vs.enable_skew_diffusion = True
+
+        vs.enable_hor_friction = True
+        vs.A_h = (2 * vs.degtom)**3 * 2e-11
+        vs.enable_hor_friction_cos_scaling = True
+        vs.hor_friction_cosPower = 1
+
+        vs.enable_bottom_friction = False
+        vs.r_bot = 1e-5
+
+        vs.enable_implicit_vert_friction = True
+
+        #vs.enable_tke = True
+        vs.c_k = 0.1
+        vs.c_eps = 0.7
+        vs.alpha_tke = 30.0
+        vs.mxl_min = 1e-8
+        vs.tke_mxl_choice = 2
+        vs.kappaM_min = 2e-4
+        vs.kappaH_min = 2e-5
+        #vs.enable_kappaH_profile = True
+        # vs.enable_tke_superbee_advection = True
+
+        vs.K_gm_0 = 1000.0
+       # vs.enable_eke = True
+        vs.eke_k_max = 1e4
+        vs.eke_c_k = 0.4
+        vs.eke_c_eps = 0.5
+        vs.eke_cross = 2.
+        vs.eke_crhin = 1.0
+        vs.eke_lmin = 100.0
+        vs.enable_eke_superbee_advection = True
+        vs.enable_eke_isopycnal_diffusion = True
+
+        vs.enable_idemix = False
+
+        vs.eq_of_state_type = 3
+        
+        vs.force_overwrite = 'True' # added for easier configuring the setting
+
+    @veros_method
+    def set_grid(self, vs):
+        #ddz = np.array([50., 100., 150., 200., 250., 300., 350., 400.,
+        #                450., 500., 550., 600., 650., 700., 750.])
+        vs.dxt[...] = 2.0
+        vs.dyt[...] = 2.0
+        vs.x_origin = 0.0
+        vs.y_origin = -40.0
+        #vs.dzt[...] = ddz[::-1] / 2.5
+        vs.dzt[...] = 20
+
+    @veros_method
+    def set_coriolis(self, vs):
+        vs.coriolis_t[:, :] = 2 * vs.omega * np.sin(vs.yt[None, :] / 180. * vs.pi)
+
+    @veros_method
+    def set_topography(self, vs):
+        x, y = np.meshgrid(vs.xt, vs.yt, indexing='ij')
+        vs.kbot[...] = np.logical_or(x > 100., y < 20).astype(np.int)
+
+    @veros_method
+    def set_initial_conditions(self, vs):
+        # initial conditions
+        
+        # design initial profile for u that is also used for restoring
+        zscale = 50
+        yvariance = 5e1
+        ycenter = vs.yt[0] + ( vs.yt[-1] - vs.yt[0] ) / 2
+        yy,zz = np.meshgrid(vs.yt,vs.zt,indexing='ij')
+        uz = np.exp(zz * (1/zscale) ) # exponential decay in z
+        uy = np.exp(- ( (yy-(ycenter) )**2 / yvariance ) ) # gaussian in y
+        uu = uz * uy
+        vs.u_ini = allocate(vs, ('xu','yt','zt'))
+        vs.u_ini = 2 * np.tile(uu[np.newaxis,:,:],[vs.xu.shape[0] , 1 , 1])
+
+        vs.u[...,0:2] = np.tile(vs.u_ini[...,np.newaxis],[1,1,1,2])
+        
+        #density
+        vs.temp[:, :, :, 0:2] = ((1 - vs.zt[None, None, :] / vs.zw[0]) * 10 * vs.maskT)[..., None]
+        vs.salt[:, :, :, 0:2] = 35.0 * vs.maskT[..., None]
+        
+        rho_background = density.get_rho(vs, vs.salt[...,0], vs.temp[...,0], np.abs(vs.zt)) * vs.maskT
+        
+        dzu = np.gradient(vs.u_ini,20)[2] * np.tile(vs.coriolis_t[...,np.newaxis],[1,1,vs.nz]) * (1025-1024) / 10
+        # cumsum to get density field
+        
+        #print(vs.yt[-1]-vs.yt[0] * )
+        sumi = np.cumsum(dzu * 2 * 111000 ,axis=1)
+        tmp = rho_background[:,2,:].copy()
+        
+        vs.rho_ini = allocate(vs, ('xt', 'yt', 'zt'))
+        
+        vs.rho_ini = np.tile(tmp[:,np.newaxis,:],[1,vs.ny+4,1]) + sumi
+        
+        vs.rho[...,0] = vs.rho_ini[...]
+        vs.rho[...,1] = vs.rho_ini[...]
+        vs.rho[...,2] = vs.rho_ini[...]
+        
+        #import matplotlib.pyplot as plt
+        #plt.contourf(yy,zz,vs.rho_ini[5,:,:],100)
+        #plt.colorbar()
+        #raise SystemExit(0)
+        
+        
+        # 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',))
+        vs.surface_taux[:, :] = 0
+
+        # surface heatflux forcing
+        vs._t_star = allocate(vs, ('yt',), fill=15)
+        vs._t_star = 0
+        vs._t_rest = 0
+
+        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]
+
+    @veros_method
+    def set_forcing(self, vs):
+        #restoring timescale only 
+        xtmp = np.zeros(vs.xu.shape[0])
+        xtmp[3:-2] = vs.xu[3:-2]
+        
+        t_tmp = 1 / (15 * vs.dt_mom) * ( 1-np.tanh((xtmp[-3]-xtmp)/10) )
+        
+        t_restoring = np.tile((t_tmp[...,np.newaxis,np.newaxis]),[1,vs.ny+4,vs.nz])
+        
+        
+        t_restoring = 1 / (vs.dt_mom * 15)
+        
+        t_restoring_tr = vs.dt_tracer * 1e-6
+        
+        #import matplotlib.pyplot as plt
+        #plt.pcolormesh(t_restoring[:,:,0]);plt.colorbar()
+        #raise SystemExit(0)
+        
+        vs.dv_rest = 0.
+        #raise SystemExit(0)
+        #print(t_restoring.shape); raise SystemExit(0)
+        
+        
+        
+        vs.du_rest[...] = vs.maskU * t_restoring * (vs.u_ini - vs.u[...,vs.taum1])
+        #vs.du[...,vs.taup1] = vs.du[...,vs.taup1] 
+        #+ vs.maskU * t_restoring * (vs.u_ini - vs.u[...,vs.tau])
+        
+        
+        #vs.rho[...,vs.tau] = vs.rho[...,vs.tau] 
+        #+ vs.maskT * t_restoring_tr * (vs.rho_ini - vs.rho[...,vs.taum1])
+        #vs.u[0:4,:,:,vs.taup1] = vs.u_restore[0:4,:,:]
+        #+ (vs.u[0:4,:,:,vs.taup1] - vs.u_restore[0:4,:,:]) * np.exp(-1e-10 * vs.dt_mom)
+        #vs.u[0:4,:,:,vs.taum1] = vs.u_restore[0:4,:,:]
+        #+ (vs.u[0:4,:,:,vs.taum1] - vs.u_restore[0:4,:,:]) * np.exp(-1e-10 * vs.dt_mom)
+        #vs.u[0:4,:,:,vs.tau] = vs.u_restore[0:4,:,:]
+        #+ (vs.u[0:4,:,:,vs.tau] - vs.u_restore[0:4,:,:]) * np.exp(-1e-10 * vs.dt_mom)
+        
+       # bc[:] = m.bic  .+ (bc .-m.bic) .*exp.(-m.cc .*m.dt)
+       # vs.u[-5:-3,15:18,8:12,vs.taup1] = 0.2
+        
+    # define velocity in restoring zone
+        
+        
+
+    @veros_method
+    def set_diagnostics(self, vs):
+        vs.diagnostics['snapshot'].output_frequency = 86400.
+        #vs.diagnostics['snapshot'].output_frequency = 1.        
+        vs.diagnostics['snapshot'].output_variables = (
+            'temp', 'psi', 'u', 'v', 'w', 'kbot','rho','salt'
+        )    
+        vs.diagnostics['averages'].output_variables = (
+            'salt', 'temp', 'u', 'v', 'w', 'psi', 'surface_taux', 'surface_tauy'
+        )
+        vs.diagnostics['averages'].output_frequency = 50 * 86400.
+        vs.diagnostics['averages'].sampling_frequency = vs.dt_tracer * 10
+        vs.diagnostics['overturning'].output_frequency = 365 * 86400. / 48.
+        vs.diagnostics['overturning'].sampling_frequency = vs.dt_tracer * 10
+        vs.diagnostics['tracer_monitor'].output_frequency = 365 * 86400. / 12.
+        vs.diagnostics['energy'].output_frequency = 365 * 86400. / 48
+        vs.diagnostics['energy'].sampling_frequency = vs.dt_tracer * 10
+        #vs.diagnostics['snapshot'].output_variables = (
+        #        'kbot')
+
+    def after_timestep(self, vs):
+        pass
+
+
+@cli
+def run(*args, **kwargs):
+    simulation = ZONJETSetup(*args, **kwargs)              
+    simulation.setup()
+    simulation.run()
+
+
+if __name__ == '__main__':
+    # Execute only if run as a script
+    run()
+    
+#xr.open_dataset('zonjet.snapshot.nc',autoclose=True).temp.isel(zt=10,Time=0).plot()
+    
+    

veros/variables.py

@@ -382,6 +382,14 @@ MAIN_VARIABLES = OrderedDict([
         'Change of v by advection', V_GRID, 'm/s^2',
         'Change of v due to advection'
     )),
+    ('du_rest', Variable(
+        'Change of u by restoring', U_GRID, 'm/s^2',
+        'Change of u due to restoring'
+    )),
+    ('dv_rest', Variable(
+        'Change of v by restoring', V_GRID, 'm/s^2',
+        'Change of v due to restoring'
+    )),
     ('p_hydro', Variable(
         'Hydrostatic pressure', T_GRID, 'm^2/s^2', 'Hydrostatic pressure', output=True
     )),
@@ -514,6 +522,10 @@ MAIN_VARIABLES = OrderedDict([
     )),
     ('w_wgrid', Variable(
         'W on W grid', W_GRID, 'm/s', 'Vertical velocity interpolated to W grid points'
+    )),
+    ('u_restore', Variable(
+        'Zonal velocity restore field', U_GRID, 'm/s', 'Zonal velocity restore field',
+        output=True, write_to_restart=True
     ))
 ])