slim.py

"""Second-generation Louvain-la-Neuve Ice-ocean model - User interface

    Contact: jonathan.lambrechts@uclouvain.be
    webpage: www.climate.be/slim

    SLIM solves the governing equations on adaptative unstructured meshes using the discontinuous Galerkin finite element method.

    To simulate the model, run the script my_run_script.py:
        - On a single core: rundgpy my_run_script.py
        - On multiple cores: mpirun -np N rundgpy my_run_script.py (replace N by the number of cores)

    Equations: 
        - 2D:
            * Shallow water equations (see :class:`~ShallowWater2d`)
            * Tracer equation (see :class:`~slim.ShallowWaterTracer2d`)
            * Particle Transport (see class TODO)
        - 3D:
            * Shallow water equations (see class TODO)
            * Tracer equation (see class TODO)
            
    Domain: 
        The domain must be represented by a mesh using the .msh format (generated by gmsh or unref)
        Various forms of physical domain can be used:
        
        - 2D domain:
            * domain which are projected by mean of cartographic transformation (such as utm projection, stereographic projection)
            * sphere. In this case, the local poblems are written in a local orthonormal curvilinear system. All the local problems are assembled in the global discrete algebraic system (for further details, see Comblen et al., 2009, A finite element method for solving the shallow water equations on the sphere).
        - 3D domain:
            * domain which are projected by mean of cartographic transformation (such as utm projection, stereographic projection)

    Data: 
        Various file format can be used: 
            * netcdf file (data.nc) 
                Those files have a special structure: data are given to each nodes at one or several time(s). The file 'slimPre' helps generating those files.
            * msh file (data_COMP_0.msh)
                 
            * idx file (data.idx or data_COMP_0.idx) 
                 
        Plane domain: In this case, data such as initial condition, wind stress, forcing ... must be expressed in the local basis (x,y) which is defined by the projection.
        
        Sphere: In this case, data such as initial condition, wind stress, forcing ... must be expressed in the global orthonormal basis (x,y,z) with z pointing upwards.
    
    partition_id:
        This function returns the number of the current partition as a string
    
    partition_nb:
        This function returns the number of partitions used as a string
"""

from dgpy.scripts import slim_private
import dgpy
 
class Domain:
    """Create the numerical domain"""
    
    def __init__(self, mesh, bathy, g=None, density=1025, solve_on_sphere=False): 
        """keyword arguments:
        
        * mesh            
            path to the mesh file (.msh format). The partitioned mesh will automatically be loaded in multiprocessing
        * bathy           
            bathymetric file [in meters, positive] (.msh, .idx or .nc format)
        * g               
            map of the mean gravitational acceleration (.msh, .idx or .nc format) (default: 9.81 m/s^2)
        * density         
            density of the liquid (default: 1025 kg/m^3)
        * solve_on_sphere 
            boolean stating if you want to solve the shallow water equation on a sphere or on any 2d manifold instead of solving it on a plane (default: False)
        """
        self._solve_on_manifold = solve_on_sphere
        if (dgpy.Msg.GetCommSize() > 1) and (not slim_private._check_if_partionned(mesh)):
            dgpy.dgMeshPartition(mesh, dgpy.Msg.GetCommSize())
            self._mesh = dgpy.dgMesh(mesh[:-4] + '_' + str(dgpy.Msg.GetCommSize()) + '.msh')
        else:
            self._mesh = dgpy.dgMesh(mesh)
        #TODO add option to run in high order
        order = 1
        self._groups = dgpy.dgGroupCollection(self._mesh, order)
        if self._solve_on_manifold:
            self._space_transform = dgpy.dgSpaceTransformFlatTriangle(self._mesh); 
            self._mesh.setSpaceTransform(self._space_transform, self._groups)
        self._bath = dgpy.dgDofContainer(self._groups, 1)
        slim_private._load(self._bath, bathy)
        self._bath_PC = dgpy.functionPrecomputed(self._groups, 3)
        self._bath_PC.compute(self._bath.getFunction())
        self._bath_PC.compute(self._bath.getFunction(),dgpy.dataCacheMap.NODE_GROUP_MODE)
        self._bath_gradient = self._bath.getFunctionGradient()
        self._bath_gradient_PC = dgpy.functionPrecomputed(self._groups, 3)
        self._bath_gradient_PC.compute(self._bath.getFunctionGradient())
        self._bath_function = self._bath_PC
        self._bath_gradient_function = self._bath_gradient_PC
        if g is None:
            self._g = dgpy.functionConstant(9.81)
        else:
            self._g = slim_private._load_function(g,self._groups)
        self._density = density
 
