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

# # CSU_RadarTools Demonstration

# <hr>
# <i>Notebook Author Info</i><br>
# Timothy Lang<br>
# tjlangco@gmail.com
# <hr>
# 
# Welcome to the CSU_RadarTools demonstration IPython notebook. CSU_RadarTools is a collection of open-source tools for weather radar data quality control and analysis, written in the Python language. The package collates a disparate number of tools that have been developed over many years at Colorado State University. The purpose of this notebook is to demonstrate how to use the package to accomplish various tasks.
# 
# In order to get started, in addition to making sure you have a robust Python installation with most major tools (e.g., numpy, matplotlib, pandas, etc.) please download and install the following modules: 
# <ul>
# <li>CSU_RadarTools (https://github.com/CSU-RadarMet/CSU_RadarTools)
# <li>Py-ART (https://github.com/ARM-DOE/pyart)
# <li>SkewT (https://pypi.python.org/pypi/SkewT); Python 3 version available at https://github.com/tjlang/SkewT
# </ul>

# In[1]:


from __future__ import print_function
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as colors
import pyart
import glob
from pyart.io.common import radar_coords_to_cart
from skewt import SkewT
from csu_radartools import (csu_fhc, csu_liquid_ice_mass, csu_blended_rain, 
                            csu_dsd, csu_kdp, csu_misc)
get_ipython().run_line_magic('matplotlib', 'inline')


# First things first, you see that csu_radartools is actually several different sub-modules: 
# <ul>
# <li><b>csu_fhc</b> - fuzzy-logic hydrometeor identification</li>
# <li><b>csu_liquid_ice_mass</b> - liquid/ice water mass calculations</li>
# <li><b>csu_blended_rain</b> - rainfall estimation via the famous CSU blended algorithm</li> 
# <li><b>csu_dsd</b> - DSD parameter estimation via several different methodologies</li> 
# <li><b>csu_kdp</b> - An FIR-based KDP estimation algorithm</li>
# <li><b>csu_misc</b> - A loose collection of miscellaneous tools, mainly focused on QC</li>
# </ul>
# 
# Let's examine each sub-module in turn. First, however, let's define a simple function to do side-by-side PPI radar plots, powered by Py-ART.

# In[2]:


def two_panel_plot(radar, sweep=0, var1='reflectivity', vmin1=0, vmax1=65, cmap1='RdYlBu_r', 
                   units1='dBZ', var2='differential_reflectivity', vmin2=-5, vmax2=5, 
                   cmap2='RdYlBu_r', units2='dB', return_flag=False, xlim=[-150,150],
                   ylim=[-150,150]):
    display = pyart.graph.RadarDisplay(radar)
    fig = plt.figure(figsize=(13, 5))
    ax1 = fig.add_subplot(121)
    display.plot_ppi(var1, sweep=sweep, vmin=vmin1, vmax=vmax1, cmap=cmap1, 
                     colorbar_label=units1, mask_outside=True)
    display.set_limits(xlim=xlim, ylim=ylim)
    ax2 = fig.add_subplot(122)
    display.plot_ppi(var2, sweep=sweep, vmin=vmin2, vmax=vmax2, cmap=cmap2, 
                     colorbar_label=units2, mask_outside=True)
    display.set_limits(xlim=xlim, ylim=ylim)
    if return_flag:
        return fig, ax1, ax2, display


# <hr>

# ### csu_fhc

# This package currently works with X-, C-, and S-band polarimetric radar data. Best results are obtained if you use a temperature sounding along with polarimetric radar data.

# In[3]:


#Read in the data
sndfile = '/Users/tjlang/Documents/OVWST/CPOL/soundings/snd_Darwin.txt'
radarfile = '/Users/tjlang/Documents/OVWST/CPOL/output/20060119/' + \
    'cfrad.20060119_170029.000_to_20060121_020810.000_CPOL_v1_PPI.nc'
radar = pyart.io.read(radarfile)
print(radar.fields.keys())
sounding = SkewT.Sounding(sndfile)


