r.drain.py

#!/usr/bin/env python3
############################################################################
#
# MODULE:       r.drain
# AUTHOR(S):    Markus Metz
# PURPOSE:      Wrapper for r.path. If no input directions are provided,
#               directions are determined with r.fill.dir
# COPYRIGHT:    (C) 2017 by metz, and the GRASS Development Team
#
#  This program is free software; you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation; either version 2 of the License, or
#  (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
############################################################################

# %module
# % description: Traces a flow through an elevation model or cost surface on a raster map.
# % keyword: raster
# % keyword: hydrology
# % keyword: cost surface
# %end
# %option
# % key: input
# % type: string
# % required: yes
# % multiple: no
# % key_desc: name
# % description: Name of input elevation or cost surface raster map
# % gisprompt: old,cell,raster
# %end
# %option
# % key: direction
# % type: string
# % required: no
# % multiple: no
# % key_desc: name
# % label: Name of input movement direction map associated with the cost surface
# % description: Direction in degrees CCW from east
# % gisprompt: old,cell,raster
# % guisection: Cost surface
# %end
# %option
# % key: output
# % type: string
# % required: yes
# % multiple: no
# % key_desc: name
# % description: Name for output raster map
# % gisprompt: new,cell,raster
# %end
# %option
# % key: drain
# % type: string
# % required: no
# % multiple: no
# % key_desc: name
# % label: Name for output drain vector map
# % description: Recommended for cost surface made using knight's move
# % gisprompt: new,vector,vector
# %end
# %option
# % key: start_coordinates
# % type: double
# % required: no
# % multiple: yes
# % key_desc: east,north
# % description: Coordinates of starting point(s) (E,N)
# % gisprompt: old,coords,coords
# % guisection: Start
# %end
# %option
# % key: start_points
# % type: string
# % required: no
# % multiple: yes
# % key_desc: name
# % label: Name of starting vector points map(s)
# % gisprompt: old,vector,vector
# % guisection: Start
# %end
# %flag
# % key: c
# % description: Copy input cell values on output
# % guisection: Path settings
# %end
# %flag
# % key: a
# % description: Accumulate input values along the path
# % guisection: Path settings
# %end
# %flag
# % key: n
# % description: Count cell numbers along the path
# % guisection: Path settings
# %end
# %flag
# % key: d
# % description: The input raster map is a cost surface (direction surface must also be specified)
# % guisection: Cost surface
# %end
# %rules
# % required: start_coordinates, start_points
# %end

import os
import sys
import atexit

import grass.script as grass


def cleanup():
    """Delete temporary direction map."""
    if tmp_maps:
        grass.run_command(
            "g.remove",
            flags="f",
            quiet=True,
            type="raster",
            name=tmp_maps,
            errors="ignore",
        )


def main():
    valmap = options["input"]
    dirmap = options["direction"]
    rpathmap = options["output"]
    vpathmap = options["drain"]
    start_coords = options["start_coordinates"]
    start_pnts = options["start_points"]

    global tmp_maps
    tmp_maps = False
    atexit.register(cleanup)

    if not dirmap:
        # get directions with r.fill.dir, without sink filling
        dirmap = "%s_tmp_%d" % (valmap, os.getpid())
        fill_map = "%s_fill_%d" % (valmap, os.getpid())
        area_map = "%s_area_%d" % (valmap, os.getpid())
        tmp_maps = dirmap + "," + fill_map + "," + area_map
        grass.run_command(
            "r.fill.dir",
            input=valmap,
            output=fill_map,
            direction=dirmap,
            areas=area_map,
            flags="f",
            format="grass",
        )

    # create args for r.path
    kwargs = {}
    kwargs["input"] = dirmap
    if flags["c"] or flags["a"]:
        kwargs["values"] = valmap
    kwargs["format"] = "degree"
    if start_coords:
        kwargs["start_coordinates"] = start_coords
    if start_pnts:
        kwargs["start_points"] = start_pnts
    if rpathmap:
        kwargs["raster_path"] = rpathmap
    if vpathmap:
        kwargs["vector_path"] = vpathmap

    pathflags = ""
    if flags["c"]:
        pathflags += "c"
    if flags["a"]:
        pathflags += "a"
    if flags["n"]:
        pathflags += "n"

    grass.run_command("r.path", flags=pathflags, **kwargs)

    return 0


if __name__ == "__main__":
    options, flags = grass.parser()
    sys.exit(main())