mgpro.py 8.0 KB
Newer Older
Mijian Xu's avatar
Mijian Xu 已提交
1
import numpy as np
Mijian Xu's avatar
Mijian Xu 已提交
2
3
4
5
# from mgpro import expand
import expand
# from mgpro.proj import *
from proj import *
Mijian Xu's avatar
Mijian Xu 已提交
6
import matplotlib.pyplot as plt
Mijian Xu's avatar
init UI    
Mijian Xu 已提交
7
import matplotlib.colors as colors
Mijian Xu's avatar
Mijian Xu 已提交
8
from scipy.fftpack import fft2, fftshift, ifft2, ifftshift
Mijian Xu's avatar
Mijian Xu 已提交
9
from scipy.interpolate import griddata
Mijian Xu's avatar
Mijian Xu 已提交
10
import argparse
Mijian Xu's avatar
Mijian Xu 已提交
11

Mijian Xu's avatar
Mijian Xu 已提交
12
np.warnings.filterwarnings('ignore')
Mijian Xu's avatar
Mijian Xu 已提交
13

Mijian Xu's avatar
update    
Mijian Xu 已提交
14
15
16
17
def msg():
    return '''
    Usage: mgpro.py -c<conti>/<diff> -d <dx>/<dy> -o <out_path> [-R <xmin/<xmax>/<ymix>/<ymax>]
    '''
Mijian Xu's avatar
Mijian Xu 已提交
18

Mijian Xu's avatar
Mijian Xu 已提交
19

Mijian Xu's avatar
Mijian Xu 已提交
20
21
22
23
24
def tick2label(ticks, lim):
    labels = ['{:.0f}'.format(tick) for tick in np.linspace(lim[0], lim[1], len(ticks))]
    return labels


Mijian Xu's avatar
Mijian Xu 已提交
25
26
27
28
29
30
31
32
33
34
35
def cal_pos(max_len, len_x, len_y):
    if max_len > 1 or max_len < 0:
        raise ValueError('Input arg of max_len must be in the range of 0 to 1')
    if len_x > len_y:
        width = max_len
        hight = max_len * (len_y / len_x)
    else:
        hight = max_len
        width = max_len * (len_x / len_y)
    return width, hight

Mijian Xu's avatar
Mijian Xu 已提交
36

Mijian Xu's avatar
Mijian Xu 已提交
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
def xyz2grd(raw_data, dx, dy, xy_limit=None):
    data = raw_data.copy()
    data[:, [0, 1]] = data[:, [1, 0]]
    if xy_limit is None:
        xmin = np.min(raw_data[:, 0])
        xmax = np.max(raw_data[:, 0])
        ymin = np.min(raw_data[:, 1])
        ymax = np.max(raw_data[:, 1])
    else:
        xmin, xmax, ymin, ymax = xy_limit
    xrange = np.arange(xmin, xmax, dx)
    yrange = np.arange(ymin, ymax, dy)
    ymesh, xmesh = np.meshgrid(yrange, xrange, indexing='ij')
    grid_data = griddata(data[:, 0:2], data[:, 2], (ymesh, xmesh))
    return xrange, yrange, grid_data

Mijian Xu's avatar
Mijian Xu 已提交
53
54
55
56
57
58
def direction_cos(i0, d0):
    P0 = np.cos(np.deg2rad(i0)) * np.cos(np.deg2rad(d0))
    Q0 = np.cos(np.deg2rad(i0)) * np.sin(np.deg2rad(d0))
    R0 = np.sin(np.deg2rad(i0))
    return P0, Q0, R0

Mijian Xu's avatar
Mijian Xu 已提交
59

Mijian Xu's avatar
dt2za    
Mijian Xu 已提交
60
def norm_uv(data, dx, dy):
Mijian Xu's avatar
Mijian Xu 已提交
61
    (len_row, len_col) = data.shape
Mijian Xu's avatar
dt2za    
Mijian Xu 已提交
62
63
64
65
66
67
68
    dom_row = 2 * np.pi / len_row / dy
    dom_col = 2 * np.pi / len_col / dx
    row = np.arange(1, len_row+1)
    col = np.arange(1, len_col+1)
    row0 = int(len_row / 2) + 1
    col0 = int(len_col / 2) + 1
    (col_mesh, row_mesh) = np.meshgrid(col, row)
Mijian Xu's avatar
Mijian Xu 已提交
69
70
    u = (col_mesh - col0)*dom_col
    v = (row_mesh - row0)*dom_row
Mijian Xu's avatar
dt2za    
Mijian Xu 已提交
71
72
73
    return u, v


Mijian Xu's avatar
init UI    
Mijian Xu 已提交
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
class JetNormalize(colors.Normalize):
    def __init__(self, vmin=None, vmax=None, midpoint=None, clip=False):
        if midpoint is None:
            self.midpoint = [0.2, 0.4, 0.6, 0.8]
        else:
            try:
                self.midpoint = list(midpoint)
            except Exception as e:
                raise ValueError(e)
        colors.Normalize.__init__(self, vmin, vmax, clip)

    def __call__(self, value, clip=None):
        # I'm ignoring masked values and all kinds of edge cases to make a
        # simple example...
        x, y = [self.vmin] + self.midpoint + [self.vmax], np.linspace(0, 1, len(self.midpoint)+2)
        return np.ma.masked_array(np.interp(value, x, y))


Mijian Xu's avatar
Mijian Xu 已提交
92
class mgmat(object):
Mijian Xu's avatar
Mijian Xu 已提交
93
    def __init__(self, file, dx, dy, xy_limit=None, to_geo=False):
Mijian Xu's avatar
Mijian Xu 已提交
94
95
        self.dx = dx
        self.dy = dy
Mijian Xu's avatar
Mijian Xu 已提交
96
        in_data = np.loadtxt(file)
Mijian Xu's avatar
Mijian Xu 已提交
97
98
        if to_geo:
            in_data = latlon2geo(in_data)
Mijian Xu's avatar
Mijian Xu 已提交
99
        self.x, self.y, self.data = xyz2grd(in_data, dx, dy, xy_limit=xy_limit)
Mijian Xu's avatar
Mijian Xu 已提交
100
101
        self.data_expand, self.row_begin, self.row_end, self.col_begin, self.col_end = expand.expand(self.data)
        self.data_sf = fftshift(fft2(self.data_expand))
Mijian Xu's avatar
init UI    
Mijian Xu 已提交
102
103
        self.h = None
        self.order = None
Mijian Xu's avatar
Mijian Xu 已提交
104
        self.result = np.array([])
Mijian Xu's avatar
Mijian Xu 已提交
105
106

    def continuation(self, h, order):
Mijian Xu's avatar
init UI    
Mijian Xu 已提交
107
108
        self.h = h
        self.order = order
Mijian Xu's avatar
dt2za    
Mijian Xu 已提交
109
110
111
112
        u, v = norm_uv(self.data_sf, self.dx, self.dy)
        H = np.sqrt(u**2+v**2) ** order * np.exp(h * np.sqrt(u**2+v**2))
        result = np.real(ifft2(ifftshift(H * self.data_sf)))[self.row_begin: self.row_end + 1, self.col_begin: self.col_end + 1]
        return result
Mijian Xu's avatar
Mijian Xu 已提交
113

Mijian Xu's avatar
Mijian Xu 已提交
114
    def gradient(self, degree):
Mijian Xu's avatar
dt2za    
Mijian Xu 已提交
115
        data_expand = np.flipud(self.data_expand)
Mijian Xu's avatar
Mijian Xu 已提交
116
117
        if degree == '0':
            grad = -np.diff(data_expand, axis=0)[self.row_begin: self.row_end + 1, self.col_begin: self.col_end + 1]
Mijian Xu's avatar
Mijian Xu 已提交
118
            grad /= self.dy
Mijian Xu's avatar
Mijian Xu 已提交
119
120
        elif degree == '90':
            grad = np.diff(data_expand, axis=1)[self.row_begin: self.row_end + 1, self.col_begin: self.col_end + 1]
Mijian Xu's avatar
Mijian Xu 已提交
121
            grad /= self.dx
Mijian Xu's avatar
Mijian Xu 已提交
122
123
124
125
126
127
128
129
130
        elif degree == 'mod':
            ns = -np.diff(data_expand, axis=0)[self.row_begin: self.row_end + 1, self.col_begin: self.col_end + 1]
            we = np.diff(data_expand, axis=1)[self.row_begin: self.row_end + 1, self.col_begin: self.col_end + 1]
            grad = np.sqrt(ns**2 + we**2)
        elif degree == '45':
            grad = np.zeros_like(data_expand)
            for _i in range(self.row_begin, self.row_end + 1):
                for _j in range(self.col_begin, self.col_end + 1):
                    grad[_i, _j] = data_expand[_i, _j] - data_expand[_i-1, _j+1]
Mijian Xu's avatar
Mijian Xu 已提交
131
            grad = grad[self.row_begin: self.row_end + 1, self.col_begin: self.col_end + 1] / np.sqrt(self.dx**2 + self.dy**2)
Mijian Xu's avatar
Mijian Xu 已提交
132
133
134
135
136
        elif degree == '135':
            grad = np.zeros_like(data_expand)
            for _i in range(self.row_begin, self.row_end + 1):
                for _j in range(self.col_begin, self.col_end + 1):
                    grad[_i, _j] = data_expand[_i, _j] - data_expand[_i-1, _j-1]
Mijian Xu's avatar
Mijian Xu 已提交
137
            grad = grad[self.row_begin: self.row_end + 1, self.col_begin: self.col_end + 1] / np.sqrt(self.dx**2 + self.dy**2)
Mijian Xu's avatar
Mijian Xu 已提交
138
        grad = np.flipud(grad)*1000
Mijian Xu's avatar
Mijian Xu 已提交
139
        return grad
Mijian Xu's avatar
::    
Mijian Xu 已提交
140

