Sam_Output_Anls/Pcolor_Peaks.py at master · nchaparr/Sam_Output_Anls · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
from __future__ import division
from netCDF4 import Dataset
import glob,os.path
import numpy as np
import numpy.ma as ma
from scipy.interpolate import UnivariateSpline
from matplotlib import cm
from matplotlib import ticker
import matplotlib.pyplot as plt
#import site
#site.addsitedir('/tera/phil/nchaparr/SAM2/sam_main/python')
#from Percentiles import *
from matplotlib.patches import Patch
import sys
#sys.path.insert(0, '/tera/phil/nchaparr/python')
import nchap_fun as nc
from Make_Timelist import *
import warnings
warnings.simplefilter('ignore', np.RankWarning)
#import pywt
from scipy import stats
from datetime import datetime
import fastfit as fsft

"""
    In testing phase -- get_fit() for identifying ML top
    To plot gradient maxima ie BL heights, and w on a 2d horizontal domain,
    and get a histogram or contour plot of BL heigths
    for an individual case
    added function to get ticks and labels based on mean and standard deviation

"""
#TODO: a mess right now.  but can be tidied up once regression code is included

def get_ticks(mean, stddev, max, min):

     """

     gets ticks and tick lavels for contour plot based on mean and standard deviation

     Arguments:
     mean, stddev, max, min

     Returns:
     ticks, tick_labels

     """
     tick_list = []
     label_list = []
     int1=int(np.ceil((mean-min)/stddev))
     int2=int(np.ceil((max-mean)/stddev))


     for i in range(int1):
         if int1==1:
             tick_list.append(min)
             label_list.append(r'$\mu - %.1f \sigma$' %((mean-min)/stddev))

         elif i > 0:
             tick_list.append(mean - (int1-i)*stddev)

             label_list.append(r'$\mu - %.1f \sigma$' %(int1-i))

         #else:
             #tick_list.append(min)

             #label_list.append(r'$\mu - %.1f \sigma$' %((mean-min)/stddev))

     tick_list.append(mean)
     label_list.append(r'$\mu$')


     for i in range(int2):

         if int2==1:
             tick_list.append(max)

             label_list.append(r'$\mu + %.1f \sigma$' %((max-mean)/stddev))

         elif i< int2-1:
             tick_list.append(mean + (i+1)*stddev)

             label_list.append(r'$\mu + %.1f \sigma$' %(i+1))

         #else:
             #tick_list.append(max)

             #label_list.append(r'$\mu + %.1f \sigma$' %((max-mean)/stddev))

     return label_list, tick_list


def get_fit(theta, height):
     """
        Fitting the local theta profile with three lines

     """

     fitvals = np.zeros_like(theta)
     RSS = np.empty((290, 290))+ np.nan
     print RSS[0,0]
     for j in range(290):
          if j > 2:
               for k in range(290):
                    if k>j+1 and k<289:
                         b_1 = (np.sum(np.multiply(height[:j], theta[:j])) - 1/j*np.sum(height[:j])*np.sum(theta[:j]))/(np.sum(height[:j]**2) - 1/j*np.sum(height[:j])**2)
                         a_1 = np.sum(np.multiply(height[:j], theta[:j]))/np.sum(height[:j]) - b_1*np.sum(height[:j]**2)/np.sum(height[:j])

                         b_2 = (np.sum(theta[j:k]) - (k-j)*(a_1+b_1*height[j]))/(np.sum(height[j:k]) - (k-j)*height[j])

                         a_2 = np.sum(np.multiply(height[j:k], theta[j:k]))/np.sum(height[j:k]) - b_2*np.sum(height[j:k]**2)/np.sum(height[j:k])

                         b_3 = (np.sum(theta[k:290]) - (290-k)*(a_2+b_2*height[k]))/(np.sum(height[k:290]) - (290-k)*height[k])
                         a_3 = np.sum(np.multiply(height[k:290], theta[k:290]))/np.sum(height[k:290]) - b_3*np.sum(height[k:290]**2)/np.sum(height[k:290])

                         RSS[j, k] = np.sum(np.add(theta[2:j], -(a_1+ b_1*height[2:j]))**2) + np.sum(np.add(theta[j:k], -(a_2+ b_2*height[j:k]))**2) + np.sum(np.add(theta[k:290], -(a_3+ b_3*height[k:290]))**2)


     RSS = ma.masked_where(np.isnan(RSS), RSS)
     [j, k] = np.unravel_index(ma.argmin(RSS), RSS.shape)

     b_1 = (np.sum(np.multiply(height[:j], theta[:j])) - 1/j*np.sum(height[:j]*np.sum(theta[:j])))/(np.sum(height[:j]**2) - 1/j*np.sum(height[2:j])**2)
     a_1 = np.sum(np.multiply(height[:j], theta[:j]))/np.sum(height[:j]) - b_1*np.sum(height[:j]**2)/np.sum(height[:j])

     b_2 = (np.sum(theta[j:k]) - (k-j)*(a_1+b_1*height[j]))/(np.sum(height[j:k]) - (k-j)*height[j])
     a_2 = np.sum(np.multiply(height[j:k], theta[j:k]))/np.sum(height[j:k]) - b_2*np.sum(height[j:k]**2)/np.sum(height[j:k])

     b_3 = (np.sum(theta[k:290]) - (290-k)*(a_2+b_2*height[k]))/(np.sum(height[k:290]) - (290-k)*height[k])
     a_3 = np.sum(np.multiply(height[k:290], theta[k:290]))/np.sum(height[k:290]) - b_3*np.sum(height[k:290]**2)/np.sum(height[k:290])


     fitvals[:j] = b_1*height[:j] + a_1
     fitvals[j:k] = b_2*height[j:k] + a_2
     fitvals[k:290] = b_3*height[k:290] + a_3


     return fitvals, RSS, j, k


#Lists of times relating to output (nc) files
dump_time_list, time_hrs = Make_Timelists(1, 600, 28800)
dump_time = dump_time_list[11]
print dump_time

for k in range(1):
     #getting variables from nc files
     [wvels, theta, tracer, height] = nc.Get_Var_Arrays("/tera2/nchaparr/Mar52014/runs/sam_case", "/OUT_3D/keep/NCHAPP1_testing_doscamiopdata_24_", dump_time, k+1)

     #getting points of maximum theta gradient, getting rid of this soon
     #[dvardz, grad_peaks] = nc.Domain_Grad(theta, height)
     #tops_indices=np.where(np.abs(grad_peaks - 1400)<10)

     #choosing one horizontal point
     for i in range(1):
          #top_index = [tops_indices[0][i], tops_indices[1][i]]
          #[i, j] = top_index
          [i, j] = [50, 50]
          thetavals = theta[:, i, j]

          startTime = datetime.now()
          #print 'Start', startTime#1
          top = np.where(np.abs(height-2300)<100)[0][0]
          print top, height[top]
          RSS, J, K = fsft.get_fit(thetavals, height, top)
          #print J, height[J]
          #print 'RSS time', (datetime.now()-startTime)
          fitvals = np.zeros_like(thetavals[:top])
          b_1 = (np.sum(np.multiply(height[9:J], thetavals[9:J])) - 1.0/(J-9)*np.sum(height[9:J]*np.sum(thetavals[9:J])))/(np.sum(height[9:J]**2) - 1.0/(J-9)*np.sum(height[9:J])**2)
          #print np.sum(np.multiply(height[9:J], thetavals[9:J])),  - 1.0/(J-9)*np.sum(height[9:J]*np.sum(thetavals[9:J])), np.sum(height[9:J]**2), - 1.0/(J-9)*np.sum(height[9:J])**2

          a_1 = np.sum(np.multiply(height[9:J], thetavals[9:J]))/np.sum(height[9:J]) - b_1*np.sum(height[9:J]**2)/np.sum(height[9:J])

          b_2 = (np.sum(thetavals[J:K]) - (K-J)*(a_1+b_1*height[J]))/(np.sum(height[J:K]) - (K-J)*height[J])
          a_2 = np.sum(np.multiply(height[J:K], thetavals[J:K]))/np.sum(height[J:K]) - b_2*np.sum(height[J:K]**2)/np.sum(height[J:K])

          b_3 = (np.sum(thetavals[K:top]) - (top-K)*(a_2+b_2*height[K]))/(np.sum(height[K:top]) - (top-K)*height[K])
          a_3 = np.sum(np.multiply(height[K:top], thetavals[K:top]))/np.sum(height[K:top]) - b_3*np.sum(height[K:top]**2)/np.sum(height[K:top])
          #print b_2, b_3

          fitvals[:J] = b_1*height[:J] + a_1
          fitvals[J:K] = b_2*height[J:K] + a_2
          fitvals[K:top] = b_3*height[K:top] + a_3

          #set up plot
          theFig = plt.figure(i)
          theFig.clf()
          theAx = theFig.add_subplot(121)
          theAx.set_title('Fit')
          theAx.set_xlabel(r'$\overline{\theta} (K)$')
          theAx.set_ylabel('z (m)')

          theAx1 = theFig.add_subplot(122)
          theAx1.set_title('Profile and Fit')
          theAx1.set_xlabel(r'$\overline{\theta} (K) $')
          theAx1.set_ylabel('z (m)')

          theAx1.plot(thetavals, height[:], 'wo')
          theAx.plot(fitvals[:J], height[:J], 'r-')
          theAx.plot(fitvals[J:K], height[J:K], 'b-')
          theAx.plot(fitvals[K:top], height[K:top], 'g-')
          theAx1.plot(fitvals[:top], height[:top], 'r-')

theAx1.set_xlim(300, 320)
theAx1.set_ylim(0, 2000)
theAx.set_ylim(0, 2000)
theAx.set_xlim(300, 320)
plt.show()