from cdo import *

cdo = Cdo()

"""
Explanation: Data Analysis with CDO
CDO User Guide https://code.mpimet.mpg.de/projects/cdo/embedded/cdo.pdf
Table of contents

About CDO
Import CDO module
Show CDO version
List CDO operators
List all operators starting with 'sel'
Show information about an operator
Turn debugging on
Show information about file content
Copy file
Select variables
Select timesteps
Select a sub-region



About CDO
See also https://code.mpimet.mpg.de/projects/cdo/wiki
CDO is a large tool set for working on climate and NWP model data. NetCDF 3/4, GRIB 1/2 including SZIP (or AEC) and JPEG compression, EXTRA, SERVICE and IEG are supported as IO-formats. Apart from that CDO can be used to analyse any kind of gridded data not related to climate science. CDO has very small memory requirements and can process files larger than the physical memory.
CDO is open source and released under the terms of the GNU General Public License v2 (GPL).
Import CDO module
Import the CDO module and set cdo to Cdo() which makes writing a little easier.
End of explanation
"""


print(cdo.__version__())


"""
Explanation: Show CDO version
python-cdo version:
End of explanation
"""


print(cdo.version())

"""
Explanation: Based on CDO version:
End of explanation
"""


tempPath = './tmp/'
cdo = Cdo(tempdir=tempPath)

"""
Explanation: Set temporary directory
Set another directory for storing tempfiles with a constructor option and remove anything left in there when you experienced a crash or something like this
End of explanation
"""


cdo.cleanTempDir()

"""
Explanation: Delete temporary files
End of explanation
"""


print(len(dict(list(cdo.operators.items()))))

"""
Explanation: List CDO operators
More than 800 operators are available.
End of explanation
"""


print(dict(list(cdo.operators.items())[0:50]))

"""
Explanation: To get the first 50 operators you can do
End of explanation
"""


[key for key, value in cdo.operators.items() if key.startswith('sel')]

"""
Explanation: List all operators starting with 'sel'
Use list comprehension
End of explanation
"""


help(cdo.selname)

"""
Explanation: Show information about an operator
End of explanation
"""


cdo.debug = True

"""
Explanation: Turn debugging on
End of explanation
"""


cdo.sinfon(input='data/rectilinear_grid_2D.nc')

"""
Explanation: Show information about file content
End of explanation
"""


cdo.debug = False

infile = 'data/rectilinear_grid_2D.nc'

cdo.sinfon(input=infile)

"""
Explanation: Turn off debugging
End of explanation
"""


import xarray as xr

xr.open_dataset(infile)

"""
Explanation: In comparison to xarray's dataset information.
End of explanation
"""


cdo.infon(input=infile)

"""
Explanation: Show information about the variables value range, times, levels in more detail
End of explanation
"""


outfile = 'outfile.nc'

"""
Explanation: Set output file.
End of explanation
"""


cdo.copy(input=infile, options='-b F64', output=outfile)
cdo.sinfon(input=outfile)

"""
Explanation: Copy file
Copy the input file to outfile.nc and change the data precision from float32 to float64.
End of explanation
"""


cdo.selvar('tsurf', input=infile, output=outfile)
cdo.sinfon(input=outfile)

cdo.selvar('u10,v10', input=infile, output=outfile)
cdo.sinfon(input=outfile)

"""
Explanation: Select variables
End of explanation
"""


cdo.seltimestep('1,10', input=infile, output=outfile)
cdo.sinfon(input=outfile)

"""
Explanation: Select timesteps
Select timestep 1 and 10:
End of explanation
"""


cdo.seltimestep('1/10', input=infile, output=outfile)
cdo.sinfon(input=outfile)

"""
Explanation: Select timestep 1 to 10:
End of explanation
"""


cdo.seltimestep('1/10', input='-selvar,u10,v10 '+infile, output=outfile)
cdo.sinfon(input=outfile)

"""
Explanation: Operator chaining
Operators with a fixed number of input files (streams) and only one output file can be combined. The input parameters must begin with an '-' and they will be executed from right to left.
Select the variables u10 and v10 and then select the first 10 timesteps:
End of explanation
"""


cdo.seltimestep('1/10', input='-selvar,u10,v10 '+infile, options='-b F64', output=outfile)
cdo.sinfon(input=outfile)

"""
Explanation: Use operators and options at once to do the above sections and change the output precision:
End of explanation
"""


cdo.sellonlatbox('20,30,70,80', input='-seltimestep,1 '+infile, output=outfile)
cdo.sinfon(input=outfile)

"""
Explanation: Select a sub-region
End of explanation
"""


cdo.fldmean(input="-selname,tsurf "+infile, output='outfile.nc')
cdo.sinfon(input=outfile)

"""
Explanation: Compute the field mean
Compute the mean of the horizontal field for each timestep (-> time series).
End of explanation
"""


cdo.fldmean(input="-selname,tsurf "+infile, returnXArray='tsurf').plot()

"""
Explanation: Plot the fieldmean data:
End of explanation
"""