Mijian Xu's avatar
dt2za    
Mijian Xu 已提交
141
    def dt2za(self, i0, d0):
Mijian Xu's avatar
Mijian Xu 已提交
142
        P0, Q0, R0 = direction_cos(i0, d0)
Mijian Xu's avatar
dt2za    
Mijian Xu 已提交
143
144
        u, v = norm_uv(self.data_sf, self.dx, self.dy)
        phi = np.sqrt(u**2 + v**2) / ((P0*u + Q0*v)*1j + R0*np.sqrt(u**2 + v**2))
Mijian Xu's avatar
Mijian Xu 已提交
145
        phi[np.where(np.isnan(phi))] = 0 + 0j
Mijian Xu's avatar
dt2za    
Mijian Xu 已提交
146
147
148
        result = np.real(ifft2(ifftshift(phi * self.data_sf)))[self.row_begin: self.row_end + 1, self.col_begin: self.col_end + 1]
        return result

Mijian Xu's avatar
Mijian Xu 已提交
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
    def dt2xa(self, i0, d0):
        P0, Q0, R0 = direction_cos(i0, d0)
        u, v = norm_uv(self.data_sf, self.dx, self.dy)
        trans_func = u*1j / ((P0*u + Q0*v)*1j + R0*np.sqrt(u**2 + v**2))
        trans_func[np.where(np.isnan(trans_func))] = 0 + 0j
        result = np.real(ifft2(ifftshift(trans_func * self.data_sf)))[self.row_begin: self.row_end + 1, self.col_begin: self.col_end + 1]
        return result

    
    def dt2ya(self, i0, d0):
        P0, Q0, R0 = direction_cos(i0, d0)
        u, v = norm_uv(self.data_sf, self.dx, self.dy)
        trans_func = v*1j / ((P0*u + Q0*v)*1j + R0*np.sqrt(u**2 + v**2))
        trans_func[np.where(np.isnan(trans_func))] = 0 + 0j
        result = np.real(ifft2(ifftshift(trans_func * self.data_sf)))[self.row_begin: self.row_end + 1, self.col_begin: self.col_end + 1]
        return result


    def pltmap(self, data, breakpoint=None):
        fig = plt.figure()
Mijian Xu's avatar
Mijian Xu 已提交
169
170
171
172
173
        f_width = fig.get_figwidth()
        f_height = fig.get_figheight()
        frac_w = min([f_height, f_width])/max([f_height, f_width])
        real_width, real_height = cal_pos(0.7, data.shape[1], data.shape[0])
        real_width *= frac_w
xumj's avatar
xumj 已提交
174
        fig.clf()
Mijian Xu's avatar
Mijian Xu 已提交
175
        ax_raw = fig.gca()
Mijian Xu's avatar
Mijian Xu 已提交
176
177
178
179
180
181
182
        divnorm = colors.DivergingNorm(vcenter=0)
        pcm = ax_raw.pcolor(data, cmap='jet', norm=divnorm)
        xticklabels = tick2label(ax_raw.get_xticks(), [self.x[0], self.x[-1]])
        yticklabels = tick2label(ax_raw.get_yticks(), [self.y[0], self.y[-1]])
        ax_raw.set_xticklabels(xticklabels)
        ax_raw.set_yticklabels(yticklabels)
        plt.setp(ax_raw.get_xticklabels(), rotation=-45, ha="left")
Mijian Xu's avatar
Mijian Xu 已提交
183
        cb = fig.colorbar(pcm, extend='both')
Mijian Xu's avatar
Mijian Xu 已提交
184
185
        # ax_raw.set_position([.1, .125, real_width, real_height], which='original')
        # cb.ax.set_position([.8, .1, real_height/6.27, real_height])
Mijian Xu's avatar
Mijian Xu 已提交
186
        fig.canvas.draw()
Mijian Xu's avatar
Mijian Xu 已提交
187
        plt.show()
Mijian Xu's avatar
Mijian Xu 已提交
188

Mijian Xu's avatar
Mijian Xu 已提交
189
190
191
192
193
194
195
196
197
198
199
200
    def savetxt(self, filename, result, to_latlon=False):
        points = np.zeros([result.size, 3])
        m = 0
        for i, x in enumerate(self.x):
            for j, y in enumerate(self.y):
                points[m, 0] = x
                points[m, 1] = y
                points[m, 2] = result[j, i]
                m += 1
        if to_latlon:
            points = geo2latlon(points)
        np.savetxt(filename, points, fmt='%.4f %.4f %.6f')
Mijian Xu's avatar
Mijian Xu 已提交
201
202


Mijian Xu's avatar
Mijian Xu 已提交
203
if __name__ == '__main__':
Mijian Xu's avatar
Mijian Xu 已提交
204
    mg = mgmat('/Users/xumj/Codes/MGPro/example/mag_proj.dat', 2000, 2000)
Mijian Xu's avatar
Mijian Xu 已提交
205
    mg.pltmap(mg.dt2ya(46, -1))