class ShallowWater2dWD:
    """Create the shallow water equation with various options"""

    def __init__(self, domain, temporal_scheme, export_every_sub_time_step=False, initial_time=None, final_time=None, hlimt=1, hlimd=0.1, linDrag=0.1): 
        """keyword arguments:
        
        * domain                     
            object of class Domain
        * temporal_scheme 
            scheme for the temporal integration
            
            * "explicit" 
                explicit Euler scheme
            * "implicit" 
                implicit order 2 Runge-Kutta
            * "multirate"
                multirate scheme (see Seny et al. 2012, 2014)
                
                Warning: This temporal scheme is not yet implemented for tracers
        * export_every_sub_time_step 
            boolean stating if you want to export the hydrodynamics every sub time step for further use of offline equation (Default: False)
        * initial_time               
            initial time of the simulation (specific to each equation) (Default: 0)
            
            Format: 

            - float [in seconds]
            - date string in format "year-month-day hours:minutes:seconds" (e.g. "2000-02-15 10:05:00")
        * final_time
            final time of the simulation (default: 0)
            
            Format:
                
            - float [in seconds]
            - date string in format "year-month-day hours:minutes:seconds" (e.g. "2007-06-25 22:42:17")                              
        """     
        fmt = '%Y-%m-%d %H:%M:%S'   
        if initial_time is None:
            self._initial_time = 0.
        elif slim_private._is_string(initial_time):
            date = slim_private.datetime.datetime.strptime(initial_time, fmt)
            self._initial_time = slim_private.calendar.timegm([date.year, date.month, date.day, date.hour, date.minute, date.second])
        else:
            self._initial_time = float(initial_time)
        if final_time is None:
            self._final_time = 0.
        elif slim_private._is_string(final_time):
            date = slim_private.datetime.datetime.strptime(final_time, fmt)
            self._final_time = slim_private.calendar.timegm([date.year, date.month, date.day, date.hour, date.minute, date.second])
        else:
            self._final_time = float(final_time)
        self._domain = domain
        self._equation = dgpy.dgConservationLawShallowWater2dC(self._domain._bath_function, self._domain._bath_gradient_function, hlimt, hlimd, linDrag)
        self._wetting_drying = None
        self._solution = dgpy.dgDofContainer(self._domain._groups, self._equation.getNbFields())
        self._solution.setAll(0.0)
        self._solution.setFieldName(0, 'H')
        self._solution.setFieldName(1, 'Hu')
        self._solution.setFieldName(2, 'Hv')
        self._name = "sw2dC"
        if temporal_scheme in ["explicit","implicit","multirate","DIRK1","DIRK2","DIRK3"]:
            self._temporal_scheme = temporal_scheme
        else:
            dgpy.Msg.Fatal("Unknown temporal scheme "+temporal_scheme+" for shallow water equation !")
        self._export_every_sub_time_step = export_every_sub_time_step
        self._evaluator = None
        self._ref = []
        self._compute_mass = False
        self._export_on_structured_grid =False
  
    def set_initial_condition(self, sse, ux, uy, uz=None):
        """
        Initial conditions

        keyword arguments:
        
        * sse 
            sea surface elevation [in m] (.msh or .nc format) 
        * ux  
            velocity along the x-axis in the local basis or in the global basis (x,y,z) if you want to solve on a sphere [in m/s] (.idx, .msh or .nc format) 
        * uy  
            velocity along the y-axis in the local basis or in the global basis (x,y,z) if you want to solve on a sphere [in m/s] (.idx, .msh or .nc format) 
        * uz  
            velocity along the z-axis in the global basis (x,y,z) if you want to solve on a sphere [in m/s] (.idx, .msh or .nc format) (default: None)
        """
        _sse = _ux = _uy = _uz = None
        if sse is not None:
            _sse = slim_private._load_function(sse, self._domain._groups)
        if ux is not None:
            _ux = slim_private._load_function(ux, self._domain._groups)
        if uy is not None:
            _uy = slim_private._load_function(uy, self._domain._groups)
        if uz is not None:
            _uz = slim_private._load_function(uz, self._domain._groups)
        velocity_local = slim_private._check_vector(_ux,_uy,_uz,self._domain)
        def nonConservativeToConservative(cm, etaF, uvF, bathF) :
            H = cm.get(bathF)
            if etaF :
                H = H+cm.get(etaF)
            if uvF :
                return slim_private.np.hstack([H,H*cm.get(uvF)])
            else :
                return slim_private.np.hstack([H,H*0,H*0])
        self._solution.interpolate(dgpy.functorNumpy(lambda cm : nonConservativeToConservative(cm,_sse,velocity_local,self._domain._bath_function)))

    def set_coriolis(self, coriolis):
        """Add a coriolis term (2*Omega*cos(latitude)) in the shallow water equation 
        
        keyword argument: 
        
       * coriolis 
            netcdf or .msh file containing the coriolis term for the whole domain. 
        """
        self._coriolis = slim_private._load_function(coriolis, self._domain._groups)
        self._coriolis_PC = dgpy.functionPrecomputed(self._domain._groups, 3)
        self._coriolis_PC.compute(self._coriolis)
        self._equation.setCoriolisFactor(self._coriolis_PC)
    
        
    def set_boundary_coast(self, physical_tag):
        """Boundary condition: normal speed set to zero and tangential speed set to zero or free (no slip or free slip condition)
        
        keyword arguments:
        
        * physical_tag 
            tag of the part of the boundary on which this boundary condition is applied.
        """
        self._equation.addBoundaryCondition(physical_tag, self._equation.newBoundaryWall())

    def set_boundary_open(self, physical_tag, discharge=None, sse=None, ux=None, uy=None, uz=None, transport_flux=False, ramp_period=0):
        """Open boundary condition: one can impose the discharge or the sea surface elevation and/or the velocity (ux and uy have to be defined at the same time but they can be zero).
        
        keyword arguments:
        
        * physical_tag  
            tag of the part of the boundary on which this boundary condition is applied.
        * discharge    
            netcdf or .msh file containing the discharge which will be applied at the boundary [in m^3s^-1] (default: None)
        * sse            
            netcdf or .msh file containing the sea surface elevation which will be applied at the boundary [in m] (default: None, i.e. the sea surface elevation computed by the model is set to this boundary)
        * ux             
            netcdf or .msh file containing the velocity along the x axis of the local basis  or the velocity along the x axis of the global basis if you solve the equation on a sphere. This velocity will be applied at the boundary [in m/s] (default: None, i.e. the velocity computed by the model is set to this boundary)
        * uy             
            netcdf or .msh file containing the velocity along the y axis of the local basis  or the velocity along the y axis of the global basis if you solve the equation on a sphere. This velocity will be applied at the boundary [in m/s] (default: None, i.e. the velocity computed by the model is set to this boundary)
        * uz             
            netcdf or .msh file containing the velocity along the z axis of the global basis if you solve the equation on a sphere. This velocity will be applied at the boundary [in m/s] (default: None, i.e. the velocity computed by the model is set to this boundary)
        * transport_flux 
            set whether ux and uy (and uz) are velocities (False) or transport (True) (Default: False) 
        * ramp_period  
            period for the linear ramp before applying the nudging [in s] (Default: 0s, no ramp) 
        """
        #TODO add option for external gradient
        _uz_= _sse = _ux = _uy = _velocity_local = _discharge = _depth_integral = None        
        if (ux is None and uy is not None) or (ux is not None and uy is None):
            dgpy.Msg.Fatal("Both ux and uy should be given for boundary condition for tag: "+physical_tag)
        if ux is not None and discharge is not None:
            dgpy.Msg.Fatal("Impossible to impose the velocity and the discharge for boundary condition for tag: "+physical_tag)
        if sse is not None and discharge is not None:
            dgpy.Msg.Fatal("Impossible to impose the sea surface elevation and the discharge for boundary condition for tag: "+physical_tag)
            #TODO add option to prescribe etaExt
        if sse is None and ux is None and discharge is None:
            _external_value = dgpy.outsideValueFree(self._solution.getFunction(), self._domain._bath_function, self._domain._g)
        elif discharge is not None:
            _discharge = slim_private._load_function(discharge, self._domain._groups)
            _depth_integral = dgpy.fullMatrixDouble(1, 1)
            dgpy.dgFunctionIntegratorInterface(self._domain._groups, self._domain._bath_PC).compute(physical_tag, _depth_integral)
            _external_value = dgpy.dischargeToBoundaryVelocity(self._solution.getFunction(), _discharge, _depth_integral(0, 0), self._domain._bath_function)
            if (ramp_period > 0.):
                _external_value_ramp = dgpy.ramp3d(_external_value, dgpy.function.getTime(), self._initial_time, ramp_period)
            else:
                _external_value_ramp = _external_value
        elif ux is None:
            _sse = slim_private._load_function(sse, self._domain._groups)
            _external_value = dgpy.outsideValueEta(self._solution.getFunction(), _sse)
        else:
            _ux_= slim_private._load_function(ux, self._domain._groups)
            _uy_ = slim_private._load_function(uy, self._domain._groups)
            if uz is not None:
                _uz_ = slim_private._load_function(uz, self._domain._groups)
            _velocity_local = slim_private._check_vector(_ux_, _uy_, _uz_, self._domain)    
            _ux = slim_private._get_x(_velocity_local).functor
            _uy = slim_private._get_y(_velocity_local).functor
            if transport_flux:
                if sse is not None:
                    _sse = slim_private._load_function(sse, self._domain._groups)    
                    _external_value = dgpy.transportToBoundaryVelocity(_sse, _ux, _uy, self._domain._bath_function)
                else:
                    _external_value = dgpy.transportToBoundaryVelocity(self._solution.getFunction(), _ux, _uy, self._domain._bath_function)
            else:
                if sse is not None:
                    _sse = slim_private._load_function(sse, self._domain._groups)    
                _external_value = dgpy.mergeVector(_sse, _ux, _uy, self._solution.getFunction())
        if (ramp_period > 0.):
            _external_value_ramp = dgpy.ramp3d(_external_value, dgpy.function.getTime(), self._initial_time, ramp_period)
        else:
            _external_value_ramp = _external_value
        def nonConservativeToConservative(cm, solF, bathF) :
            sol = cm.get(solF)
            n = cm.get(dgpy.function.getNormals())
            solIn = cm.get(cm.solution())
            H = cm.get(bathF) + sol[:,[0]]
            solExt = slim_private.np.hstack([H,H*sol[:,[1,2]]])
            solExt[:,[1,2]] = solIn[:,[1,2]]- n[:,[0,1]] * (9.81/H)**0.5*(solIn[:,[0]]-solExt[:,[0]])/2
            return solExt
        conservative = dgpy.functorNumpy(lambda cm : nonConservativeToConservative(cm,_external_value_ramp, self._domain._bath_function))
        self._equation.addBoundaryCondition(physical_tag, self._equation.newOutsideValueBoundary(conservative))
        self._ref += [_sse, _ux, _uy, _velocity_local, _external_value, _external_value_ramp, ramp_period, _discharge, _depth_integral, conservative] 
 
    def export_value_at_point(self, output_name, x, y, z=None):
        """set the coordinates where a time serie of the solution is exported

        keyword arguments:
        
        * output_name 
            name of the output file
        * coord_x     
            vector containing the x coordinates where the solution is evaluated
        * coord_y     
            vector containing the y coordinates where the solution is evaluated
        * coord_z     
            vector containing the z coordinates where the solution is evaluated (Default: None, replaced by vector of 0)
        """
        if self._evaluator is None:
            self._evaluator = dgpy.dgFunctionEvaluator(self._domain._groups, self._solution.getFunction())
        self._evaluator_points_x = x
        self._evaluator_points_y = y
        if z is None:
            z = [0.]*len(x)
        self._evaluator_points_z = z
        self._evaluator_output_name = output_name
        f = open(output_name, 'w')
        sx = "x,"
        sy = "y,"
        sz = "z,"
        for i in range(len(x)):
            sx += str(x[i])+","
            sy += str(y[i])+","
            sz += str(z[i])+","
        f.write(sx[:-1]+'\n')
        f.write(sy[:-1]+'\n')
        f.write(sz[:-1]+'\n')
        f.close()
    
    def compute_mass(self,output_name):
        """compute the total mass of the domain
        
        keyword arguments:
        
        * output_name 
            name of the output file
        """
        self._output_mass = output_name
        self._compute_mass = True
        self._height = dgpy.functorNumpy(lambda cm : (cm.get(self._solution.getFunction())[:,[0]])+cm.get(self._domain._bath_function))
        self._integrator_mass = dgpy.dgFunctionIntegrator(self._domain._groups, self._height)
        self._integrator_mass.setSolution(self._solution.getFunction(), self._solution.getFunctionGradient())
        self._mass = dgpy.fullMatrixDouble(1, 1)
        self._integrator_mass.compute(self._mass,"")

