import pyart
from matplotlib import pyplot as plt
import numpy as np
%matplotlib inline
print pyart.__version__

filename ='data/sgpcsaprppi_20110520095101.nc'
radar = pyart.io.read(filename)

print radar.fields.keys()

display = pyart.graph.RadarDisplay(radar)
fig = plt.figure(figsize = [15,10])
display.plot_ppi('reflectivity', vmin = -8, vmax =64)

fig = plt.figure(figsize = [15,10])
display.plot_ppi('differential_phase', vmin = -180, vmax =180)

# Perform a Linear programming phase processing to get the conditionally
# fitted phase profile and sobel filtered KDP field

c_offset = - 3.42

gates = radar.range['data'][1] - radar.range['data'][0]
rge = 30.0 * 1000.0
ng = rge / gates


reproc_phase, sob_kdp = pyart.correct.phase_proc.phase_proc_lp(
    radar, c_offset, debug=True, nowrap=ng, min_ncp=0.6, min_rhv=0.8)

# in some sectors where there is not enough data the LP method fails.. find these and mask

reproc_phase['data'][np.where(reproc_phase['data'] < 0.0)] = 0.0
radar.add_field('recalculated_diff_phase', sob_kdp)
radar.add_field('proc_dp_phase_shift', reproc_phase)

fig = plt.figure(figsize = [15,10])
display.plot_ppi('proc_dp_phase_shift', vmin = 0, vmax =400)

fig = plt.figure(figsize = [15,10])
display.plot_ppi('recalculated_diff_phase', vmin = -1, vmax =10)

spec_at, cor_z = pyart.correct.attenuation.calculate_attenuation(
    radar, c_offset, debug=True, ncp_min=0.4)

# again filter where the algorithm has failed due to lack of data
spec_at['data'][np.where(spec_at['data'] < 0.0)] = 0.0

radar.add_field('specific_attenuation', spec_at)
radar.add_field('corrected_reflectivity', cor_z)

fig = plt.figure(figsize = [15,10])
display.plot_ppi('specific_attenuation', vmin = 0, vmax =1)

R = 300.0 * (radar.fields['specific_attenuation']['data']) ** 0.89
rain_rate_data = np.ma.masked_where(np.logical_or(
    (radar.fields['normalized_coherent_power']['data'] < 0.4),
    (radar.fields['cross_correlation_ratio']['data'] < 0.8)), R)
radar.add_field_like('unfolded_differential_phase', 
                     'rain_rate_A', rain_rate_data)
rainrate = radar.fields['rain_rate_A']
rainrate['valid_min'] = 0.0
rainrate['valid_max'] = 400.0
rainrate['standard_name'] = 'rainfall_rate'
rainrate['long_name'] = 'rainfall_rate'
rainrate['least_significant_digit'] = 1
rainrate['units'] = 'mm/hr'
rainrate['comment'] = 'Rain rate calculated from specific_attenuation, R=300.0*specific_attenuation**0.89, note R=0.0 where norm coherent power < 0.4 or rhohv < 0.8'
mask = radar.fields['corrected_reflectivity']['data'].mask
# set masked values to 0 for a nicer plot
radar.fields['rain_rate_A']['data'][np.where(mask)]=0.0

fig = plt.figure(figsize = [15,10])
display.plot_ppi('rain_rate_A', vmin = 0, vmax =150)

fig = plt.figure(figsize = [15,10])
display.plot_ppi('rain_rate_A',sweep=1, vmin = 0, vmax =150)