Add functions for saving Za to memory

mgpro/client.py

@@ -102,41 +102,65 @@ horizontal gradient in azimuth of 0, 45, 90 or 135 or module of horizontal gradi
         '''
 dt2za (z): Vertical magnetization
 
-Syntax: dt2za i0 d0
+Syntax: dt2za i0 d0 ['mem']
     i0: Magnetic dip of the region
     d0: Magnetic declination of the region
+    mem: Switch for saving Za to memory and replacing the data read before
         '''
         if len(arg.split()) == 2:
             i0, d0 = [float(value) for value in arg.split()]
             self.result = self.mg.dt2za(i0, d0)
+        elif len(arg.split()) == 3:
+            if arg.split()[2].lower() == 'mem':
+                i0, d0 = [float(value) for value in arg.split()[0:2]]
+                self.result = self.mg.dt2za(i0, d0, replace=True)
             else:
-            print('Two argumrnts required')
+                print('Error: \'mem\' should be specified to save Za to memory')
+                return
+        else:
+            print('Two or three argumrnts required')
 
     def do_dt2xa(self, arg):
         '''
 dt2xa (x): Horizontal magnetization along x direction
 
-Syntax: dt2xa i0 d0
+Syntax: dt2xa i0 d0 ['mem']
     i0: Magnetic dip of the region
     d0: Magnetic declination of the region
+    mem: Switch for saving Xa to memory and replacing the data read before
         '''
         if len(arg.split()) == 2:
             i0, d0 = [float(value) for value in arg.split()]
             self.result = self.mg.dt2xa(i0, d0)
+        elif len(arg.split()) == 3:
+            if arg.split()[2].lower() == 'mem':
+                i0, d0 = [float(value) for value in arg.split()[0:2]]
+                self.result = self.mg.dt2xa(i0, d0, replace=True)
             else:
-            print('Two argumrnts required')
+                print('Error: \'mem\' should be specified to save Xa to memory')
+                return
+        else:
+            print('Two or three argumrnts required')
 
     def do_dt2ya(self, arg):
         '''
 dt2xa (y): Horizontal magnetization along y direction
 
-Syntax: dt2ya i0 d0
+Syntax: dt2ya i0 d0 ['mem']
     i0: Magnetic dip of the region
     d0: Magnetic declination of the region
+    mem: Switch for saving Ya to memory and replacing the data read before
         '''
         if len(arg.split()) == 2:
             i0, d0 = [float(value) for value in arg.split()]
             self.result = self.mg.dt2ya(i0, d0)
+        elif len(arg.split()) == 3:
+            if arg.split()[2].lower() == 'mem':
+                i0, d0 = [float(value) for value in arg.split()[0:2]]
+                self.result = self.mg.dt2ya(i0, d0, replace=True)
+            else:
+                print('Error: \'mem\' should be specified to save Ya to memory')
+                return
         else:
             print('Two argumrnts required')
 
@@ -172,9 +196,15 @@ Syntax: write filename [\'lalo\']
         '''
 plot (p): Preview the result.
 
-Syntax: plot
+Syntax: plot [data|result]
+    Specify 'data' or 'result' to draw map view of the raw data or the result (default)
         ''' 
+        if arg.lower() == 'data':
+            self.mg.pltmap(self.mg.data_expand[self.mg.row_begin: self.mg.row_end + 1, self.mg.col_begin: self.mg.col_end + 1])
+        elif  arg.lower() == 'result' or arg == '':
             self.mg.pltmap(self.result)
+        else:
+            print('Error: argument should be in \'data\' and \'result\'')
 
     def completedefault(self, *args):
         readline.set_completer_delims(' \t\n;')

mgpro/mgpro.py

@@ -138,29 +138,44 @@ class mgmat(object):
         grad = np.flipud(grad)*1000
         return grad
 
-    def dt2za(self, i0, d0):
+    def dt2za(self, i0, d0, replace=False):
         P0, Q0, R0 = direction_cos(i0, d0)
         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))
         phi[np.where(np.isnan(phi))] = 0 + 0j
+        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]
         return result
 
-    def dt2xa(self, i0, d0):
+    def dt2xa(self, i0, d0, replace=False):
         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
+        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]
         return result
 
     
-    def dt2ya(self, i0, d0):
+    def dt2ya(self, i0, d0, replace=False):
         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
+        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]
         return result
 
 

setup.py

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