提交 307b9908 编辑于 作者: Mijian Xu's avatar Mijian Xu 😷
浏览文件

add function for power spectrum

上级 888b1e09
流水线 #1152 已通过 ,包含阶段
in 1 minute 和 19 second
#!/bin/bash
mgpro<< eof
r mag_proj.dat 2000 2000
s
s 0.00015 0.0014
s 0.00015 0.0014
s afdsf sdfd
q
eof
...@@ -19,6 +19,7 @@ class Client(Cmd): ...@@ -19,6 +19,7 @@ class Client(Cmd):
'p': self.do_plot, 'p': self.do_plot,
'q': self.do_quit, 'q': self.do_quit,
'r': self.do_read, 'r': self.do_read,
's': self.do_powspec,
'w': self.do_write, 'w': self.do_write,
'x': self.do_dt2xa, 'x': self.do_dt2xa,
'y': self.do_dt2ya, 'y': self.do_dt2ya,
...@@ -205,6 +206,30 @@ Syntax: plot [data|result] ...@@ -205,6 +206,30 @@ Syntax: plot [data|result]
self.mg.pltmap(self.result) self.mg.pltmap(self.result)
else: else:
print('Error: argument should be in \'data\' and \'result\'') print('Error: argument should be in \'data\' and \'result\'')
def do_powspec(self, arg):
'''
powspec (s): compute power spectrum of the data.
Syntax: powspec [w1, w2]
Specify begining and ending frequency
'''
if self.mg.power is None:
self.mg.power_specf()
if arg == '':
self.mg.plotpower()
elif len(arg.split()) == 2:
try:
w1, w2 = [float(w) for w in arg.split()]
except:
print('arguments should be in float type')
return
if w1 >= w2:
print('w1 should be less than w2')
self.mg.power_fit(w1, w2)
self.mg.plotpower()
else:
print('Only two arguments can be accepted')
def completedefault(self, *args): def completedefault(self, *args):
readline.set_completer_delims(' \t\n;') readline.set_completer_delims(' \t\n;')
......
...@@ -7,6 +7,7 @@ import matplotlib.pyplot as plt ...@@ -7,6 +7,7 @@ import matplotlib.pyplot as plt
import matplotlib.colors as colors import matplotlib.colors as colors
from scipy.fftpack import fft2, fftshift, ifft2, ifftshift from scipy.fftpack import fft2, fftshift, ifft2, ifftshift
from scipy.interpolate import griddata from scipy.interpolate import griddata
from scipy.optimize import curve_fit
import argparse import argparse
np.warnings.filterwarnings('ignore') np.warnings.filterwarnings('ignore')
...@@ -102,6 +103,7 @@ class mgmat(object): ...@@ -102,6 +103,7 @@ class mgmat(object):
self.h = None self.h = None
self.order = None self.order = None
self.result = np.array([]) self.result = np.array([])
self.power = None
def continuation(self, h, order): def continuation(self, h, order):
self.h = h self.h = h
...@@ -178,6 +180,36 @@ class mgmat(object): ...@@ -178,6 +180,36 @@ class mgmat(object):
result = self.data_expand[self.row_begin: self.row_end + 1, self.col_begin: self.col_end + 1] result = self.data_expand[self.row_begin: self.row_end + 1, self.col_begin: self.col_end + 1]
return result return result
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)])
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()
def pltmap(self, data, breakpoint=None): def pltmap(self, data, breakpoint=None):
fig = plt.figure() fig = plt.figure()
...@@ -188,8 +220,8 @@ class mgmat(object): ...@@ -188,8 +220,8 @@ class mgmat(object):
real_width *= frac_w real_width *= frac_w
fig.clf() fig.clf()
ax_raw = fig.gca() ax_raw = fig.gca()
divnorm = colors.DivergingNorm(vcenter=0) # divnorm = colors.DivergingNorm(vcenter=0)
pcm = ax_raw.pcolor(data, cmap='jet', norm=divnorm) pcm = ax_raw.pcolor(data, cmap='jet')
xticklabels = tick2label(ax_raw.get_xticks(), [self.x[0], self.x[-1]]) xticklabels = tick2label(ax_raw.get_xticks(), [self.x[0], self.x[-1]])
yticklabels = tick2label(ax_raw.get_yticks(), [self.y[0], self.y[-1]]) yticklabels = tick2label(ax_raw.get_yticks(), [self.y[0], self.y[-1]])
ax_raw.set_xticklabels(xticklabels) ax_raw.set_xticklabels(xticklabels)
...@@ -217,4 +249,6 @@ class mgmat(object): ...@@ -217,4 +249,6 @@ class mgmat(object):
if __name__ == '__main__': if __name__ == '__main__':
mg = mgmat('/Users/xumj/Codes/MGPro/example/mag_proj.dat', 2000, 2000) mg = mgmat('/Users/xumj/Codes/MGPro/example/mag_proj.dat', 2000, 2000)
mg.pltmap(mg.dt2ya(46, -1)) mg.power_specf()
mg.power_fit(0., 0.00015)
mg.plotpower()
...@@ -5,7 +5,7 @@ packages = find_packages() ...@@ -5,7 +5,7 @@ packages = find_packages()
with open("README.md", "r") as fh: with open("README.md", "r") as fh:
long_description = fh.read() long_description = fh.read()
VERSION = "0.1.13" VERSION = "0.1.14"
setup(name='mgpro', setup(name='mgpro',
version=VERSION, version=VERSION,
author='Mijian Xu', author='Mijian Xu',
......
Supports Markdown
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册