# This volume is from the CPOL C-band polarimteric Doppler radar. The fields of note are corrected reflectivity (ZC), differential reflectivity (ZD), specific differential phase (KD), and correlation coefficient (RH). The sounding is a nearby sounding in the UWyo format (i.e., from http://weather.uwyo.edu/upperair/sounding.html).
# 
# The way CSU_RadarTools is designed is it works on Python arrays or scalars. This allows you to use it with any kind of radar data, whether it was ingested via Py-ART or something else, or is in polar coordinate or gridded form. Thus, we first need to extract the relevant fields from the Py-ART radar object.

# In[4]:


dz = radar.fields['ZC']['data']
dr = radar.fields['ZD']['data']
kd = radar.fields['KD']['data']
rh = radar.fields['RH']['data']


# But we also need to get the sounding data interpolated onto the same structure as the radar data. Here is a way to do that via Py-ART and numpy.interp().

# In[5]:


def get_z_from_radar(radar):
    """Input radar object, return z from radar (km, 2D)"""
    azimuth_1D = radar.azimuth['data']
    elevation_1D = radar.elevation['data']
    srange_1D = radar.range['data']
    sr_2d, az_2d = np.meshgrid(srange_1D, azimuth_1D)
    el_2d = np.meshgrid(srange_1D, elevation_1D)[1]
    xx, yy, zz = radar_coords_to_cart(sr_2d/1000.0, az_2d, el_2d)
    return zz + radar.altitude['data']


def check_sounding_for_montonic(sounding):
    """
    So the sounding interpolation doesn't fail, force the sounding to behave
    monotonically so that z always increases. This eliminates data from
    descending balloons.
    """
    snd_T = sounding.soundingdata['temp']  # In old SkewT, was sounding.data
    snd_z = sounding.soundingdata['hght']  # In old SkewT, was sounding.data
    dummy_z = []
    dummy_T = []
    if not snd_T.mask[0]:  # May cause issue for specific soundings
        dummy_z.append(snd_z[0])
        dummy_T.append(snd_T[0])
        for i, height in enumerate(snd_z):
            if i > 0:
                if snd_z[i] > snd_z[i-1] and not snd_T.mask[i]:
                    dummy_z.append(snd_z[i])
                    dummy_T.append(snd_T[i])
        snd_z = np.array(dummy_z)
        snd_T = np.array(dummy_T)
    return snd_T, snd_z


def interpolate_sounding_to_radar(sounding, radar):
    """Takes sounding data and interpolates it to every radar gate."""
    radar_z = get_z_from_radar(radar)
    radar_T = None
    snd_T, snd_z = check_sounding_for_montonic(sounding)
    shape = np.shape(radar_z)
    rad_z1d = radar_z.ravel()
    rad_T1d = np.interp(rad_z1d, snd_z, snd_T)
    return np.reshape(rad_T1d, shape), radar_z


# In[6]:


radar_T, radar_z = interpolate_sounding_to_radar(sounding, radar)


# And now we are ready to run the HID algorithm. Behold ...

# In[7]:


scores = csu_fhc.csu_fhc_summer(dz=dz, zdr=dr, rho=rh, kdp=kd, use_temp=True, band='C',
                                T=radar_T)
fh = np.argmax(scores, axis=0) + 1


# To enable the ability to find out the second-ranked (or third, etc.) species, csu_fhc_summer() returns the scores for all the different categories, not just the max. So to get the traditional HID category number you have to use numpy.argmax() as above. The summer HID from CSU returns 10 possible categories:
# <ul>
# <li>1. Drizzle</li>
# <li>2. Rain</li>
# <li>3. Ice Crystals</li>
# <li>4. Aggregates</li>
# <li>5. Wet/Melting Snow</li>
# <li>6. Vertically Aligned Ice</li>
# <li>7. Low-Density Graupel</li>
# <li>8. High-Density Graupel</li>
# <li>9. Hail</li>
# <li>10. Big Drops</li>
# </ul>
# 
# And these are represented as integers in the newly created fh array, which as the same structure as dz, dr, etc. We'd like to plot these data using Py-ART, which means we need to turn fh in a radar object field. Let's do that.

# In[8]:


def add_field_to_radar_object(field, radar, field_name='FH', units='unitless', 
                              long_name='Hydrometeor ID', standard_name='Hydrometeor ID',
                              dz_field='ZC'):
    """
    Adds a newly created field to the Py-ART radar object. If reflectivity is a masked array,
    make the new field masked the same as reflectivity.
    """
    masked_field = np.ma.asanyarray(field)
    fill_value = -32768
    if hasattr(radar.fields[dz_field]['data'], 'mask'):
        setattr(masked_field, 'mask', radar.fields[dz_field]['data'].mask)
        fill_value = radar.fields[dz_field]['_FillValue']
    field_dict = {'data': masked_field,
                  'units': units,
                  'long_name': long_name,
                  'standard_name': standard_name,
                  '_FillValue': fill_value}
    radar.add_field(field_name, field_dict, replace_existing=True)
    return radar


# In[9]:


radar = add_field_to_radar_object(fh, radar)


# And now let's plot it up! First, however, let's fix Py-ART's colorbars to look nice for HID and other category-style fields.

# In[10]:


hid_colors = ['White', 'LightBlue', 'MediumBlue', 'DarkOrange', 'LightPink',
              'Cyan', 'DarkGray', 'Lime', 'Yellow', 'Red', 'Fuchsia']
cmaphid = colors.ListedColormap(hid_colors)
cmapmeth = colors.ListedColormap(hid_colors[0:6])

def adjust_fhc_colorbar_for_pyart(cb):
    cb.set_ticks(np.arange(1.4, 10, 0.9))
    cb.ax.set_yticklabels(['Drizzle', 'Rain', 'Ice Crystals', 'Aggregates',
                           'Wet Snow', 'Vertical Ice', 'LD Graupel',
                           'HD Graupel', 'Hail', 'Big Drops'])
    cb.ax.set_ylabel('')
    cb.ax.tick_params(length=0)
    return cb

def adjust_meth_colorbar_for_pyart(cb):
    cb.set_ticks(np.arange(1.25, 5, 0.833))
    cb.ax.set_yticklabels(['R(Kdp, Zdr)', 'R(Kdp)', 'R(Z, Zdr)', 'R(Z)', 'R(Zrain)'])
    cb.ax.set_ylabel('')
    cb.ax.tick_params(length=0)
    return cb


# In[11]:


# Actual plotting done here
lim = [-80, 80]
fig, ax1, ax2, display = two_panel_plot(radar, sweep=5, var1='ZC', var2='FH', vmin2=0, 
                                        vmax2=10, cmap2=cmaphid, units2='', return_flag=True, 
                                        xlim=lim, ylim=lim)
display.cbs[1] = adjust_fhc_colorbar_for_pyart(display.cbs[1])


# The HID algorithm, like most algorithms in CSU_RadarTools, works on scalars too. So if you are just curious what category a set of polarimetric values will get, try the following:

# In[12]:


scores = csu_fhc.csu_fhc_summer(dz=45.0, zdr=0.0, kdp=-0.2, rho=0.95, T=-1) 
print(np.argmax(scores, axis=0) + 1)


# ... which is high-density graupel.

# In[13]:


#Here's help to see what else you can do with csu_fhc.csu_fhc_summer
help(csu_fhc.csu_fhc_summer)


# <hr>

# ### csu_blended_rain

# This module provides access to two distinct CSU blended rainfall algorithm methodologies. The first function is calc_blended_rain(), which is the classic version that utilizes ZDP to estimate ice fraction, and the second function is calc_hidro_rain(), which is newer and uses HID to help distinguish regions of ice.
# 
# There are also a number of individual rain rate calculation functions (e.g., Z-R, R(Kdp), R(Z,Zdr), etc.) provided with this module. The user can specify the coefficients or just use the defaults, which come from standard literature references.
# 
# Output from the main blended functions include the method used at each gate, which the user can also plot up. For the example below we show the results from calc_hidro_rain(). Since we already calculated HID, we can use our previous output from csu_fhc to supply that field.

# In[14]:


rain, method = csu_blended_rain.csu_hidro_rain(dz=dz, zdr=dr, kdp=kd, fhc=fh)
radar = add_field_to_radar_object(rain, radar, field_name='rain', units='mm h-1',
                                  long_name='HIDRO Rainfall Rate', 
                                  standard_name='Rainfall Rate')
