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

# #PyTDA Demo
# 
# <b>Author</b><br>
# Timothy Lang, NASA MSFC<br>
# timothy.j.lang@nasa.gov
# 
# <b>Overview</b><br>
# PyTDA is a Python module that allows the use to estimate eddy dissipation rate (a measure of turbulence) from Doppler weather radar data. It interfaces seamlessly with Py-ART to make this calculation as simple as one line of code.
# 
# To get started, you need the following to be installed:
# <ul>
# <li>Py-ART (https://github.com/ARM-DOE/pyart)
# <li>Standard python libraries like numpy, sklearn, scipy, Cython, etc.
# </ul>
# 
# PyTDA is currently a pseudo-package, so you need to add the code location to your PYTHONPATH. Then you need to compile the Cython code using the provided compile_pytda_cython_code.sh script. Then you should be good to go for using this demo. PyTDA is tested and works under Python 2.7 and Python 3.4. If you switch between the two, you need to recompile the Cython code and recreate the shared object file after doing so.

# In[1]:


from __future__ import division, print_function
import numpy as np
import matplotlib.pyplot as plt
import os
import glob
import pyart
import pytda
get_ipython().run_line_magic('matplotlib', 'inline')


# So PyTDA is very simple to use. First you need a Py-ART radar object:

# In[13]:


files = glob.glob('*.nc')
print(files)
radar = pyart.io.read(files[0])


# In[14]:


print(radar.fields.keys())


# In[15]:


print(radar.instrument_parameters['radar_beam_width_h']['data'][0])


# Once you have that, send the radar object as an argument to the desired PyTDA function. There are two: calc_turb_sweep and calc_turb_vol. The former returns the turbulence field as an ndarray. The latter just successively calls the former to process a volume and modify the radar object and add the turbulence field. Currently, PyTDA only works on PPI sweeps/volumes.

# In[16]:


pytda.calc_turb_vol(radar, name_sw='SW', name_dz='DZ', verbose=False,
                    gate_spacing=250.0/1000.0, use_ntda=False,
                    beamwidth=radar.instrument_parameters['radar_beam_width_h']['data'][0])


# The default is to do fuzzy-logic-based QC on the turbulence retrievals, to ensure better quality. This is like what the NCAR Turbulence Detection Algorithm (NTDA) does. However, there is an option to turn that off and just have straight inversion of spectrum width to EDR. You can use the verbose keyword to turn off the text, I just turned it on to show more of what was going on. See the documentation below to find out all the different keyword options.
# 
# That's about it to do the processing. Usually takes about a minute or two per volume. Now let's use Py-ART to plot the results! First, let's define a function to do easy multi-panel plots in Py-ART.

# In[17]:


def plot_list_of_fields(radar, sweep=0, fields=['reflectivity'], vmins=[0],
                        vmaxs=[65], units=['dBZ'], cmaps=['RdYlBu_r'],
                        return_flag=False, xlim=[-150, 150], ylim=[-150, 150],
                        mask_tuple=None):
    num_fields = len(fields)
    if mask_tuple is None:
        mask_tuple = []
        for i in np.arange(num_fields):
            mask_tuple.append(None)
    nrows = (num_fields + 1) // 2
    ncols = (num_fields + 1) % 2 + 1
    fig = plt.figure(figsize=(14.0, float(nrows)*5.5))
    display = pyart.graph.RadarDisplay(radar)
    for index, field in enumerate(fields):
        ax = fig.add_subplot(nrows, 2, index+1)
        display.plot_ppi(field, sweep=sweep, vmin=vmins[index],
                         vmax=vmaxs[index],
                         colorbar_label=units[index], cmap=cmaps[index],
                         mask_tuple=mask_tuple[index])
        display.set_limits(xlim=xlim, ylim=ylim)
    plt.tight_layout()
    if return_flag:
        return display


# And now let's plot.

# In[19]:


plot_list_of_fields(radar, sweep=4, xlim=[-50, 50], ylim=[50, 150],
                    fields=['DZ', 'VR', 'SW', 'turbulence'],
                    vmins=[0, -30, 0, 0], vmaxs=[60, 30, 10, 1.0], 
                    units=['dBZ', 'm/s', 'm/s', 'EDR^0.33'],
                    cmaps=['pyart_LangRainbow12', 'seismic', 'YlOrRd', 'cubehelix'])


# In[21]:


pyart.io.write_cfradial('file_with_turbulence.nc', radar)


# In[20]:


help(pytda)


# In[ ]: