Gradient plotting worked. Problems with adding scalar fieldplotting

Lewis-Hou · May 17, 2018 · afba80e · afba80e
1 parent 5353a9f
commit afba80e
Show file tree

Hide file tree

Showing 3 changed files with 1,088 additions and 88 deletions.
diff --git a/gempy/gempy_front.py b/gempy/gempy_front.py
@@ -892,6 +892,13 @@ def plot_scalar_field(geo_data, potential_field, cell_number, N=20,
                               direction=direction,  plot_data=plot_data, series=series,
                               *args, **kwargs)
 
+def plot_gradient(geo_data, scalar_field, gx, gy, gz, cell_number, r=5,
+                      direction="y", plot_scalar=True, **kwargs):
+
+    plot = PlotData2D(geo_data)
+    plot.plot_gradient(scalar_field, gx, gy, gz, cell_number, r=r,
+                           direction=direction, plot_scalar=plot_scalar,
+                           **kwargs)
 
 def plot_data_3D(geo_data, **kwargs):
     """

diff --git a/gempy/plotting/visualization.py b/gempy/plotting/visualization.py
@@ -377,15 +377,48 @@ def plot_topo_g(geo_data, G, centroids, direction="y"):
                          centroids[node][c2] * e2 / r2 + d2, str(node), color="white", size=6, ha="center", va="center",
                          weight="ultralight", family="monospace")
 
-    def plot_gradient(self, geo_data, cell_number, direction="y", r = 5):
+    def plot_gradient(self, scalar_field, gx, gy, gz, cell_number, r=5,
+                      direction="y", plot_scalar = True, *args, **kwargs):
         _a, _b, _c, extent_val, x, y, Gx, Gy = self._slice(direction, cell_number)
 
-
-        plt.quiver(geo_data.grid.values[:, 0].reshape(50, 50, 50)[::r, cell_number, ::r].T,
-                   geo_data.grid.values[:, 2].reshape(50, 50, 50)[::r, cell_number, ::r].T, s2, V2, pivot="tail",
+        if direction == "y":
+            if plot_scalar:
+                self.plot_scalar_field(scalar_field, cell_number, direction=direction, plot_data=False)
+            s2 = gx.reshape(self._data.resolution[0], self._data.resolution[1], self._data.resolution[2])[::r,
+                 cell_number, ::r].T
+            V2 = gz.reshape(self._data.resolution[0], self._data.resolution[1], self._data.resolution[2])[::r,
+                 cell_number, ::r].T
+            plt.quiver(self._data.grid.values[:, 0].reshape(self._data.resolution[0], self._data.resolution[1], self._data.resolution[2])[::r, cell_number, ::r].T,
+                   self._data.grid.values[:, 2].reshape(self._data.resolution[0], self._data.resolution[1], self._data.resolution[2])[::r, cell_number, ::r].T, s2, V2, pivot="tail",
                    color='blue', alpha=.6)
-
-        print(Gx, Gy)
+        elif direction == "x":
+            if plot_scalar:
+                self.plot_scalar_field(scalar_field, cell_number, direction=direction, plot_data=False)
+            s2 = gy.reshape(self._data.resolution[0], self._data.resolution[1], self._data.resolution[2])[::r,
+                 cell_number, ::r].T
+            V2 = gz.reshape(self._data.resolution[0], self._data.resolution[1], self._data.resolution[2])[::r,
+                 cell_number, ::r].T
+            plt.quiver(self._data.grid.values[:, 0].reshape(self._data.resolution[0], self._data.resolution[1],
+                                                            self._data.resolution[2])[::r, cell_number, ::r].T,
+                       self._data.grid.values[:, 2].reshape(self._data.resolution[0], self._data.resolution[1],
+                                                            self._data.resolution[2])[::r, cell_number, ::r].T, s2, V2,
+                       pivot="tail",
+                       color='blue', alpha=.6)
+        elif direction== "z":
+            if plot_scalar:
+                self.plot_scalar_field(scalar_field, cell_number, direction=direction, plot_data=False)
+            s2 = gx.reshape(self._data.resolution[0], self._data.resolution[1], self._data.resolution[2])[::r,
+                 cell_number, ::r].T
+            V2 = gy.reshape(self._data.resolution[0], self._data.resolution[1], self._data.resolution[2])[::r,
+                 cell_number, ::r].T
+            plt.quiver(self._data.grid.values[:, 0].reshape(self._data.resolution[0], self._data.resolution[1],
+                                                            self._data.resolution[2])[::r, cell_number, ::r].T,
+                       self._data.grid.values[:, 2].reshape(self._data.resolution[0], self._data.resolution[1],
+                                                            self._data.resolution[2])[::r, cell_number, ::r].T, s2, V2,
+                       pivot="tail",
+                       color='blue', alpha=.6)
+        else:
+            raise AttributeError(str(direction) + "must be a cartesian direction, i.e. xyz")
 
     # TODO: Incorporate to the class
     @staticmethod