precipitation = cdo.selvar('precip', input=infile, returnXArray='precip')
print(precipitation)

"""
Explanation: Assign data to variable
Assign the file variable precip to the python variable precipitation:
End of explanation
"""


print(precipitation.values[0,0,0])

"""
Explanation: Print the value of the first timestep, first latitude, first longitude:
End of explanation
"""


tsurf_fldmean = cdo.fldmean(input=infile, returnXArray='tsurf')
print(tsurf_fldmean.values[0:10,0,0])

"""
Explanation: Compute the fieldmean:
End of explanation
"""


cdo.remapbil('r360x180', input='-seltimestep,1 -selvar,tsurf '+infile, output=outfile)

"""
Explanation: Remapping
Interpolate the data of input file to a new grid using the bilinear interpolation method.
For the ease of use, we select the variable tsurf and only the first timestep to interpolate the data to a longitude 1 deg x latitude 1 deg (r360x180) grid.
End of explanation
"""


import cartopy.crs as ccrs

tsurf_orig = xr.open_dataset(infile).tsurf[0,:,:]

data = tsurf_orig.plot(subplot_kws=dict(projection=ccrs.PlateCarree(), facecolor='gray'), transform=ccrs.PlateCarree(),)

data.axes.set_extent([-10.,20.,30.,60.])
data.axes.coastlines()

tsurf = cdo.remapbil('r720x360', input='-seltimestep,1 '+infile, returnXArray='tsurf')

data = tsurf.plot(subplot_kws=dict(projection=ccrs.PlateCarree(), facecolor='gray'), transform=ccrs.PlateCarree(),)

data.axes.set_extent([-10.,20.,30.,60.])
data.axes.coastlines()

"""
Explanation: To demonstrate the functionality we increase the resolution of the input data to a 0.5 deg grid and plot the original and the interpolated data.
Original input data:
End of explanation
"""


cdo.topo(options='-f nc', output='topo.nc')

tsurf_2 = cdo.remapbil('topo.nc', input="-seltimestep,1 "+infile, returnXArray='tsurf')

"""
Explanation: Remap data to a grid of another file
End of explanation
"""


lsm = cdo.gtc(0.5, input='-remapbil,'+infile+' -topo', returnXArray='topo')
print(lsm)

p = lsm.plot(subplot_kws=dict(projection=ccrs.PlateCarree(), facecolor='gray'), transform=ccrs.PlateCarree(),)
p.axes.set_extent([-20.,20.,30.,60.])
p.axes.coastlines()

"""
Explanation: Create a land sea mask
CDO provides a global 0.5 degree topography dataset that can be used to generate a land sea mask. First, we interpolate the topography data to the same grid as the data file. With the operator gtc we set all values greater than 0.5 m to 1 and all other values to 0.
End of explanation
"""


cdo.setname('lsm', input='-gtc,0.5 -remapbil,'+infile+' -topo', options='-f nc', output='lsm.nc')

"""
Explanation: Masking data
Before we can use a mask on a data variable we need to create a mask file using the same grid as the data variable.
End of explanation
"""


masked = cdo.setctomiss(0, input='-mul lsm.nc -seltimestep,1 '+infile, returnXArray='tsurf')

"""
Explanation: Now, we want to get only tsurf values over land for timestep 1.
End of explanation
"""


p = masked.plot(subplot_kws=dict(projection=ccrs.PlateCarree(), facecolor='gray'), transform=ccrs.PlateCarree(),)
p.axes.set_extent([-20.,20.,30.,60.])
p.axes.coastlines()

"""
Explanation: Let's see how it looks like.
End of explanation
"""


lsm = cdo.setname('lsm', input='-gtc,0.5 -remapbil,'+infile+' -topo', options='-f nc', returnXArray='lsm')

tsurf = cdo.seltimestep(1, input='-selvar,tsurf '+infile, returnXArray='tsurf')

masked = tsurf * lsm
print(masked)

"""
Explanation: Use a variables instead of output files.
End of explanation
"""


masked.plot(subplot_kws=dict(projection=ccrs.PlateCarree(), facecolor='gray'), transform=ccrs.PlateCarree(),
           cbar_kwargs={'orientation': 'horizontal'})
p.axes.set_extent([-20.,20.,30.,60.])
p.axes.coastlines()  

"""
Explanation: Plot the masked data.
End of explanation
"""


import numpy as np
import matplotlib as mpl
import copy

lsm = np.where(lsm==0, 1, np.nan)

masked = tsurf * lsm

masked.plot(subplot_kws=dict(projection=ccrs.PlateCarree(), facecolor='gray'), transform=ccrs.PlateCarree(),
           cbar_kwargs={'orientation': 'horizontal'})
p.axes.set_extent([-20.,20.,30.,60.])
p.axes.coastlines()

"""
Explanation: And now we want to mask the land part of the data. The easiest way is to set all zeros to one and all ones to NaN. Matplotlib will automatically use the color grey for the missing values.
End of explanation
"""


cdo.cleanTempDir()

"""
Explanation: Delete temporary files
End of explanation
"""


