Merge pull request #95 from simpeg/StreamThickness

Stream thickness

Author: lheagy

.bumpversion.cfg

@@ -1,4 +1,4 @@
 [bumpversion]
-current_version = 0.3.6
+current_version = 0.3.7
 files = setup.py discretize/__init__.py docs/conf.py
 

discretize/View.py

@@ -206,7 +206,8 @@ def plotSlice(
         range_x=None,
         range_y=None,
         sample_grid=None,
-        stream_threshold=None
+        stream_threshold=None,
+        stream_thickness=None
     ):
 
         """
@@ -249,7 +250,8 @@ def plotSlice(
                         v, vType=vTypeI, normal=normal, ind=ind, grid=grid,
                         view=view, ax=ax, clim=clim, showIt=False,
                         pcolorOpts=pcolorOpts, streamOpts=streamOpts,
-                        gridOpts=gridOpts, stream_threshold=stream_threshold
+                        gridOpts=gridOpts, stream_threshold=stream_threshold,
+                        stream_thickness=stream_thickness
                     )
                 ]
             return out
@@ -351,7 +353,9 @@ def doSlice(v):
             range_x=range_x,
             range_y=range_y,
             sample_grid=sample_grid,
-            stream_threshold=stream_threshold
+            stream_threshold=stream_threshold,
+            stream_thickness=stream_thickness
+
         )
 
         ax.set_xlabel('y' if normal == 'X' else 'x')
@@ -368,7 +372,8 @@ def _plotImage2D(
         range_x=None,
         range_y=None,
         sample_grid=None,
-        stream_threshold=None
+        stream_threshold=None,
+        stream_thickness=None
     ):
 
         if pcolorOpts is None:
@@ -494,6 +499,36 @@ def _plotImage2D(
                 Ui = np.ma.masked_where(mask_me, Ui)
                 Vi = np.ma.masked_where(mask_me, Vi)
 
+
+            if stream_thickness is not None:
+                scaleFact = np.copy(stream_thickness)
+
+                # Calculate vector amplitude
+                vecAmp = np.sqrt(U**2 + V**2).T
+
+                # Form bounds to knockout the top and bottom 10%
+                vecAmp_sort = np.sort(vecAmp.ravel())
+                nVecAmp = vecAmp.size
+                tenPercInd = int(np.ceil(0.1*nVecAmp))
+                lowerBound = vecAmp_sort[tenPercInd]
+                upperBound = vecAmp_sort[-tenPercInd]
+
+                lowInds = np.where(vecAmp < lowerBound)
+                vecAmp[lowInds] = lowerBound
+
+                highInds = np.where(vecAmp > upperBound)
+                vecAmp[highInds] = upperBound
+
+                # Normalize amplitudes 0-1
+                norm_thickness = vecAmp/vecAmp.max()
+
+                # Scale by user defined thickness factor
+                stream_thickness = scaleFact*norm_thickness
+
+                # Add linewidth to streamOpts
+                streamOpts.update({'linewidth':stream_thickness})
+
+
             out += (
                 ax.pcolormesh(
                     x, y, np.sqrt(U**2+V**2).T, vmin=clim[0], vmax=clim[1],

discretize/__init__.py

@@ -27,7 +27,7 @@
             """
             )
 
-__version__   = '0.3.6'
+__version__   = '0.3.7'
 __author__    = 'SimPEG Team'
 __license__   = 'MIT'
 __copyright__ = '2013 - 2017, SimPEG Developers, http://simpeg.xyz'

docs/conf.py

@@ -60,9 +60,9 @@
 # built documents.
 #
 # The short X.Y version.
-version = '0.3.6'
+version = '0.3.7'
 # The full version, including alpha/beta/rc tags.
-release = '0.3.6'
+release = '0.3.7'
 
 # The language for content autogenerated by Sphinx. Refer to documentation
 # for a list of supported languages.

docs/examples/plot_streamThickness.py

@@ -0,0 +1,85 @@
+"""
+Plotting: Streamline thickness
+==============================
+
+A simple example to vary streamline thickness based on the vector amplitudes
+
+Author: `@micmitch <https://github.com/micmitch>`
+"""
+import numpy as np
+import matplotlib.pyplot as plt
+
+from discretize import TensorMesh
+
+###############################################################################
+# Create mesh
+# -----------
+#
+# Minimum cell size in each direction
+dx = 1.
+dy = 1.
+dz = 1.
+
+# Number of core cells in each direction
+nCoreX = 43.
+nCoreY = 43.
+nCoreZ = 43.
+
+# Cell widths
+hx = [(dx,nCoreX)]
+hy = [(dy,nCoreY)]
+hz = [(dz,nCoreZ)]
+
+# Desired Core mesh origin (Bottom SW corner)
+x0 = -21.5
+y0 = -21.5
+z0 = -21.5
+
+mesh = TensorMesh([hx, hy, hz],[x0,y0,z0])
+print(mesh.nC)
+
+###############################################################################
+# Define arbitrary function to plot
+# ---------------------------------
+#
+
+X = mesh.gridCC[:,0]
+Y = mesh.gridCC[:,1]
+Z = mesh.gridCC[:,2]
+
+U = -1 - X**2 + Y + Z
+V = 1 + X - Y**2 + Z
+W = 1 + X + Y - Z**2
+
+###############################################################################
+# Plot streamlines
+# ----------------
+#
+# Create figure
+fig = plt.figure()
+ax = plt.subplot(111)
+fig.set_figheight(15)
+fig.set_figwidth(15)
+labelsize = 30.
+ticksize = 30.
+
+# Create data vector
+dataVec = np.hstack([U,V,W])
+print(dataVec.shape)
+
+# Set streamline plotting options
+streamOpts = {'color':'w', 'density':2.0}
+pcolorOpts = {"cmap":"viridis"}
+
+dat = mesh.plotSlice(
+    dataVec, ax=ax, normal='Z', ind=5, vType='CCv', view='vec',
+    streamOpts=streamOpts, gridOpts={"color":"k", "alpha":0.1}, grid=True,
+    clim=None, stream_thickness=3
+)
+
+###############################################################################
+# Moving Forward
+# --------------
+#
+# If you have suggestions for improving this example, please create a
+# pull request on the example in discretize

setup.py

@@ -56,7 +56,7 @@ def configuration(parent_package='', top_path=None):
 
 setup(
     name="discretize",
-    version="0.3.6",
+    version="0.3.7",
     install_requires=[
         'numpy>=1.7',
         'scipy>=0.13',