mgpro.py 10.3 KB
Newer Older
Mijian Xu's avatar
Mijian Xu 已提交
1
import numpy as np
Mijian Xu's avatar
bug fix    
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
from scipy.optimize import curve_fit
Mijian Xu's avatar
Mijian Xu 已提交
11
import argparse
Mijian Xu's avatar
Mijian Xu 已提交
12

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

Mijian Xu's avatar
update    
Mijian Xu 已提交
15
16
17
18
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 已提交
19

Mijian Xu's avatar
Mijian Xu 已提交
20

Mijian Xu's avatar
Mijian Xu 已提交
21
22
23
24
25
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 已提交
26
27
28
29
30
31
32
33
34
35
36
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 已提交
37

Mijian Xu's avatar
Mijian Xu 已提交
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
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 已提交
54
55
56
57
58
59
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 已提交
60

Mijian Xu's avatar
dt2za    
Mijian Xu 已提交
61
def norm_uv(data, dx, dy):
Mijian Xu's avatar
Mijian Xu 已提交
62
    (len_row, len_col) = data.shape
Mijian Xu's avatar
dt2za    
Mijian Xu 已提交
63
64
65
66
67
68
69
    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 已提交
70
71
    u = (col_mesh - col0)*dom_col
    v = (row_mesh - row0)*dom_row
Mijian Xu's avatar
dt2za    
Mijian Xu 已提交
72
73
74
    return u, v


Mijian Xu's avatar
init UI    
Mijian Xu 已提交
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
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 已提交
93
class mgmat(object):
Mijian Xu's avatar
Mijian Xu 已提交
94
    def __init__(self, file, dx, dy, xy_limit=None, to_geo=False):
Mijian Xu's avatar
Mijian Xu 已提交
95
96
        self.dx = dx
        self.dy = dy
Mijian Xu's avatar
Mijian Xu 已提交
97
        in_data = np.loadtxt(file)
Mijian Xu's avatar
Mijian Xu 已提交
98
99
        if to_geo:
            in_data = latlon2geo(in_data)
Mijian Xu's avatar
Mijian Xu 已提交
100
        self.x, self.y, self.data = xyz2grd(in_data, dx, dy, xy_limit=xy_limit)
Mijian Xu's avatar
Mijian Xu 已提交
101
102
        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 已提交
103
104
        self.h = None
        self.order = None
Mijian Xu's avatar
Mijian Xu 已提交
105
        self.result = np.array([])
Mijian Xu's avatar
Mijian Xu 已提交
106
        self.power = None
Mijian Xu's avatar
Mijian Xu 已提交
107
108

    def continuation(self, h, order):
Mijian Xu's avatar
init UI    
Mijian Xu 已提交
109
110
        self.h = h
        self.order = order
Mijian Xu's avatar
dt2za    
Mijian Xu 已提交
111
112
113
114
        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 已提交
115

Mijian Xu's avatar
Mijian Xu 已提交
116
    def gradient(self, degree):
Mijian Xu's avatar
dt2za    
Mijian Xu 已提交
117
        data_expand = np.flipud(self.data_expand)
Mijian Xu's avatar
Mijian Xu 已提交
118
119
        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 已提交
120
            grad /= self.dy
Mijian Xu's avatar
Mijian Xu 已提交
121
122
        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 已提交
123
            grad /= self.dx
Mijian Xu's avatar
Mijian Xu 已提交
124
        elif degree == 'mod':
Mijian Xu's avatar
Mijian Xu 已提交
125
126
            ns = -np.diff(data_expand, axis=0)[self.row_begin: self.row_end + 1, self.col_begin: self.col_end + 1] / self.dy
            we = np.diff(data_expand, axis=1)[self.row_begin: self.row_end + 1, self.col_begin: self.col_end + 1] / self.dx
Mijian Xu's avatar
Mijian Xu 已提交
127
128
129
130
131
132
            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 已提交
133
            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 已提交
134
135
136
137
138
        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 已提交
139
            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 已提交
140
        grad = np.flipud(grad)*1000
Mijian Xu's avatar
Mijian Xu 已提交
141
        return grad
Mijian Xu's avatar
::    
Mijian Xu 已提交
142

143
    def dt2za(self, i0, d0, replace=False):
Mijian Xu's avatar
Mijian Xu 已提交
144
        P0, Q0, R0 = direction_cos(i0, d0)
Mijian Xu's avatar
dt2za    
Mijian Xu 已提交
145
146
        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 已提交
147
        phi[np.where(np.isnan(phi))] = 0 + 0j
148
149
150
151
152
153
        if not replace:
            result = np.real(ifft2(ifftshift(phi * self.data_sf)))[self.row_begin: self.row_end + 1, self.col_begin: self.col_end + 1]
        else:
            self.data_sf = phi * self.data_sf
            self.data_expand = np.real(ifft2(ifftshift(self.data_sf)))
            result = self.data_expand[self.row_begin: self.row_end + 1, self.col_begin: self.col_end + 1]
