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

Mijian Xu's avatar
Mijian Xu committed
13
np.warnings.filterwarnings('ignore')
Mijian Xu's avatar
Mijian Xu committed
14

Mijian Xu's avatar
update  
Mijian Xu committed
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 committed
19

Mijian Xu's avatar
Mijian Xu committed
20

Mijian Xu's avatar
Mijian Xu committed
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 committed
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 committed
37

Mijian Xu's avatar
Mijian Xu committed
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 committed
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 committed
60

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


Mijian Xu's avatar
init UI  
Mijian Xu committed
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 committed
93
class mgmat(object):
Mijian Xu's avatar
Mijian Xu committed
94
    def __init__(self, file, dx, dy, xy_limit=None, to_geo=False):
Mijian Xu's avatar
Mijian Xu committed
95 96
        self.dx = dx
        self.dy = dy
Mijian Xu's avatar
Mijian Xu committed
97
        in_data = np.loadtxt(file)
Mijian Xu's avatar
Mijian Xu committed
98 99
        if to_geo:
            in_data = latlon2geo(in_data)
Mijian Xu's avatar
Mijian Xu committed
100
        self.x, self.y, self.data = xyz2grd(in_data, dx, dy, xy_limit=xy_limit)
Mijian Xu's avatar
Mijian Xu committed
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 committed
103 104
        self.h = None
        self.order = None
Mijian Xu's avatar
Mijian Xu committed
105
        self.result = np.array([])
Mijian Xu's avatar
Mijian Xu committed
106
        self.power = None
Mijian Xu's avatar
Mijian Xu committed
107 108

    def continuation(self, h, order):
Mijian Xu's avatar
init UI  
Mijian Xu committed
109 110
        self.h = h
        self.order = order
Mijian Xu's avatar
dt2za  
Mijian Xu committed
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 committed
115

Mijian Xu's avatar
Mijian Xu committed
116
    def gradient(self, degree):
Mijian Xu's avatar
dt2za  
Mijian Xu committed
117
        data_expand = np.flipud(self.data_expand)
Mijian Xu's avatar
Mijian Xu committed
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 committed
120
            grad /= self.dy
Mijian Xu's avatar
Mijian Xu committed
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 committed
123
            grad /= self.dx
Mijian Xu's avatar
Mijian Xu committed
124
        elif degree == 'mod':
Mijian Xu's avatar
Mijian Xu committed
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 committed
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 committed
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 committed
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 committed
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 committed
140
        grad = np.flipud(grad)*1000
Mijian Xu's avatar
Mijian Xu committed
141
        return grad
Mijian Xu's avatar
::  
Mijian Xu committed
142

143
    def dt2za(self, i0, d0, replace=False):
Mijian Xu's avatar
Mijian Xu committed
144
        P0, Q0, R0 = direction_cos(i0, d0)
Mijian Xu's avatar
dt2za  
Mijian Xu committed
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 committed
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 committed
154 155
        return result

156
    def dt2xa(self, i0, d0, replace=False):
Mijian Xu's avatar
Mijian Xu committed
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 committed
167 168 169
        return result

    
170
    def dt2ya(self, i0, d0, replace=False):
Mijian Xu's avatar
Mijian Xu committed
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 committed
181 182
        return result

Mijian Xu's avatar
Mijian Xu committed
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 committed
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 committed
196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211
        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()
Mijian Xu's avatar
Mijian Xu committed
212
        plt.ylabel('ln(Amp)')
Mijian Xu's avatar
Mijian Xu committed
213 214
        plt.xlabel('Angular frequency')
        plt.show()
Mijian Xu's avatar
Mijian Xu committed
215 216 217

    def pltmap(self, data, breakpoint=None):
        fig = plt.figure()
Mijian Xu's avatar
Mijian Xu committed
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 committed
223
        fig.clf()
Mijian Xu's avatar
Mijian Xu committed
224
        ax_raw = fig.gca()
Mijian Xu's avatar
Mijian Xu committed
225 226
        # divnorm = colors.DivergingNorm(vcenter=0)
        pcm = ax_raw.pcolor(data, cmap='jet')
Mijian Xu's avatar
Mijian Xu committed
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 committed
232
        cb = fig.colorbar(pcm, extend='both')
Mijian Xu's avatar
Mijian Xu committed
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 committed
235
        fig.canvas.draw()
Mijian Xu's avatar
Mijian Xu committed
236
        plt.show()
Mijian Xu's avatar
Mijian Xu committed
237

Mijian Xu's avatar
Mijian Xu committed
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 committed
250 251


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