class ShallowWater2d:
    """Create the shallow water equation with various options"""

    def __init__(self, domain, temporal_scheme, linear_equation=False, wetting_drying=None, export_every_sub_time_step=False, initial_time=None, final_time=None): 
        """keyword arguments:
        
        * domain                     
            object of class Domain
        * temporal_scheme 
            scheme for the temporal integration
            
            * "explicit" 
                explicit Euler scheme
            * "implicit" 
                implicit order 2 Runge-Kutta
            * "multirate"
                multirate scheme (see Seny et al. 2012, 2014)
                
                Warning: This temporal scheme is not yet implemented for tracers
        * linear_equation            
            boolean stating if you want tos solve the linear version of the  shallow water equation (Default: False)
        * wetting_drying             
            value for the wetting & drying algorithm (see Karna et al. 2011) (Default: None)
        * export_every_sub_time_step 
            boolean stating if you want to export the hydrodynamics every sub time step for further use of offline equation (Default: False)
        * initial_time               
            initial time of the simulation (specific to each equation) (Default: 0)
            
            Format: 

            - float [in seconds]
            - date string in format "year-month-day hours:minutes:seconds" (e.g. "2000-02-15 10:05:00")
        * final_time
            final time of the simulation (default: 0)
            
            Format:
                
            - float [in seconds]
            - date string in format "year-month-day hours:minutes:seconds" (e.g. "2007-06-25 22:42:17")                              
        """     
        fmt = '%Y-%m-%d %H:%M:%S'   
        if initial_time is None:
            self._initial_time = 0.
        elif slim_private._is_string(initial_time):
            date = slim_private.datetime.datetime.strptime(initial_time, fmt)
            self._initial_time = slim_private.calendar.timegm([date.year, date.month, date.day, date.hour, date.minute, date.second])
        else:
            self._initial_time = float(initial_time)
        if final_time is None:
            self._final_time = 0.
        elif slim_private._is_string(final_time):
            date = slim_private.datetime.datetime.strptime(final_time, fmt)
            self._final_time = slim_private.calendar.timegm([date.year, date.month, date.day, date.hour, date.minute, date.second])
        else:
            self._final_time = float(final_time)
        self._domain = domain
        self._equation = dgpy.dgConservationLawShallowWater2d()
        self._equation.setIsLinear(linear_equation)
        self._equation.setLaxFriedrichs(True)
        self._wetting_drying = wetting_drying
        self._equation.setBathymetry(self._domain._bath_function)
        self._equation.setBathymetryGradient(self._domain._bath_gradient_function)
        self._equation.setGravity(self._domain._g)
        self._coriolis = None
        self._solution = dgpy.dgDofContainer(self._domain._groups, self._equation.getNbFields())
        self._solution.setAll(0.0)
        self._solution.setFieldName(0, 'eta')
        self._solution.setFieldName(1, 'u')
        self._solution.setFieldName(2, 'v')
        if self._wetting_drying is not None:
            self._equation.setMovingBathWettingDrying(self._solution.getFunction(), self._solution.getFunctionGradient(), self._wetting_drying)
            self._domain._bath_function = self._equation.getBathymetry()
            self._domain._bath_gradient_function = self._equation.getBathymetryGradient()
        self._name = "sw2d"
        if temporal_scheme == "explicit" or temporal_scheme == "implicit" or temporal_scheme == "multirate":
            self._temporal_scheme = temporal_scheme
        else:
            dgpy.Msg.Fatal("Unknown temporal scheme "+temporal_scheme+" for shallow water equation !")
        self._export_every_sub_time_step = export_every_sub_time_step
        self._evaluator = None
        self._ref = []
        self._compute_mass = False
        self._export_on_structured_grid =False
  
    def set_initial_condition(self, sse, ux, uy, uz=None):
        """
        Initial conditions

        keyword arguments:
        
        * sse 
            sea surface elevation [in m] (.msh or .nc format) 
        * ux  
            velocity along the x-axis in the local basis or in the global basis (x,y,z) if you want to solve on a sphere [in m/s] (.idx, .msh or .nc format) 
        * uy  
            velocity along the y-axis in the local basis or in the global basis (x,y,z) if you want to solve on a sphere [in m/s] (.idx, .msh or .nc format) 
        * uz  
            velocity along the z-axis in the global basis (x,y,z) if you want to solve on a sphere [in m/s] (.idx, .msh or .nc format) (default: None)
        """
        _sse = _ux = _uy = _uz = _ux_ = _uy_ = None
        if sse is not None:
            _sse = slim_private._load_function(sse, self._domain._groups)
        if ux is not None:
            _ux = slim_private._load_function(ux, self._domain._groups)
        if uy is not None:
            _uy = slim_private._load_function(uy, self._domain._groups)
        if uz is not None:
            _uz = slim_private._load_function(uz, self._domain._groups)
        _velocity_local = slim_private._check_vector(_ux,_uy,_uz,self._domain)
        _ux_ = slim_private._get_x(_velocity_local).functor
        _uy_ = slim_private._get_y(_velocity_local).functor
        _f0 = dgpy.functionConstant([0.,0.,0.])
        _initial_condition = dgpy.mergeVector(_sse, _ux_, _uy_, _f0)
        self._solution.interpolate(_initial_condition)

    def set_viscosity(self, mode, smagorinsky_coefficient=0.1, constant_viscosity=1e-6):
        """Turbulent viscosity : Smagorinsky scheme or constant value
        
        keyword arguments:
        
        * mode       
            type of turbulent viscosity
            
            * "smagorinsky" 
                Smagorinsky scheme (see Smagorinsky 1962)
            * "constant"    
                constant value 
        * smagorinsky_coefficient 
            coefficient value for Smagorinsky scheme (default: 0.1)
        * constant_viscosity      
            constant value [in m^2/s] (default: 1e-6 m^2/s)
        """
        if mode == "smagorinsky":
            self._smagorinsky_coefficient = smagorinsky_coefficient
            self._nu = dgpy.smagorinsky2d(self._smagorinsky_coefficient, self._solution.getFunctionGradient(), dgpy.function.getElementVolume())
            self._equation.setDiffusivity(self._nu)
        elif mode == "constant":
            self._constant_viscosity = constant_viscosity
            self._nu = dgpy.functionConstant(self._constant_viscosity)
            self._equation.setDiffusivity(self._nu)
        else: 
            dgpy.Msg.Fatal("Viscosity scheme not implemented, use smagorinsky or a constant viscosity scheme")

    def set_dissipation(self, mode, coefficient=None):
        """Dissipation due to bottom friction: linear, basic quadratic scheme or Chezy-Manning scheme
        
        keyword arguments:
        
        * mode
            scheme used for the bottom friction
                                      
            * "linear"    
                linear dissipation tau_b/(rho*H) = gamma*u (default: gamma = 1e-6 s^-1)
            * "quadratic" 
                basic quadratic dissipation tau_b/(rho*H) = (C_d/H)*|u|*u (default: C_d = 2.5e-3)
            * "manning"   
                Chezy-Manning-Strickler formulation tau_b/(rho*H) = (n*n*g/(H^4/3))*|u|*u (default: n = 0.03 sm^(-1/3))
        * coefficient
            netcdf or .msh file containing the coefficient value (gamma, C_d or n depending on the mode) over the whole domain (default: None, the default value of the mode will be used)
        """
        if coefficient is not None:
            self._coefficient = slim_private._load_function(coefficient, self._domain._groups)
        if mode == "linear":
            if coefficient is None:
                self._coefficient = dgpy.functionConstant(1e-6)
            self._equation.setLinearDissipation(self._coefficient)
        elif mode == "quadratic":
            if coefficient is None:
                self._coefficient = dgpy.functionConstant(2.5e-3)
            self._cd = dgpy.quadraticBottomFriction(self._coefficient, self._domain._bath_function, self._solution.getFunction())
            self._equation.setQuadraticDissipation(self._cd)
        elif mode == "manning":
            if coefficient is None:
                self._coefficient = dgpy.functionConstant(0.03)
            self._cd = dgpy.manningBottomFriction(self._coefficient, self._domain._g, self._domain._bath_function, self._solution.getFunction())
            self._equation.setQuadraticDissipation(self._cd)
        else: 
            dgpy.Msg.Fatal("Dissipation scheme not implemented!")
    
    def set_coriolis(self, coriolis):
        """Add a coriolis term (2*Omega*cos(latitude)) in the shallow water equation 
        
        keyword argument: 
        
       * coriolis 
            netcdf or .msh file containing the coriolis term for the whole domain. 
        """
        self._coriolis = slim_private._load_function(coriolis, self._domain._groups)
        self._coriolis_PC = dgpy.functionPrecomputed(self._domain._groups, 3)
        self._coriolis_PC.compute(self._coriolis)
        self._equation.setCoriolisFactor(self._coriolis_PC)
    
    def set_wind_stress(self, mode, wind_x, wind_y, density_air=-1, wind_z= None):
        """Add a wind stress term in the shallow water equation.
        
        keyword arguments:
        
        * mode                  
            type of wind forcing
            
            * "speed"   
                wind speed given in m s^-1. wind_stress will be computed with Smith-Banke formula.
                  stress = C_D(speed)*density_air*speed
                Setting the air density is mandatory 
            * "stress" 
                wind stress given in kg m^-1 s^-1. (air density will not be used)
        * wind_x 
            netcdf or .msh file containing the wind term along the x-axis in the local basis 
            
            (except for equation on a sphere for which it has to be expressed in the global basis (x,y,z)) for the whole domain.  
        * wind_y 
            netcdf or .msh file containing the wind term along the y-axis in the local basis 
            
            (except for equation on a sphere for which it has to be expressed in the global basis (x,y,z)) for the whole domain.  
        * density_air  
            density of the ambiant air (only necessary for "speed" mode) (default: -1)
        * wind_z 
            netcdf or .msh file containing the wind term along the z-axis in the global basis (x,y,z) for the whole domain (only for equation on a sphere). (default: None)
        """
        if mode == 'speed':
            if density_air < 0:
                dgpy.Msg.Fatal("Setting wind stress. in 'speed' mode, density_air must be set")
            _wind_speed_x = slim_private._load_function(wind_x, self._domain._groups)
            _wind_speed_y = slim_private._load_function(wind_y, self._domain._groups)
            if wind_z is not None:
                _wind_speed_z = slim_private._load_function(wind_z, self._domain._groups)
            else:
                _wind_speed_z = None
            self._wind_speed_local = slim_private._check_vector(_wind_speed_x, _wind_speed_y, _wind_speed_z, self._domain)
            self._wind_stress_local = dgpy.windCubicStress(self._wind_speed_local, density_air)
            self._ref += [density_air]    
            self._equation.setWindStress(self._wind_stress_local)   
        elif mode == 'stress':
            _wind_stress_x = slim_private._load_function(wind_x, self._domain._groups)
            _wind_stress_y = slim_private._load_function(wind_y, self._domain._groups)
            if wind_z is not None:
                _wind_stress_z = slim_private._load_function(wind_z, self._domain._groups)
            else:
                _wind_stress_z = None
            self._wind_stress_local = slim_private._check_vector(_wind_stress_x, _wind_stress_y, _wind_stress_z, self._domain)
            self._equation.setWindStress(self._wind_stress_local)
        
    def set_forcing(self, forcing_x, forcing_y, forcing_z = None):
        """Add a momentum source term derived from a gravitational potential in the shallow water equation.
        
        keyword arguments: 
        
        * forcing_x 
            netcdf or .msh file containing the forcing term along the x-axis in the local basis 
            
            (except for equation on a sphere for which it has to be expressed in the global basis (x,y,z)) for the whole domain [in ms^-2].  
        * forcing_y 
            netcdf or .msh file containing the forcing term along the y-axis in the local basis 
            
            (except for equation on a sphere for which it has to be expressed in the global basis (x,y,z)) for the whole domain [in ms^-2].
        * forcing_z 
            netcdf or .msh file containing the forcing term along the z-axis in the local basis 
            
            (except for equation on a sphere for which it has to be expressed in the global basis (x,y,z)) for the whole domain [in ms^-2].
        """
        _forcing_x = slim_private._load_function(forcing_x, self._domain._groups)
        _forcing_y = slim_private._load_function(forcing_y, self._domain._groups)
        _forcing_z = None
        if forcing_z is not None:
            _forcing_z = slim_private._load_function(forcing_z, self._domain._groups)
        self._forcing_local = slim_private._check_vector(_forcing_x, _forcing_y, _forcing_z, self._domain)
        self._equation.setSource(self._forcing_local)

    def set_nudging(self, coefficient, external_sse, external_ux, external_uy, external_uz=None, ramp_period=0, transport_flux=False):
        """Add a nudging term as a source term in the shallow water equation.
        
        keyword arguments: 
        
        * coefficient    
            netcdf or .msh file containing the nudging coefficient for the whole domain [in s^-1]
        * external_sse   
            netcdf or .msh file containing the external sea surface elevation for the whole domain [in m]
        * external_ux    
            netcdf or .msh file containing the external velocity along the x-axis expressed in the local basis (except for equation on a sphere for which it has to be expressed in the global basis (x,y,z)) for the whole domain [in m/s]
        * external_uy    
            netcdf or .msh file containing the external velocity along the y-axis expressed in the local basis (except for equation on a sphere for which it has to be expressed in the global basis (x,y,z)) for the whole domain [in m/s]
        * external_uz    
            netcdf or .msh file containing the external velocity along the z-axis expressed in the global basis (x,y,z) for the whole domain [in m/s]
        * ramp_period    
            period for the linear ramp before applying the nudging [in s] (Default: 0s, no ramp) 
        * transport_flux 
            set whether ux and uy (and uz) are velocities (False) or transport (True) (Default: False) 
        """
        _coeff_nudging = slim_private._load_function(coefficient, self._domain._groups)
        _external_sse = slim_private._load_function(external_sse, self._domain._groups)
        _external_ux = slim_private._load_function(external_ux, self._domain._groups)
        _external_uy = slim_private._load_function(external_uy, self._domain._groups)
        _external_uz = None
        if external_uz is not None:
            _external_uz = slim_private._load_function(external_uz, self._domain._groups)
        _external_vel_local_ = slim_private._check_vector(_external_ux, _external_uy, _external_uz, self._domain)
        _external_vel_local_x = slim_private._get_x(_external_vel_local_).functor
        _external_vel_local_y = slim_private._get_y(_external_vel_local_).functor
        _external_sol_local = slim_private._merger_3d(_external_sse, _external_vel_local_x, _external_vel_local_y).functor
        self._ref += [_coeff_nudging, ramp_period, transport_flux]    
        self._equation.setNudgingCoefficient(_coeff_nudging)
        if (ramp_period > 0.):
            self._external_sol_local = dgpy.ramp3d(_external_sol_local, dgpy.function.getTime(), self._initial_time, ramp_period)
        else:
            self._external_sol_local = _external_sol_local
        self._equation.setNudgingVelocity(self._external_sol_local, transport_flux)

    def set_boundary_coast(self, physical_tag, slip=False):
        """Boundary condition: normal speed set to zero and tangential speed set to zero or free (no slip or free slip condition)
        
        keyword arguments:
        
        * physical_tag 
            tag of the part of the boundary on which this boundary condition is applied.
        * slip         
            flag wheter applying slip (True) or no slip (False) condition at the boundary (Default: False)
        """
        self._equation.addBoundaryCondition(physical_tag, self._equation.newBoundaryWall(slip))

    def set_boundary_open(self, physical_tag, discharge=None, sse=None, ux=None, uy=None, uz=None, transport_flux=False, ramp_period=0):
        """Open boundary condition: one can impose the discharge or the sea surface elevation and/or the velocity (ux and uy have to be defined at the same time but they can be zero).
        
        keyword arguments:
        
        * physical_tag  
            tag of the part of the boundary on which this boundary condition is applied.
        * discharge    
            netcdf or .msh file containing the discharge which will be applied at the boundary [in m^3s^-1] (default: None)
        * sse            
            netcdf or .msh file containing the sea surface elevation which will be applied at the boundary [in m] (default: None, i.e. the sea surface elevation computed by the model is set to this boundary)
        * ux             
            netcdf or .msh file containing the velocity along the x axis of the local basis  or the velocity along the x axis of the global basis if you solve the equation on a sphere. This velocity will be applied at the boundary [in m/s] (default: None, i.e. the velocity computed by the model is set to this boundary)
        * uy             
            netcdf or .msh file containing the velocity along the y axis of the local basis  or the velocity along the y axis of the global basis if you solve the equation on a sphere. This velocity will be applied at the boundary [in m/s] (default: None, i.e. the velocity computed by the model is set to this boundary)
        * uz             
            netcdf or .msh file containing the velocity along the z axis of the global basis if you solve the equation on a sphere. This velocity will be applied at the boundary [in m/s] (default: None, i.e. the velocity computed by the model is set to this boundary)
        * transport_flux 
            set whether ux and uy (and uz) are velocities (False) or transport (True) (Default: False) 
        * ramp_period  
            period for the linear ramp before applying the nudging [in s] (Default: 0s, no ramp) 
        """
        #TODO add option for external gradient
        _uz_= _sse = _ux = _uy = _velocity_local = _discharge = _depth_integral = None        
        if (ux is None and uy is not None) or (ux is not None and uy is None):
            dgpy.Msg.Fatal("Both ux and uy should be given for boundary condition for tag: "+physical_tag)
        if ux is not None and discharge is not None:
            dgpy.Msg.Fatal("Impossible to impose the velocity and the discharge for boundary condition for tag: "+physical_tag)
        if sse is not None and discharge is not None:
            dgpy.Msg.Fatal("Impossible to impose the sea surface elevation and the discharge for boundary condition for tag: "+physical_tag)
            #TODO add option to prescribe etaExt
        if sse is None and ux is None and discharge is None:
            _external_value = dgpy.outsideValueFree(self._solution.getFunction(), self._domain._bath_function, self._domain._g)
        elif discharge is not None:
            _discharge = slim_private._load_function(discharge, self._domain._groups)
            _depth_integral = dgpy.fullMatrixDouble(1, 1)
            dgpy.dgFunctionIntegratorInterface(self._domain._groups, self._domain._bath_PC).compute(physical_tag, _depth_integral)
            _external_value = dgpy.dischargeToBoundaryVelocity(self._solution.getFunction(), _discharge, _depth_integral(0, 0), self._domain._bath_function)
            if (ramp_period > 0.):
                _external_value_ramp = dgpy.ramp3d(_external_value, dgpy.function.getTime(), self._initial_time, ramp_period)
            else:
                _external_value_ramp = _external_value
        elif ux is None:
            _sse = slim_private._load_function(sse, self._domain._groups)
            _external_value = dgpy.outsideValueEta(self._solution.getFunction(), _sse)
        else:
            _ux_= slim_private._load_function(ux, self._domain._groups)
            _uy_ = slim_private._load_function(uy, self._domain._groups)
            if uz is not None:
                _uz_ = slim_private._load_function(uz, self._domain._groups)
            _velocity_local = slim_private._check_vector(_ux_, _uy_, _uz_, self._domain)    
            _ux = slim_private._get_x(_velocity_local).functor
            _uy = slim_private._get_y(_velocity_local).functor
            if transport_flux:
                if sse is not None:
                    _sse = slim_private._load_function(sse, self._domain._groups)    
                    _external_value = dgpy.transportToBoundaryVelocity(_sse, _ux, _uy, self._domain._bath_function)
                else:
                    _external_value = dgpy.transportToBoundaryVelocity(self._solution.getFunction(), _ux, _uy, self._domain._bath_function)
            else:
                if sse is not None:
                    _sse = slim_private._load_function(sse, self._domain._groups)    
                _external_value = dgpy.mergeVector(_sse, _ux, _uy, self._solution.getFunction())
        if (ramp_period > 0.):
            _external_value_ramp = dgpy.ramp3d(_external_value, dgpy.function.getTime(), self._initial_time, ramp_period)
        else:
            _external_value_ramp = _external_value
        self._equation.addBoundaryCondition(physical_tag, self._equation.newOutsideValueBoundary(_external_value_ramp))
        self._ref += [_sse, _ux, _uy, _velocity_local, _external_value, _external_value_ramp, ramp_period, _discharge, _depth_integral] 
 
    def export_value_at_point(self, output_name, x, y, z=None):
        """set the coordinates where a time serie of the solution is exported

        keyword arguments:
        
        * output_name 
            name of the output file
        * coord_x     
            vector containing the x coordinates where the solution is evaluated
        * coord_y     
            vector containing the y coordinates where the solution is evaluated
        * coord_z     
            vector containing the z coordinates where the solution is evaluated (Default: None, replaced by vector of 0)
        """
        if self._evaluator is None:
            self._evaluator = dgpy.dgFunctionEvaluator(self._domain._groups, self._solution.getFunction())
        self._evaluator_points_x = x
        self._evaluator_points_y = y
        if z is None:
            z = [0.]*len(x)
        self._evaluator_points_z = z
        self._evaluator_output_name = output_name
        f = open(output_name, 'w')
        sx = "x,"
        sy = "y,"
        sz = "z,"
        for i in range(len(x)):
            sx += str(x[i])+","
            sy += str(y[i])+","
            sz += str(z[i])+","
        f.write(sx[:-1]+'\n')
        f.write(sy[:-1]+'\n')
        f.write(sz[:-1]+'\n')
        f.close()
    
    def compute_mass(self,output_name):
        """compute the total mass of the domain
        
        keyword arguments:
        
        * output_name 
            name of the output file
        """
        self._output_mass = output_name
        self._compute_mass = True
        self._height = dgpy.functorNumpy(lambda cm : (cm.get(self._solution.getFunction())[:,[0]])+cm.get(self._domain._bath_function))
        self._integrator_mass = dgpy.dgFunctionIntegrator(self._domain._groups, self._height)
        self._integrator_mass.setSolution(self._solution.getFunction(), self._solution.getFunctionGradient())
        self._mass = dgpy.fullMatrixDouble(1, 1)
        self._integrator_mass.compute(self._mass,"")


