#!/usr/bin/env python
# coding: utf-8

# ## Tutorial on how to use Parcels on 3D C grids

# One of the features of Parcels is that it can directly and natively work with `Field` data discretised on C-grids. These C grids are very popular in OGCMs, so velocity fields outputted by OGCMs are often provided on such grids, except if they have been firstly re-interpolated on a A grid.
# 
# More information about C-grid interpolation can be found in [Delandmeter et al., 2019](https://www.geosci-model-dev-discuss.net/gmd-2018-339/).
# An example of such a discretisation is the NEMO model, which is one of the models supported in Parcels. A tutorial teaching how to [interpolate 2D data on a NEMO grid](https://nbviewer.jupyter.org/github/OceanParcels/parcels/blob/master/parcels/examples/tutorial_nemo_curvilinear.ipynb) is available within Parcels.
# 
# Here, we focus on 3D fields. Basically, it is a straightforward extension of the 2D example, but it is very easy to do a mistake in the setup of the vertical discretisation that would affect the interpolation scheme.
# 
# ### Preliminary comments
# 
# *How to know if your data is discretised on a C grid?* The best way is to read the documentation that comes with the data. Alternatively, an easy check is to assess the coordinates of the U, V and W fields: for an A grid, U, V and W are distributed on the same nodes, such that the coordinates are the same. For a C grid, there is a shift of half a cell between the different variables.
# 
# *What about grid indexing?* Since the C-grid variables are not located on the same nodes, there is not one obvious way to define the indexing, i.e. where is `u[k,j,i]` compared to `v[k,j,i]` and `w[k,j,i]`. In Parcels, we use the same notation as in NEMO: see [horizontal indexing](https://www.nemo-ocean.eu/doc/img360.png) and [vertical indexing](https://www.nemo-ocean.eu/doc/img362.png).
# It is important that you check if your data is following the same notation. Otherwise, you should re-index your data properly (this can be done within Parcels, there is no need to regenerate new netcdf files).
# 
# *What about the accuracy?* By default in Parcels, particle coordinates (i.e. longitude, latitude and depth) are stored using single-precision `np.float32` numbers. The advantage of this is that it saves some memory ressources for the computation. In some applications, especially where particles travel very close to the coast, the single precision accuracy can lead to uncontrolled particle beaching, due to numerical rounding errors. In such case, you may want to double the coordinate precision to `np.float64`. This can be done by adding the parameter `lonlatdepth_dtype=np.float64` to the ParticleSet constructor. Note that for C-grid fieldsets such as in NEMO, the coordinates precision is set by default to `np.float64`.
# 
# ### How to create a 3D NEMO `dimensions` dictionary?
# 
# In the following, we will show how to create the `dimensions` dictionary for 3D NEMO simulations. What you require is a 'mesh_mask' file, which in our case is called `coordinates.nc` but in some other versions of NEMO has a different name. In any case, it will have to contain the variables `glamf`, `gphif` and `depthw`, which are the longitude, latitude and depth of the mesh nodes, respectively. Note that `depthw` is not part of the mesh_mask file, but is in the same file as the w data (`wfiles[0]`).
# 
# For the C-grid interpolation in Parcels to work properly, it is important that `U`, `V` and `W` are on the same grid.
# 
# The code below is an example of how to create a 3D simulation with particles, starting in the mouth of the river Rhine at 1m depth, and advecting them through the North Sea using the `AdvectionRK4_3D`

# In[1]:


from parcels import FieldSet, ParticleSet, JITParticle, AdvectionRK4_3D
from glob import glob
import numpy as np
from datetime import timedelta as delta
from os import path

data_path = 'NemoNorthSeaORCA025-N006_data/'
ufiles = sorted(glob(data_path+'ORCA*U.nc'))
vfiles = sorted(glob(data_path+'ORCA*V.nc'))
wfiles = sorted(glob(data_path+'ORCA*W.nc'))
mesh_mask = data_path + 'coordinates.nc'

filenames = {'U': {'lon': mesh_mask, 'lat': mesh_mask, 'depth': wfiles[0], 'data': ufiles},
             'V': {'lon': mesh_mask, 'lat': mesh_mask, 'depth': wfiles[0], 'data': vfiles},
             'W': {'lon': mesh_mask, 'lat': mesh_mask, 'depth': wfiles[0], 'data': wfiles}}

variables = {'U': 'uo',
             'V': 'vo',
             'W': 'wo'}
dimensions = {'U': {'lon': 'glamf', 'lat': 'gphif', 'depth': 'depthw', 'time': 'time_counter'},
              'V': {'lon': 'glamf', 'lat': 'gphif', 'depth': 'depthw', 'time': 'time_counter'},
              'W': {'lon': 'glamf', 'lat': 'gphif', 'depth': 'depthw', 'time': 'time_counter'}}

fieldset = FieldSet.from_nemo(filenames, variables, dimensions)

pset = ParticleSet.from_line(fieldset=fieldset, pclass=JITParticle,
                             size=10,
                             start=(1.9, 52.5),
                             finish=(3.4, 51.6),
                             depth=1)

kernels = pset.Kernel(AdvectionRK4_3D)
pset.execute(kernels, runtime=delta(days=4), dt=delta(hours=6))


# In[2]:


get_ipython().run_line_magic('matplotlib', 'inline')

depth_level = 8
print("Level[%d] depth is: [%g %g]" % (depth_level, fieldset.W.grid.depth[depth_level], fieldset.W.grid.depth[depth_level+1]))

pset.show(field=fieldset.W, domain={'N':60, 'S':49, 'E':15 ,'W':0}, depth_level=depth_level)