[GRASS-SVN] r72661 - grass-addons/grass7/raster/r.tri

[svn_grass at osgeo.org]

Author: spawley
Date: 2018-05-01 22:19:17 -0700 (Tue, 01 May 2018)
New Revision: 72661

Modified:
   grass-addons/grass7/raster/r.tri/r.tri.py
Log:
r.tri deal with null values as part of mapcalc expression

Modified: grass-addons/grass7/raster/r.tri/r.tri.py
===================================================================
--- grass-addons/grass7/raster/r.tri/r.tri.py	2018-05-02 04:24:16 UTC (rev 72660)
+++ grass-addons/grass7/raster/r.tri/r.tri.py	2018-05-02 05:19:17 UTC (rev 72661)
@@ -38,22 +38,19 @@
 #%end
 
 #%flag
-#% key: p
-#% label: Cell padding
-#% description: Use cell padding to reduce edge effect
-#% guisection: Optional
+#% key: c
+#% description: Use circular neighborhood
 #%end
 
-
+from grass.pygrass.modules.shortcuts import general as g
+from grass.pygrass.modules.shortcuts import raster as r
+from grass.pygrass.gis.region import Region
+from grass.pygrass.raster import RasterRow
 import sys
 import atexit
 import random
 import string
 import grass.script as gs
-from grass.pygrass.modules.shortcuts import general as g
-from grass.pygrass.modules.shortcuts import raster as r
-from grass.pygrass.gis.region import Region
-from grass.pygrass.raster import RasterRow
 
 TMP_RAST = []
 
@@ -73,60 +70,78 @@
     return (tmp)
 
 
+def focal_expr(radius, window_square=False):
+    """Returns array offsets relative to centre pixel (0,0) for a matrix of
+    size radius
+
+    Args
+    ----
+    radius : int
+        Radius of the focal function
+    window_square : bool. Optional (default is False)
+        Boolean to use a circular or square window
+
+    Returns
+    -------
+    offsets : list
+        List of pixel positions (row, col) relative to the center pixel
+        ( 1, -1)  ( 1, 0)  ( 1, 1)
+        ( 0, -1)  ( 0, 0)  ( 0, 1)
+        (-1, -1)  (-1, 0)  (-1, 1)"""
+
+    offsets = []
+
+    # generate a list of spatial neighbourhood offsets for the chosen radius
+    # ignoring the centre cell
+    if window_square:
+        
+        for i in range(-radius, radius+1):
+            for j in range(-radius, radius+1):
+                if (i,j) != (0,0):
+                    offsets.append((i, j))
+    
+    else:
+
+        for i in range(-radius, radius+1):
+            for j in range(-radius, radius+1):
+                row = i + radius
+                col = j + radius
+
+                if pow(row - radius, 2) + pow(col - radius, 2) <= \
+                    pow(radius, 2) and (i, j) != (0,0):
+                    offsets.append((j, i))
+
+    return offsets
+
+
 def main():
     dem = options['input']
     tri = options['output']
     size = int(options['size'])
-    p = flags['p']
+    circular = flags['c']
 
     radius = (size-1)/2
-    ncells = size**2
 
-    # store current region settings
-    current_reg = Region()
-
-    if p:
-        # check for existing mask
-        if RasterRow('MASK').exist():
-            gs.fatal('Cell padding option cannot be used with existing mask...please remove first')
-
-        # grow the map to remove border effects
-        gs.message('Growing DEM to avoid edge effects...')
-        dem_grown = temp_map('tmp_elevation_grown')
-        g.region(n=current_reg.north + (current_reg.nsres * (radius+2)),
-                 s=current_reg.south - (current_reg.nsres * (radius+2)),
-                 w=current_reg.west - (current_reg.ewres * (radius+2)),
-                 e=current_reg.east + (current_reg.ewres * (radius+2)))
-        r.grow(input=dem, output=dem_grown, radius=radius+2, quiet=True)
-        dem = dem_grown
-
     # calculate TRI based on map calc statements
     gs.message("Calculating the Topographic Ruggedness Index...")
 
-    # generate a list of spatial neighbourhood indexs for the chosen radius
+    # generate a list of spatial neighbourhood offsets for the chosen radius
     # ignoring the center cell
-    offsets = []
-    for j in range(-radius, radius+1):
-        for i in range(-radius, radius+1):
-            if (j,i) != (0,0):
-                offsets.append((j,i))
+    offsets = focal_expr(radius = radius, window_square=not circular)
+    terms = []
+    for d in offsets:
+        valid = ','.join(map(str, d))        
+        invalid = ','.join([str(-d[0]), str(-d[1])])
+        terms.append(
+            "if(isnull({dem}[{d}]), if(isnull({dem}[{e}]), (median({dem})-{dem})^2, ({dem}[{e}]-{dem}^2)), ({dem}[{d}]-{dem})^2)".format(
+            dem=dem, d=valid, e=invalid))
 
-    # define the calculation term
-    terms = ["($dem - $dem[%d,%d])^2" % d for d in offsets]
-
     # define the calculation expression
+    ncells = len(offsets)+1
     expr = "$tri = sqrt((%s" % " + ".join(terms) + ") / $ncells)"
 
     # perform the r.mapcalc calculation with the moving window
-    if p:
-        output_tmp = temp_map('tmp')
-        gs.mapcalc(expr, tri=output_tmp, dem=dem, ncells=ncells)
-        r.mask(raster=options['input'], quiet=True)
-        r.mapcalc('{x}={y}'.format(x=tri, y=output_tmp))
-        r.mask(flags='r', quiet=True)
-        Region.write(current_reg)
-    else:
-        gs.mapcalc(expr, tri=tri, dem=dem, ncells=ncells)
+    gs.mapcalc(expr, tri=tri, dem=dem, ncells=ncells)
 
     return 0