Convert v.rast.ststs to Python

git-svn-id: https://svn.osgeo.org/grass/grass/trunk@33522 15284696-431f-4ddb-bdfa-cd5b030d7da7

scripts/v.rast.stats/v.rast.stats.py

@@ -0,0 +1,294 @@
+#!/usr/bin/env python
+
+############################################################################
+#
+# MODULE:	v.rast.stats
+# AUTHOR(S):	Markus Neteler, converted to Python by Glynn Clements
+# PURPOSE:	Calculates univariate statistics from a GRASS raster map
+#		only for areas covered by vector objects on a per-category base
+# COPYRIGHT:	(C) 2005-2008 by the GRASS Development Team
+#
+#		This program is free software under the GNU General Public
+#		License (>=v2). Read the file COPYING that comes with GRASS
+#		for details.
+#
+# TODO: do we need layer= management?
+#############################################################################
+
+#%Module
+#%  description: Calculates univariate statistics from a GRASS raster map based on vector polygons and uploads statistics to new attribute columns.
+#%  keywords: vector, raster, statistics
+#%End
+#%flag
+#%  key: c
+#%  description: Continue if upload column(s) already exist
+#%END
+#%flag
+#%  key: e
+#%  description: Calculate extended statistics
+#%END
+#%option
+#% key: vector
+#% type: string
+#% key_desc: name
+#% gisprompt: old,vector,vector
+#% description: Name of vector polygon map
+#% required : yes
+#%End
+#%option
+#% key: layer
+#% type: integer
+#% description: Layer to which the table to be changed is connected
+#% answer: 1
+#% required : no
+#%end
+#%option
+#% key: raster
+#% type: string
+#% key_desc: name
+#% gisprompt: old,cell,raster
+#% description: Name of raster map to calculate statistics from
+#% required : yes
+#%END
+#%option
+#% key: colprefix
+#% type: string
+#% description: Column prefix for new attribute columns
+#% required : yes
+#%end
+#%option
+#% key: percentile
+#% type: integer
+#% description: Percentile to calculate (requires extended statistics flag)
+#% options: 0-100
+#% answer: 90
+#% required : no
+#%end
+
+import sys
+import os
+import atexit
+import grass
+
+def has_column(vector, col):
+    s = grass.read_command('v.info', flags = 'c', map = vector, quiet = True)
+    for l in s.splitlines():
+	f = l.split('|')
+	if len(f) < 2:
+	    continue
+	if f[1] == col:
+	    return True
+    return False
+
+def cleanup():
+    grass.run_command('g.remove', rast = '%s_%s' % (vector, tmpname), quiet = True)
+    grass.run_command('g.remove', rast = 'MASK', quiet = True, stderr = nuldev)
+    if mask_found:
+	grass.message("Restoring previous MASK...")
+	grass.run_command('g.rename', rast = (tmpname + "_origmask", 'MASK'), quiet = True)
+    for f in [tmp, tmpname, sqltmp]:
+	try:
+	    os.remove(f)
+	except:
+	    pass
+
+def main():
+    global tmp, sqltmp, tmpname, nuldev, vector, mask_found
+    #### setup temporary files
+    tmp = grass.tempfile()
+    sqltmp = tmp + ".sql"
+    # we need a random name
+    tmpname = os.path.basename(tmp)
+
+    nuldev = file(os.devnull, 'w')
+
+    raster = options['raster']
+    colprefix = options['colprefix']
+    vector = options['vector']
+    layer = options['layer']
+    percentile = options['percentile']
+
+    ### setup enviro vars ###
+    env = grass.gisenv()
+    mapset = env['MAPSET']
+
+    vs = vector.split('@')
+    if len(vs) > 1:
+	vect_mapset = vs[1]
+    else:
+	vect_mapset = mapset
+
+    # does map exist in CURRENT mapset?
+    if vect_mapset != mapset or not grass.find_file(vector, 'vector', mapset)['file']:
+	grass.fatal("Vector map <%s> not found in current mapset" % vector)
+
+    vector = vs[0]
+
+    #check the input raster map
+    if not grass.find_file(raster, 'cell')['file']:
+	grass.fatal("Raster map <%s> not found" % raster)
+
+    #check presence of raster MASK, put it aside
+    mask_found = bool(grass.find_file('MASK', 'cell')['file'])
+    if mask_found:
+	grass.message("Raster MASK found, temporarily disabled")
+	grass.run_command('g.rename', rast = ('MASK', tmpname + "_origmask"), quiet = True)
+
+    #get RASTER resolution of map which we want to query:
+    #fetch separated to permit for non-square cells (latlong etc)
+    s = grass.read_command('r.info', flags = 's', map = raster)
+    d = {}
+    for l in s.splitlines():
+	kv = l.split('=')
+	d[kv[0]] = float(kv[1])
+    nsres = d['nsres']
+    ewres = d['ewres']
+
+    #save current settings:
+    grass.use_temp_region()
+
+    #Temporarily setting raster resolution to $RASTER resolution
+    #keep boundary settings
+    grass.run_command('g.region', flags = 'a', nsres = nsres, ewres = ewres)
+
+    #prepare raster MASK
+    if grass.run_command('v.to.rast', input = vector, output = "%s_%s" % (vector, tmpname),
+			 use = 'cat', quiet = True) != 0:
+	grass.fatal("An error occurred while converting vector to raster")
+
+    #dump cats to file to avoid "too many argument" problem:
+    p = grass.pipe_command('r.category', map = '%s_%s' % (vector, tmpname), fs = ';', quiet = True)
+    cats = []
+    for line in p.stdout:
+	cats.append(line.rstrip('\r\n').split(';')[0])
+    p.wait()
+
+    #echo "List of categories found: $CATSLIST"
+    number = len(cats)
+    if number < 1:
+	grass.fatal("No categories found in raster map")
+
+    #check if DBF driver used, in this case cut to 10 chars col names:
+    s = grass.read_command('v.db.connect', flags = 'g', map = vector, layer = layer)
+    # we need this for non-DBF driver:
+    table = s.split()[1]
+    db_database = s.split()[3]
+    db_sqldriver = s.split()[4]
+    dbfdriver = db_sqldriver == 'dbf'
+
+    #Find out which table is linked to the vector map on the given layer
+    if not table:
+	grass.fatal('There is no table connected to this map! Run v.db.connect or v.db.addtable first.')
+
+    basecols = ['n', 'min', 'max', 'range', 'mean', 'stddev', 'variance', 'cf_var', 'sum']
+
+    # we need at least three chars to distinguish [mea]n from [med]ian
+    # so colprefix can't be longer than 6 chars with DBF driver
+    if dbfdriver:
+	colprefix = colprefix[:6]
+
+    # do extended stats?
+    if flags['e']:
+	# namespace is limited in DBF but the % value is important
+	if dbfdriver:
+	    perccol = "per" + percentile
+	else:
+	    perccol = "percentile_" + percentile
+	extracols = ['first_quartile', 'median', 'third_quartile'] + [perccol]
+    else:
+	extracols = []
+
+    addcols = []
+    for i in basecols + extracols:
+	#check if column already present
+	currcolumn = ("%s_%s" % (colprefix, i))
+	if dbfdriver:
+	    currcolumn = currcolumn[:10]
+
+	if has_column(vector, currcolumn):
+	    if not flags['c']:
+		grass.fatal(("Cannot create column <%s> (already present)." % currcolumn) +
+			    "Use -c flag to update values in this column.")
+	else:
+	    if i == "n":
+		coltype = "INTEGER"
+	    else:
+		coltype = "DOUBLE PRECISION"
+	    addcols.append(currcolumn + ' ' + coltype)
+
+    grass.verbose("Adding columns <%s>" % addcols)
+    if grass.run_command('v.db.addcol', map = vector, columns = addcols) != 0:
+	grass.fatal("Cannot continue (problem adding columns).")
+
+    #loop over cats and calculate statistics:
+    grass.verbose("Processing data ...")
+
+    # get rid of any earlier attempts
+    try:
+	os.remove(sqltmp)
+    except:
+	pass
+
+    # do extended stats?
+    if flags['e']:
+	extstat = 'e'
+    else:
+	extstat = ""
+
+    f = file(sqltmp, 'w')
+    currnum = 1
+
+    for i in cats:
+	grass.verbose("Processing category %s (%d/%d)" % (i, currnum, number))
+	grass.run_command('g.remove', rast = 'MASK', quiet = True, stderr = nuldev)
+	grass.run_command('r.mapcalc', quiet = True,
+			  expression = "MASK = if(%s_%s == %s, 1, null())" % (vector, tmpname, i))
+
+        #n, min, max, range, mean, stddev, variance, coeff_var, sum
+	# How to test r.univar $? exit status? using -e creates the real possibility of out-of-memory errors
+	s = grass.read_command('r.univar', flags = 'g' + extstat, map = raster, percentile = percentile)
+	vars = {}
+	for l in s.splitlines():
+	    kv = l.split('=')
+	    if len(kv) != 2:
+		continue
+	    var = kv[0]
+	    val = kv[1]
+	    if val.lower() == 'nan':
+		val = 'NULL'
+	    vars[var] = val
+
+	vars['cf_var'] = vars['coeff_var']
+
+	if flags['e'] and dbfdriver:
+	    percvar = 'percentile_' + percentile
+	    vars[perccol] = vars[percvar]
+
+	for var in basecols + extracols:
+	    value = vars[var]
+	    colname = '%s_%s' % (colprefix, var)
+	    if dbfdriver:
+		colname = colname[:10]
+	    f.write("UPDATE %s SET %s=%s WHERE cat=%s;\n" % (table, colname, value, i))
+
+	currnum += 1
+
+    f.close()
+
+    grass.verbose("Updating the database ...")
+    exitcode = grass.run_command('db.execute', input = sqltmp,
+				 database = db_database, driver = db_sqldriver)
+
+    grass.run_command('g.remove', rast = 'MASK', quiet = True, stderr = nuldev)
+
+    if exitcode == 0:
+	grass.message(("Statistics calculated from raster map %s" % raster) +
+		      (" and uploaded to attribute table of vector map %s." % vector))
+	grass.message("Done.")
+
+    sys.exit(exitcode)
+
+if __name__ == "__main__":
+    options, flags = grass.parser()
+    atexit.register(cleanup)
+    main()