Mijian Xu's avatar
dt2za    
Mijian Xu 已提交
154
155
        return result

156
    def dt2xa(self, i0, d0, replace=False):
Mijian Xu's avatar
Mijian Xu 已提交
157
158
159
160
        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
161
162
163
164
165
166
        if not replace:
            result = np.real(ifft2(ifftshift(trans_func * self.data_sf)))[self.row_begin: self.row_end + 1, self.col_begin: self.col_end + 1]
        else:
            self.data_sf = trans_func * self.data_sf
            self.data_expand = np.real(ifft2(ifftshift(self.data_sf)))
            result = self.data_expand[self.row_begin: self.row_end + 1, self.col_begin: self.col_end + 1]
Mijian Xu's avatar
Mijian Xu 已提交
167
168
169
        return result

    
170
    def dt2ya(self, i0, d0, replace=False):
Mijian Xu's avatar
Mijian Xu 已提交
171
172
173
174
        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
175
176
177
178
179
180
        if not replace:
            result = np.real(ifft2(ifftshift(trans_func * self.data_sf)))[self.row_begin: self.row_end + 1, self.col_begin: self.col_end + 1]
        else:
            self.data_sf = trans_func * self.data_sf
            self.data_expand = np.real(ifft2(ifftshift(self.data_sf)))
            result = self.data_expand[self.row_begin: self.row_end + 1, self.col_begin: self.col_end + 1]
Mijian Xu's avatar
Mijian Xu 已提交
181
182
        return result

Mijian Xu's avatar
Mijian Xu 已提交
183
184
185
186
187
188
189
190
191
192
193
    def power_specf(self):
        self.pcot = None
        sf_amp = np.abs(self.data_sf)
        u, v = norm_uv(self.data_sf, self.dx, self.dy)
        x_freq = u[0, :]
        y_freq = v[:, 0]
        omega = np.sqrt(u**2 + v**2)
        radfreq = np.sort(np.sqrt(x_freq**2 + y_freq**2))
        power = np.zeros_like(radfreq)
        for i, frq in enumerate(radfreq):
            power[i] = np.mean(sf_amp[np.where(omega == frq)])
Mijian Xu's avatar
Mijian Xu 已提交
194
195
        power = power[np.where(np.logical_not(np.isnan(power)))]
        radfreq = radfreq[np.where(np.logical_not(np.isnan(power)))]
Mijian Xu's avatar
Mijian Xu 已提交
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
        self.power = np.vstack((radfreq, np.log(power)))
    
    def power_fit(self, freqmin, freqmax):
        eff_idx = np.where((self.power[0] > freqmin) & (self.power[0] < freqmax))
        self.eff_freq = self.power[0][eff_idx]
        eff_amp = self.power[1][eff_idx]
        self.pcot, _ = curve_fit(lambda xdata, a, b: a*xdata+b, self.eff_freq, eff_amp)

    def plotpower(self):
        fig = plt.figure()
        fig.clf()
        plt.grid(True)
        plt.plot(self.power[0, :], self.power[1, :], 'r.', ms=4, alpha=0.7)
        if self.pcot is not None:
            line = plt.plot(self.eff_freq, self.eff_freq*self.pcot[0]+self.pcot[1], 'k', label=r"Z = {:.2f}m".format(self.pcot[0]))
            plt.legend()
        plt.ylabel('Amplitude')
        plt.xlabel('Angular frequency')
        plt.show()
Mijian Xu's avatar
Mijian Xu 已提交
215
216
217

    def pltmap(self, data, breakpoint=None):
        fig = plt.figure()
Mijian Xu's avatar
Mijian Xu 已提交
218
219
220
221
222
        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 已提交
223
        fig.clf()
Mijian Xu's avatar
Mijian Xu 已提交
224
        ax_raw = fig.gca()
Mijian Xu's avatar
Mijian Xu 已提交
225
226
        # divnorm = colors.DivergingNorm(vcenter=0)
        pcm = ax_raw.pcolor(data, cmap='jet')
Mijian Xu's avatar
Mijian Xu 已提交
227
228
229
230
231
        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 已提交
232
        cb = fig.colorbar(pcm, extend='both')
Mijian Xu's avatar
Mijian Xu 已提交
233
234
        # 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 已提交
235
        fig.canvas.draw()
Mijian Xu's avatar
Mijian Xu 已提交
236
        plt.show()
Mijian Xu's avatar
Mijian Xu 已提交
237

Mijian Xu's avatar
Mijian Xu 已提交
238
239
240
241
242
243
244
245
246
247
248
249
    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 已提交
250
251


Mijian Xu's avatar
Mijian Xu 已提交
252
if __name__ == '__main__':
Mijian Xu's avatar
Mijian Xu 已提交
253
    mg = mgmat('/Users/xumj/Codes/MGPro/example/mag_proj.dat', 2000, 2000)
Mijian Xu's avatar
Mijian Xu 已提交
254
255
256
    mg.power_specf()
    mg.power_fit(0., 0.00015)
    mg.plotpower()