radar = add_field_to_radar_object(method, radar, field_name='method', units='',
                                  long_name='HIDRO Rainfall Method', 
                                  standard_name='Rainfall Method')


# In[15]:


fig, ax1, ax2, display = two_panel_plot(radar, sweep=0, var1='rain', vmin1=0, vmax1=150,
                                        cmap1='GnBu', var2='method', vmin2=0, vmax2=5, 
                                        cmap2=cmapmeth, units2='', return_flag=True, 
                                        xlim=lim, ylim=lim, units1='mm h-1')
display.cbs[1] = adjust_meth_colorbar_for_pyart(display.cbs[1])


# The method variable is an integer output that matches the following legend:
# 
# <u>Method Legend</u>
# <ul>
# <li>1. R(Kdp, Zdr)</li>
# <li>2. R(Kdp)</li>
# <li>3. R(Z, Zdr)</li>
# <li>4. R(Z) - Z uncorrected to remove ice contamination</li>
# <li>5. R(Zrain) - Z corrected to remove ice contamination</li>
# </ul>
# 
# The last method is only used by calc_blended_rain(). Here is how to use that algorithm. It has additional outputs, which include ZDP and ice fraction (unitless), if you set ice_flag = True.

# In[16]:


rain, method, zdp, fi = csu_blended_rain.calc_blended_rain(dz=dz, zdr=dr, 
                                                           kdp=kd, ice_flag=True)
radar = add_field_to_radar_object(rain, radar, field_name='rain_blend', units='mm h-1',
                                  long_name='Blended Rainfall Rate', 
                                  standard_name='Rainfall Rate')
radar = add_field_to_radar_object(method, radar, field_name='method_blend', units='',
                                  long_name='Blended Rainfall Method', 
                                  standard_name='Rainfall Method')
radar = add_field_to_radar_object(zdp, radar, field_name='ZDP', units='dB',
                                  long_name='Difference Reflectivity',
                                  standard_name='Difference Reflectivity')
radar = add_field_to_radar_object(fi, radar, field_name='FI', units='', 
                                  long_name='Ice Fraction',
                                  standard_name='Ice Fraction')


# In[17]:


fig, ax1, ax2, display = two_panel_plot(radar, sweep=0, var1='rain_blend', vmin1=0, vmax1=150,
                                        cmap1='GnBu', var2='method_blend', vmin2=0, vmax2=5, 
                                        cmap2=cmapmeth, units2='', return_flag=True, xlim=lim, 
                                        ylim=lim, units1='mm h-1')
display.cbs[1] = adjust_meth_colorbar_for_pyart(display.cbs[1])


# We can plot ZDP and FI too, to gauge how they are influencing the rainfall method choices seen above.

# In[18]:


two_panel_plot(radar, sweep=0, var1='ZDP', units1='dB', vmin1=0, vmax1=65,
               cmap1='cubehelix', var2='FI', vmin2=0, vmax2=1, 
               cmap2='YlOrRd_r', units2='', xlim=lim, ylim=lim)


# In[19]:


#And of course, we can get help on all the other stuff!
#Original IDL notes are listed toward the top, with the Python-specific notes below
help(csu_blended_rain)


# If you dig into the docstrings above, you see individual functions for rainfall, like R(Kdp) and R(Z,Zdr). So you can just access those on your own, if you like.

# In[20]:


rkdp = csu_blended_rain.calc_rain_kdp(kd)
radar = add_field_to_radar_object(rkdp, radar, field_name='rkdp', units='mm h-1',  
                                  long_name='Rainfall Rate R(Kdp)', 
                                  standard_name='Rainfall Rate')


# In[21]:


#Compare R(Kdp) with Kdp!
two_panel_plot(radar, sweep=0, var1='rkdp', vmin1=0, vmax1=150, 
               cmap1='GnBu', units1='mm h-1',
               var2='KD', vmin2=-5, vmax2=5, cmap2='RdYlBu', 
               units2='deg km-1', xlim=lim, ylim=lim)


# <hr>

# ### csu_dsd

# This module enables retrieval of common drop-size distribution (DSD) parameters. As with most aspects of CSU_RadarTools, the end user is responsible for proper masking of their data, as DSD retrievals are only valid in pure rain. Remember, garbage in = garbage out! The module is designed to work with C-band or S-band, and the retrievals are different depending on radar frequency. The main function is calc_dsd() and it returns median volume diameter (D0), normalized intercept parameter (Nw), and mu assuming a gamma distribution DSD model.
# 
# Let's first check things out for C-band.

