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

[svn_grass at osgeo.org]

Author: annakrat
Date: 2014-07-01 10:49:08 -0700 (Tue, 01 Jul 2014)
New Revision: 61105

Modified:
   sandbox/annakrat/r3.flow/rast3d_functions.py
Log:
r3.flow: fix on-the-fly gradient computation

Modified: sandbox/annakrat/r3.flow/rast3d_functions.py
===================================================================
--- sandbox/annakrat/r3.flow/rast3d_functions.py	2014-07-01 15:36:55 UTC (rev 61104)
+++ sandbox/annakrat/r3.flow/rast3d_functions.py	2014-07-01 17:49:08 UTC (rev 61105)
@@ -37,55 +37,71 @@
     if not velocity_obj.last_neighbors_extent or neighbors_extent != velocity_obj.last_neighbors_extent:
         velocity_obj.last_neighbors_extent = neighbors_extent
 
-        if minx < 0 or maxx >= rast3d_region['cols'] or \
-           miny < 0 or maxy >= rast3d_region['rows'] or \
-           minz < 0 or maxz >= rast3d_region['depths']:
+        # just to be sure, we check that at least one voxel is inside
+        if maxx < 0 or minx >= rast3d_region['cols'] or \
+           maxy < 0 or miny >= rast3d_region['rows'] or \
+           maxz < 0 or minz >= rast3d_region['depths']:
             return None
 
-        if minx == 0:
-            maxx = minx + 3
-            xshift = 0
-        elif maxx == rast3d_region['cols'] - 1:
-            minx = maxx - 3
-            xshift = 1
+        # these if's are here to handle edge cases
+        # min, max are changed to represent the min max of the 4x4x4 array
+        # from which the gradient will be computed
+        # shift is relative position of the neighbors within this 4x4x4 array
+        if minx == 0 or minx == -1:
+            maxx = minx + 3 if minx == 0 else minx + 4
+            xshift = minx
+            minx = 0
+        elif maxx >= rast3d_region['cols'] - 1:
+            minx = maxx - 3 if maxx < rast3d_region['cols'] else maxx - 4
+            xshift = 2 if maxx < rast3d_region['cols'] else 3
+            maxx = rast3d_region['cols'] - 1
         else:
             minx -= 1
             maxx += 1
-            xshift = 2
+            xshift = 1
 
-        if miny == 0:
-            maxy = miny + 3
-            yshift = 0
-        elif maxy == rast3d_region['rows'] - 1:
-            miny = maxy - 3
-            yshift = 1
+        if miny == 0 or miny == -1:
+            maxy = miny + 3 if miny == 0 else miny + 4
+            yshift = miny
+            miny = 0
+        elif maxy >= rast3d_region['rows'] - 1:
+            miny = maxy - 3 if maxy < rast3d_region['rows'] else maxy - 4
+            yshift = 2 if maxy < rast3d_region['rows'] else 3
+            maxy = rast3d_region['rows'] - 1
         else:
             miny -= 1
             maxy += 1
-            yshift = 2
+            yshift = 1
 
-        if minz == 0:
-            maxz = minz + 3
-            zshift = 0
-        elif maxz == rast3d_region['depths'] - 1:
-            minz = maxz - 3
-            zshift = 1
+        if minz == 0 or minz == -1:
+            maxz = minz + 3 if minz == 0 else minz + 4
+            zshift = minz
+            minz = 0
+        elif maxz >= rast3d_region['depths'] - 1:
+            minz = maxz - 3 if maxz < rast3d_region['depths'] else maxz - 4
+            zshift = 2 if maxz < rast3d_region['depths'] else 3
+            maxz = rast3d_region['depths'] - 1
         else:
             minz -= 1
             maxz += 1
-            zshift = 2
+            zshift = 1
 
+        # get the 4x4x4 block of the array
         array = velocity_obj.scalar_array[minz:maxz + 1, miny:maxy + 1, minx:maxx + 1]
-#        print (minx, maxx), (miny, maxy), (minz, maxz)
         xyz_gradients = gradient(array, step=[rast3d_region['ewres'], rast3d_region['nsres'],
                                               rast3d_region['tbres']])
-
+        # go through x, y, z and all 8 neighbors and store their value
+        # if the voxel is outside, add 0 (weight)
         array_values = [[] for i in range(3)]
         for i in range(3):
             for z in range(2):
-                for y in range(2):
+                for y in range(1, -1, -1):
                     for x in range(2):
-                        array_values[i].append(xyz_gradients[i][z + zshift][y + yshift][x + xshift])
+                        if z + zshift < 0 or z + zshift > 3 or y + yshift < 0 or y + yshift > 3 \
+                           or x + xshift < 0 or x + xshift > 3:
+                            array_values[i].append(0)
+                        else:
+                            array_values[i].append(xyz_gradients[i][z + zshift][y + yshift][x + xshift])
 
         velocity_obj.last_neighbors_values = array_values
     else: