[GRASS-SVN] r60729 - sandbox/annakrat/r3.flow

[svn_grass at osgeo.org]

Author: annakrat
Date: 2014-06-06 10:46:12 -0700 (Fri, 06 Jun 2014)
New Revision: 60729

Modified:
   sandbox/annakrat/r3.flow/r3.flow.py
   sandbox/annakrat/r3.flow/rast3d_functions.py
Log:
r3.flow: support for scalar (gradient computed by numpy)

Modified: sandbox/annakrat/r3.flow/r3.flow.py
===================================================================
--- sandbox/annakrat/r3.flow/r3.flow.py	2014-06-06 16:00:50 UTC (rev 60728)
+++ sandbox/annakrat/r3.flow/r3.flow.py	2014-06-06 17:46:12 UTC (rev 60729)
@@ -25,6 +25,13 @@
 #% keywords: vector
 #%end
 #%option G_OPT_R3_INPUT
+#% key: input
+#% required: no
+#% description: Name of 3D raster maps
+#%end
+#%option G_OPT_R3_INPUTS
+#% key: vector_field
+#% required: no
 #% description: Names of three 3D raster maps describing x, y, z components of vector field
 #%end
 #%option G_OPT_V_INPUT
@@ -76,7 +83,7 @@
 from grass.pygrass.vector.geometry import Line, Point
 
 from integrate import integrate_next, get_time_step
-from rast3d_functions import get_velocity, norm
+from rast3d_functions import get_velocity, norm, compute_gradient
 
 EPSILON = 1e-7
 
@@ -84,17 +91,26 @@
 def main():
     options, flags = gcore.parser()
 
-    try:
-        v_x, v_y, v_z = options['input'].split(',')
-    except ValueError:
-        gcore.fatal(_("Please provide 3 input 3D raster maps."))
+    if options['vector_field']:
+        try:
+            v_x, v_y, v_z = options['vector_field'].split(',')
+        except ValueError:
+            gcore.fatal(_("Please provide 3 input 3D raster maps representing components of vector field."))
+        # read input 3D maps
+        velocity = [garray.array3d(), garray.array3d(), garray.array3d()]
+        for inp, vel in zip((v_x, v_y, v_z), velocity):
+            if vel.read(mapname=inp) != 0:
+                gcore.fatal(_("Error when reading input 3D raster maps"))
+        map_info = grast3d.raster3d_info(v_x)
 
-    # read input 3D maps
-    velocity = [garray.array3d(), garray.array3d(), garray.array3d()]
-    for inp, vel in zip((v_x, v_y, v_z), velocity):
-        if vel.read(mapname=inp) != 0:
-            gcore.fatal(_("Error when reading input 3D raster maps"))
-    map_info = grast3d.raster3d_info(v_x)
+    elif options['input']:
+        array = garray.array3d()
+        array.read(mapname=options['input'])
+        velocity = compute_gradient(array)        
+        map_info = grast3d.raster3d_info(options['input'])
+    else:
+        # this should be handled by parser (in the future)
+        gcore.fatal("Either input or vector_field options must be specified.")
 
     # integrating forward means upstream, backward means downstream
     if flags['u']:

Modified: sandbox/annakrat/r3.flow/rast3d_functions.py
===================================================================
--- sandbox/annakrat/r3.flow/rast3d_functions.py	2014-06-06 16:00:50 UTC (rev 60728)
+++ sandbox/annakrat/r3.flow/rast3d_functions.py	2014-06-06 17:46:12 UTC (rev 60729)
@@ -4,6 +4,7 @@
 """
 
 import math
+import numpy
 #from grass.script import raster3d as grast3d
 
 
@@ -18,6 +19,14 @@
     return math.sqrt(vector[0] * vector[0] + vector[1] * vector[1] + vector[2] * vector[2])
 
 
+def compute_gradient(array):
+    """!Computes gradient of 3D array using numpy gradient function.
+    Returns x, y, z components."""
+    vz, vy, vx = numpy.gradient(array)
+    # we have to reverse y because we have different system than numpy
+    vy = -vy
+    return vx, vy, vz
+
 # following functions are already in raster3d library, this is simpler Python version
 # except for trilinear interpolation which could be added there.