class ShallowWaterTracer2d:
    """Create the shallow water equation for a tracer"""
    def __init__(self, domain, temporal_scheme, hydro_sol, linear_equation=False, name="Tracer", offline=False, wetting_drying=None, export_every_sub_time_step=False, initial_time=None, final_time=None): 
        """keyword arguments:
        
        * domain                     
            object of class Domain
        * temporal_scheme 
            scheme for the temporal integration
                               
            * "explicit" 
                explicit Euler scheme
            * "implicit" 
                implicit order 2 Runge-Kutta
        * hydro_sol
            hydrodynamics solution
            
            Format:
            
            - equation : ShallowWater2d object (online mode)
            - string : path to the folder containing the hydrodynamics folder 'offline_sw2d' (offline mode).
        * linear_equation            
            boolean stating linear shallow water tracer equation (default: False)
        * name                       
            name of the Tracer (default: "Tracer")
        * offline                    
            resolve shallow water equation (False) or use file to provide hydrodynamics (True) (default: False)
        * wetting_drying             
            value for the wetting & drying algorithm. It has to be the same as for the shallow water equation (see Karna et al. 2011) (default: None)
        * export_every_sub_time_step 
            boolean stating if you want to export the hydrodynamics every sub time step for further use of offline tracer (default: False)
        * initial_time               
            initial time of the simulation (default: 0)
            
            Format:
            
            - float [in seconds]
            - date string in format "year-month-day hours:minutes:seconds" (e.g. "2000-06-20 10:05:00")
        * final_time                 
            final time of the simulation (default: 0)
            
            Format:
            
            - float [in seconds]
            - date string in format "year-month-day hours:minutes:seconds" (e.g. "2000-06-20 22:42:17")
        """
        fmt = '%Y-%m-%d %H:%M:%S'   
        if initial_time is None:
            self._initial_time = 0.
        elif slim_private._is_string(initial_time):
            date = slim_private.datetime.datetime.strptime(initial_time, fmt)
            self._initial_time = slim_private.calendar.timegm([date.year, date.month, date.day, date.hour, date.minute, date.second])
        else:
            self._initial_time = float(initial_time)
        if final_time is None:
            self._final_time = 0.
        elif slim_private._is_string(final_time):
            date = slim_private.datetime.datetime.strptime(final_time, fmt)
            self._final_time = slim_private.calendar.timegm([date.year, date.month, date.day, date.hour, date.minute, date.second])
        else:
            self._final_time = float(final_time)
        self._domain = domain
        self._equation = dgpy.dgConservationLawShallowWaterTracer2d()
        self._equation.setIsLinear(linear_equation)
        self._equation.hydroWithLaxFriedrichs(True)
        self._equation.setGravity(self._domain._g)
        self._wetting_drying = wetting_drying
        self._equation.setBathymetry(self._domain._bath_function)
        self._solution = dgpy.dgDofContainer(self._domain._groups, 1)
        self._solution.setAll(0.0)
        self._solution.setFieldName(0, name)
        self._hydro_sol = None
        self._name = name
        if temporal_scheme == "explicit" or temporal_scheme == "implicit":
            self._temporal_scheme = temporal_scheme
        else:
            dgpy.Msg.Fatal("Unknown temporal scheme "+temporal_scheme+" for shallow water tracer equation !")
        self._offline = offline
        if self._offline:
            if not slim_private._is_string(hydro_sol):
                dgpy.Msg.Fatal("You have to provide the sea surface elevation and the velocity to solve the tracer in offline mode!")
            self._datafile = hydro_sol+'/offline_sw2d/'
            self._hydro_sol_dof = dgpy.dgDofContainer(self._domain._groups, 3)
            self._hydro_sol_dof.importMsh(self._datafile+"offline_sw2d-000000_000000/offline_sw2d-000000_000000")
            self._hydro_sol = self._hydro_sol_dof.getFunction()
        else:
            if not isinstance(hydro_sol, ShallowWater2d):
                dgpy.Msg.Fatal("You have to provide the hydrodynamic equation to solve the tracer!")
            self._hydro_sol = hydro_sol._solution.getFunction()
        if self._wetting_drying is not None:
            self._equation.setMovingBathWettingDrying(self._hydro_sol, None, self._wetting_drying)
            self._domain._bath_function = self._equation.getBathymetry()
        self._equation.setHydroSolution(self._hydro_sol) 
        self._export_every_sub_time_step = export_every_sub_time_step
        self._evaluator = None
        self._ref = []
        self._compute_mass = False
        self._export_on_structured_grid =False
           
    def set_initial_condition(self, initial_condition):
       """Initial conditions
        
        keyword argument:
        
        * tracer 
            tracer initial condition [-] (.msh or .nc format) (default: 0 )
        """
       slim_private._load(self._solution, initial_condition)
    
    def set_diffusivity(self, mode, okubo_coefficient=0.03, constant_diffusivity=1e-6):
        """Diffusivity
        
        keyword arguments:

        * mode                  
            type of diffusivity
            
            * "okubo"   
                Okubo scheme 
            * "constant" 
                constant value 
        * okubo_coefficient   
            coefficient for Okubo scheme [in m^0.85/s] (default: 0.03 m^0.85/s)
        * constant_diffusivity 
            constant value [in m^2/s] (default: 1e-6 m^2/s)
        """
        if mode == "okubo":
            self._kappa = dgpy.okubo2d(okubo_coefficient, dgpy.function.getElementVolume())
            self._equation.setDiffusivity(self._kappa)
        elif mode == "constant":
            self._kappa = dgpy.functionConstant(constant_diffusivity)
            self._equation.setDiffusivity(self._kappa)
        else: 
            dgpy.Msg.Fatal("Diffusivity scheme not implemented, use okubo or a constant diffusivity scheme") 

    def set_source(self, tracer_source):
        """Add a tracer source term in the tracer equation.
        
        keyword argument: 
        
        * tracer_source 
            netcdf or .msh file containing the source term for the whole domain [in c*s^-1]. 
        """
        self._source = slim_private._load_function(tracer_source, self._domain._groups)
        self._equation.setSource(self._source)    

    def set_boundary_coast(self, physical_tag):
        """Boundary condition: no flux condition at the boundary
        
        keyword argument:
        
        * physical_tag 
            tag of the part of the boundary on which this boundary condition is applied.
        """
        self._equation.addBoundaryCondition(physical_tag, self._equation.newBoundaryWall())

    def set_boundary_open(self, physical_tag, concentration):
        """Open boundary condition: one can impose the external concentration
        
        keyword arguments:
        
        * physical_tag             
            tag of the part of the boundary on which this boundary condition is applied.
        * concentration            
            netcdf or .msh file containing the concentration which will be applied at the boundary [in c]
        """
        #TODO add option for external gradient
        _external_value = slim_private._load_function(concentration, self._domain._groups)
        self._equation.addBoundaryCondition(physical_tag, self._equation.newOutsideValueBoundary(_external_value))
        self._ref += [_external_value]
       
    def set_boundary_flux(self, physical_tag, flux):
        """Boundary condition for a tracer flux: one can impose the flux throw the boundary condition, in [c s^-1], (with the tracer units [c]).
        
        keyword arguments:
        
        * physical_tag             
            tag of the part of the boundary on which this boundary condition is applied.
        * flux            
            netcdf or .msh file containing the flux which will be applied at the boundary in [c s^-1]
        """       
        _flux = slim_private._load_function(flux, self._domain._groups)
        _length_interface = dgpy.fullMatrixDouble(1, 1)
        one = slim_private.dgpy.functionConstant(1)
        dgpy.dgFunctionIntegratorInterface(self._domain._groups, one).compute(physical_tag, _length_interface)
        _flux_HC = dgpy.tracerFluxToNeumann(_flux, _length_interface(0,0))