add function for power spectrum

example/run.sh

+#!/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

mgpro/client.py

@@ -19,6 +19,7 @@ class Client(Cmd):
                         'p': self.do_plot,
                         'q': self.do_quit,
                         'r': self.do_read,
+                        's': self.do_powspec,
                         'w': self.do_write,
                         'x': self.do_dt2xa,
                         'y': self.do_dt2ya,
@@ -205,6 +206,30 @@ Syntax: plot [data|result]
             self.mg.pltmap(self.result)
         else:
             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):
         readline.set_completer_delims(' \t\n;')

mgpro/mgpro.py

@@ -7,6 +7,7 @@ import matplotlib.pyplot as plt
 import matplotlib.colors as colors
 from scipy.fftpack import fft2, fftshift, ifft2, ifftshift
 from scipy.interpolate import griddata
+from scipy.optimize import curve_fit
 import argparse
 
 np.warnings.filterwarnings('ignore')
@@ -102,6 +103,7 @@ class mgmat(object):
         self.h = None
         self.order = None
         self.result = np.array([])
+        self.power = None
 
     def continuation(self, h, order):
         self.h = h
@@ -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]
         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):
         fig = plt.figure()
@@ -188,8 +220,8 @@ class mgmat(object):
         real_width *= frac_w
         fig.clf()
         ax_raw = fig.gca()
-        divnorm = colors.DivergingNorm(vcenter=0)
-        pcm = ax_raw.pcolor(data, cmap='jet', norm=divnorm)
+        # divnorm = colors.DivergingNorm(vcenter=0)
+        pcm = ax_raw.pcolor(data, cmap='jet')
         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)
@@ -217,4 +249,6 @@ class mgmat(object):
 
 if __name__ == '__main__':
     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()

setup.py

@@ -5,7 +5,7 @@ packages = find_packages()
 with open("README.md", "r") as fh:
     long_description = fh.read()
 
-VERSION = "0.1.13"
+VERSION = "0.1.14"
 setup(name='mgpro',
       version=VERSION,
       author='Mijian Xu',