# In[22]:


d0, Nw, mu = csu_dsd.calc_dsd(dz=dz, zdr=dr, kdp=kd, band='C')
radar = add_field_to_radar_object(d0, radar, field_name='D0', units='mm', 
                                  long_name='Median Volume Diameter',
                                  standard_name='Median Volume Diameter')
logNw = np.log10(Nw)
radar = add_field_to_radar_object(logNw, radar, field_name='NW', units='', 
                                  long_name='Normalized Intercept Parameter',
                                  standard_name='Normalized Intercept Parameter')
radar = add_field_to_radar_object(mu, radar, field_name='MU', units='', 
                                  long_name='Mu', standard_name='Mu')


# The C-band retrievals, based on Bringi et al. (2009), don't bother retrieving mu. So CSU_RadarTools will fix it at 3 unless the end user changes the csu_dsd.DEFAULT_MU global variable. So let's check out instead what happened with D0 and Nw.

# In[23]:


two_panel_plot(radar, sweep=0, var1='D0', vmin1=0, vmax1=2, cmap1='GnBu', units1='mm',
               var2='NW', vmin2=0, vmax2=8, cmap2='cubehelix', units2='log10(Nw)', xlim=lim, 
               ylim=lim)


# Note the differences in D0 & Nw pairs inside and outside of convection. Both here and above with ice fraction, which also depends heavily on Zdr, there are some possible differential attenuation impacts on the retrievals downrange of one of the cores just west of the radar.
# 
# What about S-band DSD retrievals? For that we need a different radar. Let's try S-PolKa during the NAME project in 2004.

# In[24]:


fdir = '/Users/tjlang/Documents/OVWST/NAME/output/20040805/'
files = sorted(glob.glob(fdir+'*nc'))
radarS = pyart.io.read(files[0])
print(radarS.fields.keys())


# In[25]:


dzS = radarS.fields['DZ']['data']
drS = radarS.fields['DR']['data']
kdS = radarS.fields['KD']['data']


# With S-band, there are two different methodologies available: Bringi et al. (2013) and Bringi et al. (2004). The latter allows mu to vary, so let's try that.

# In[26]:


d0, Nw, mu = csu_dsd.calc_dsd(dz=dzS, zdr=drS, kdp=kdS, band='S', method='2004')
radarS = add_field_to_radar_object(d0, radarS, field_name='D0', units='mm', 
                                   long_name='Median Volume Diameter',
                                   standard_name='Median Volume Diameter', dz_field='DZ')
logNw = np.log10(Nw)
radarS = add_field_to_radar_object(logNw, radarS, field_name='NW', units='', 
                                   long_name='Normalized Intercept Parameter',
                                   standard_name='Normalized Intercept Parameter', 
                                   dz_field='DZ')
radarS = add_field_to_radar_object(mu, radarS, field_name='MU', units='', long_name='Mu', 
                                   standard_name='Mu', dz_field='DZ')


# In[27]:


limS = [-100, 100]
# Changing the color scale a bit to allow for the bigger drops in this case. 
two_panel_plot(radarS, sweep=0, var1='D0', vmin1=0, vmax1=3.5, cmap1='GnBu', units1='mm',
               var2='NW', vmin2=0, vmax2=8, cmap2='cubehelix', units2='log10(Nw)', xlim=limS, 
               ylim=limS)


# In[28]:


# Now let's see what mu is up to ...
two_panel_plot(radarS, sweep=0, var1='DZ', vmin1=0, vmax1=65, cmap1='RdYlBu_r', units1='dBZ',
               var2='MU', vmin2=0, vmax2=5, cmap2='cubehelix', units2='mu', xlim=limS, 
               ylim=limS)


# So for the most part mu gets set to 3, except there is some variability in the cores and a few other places. Normally, mu is very difficult to retrieve accurately, but at least the user has the option of testing different methodologies with csu_dsd.

# In[29]:


# And you can always ...
help(csu_dsd)


# <hr>

# In[ ]: