grass.tools: Add NumPy arrays IO to Tools

python/grass/benchmark/runners.py

@@ -173,7 +173,7 @@ def benchmark_resolutions(module, resolutions, label, repeat=5, nprocs=None):
         region = gs.region()
         n_cells.append(region["cells"])
         print("\u2500" * term_size.columns)
-        print(f"Benchmark with {resolution} resolution...\n")
+        print(f"Benchmark with resolution {resolution}...\n")
         time_sum = 0
         measured_times = []
         for _ in range(repeat):

python/grass/tools/benchmark/benchmark_grass_tools_numpy.py

@@ -0,0 +1,121 @@
+import time
+import numpy as np
+
+
+from grass.tools import Tools
+from grass.benchmark import (
+    num_cells_plot,
+    benchmark_resolutions,
+    load_results,
+    save_results,
+)
+
+
+class TimeMeasurer:
+    def __init__(self):
+        self._time = None
+        self._start = None
+
+    @property
+    def time(self):
+        return self._time
+
+    def start(self):
+        self._start = time.perf_counter()
+
+    def stop(self):
+        self._time = time.perf_counter() - self._start
+
+
+class PlainNumPyBenchmark(TimeMeasurer):
+    def run(self):
+        tools = Tools()
+        region = tools.g_region(flags="p", format="json")
+        a = np.full((region["rows"], region["cols"]), 1)
+        b = np.full((region["rows"], region["cols"]), 1)
+
+        self.start()
+        c = 2 * np.sqrt(a + b) * np.sqrt(a) + np.sqrt(b) + a / 2
+        self.stop()
+
+        print(c.sum())
+        print(c.size)
+
+        del a
+        del b
+        del c
+
+
+class PlainGRASSBenchmark(TimeMeasurer):
+    def run(self):
+        tools = Tools(overwrite=True)
+        tools.r_mapcalc(expression="a = 1")
+        tools.r_mapcalc(expression="b = 1")
+
+        self.start()
+        tools.r_mapcalc(expression="c = 2 * sqrt(a + b) * sqrt(a) * sqrt(b) + a / 2")
+        self.stop()
+
+        c_stats = tools.r_univar(map="c", format="json")
+        print(c_stats["sum"])
+        print(c_stats["cells"])
+
+
+class NumPyGRASSBenchmark(TimeMeasurer):
+    def run(self):
+        tools = Tools()
+        region = tools.g_region(flags="p", format="json")
+        a = np.full((region["rows"], region["cols"]), 1)
+        b = np.full((region["rows"], region["cols"]), 1)
+
+        self.start()
+        c = tools.r_mapcalc_simple(
+            expression="2* sqrt(A + B) * sqrt(A) * sqrt(B) + A / 2",
+            a=a,
+            b=b,
+            output=np.array,
+        )
+        self.stop()
+
+        c_stats = tools.r_univar(map=c, format="json")
+        print(c_stats["sum"])
+        print(c_stats["cells"])
+
+        del a
+        del b
+        del c
+
+
+def main():
+    resolutions = [5, 2, 1, 0.5]
+    repeat = 10
+    results = [
+        benchmark_resolutions(
+            module=PlainNumPyBenchmark(),
+            label="NumPy",
+            resolutions=resolutions,
+            repeat=repeat,
+        ),
+        benchmark_resolutions(
+            module=PlainGRASSBenchmark(),
+            label="GRASS",
+            resolutions=resolutions,
+            repeat=repeat,
+        ),
+        benchmark_resolutions(
+            module=NumPyGRASSBenchmark(),
+            label="NumPy GRASS",
+            resolutions=resolutions,
+            repeat=repeat,
+        ),
+    ]
+    print(results)
+    results = load_results(save_results(results))
+    print(results)
+    plot_file = "test_res_plot.png"
+    num_cells_plot(results.results, filename=plot_file)
+    print(plot_file)
+
+
+if __name__ == "__main__":
+    main()

python/grass/tools/session_tools.py

@@ -52,6 +52,19 @@ class Tools:
     the *ToolResult* object:
 
     >>> tools.g_region(flags="p").text  # doctest: +SKIP
+
+    Raster input and outputs can be NumPy arrays:
+
+    >>> import numpy as np
+    >>> tools.g_region(rows=2, cols=3)
+    ToolResult(...)
+    >>> slope = tools.r_slope_aspect(elevation=np.ones((2, 3)), slope=np.ndarray)
+
+    When multiple outputs are returned, they are returned as a tuple:
+
+    >>> (slope, aspect) = tools.r_slope_aspect(
+    ...     elevation=np.ones((2, 3)), slope=np.array, aspect=np.array
+    ... )
     """
 
     def __init__(
@@ -167,13 +180,35 @@ def run(self, tool_name_: str, /, **kwargs):
         # Get a fixed env parameter at at the beginning of each execution,
         # but repeat it every time in case the referenced environment is modified.
         args, popen_options = gs.popen_args_command(tool_name_, **kwargs)
+
+        # Compute the environment for subprocesses and store it for later use.
+        if "env" not in popen_options:
+            popen_options["env"] = self._modified_env_if_needed()
+
+        object_parameter_handler.translate_objects_to_data(
+            kwargs, env=popen_options["env"]
+        )
+
         # We approximate original kwargs with the possibly-modified kwargs.
-        return self.run_cmd(
+        result = self.run_cmd(
             args,
             tool_kwargs=kwargs,
             input=object_parameter_handler.stdin,
             **popen_options,
         )
+        use_objects = object_parameter_handler.translate_data_to_objects(
+            kwargs, env=popen_options["env"]
+        )
+        if use_objects:
+            result = object_parameter_handler.result
+        if object_parameter_handler.temporary_rasters:
+            self.call(
+                "g.remove",
+                type="raster",
+                name=object_parameter_handler.temporary_rasters,
+                flags="f",
+            )
+        return result
 
     def run_cmd(
         self,

python/grass/tools/support.py

@@ -25,22 +25,81 @@
 import shutil
 from io import StringIO
 
+import numpy as np
+
 import grass.script as gs
+import grass.script.array as garray
 
 
 class ParameterConverter:
+    """Converts parameter values to strings and facilitates flow of the data."""
+
     def __init__(self):
         self._numpy_inputs = {}
-        self._numpy_outputs = {}
+        self._numpy_outputs = []
         self._numpy_inputs_ordered = []
         self.stdin = None
+        self.result = None
+        self.temporary_rasters = []
 
     def process_parameters(self, kwargs):
+        """Converts high level parameter values to strings.
+
+        Converts io.StringIO to dash and stores the string in the *stdin* attribute.
+        Replaces NumPy arrays by temporary raster names and stores the arrays.
+        Replaces NumPy array types by temporary raster names.
+
+        Temporary names are accessible in the *temporary_rasters* attribute and need
+        to be cleaned.
+        The functions *translate_objects_to_data* and *translate_data_to_objects*
+        need to be called before and after the computation to do the translations
+        from NumPy arrays to GRASS data and from GRASS data to NumPy arrays.
+
+        Simple type conversions from numbers and iterables to strings are expected to
+        be done by lower level code.
+        """
         for key, value in kwargs.items():
-            if isinstance(value, StringIO):
+            if isinstance(value, np.ndarray):
+                name = gs.append_uuid("tmp_serialized_input_array")
+                kwargs[key] = name
+                self._numpy_inputs[key] = (name, value)
+            elif value in (np.ndarray, np.array, garray.array):
+                # We test for class or the function.
+                name = gs.append_uuid("tmp_serialized_output_array")
+                kwargs[key] = name
+                self._numpy_outputs.append((name, value))
+            elif isinstance(value, StringIO):
                 kwargs[key] = "-"
                 self.stdin = value.getvalue()
 
+    def translate_objects_to_data(self, kwargs, env):
+        """Convert NumPy arrays to GRASS data"""
+        for name, value in self._numpy_inputs.values():
+            map2d = garray.array(env=env)
+            map2d[:] = value
+            map2d.write(name)
+            self.temporary_rasters.append(name)
+
+    def translate_data_to_objects(self, kwargs, env):
+        """Convert GRASS data to NumPy arrays
+
+        Returns True if there is one or more output arrays, False otherwise.
+        The arrays are stored in the *result* attribute.
+        """
+        output_arrays = []
+        for name, value in self._numpy_outputs:
+            output_array = garray.array(name, env=env)
+            output_arrays.append(output_array)
+            self.temporary_rasters.append(name)
+        if len(output_arrays) == 1:
+            self.result = output_arrays[0]
+            return True
+        if len(output_arrays) > 1:
+            self.result = tuple(output_arrays)
+            return True
+        self.result = None
+        return False
+
 
 class ToolFunctionResolver:
     def __init__(self, *, run_function, env):

python/grass/tools/tests/grass_tools_session_tools_numpy_test.py

@@ -0,0 +1,85 @@
+import numpy as np
+
+import grass.script as gs
+from grass.tools import Tools
+
+
+def test_numpy_one_input(xy_dataset_session):
+    """Check that global overwrite is not used when separate env is used"""
+    tools = Tools(session=xy_dataset_session)
+    tools.r_slope_aspect(elevation=np.ones((1, 1)), slope="slope")
+    assert tools.r_info(map="slope", format="json")["datatype"] == "FCELL"
+
+
+def test_numpy_one_input_one_output(xy_dataset_session):
+    """Check that a NumPy array works as input and for signaling output
+
+    It tests that the np.ndarray class is supported to signal output.
+    Return type is not strictly defined, so we are not testing for it explicitly
+    (only by actually using it as an NumPy array).
+    """
+    tools = Tools(session=xy_dataset_session)
+    tools.g_region(rows=2, cols=3)
+    slope = tools.r_slope_aspect(elevation=np.ones((2, 3)), slope=np.ndarray)
+    assert slope.shape == (2, 3)
+    assert np.all(slope == np.full((2, 3), 0))
+
+
+def test_numpy_with_name_and_parameter(xy_dataset_session):
+    """Check that a NumPy array works as input and for signaling output
+
+    It tests that the np.ndarray class is supported to signal output.
+    Return type is not strictly defined, so we are not testing for it explicitly
+    (only by actually using it as an NumPy array).
+    """
+    tools = Tools(session=xy_dataset_session)
+    tools.g_region(rows=2, cols=3)
+    slope = tools.run("r.slope.aspect", elevation=np.ones((2, 3)), slope=np.ndarray)
+    assert slope.shape == (2, 3)
+    assert np.all(slope == np.full((2, 3), 0))
+
+
+def test_numpy_one_input_multiple_outputs(xy_dataset_session):
+    """Check that a NumPy array function works for signaling multiple outputs
+
+    Besides multiple outputs it tests that np.array is supported to signal output.
+    """
+    tools = Tools(session=xy_dataset_session)
+    tools.g_region(rows=2, cols=3)
+    (slope, aspect) = tools.r_slope_aspect(
+        elevation=np.ones((2, 3)), slope=np.array, aspect=np.array
+    )
+    assert slope.shape == (2, 3)
+    assert np.all(slope == np.full((2, 3), 0))
+    assert aspect.shape == (2, 3)
+    assert np.all(aspect == np.full((2, 3), 0))
+
+
+def test_numpy_multiple_inputs_one_output(xy_dataset_session):
+    """Check that a NumPy array works for multiple inputs"""
+    tools = Tools(session=xy_dataset_session)
+    tools.g_region(rows=2, cols=3)
+    result = tools.r_mapcalc_simple(
+        expression="A + B", a=np.full((2, 3), 2), b=np.full((2, 3), 5), output=np.array
+    )
+    assert result.shape == (2, 3)
+    assert np.all(result == np.full((2, 3), 7))
+
+
+def test_numpy_grass_array_input_output(xy_dataset_session):
+    """Check that global overwrite is not used when separate env is used
+
+    When grass array output is requested, we explicitly test the return value type.
+    """
+    tools = Tools(session=xy_dataset_session)
+    rows = 2
+    cols = 3
+    tools.g_region(rows=rows, cols=cols)
+    tools.r_mapcalc_simple(expression="5", output="const_5")
+    const_5 = gs.array.array("const_5", env=xy_dataset_session.env)
+    result = tools.r_mapcalc_simple(
+        expression="2 * A", a=const_5, output=gs.array.array
+    )
+    assert result.shape == (rows, cols)
+    assert np.all(result == np.full((rows, cols), 10))
+    assert isinstance(result, gs.array.array)