Merge pull request #4843 from alexbruy/processing-heatmap

nyalldawson · web-flow · commit cfbed91103b9 · 2017-07-13T20:07:02.000+10:00
[processing] port heatmap algorithm
diff --git a/python/plugins/processing/algs/qgis/Heatmap.py b/python/plugins/processing/algs/qgis/Heatmap.py
@@ -26,24 +26,24 @@
 __revision__ = '$Format:%H$'
 
 import os
+from collections import OrderedDict
 
 from qgis.PyQt.QtGui import QIcon
 
 from qgis.core import (QgsFeatureRequest,
-                       QgsMessageLog,
-                       QgsProcessingUtils,
-                       QgsProcessingParameterDefinition)
+                       QgsProcessing,
+                       QgsProcessingException,
+                       QgsProcessingParameterVectorLayer,
+                       QgsProcessingParameterNumber,
+                       QgsProcessingParameterField,
+                       QgsProcessingParameterEnum,
+                       QgsProcessingParameterDefinition,
+                       QgsProcessingParameterRasterDestination)
+
 from qgis.analysis import QgsKernelDensityEstimation
 
 from processing.algs.qgis.QgisAlgorithm import QgisAlgorithm
-from processing.core.GeoAlgorithmExecutionException import GeoAlgorithmExecutionException
-from processing.core.parameters import ParameterVector
-from processing.core.parameters import ParameterNumber
-from processing.core.parameters import ParameterSelection
-from processing.core.parameters import ParameterTableField
-from processing.core.outputs import OutputRaster
-from processing.tools import dataobjects, raster
-from processing.algs.qgis.ui.HeatmapWidgets import HeatmapPixelSizeWidgetWrapper
+from processing.tools import raster
 
 pluginPath = os.path.split(os.path.split(os.path.dirname(__file__))[0])[0]
 
@@ -55,18 +55,8 @@ class Heatmap(QgisAlgorithm):
     RADIUS_FIELD = 'RADIUS_FIELD'
     WEIGHT_FIELD = 'WEIGHT_FIELD'
     PIXEL_SIZE = 'PIXEL_SIZE'
-
-    KERNELS = ['Quartic',
-               'Triangular',
-               'Uniform',
-               'Triweight',
-               'Epanechnikov'
-               ]
     KERNEL = 'KERNEL'
     DECAY = 'DECAY'
-    OUTPUT_VALUES = ['Raw',
-                     'Scaled'
-                     ]
     OUTPUT_VALUE = 'OUTPUT_VALUE'
     OUTPUT_LAYER = 'OUTPUT_LAYER'
 
@@ -79,74 +69,115 @@ def tags(self):
     def group(self):
         return self.tr('Interpolation')
 
+    def name(self):
+        return 'heatmapkerneldensityestimation'
+
+    def displayName(self):
+        return self.tr('Heatmap (Kernel Density Estimation)')
+
     def __init__(self):
         super().__init__()
 
     def initAlgorithm(self, config=None):
-        self.addParameter(ParameterVector(self.INPUT_LAYER,
-                                          self.tr('Point layer'), [dataobjects.TYPE_VECTOR_POINT]))
-        self.addParameter(ParameterNumber(self.RADIUS,
-                                          self.tr('Radius (layer units)'),
-                                          0.0, 9999999999, 100.0))
-
-        radius_field_param = ParameterTableField(self.RADIUS_FIELD,
-                                                 self.tr('Radius from field'), self.INPUT_LAYER, optional=True, datatype=ParameterTableField.DATA_TYPE_NUMBER)
+        self.KERNELS = OrderedDict([(self.tr('Quartic'), QgsKernelDensityEstimation.KernelQuartic),
+                                    (self.tr('Triangular'), QgsKernelDensityEstimation.KernelTriangular),
+                                    (self.tr('Uniform'), QgsKernelDensityEstimation.KernelUniform),
+                                    (self.tr('Triweight'), QgsKernelDensityEstimation.KernelTriweight),
+                                    (self.tr('Epanechnikov'), QgsKernelDensityEstimation.KernelEpanechnikov)])
+
+        self.OUTPUT_VALUES = OrderedDict([(self.tr('Raw'), QgsKernelDensityEstimation.OutputRaw),
+                                          (self.tr('Scaled'), QgsKernelDensityEstimation.OutputScaled)])
+
+        self.addParameter(QgsProcessingParameterVectorLayer(self.INPUT_LAYER,
+                                                            self.tr('Point layer'),
+                                                            [QgsProcessing.TypeVectorPoint]))
+
+        self.addParameter(QgsProcessingParameterNumber(self.RADIUS,
+                                                       self.tr('Radius (layer units)'),
+                                                       QgsProcessingParameterNumber.Double,
+                                                       100.0, False, 0.0, 9999999999.99))
+
+        radius_field_param = QgsProcessingParameterField(self.RADIUS_FIELD,
+                                                         self.tr('Radius from field'),
+                                                         None,
+                                                         self.INPUT_LAYER,
+                                                         QgsProcessingParameterField.Numeric,
+                                                         optional=True
+                                                         )
         radius_field_param.setFlags(radius_field_param.flags() | QgsProcessingParameterDefinition.FlagAdvanced)
         self.addParameter(radius_field_param)
 
-        class ParameterHeatmapPixelSize(ParameterNumber):
+        class ParameterHeatmapPixelSize(QgsProcessingParameterNumber):
 
             def __init__(self, name='', description='', parent_layer=None, radius_param=None, radius_field_param=None, minValue=None, maxValue=None,
-                         default=None, optional=False, metadata={}):
-                ParameterNumber.__init__(self, name, description, minValue, maxValue, default, optional, metadata)
+                         default=None, optional=False):
+                QgsProcessingParameterNumber.__init__(self, name, description, QgsProcessingParameterNumber.Double, default, optional, minValue, maxValue)
                 self.parent_layer = parent_layer
                 self.radius_param = radius_param
                 self.radius_field_param = radius_field_param
 
-        self.addParameter(ParameterHeatmapPixelSize(self.PIXEL_SIZE,
-                                                    self.tr('Output raster size'), parent_layer=self.INPUT_LAYER, radius_param=self.RADIUS,
-                                                    radius_field_param=self.RADIUS_FIELD,
-                                                    minValue=0.0, maxValue=9999999999, default=0.1,
-                                                    metadata={'widget_wrapper': HeatmapPixelSizeWidgetWrapper}))
-
-        weight_field_param = ParameterTableField(self.WEIGHT_FIELD,
-                                                 self.tr('Weight from field'), self.INPUT_LAYER, optional=True, datatype=ParameterTableField.DATA_TYPE_NUMBER)
+        pixel_size_param = ParameterHeatmapPixelSize(self.PIXEL_SIZE,
+                                                     self.tr('Output raster size'),
+                                                     parent_layer=self.INPUT_LAYER,
+                                                     radius_param=self.RADIUS,
+                                                     radius_field_param=self.RADIUS_FIELD,
+                                                     minValue=0.0,
+                                                     maxValue=9999999999,
+                                                     default=0.1)
+        pixel_size_param.setMetadata({
+            'widget_wrapper': {
+                'class': 'processing.algs.qgis.ui.HeatmapWidgets.HeatmapPixelSizeWidgetWrapper'}})
+        self.addParameter(pixel_size_param)
+
+        weight_field_param = QgsProcessingParameterField(self.WEIGHT_FIELD,
+                                                         self.tr('Weight from field'),
+                                                         None,
+                                                         self.INPUT_LAYER,
+                                                         QgsProcessingParameterField.Numeric,
+                                                         optional=True
+                                                         )
         weight_field_param.setFlags(weight_field_param.flags() | QgsProcessingParameterDefinition.FlagAdvanced)
         self.addParameter(weight_field_param)
-        kernel_shape_param = ParameterSelection(self.KERNEL,
-                                                self.tr('Kernel shape'), self.KERNELS)
+
+        keys = list(self.KERNELS.keys())
+        kernel_shape_param = QgsProcessingParameterEnum(self.KERNEL,
+                                                        self.tr('Kernel shape'),
+                                                        keys,
+                                                        allowMultiple=False,
+                                                        defaultValue=0)
         kernel_shape_param.setFlags(kernel_shape_param.flags() | QgsProcessingParameterDefinition.FlagAdvanced)
         self.addParameter(kernel_shape_param)
-        decay_ratio = ParameterNumber(self.DECAY,
-                                      self.tr('Decay ratio (Triangular kernels only)'),
-                                      -100.0, 100.0, 0.0)
+
+        decay_ratio = QgsProcessingParameterNumber(self.DECAY,
+                                                   self.tr('Decay ratio (Triangular kernels only)'),
+                                                   QgsProcessingParameterNumber.Double,
+                                                   0.0, True, -100.0, 100.0)
         decay_ratio.setFlags(decay_ratio.flags() | QgsProcessingParameterDefinition.FlagAdvanced)
         self.addParameter(decay_ratio)
-        output_scaling = ParameterSelection(self.OUTPUT_VALUE,
-                                            self.tr('Output value scaling'), self.OUTPUT_VALUES)
+
+        keys = list(self.OUTPUT_VALUES.keys())
+        output_scaling = QgsProcessingParameterEnum(self.OUTPUT_VALUE,
+                                                    self.tr('Output value scaling'),
+                                                    keys,
+                                                    allowMultiple=False,
+                                                    defaultValue=0)
         output_scaling.setFlags(output_scaling.flags() | QgsProcessingParameterDefinition.FlagAdvanced)
         self.addParameter(output_scaling)
-        self.addOutput(OutputRaster(self.OUTPUT_LAYER,
-                                    self.tr('Heatmap')))
 
-    def name(self):
-        return 'heatmapkerneldensityestimation'
-
-    def displayName(self):
-        return self.tr('Heatmap (Kernel Density Estimation)')
+        self.addParameter(QgsProcessingParameterRasterDestination(self.OUTPUT_LAYER, self.tr('Heatmap')))
 
     def processAlgorithm(self, parameters, context, feedback):
-        layer = QgsProcessingUtils.mapLayerFromString(self.getParameterValue(self.INPUT_LAYER), context)
-
-        radius = self.getParameterValue(self.RADIUS)
-        kernel_shape = self.getParameterValue(self.KERNEL)
-        pixel_size = self.getParameterValue(self.PIXEL_SIZE)
-        decay = self.getParameterValue(self.DECAY)
-        output_values = self.getParameterValue(self.OUTPUT_VALUE)
-        output = self.getOutputValue(self.OUTPUT_LAYER)
-        output_format = raster.formatShortNameFromFileName(output)
-        weight_field = self.getParameterValue(self.WEIGHT_FIELD)
-        radius_field = self.getParameterValue(self.RADIUS_FIELD)
+        layer = self.parameterAsVectorLayer(parameters, self.INPUT_LAYER, context)
+
+        radius = self.parameterAsDouble(parameters, self.RADIUS, context)
+        kernel_shape = self.parameterAsEnum(parameters, self.KERNEL, context)
+        pixel_size = self.parameterAsDouble(parameters, self.PIXEL_SIZE, context)
+        decay = self.parameterAsDouble(parameters, self.DECAY, context)
+        output_values = self.parameterAsEnum(parameters, self.OUTPUT_VALUE, context)
+        outputFile = self.parameterAsOutputLayer(parameters, self.OUTPUT_LAYER, context)
+        output_format = raster.formatShortNameFromFileName(outputFile)
+        weight_field = self.parameterAsString(parameters, self.WEIGHT_FIELD, context)
+        radius_field = self.parameterAsString(parameters, self.RADIUS_FIELD, context)
 
         attrs = []
 
@@ -167,22 +198,27 @@ def processAlgorithm(self, parameters, context, feedback):
         kde_params.decayRatio = decay
         kde_params.outputValues = output_values
 
-        kde = QgsKernelDensityEstimation(kde_params, output, output_format)
+        kde = QgsKernelDensityEstimation(kde_params, outputFile, output_format)
 
         if kde.prepare() != QgsKernelDensityEstimation.Success:
-            raise GeoAlgorithmExecutionException(
+            raise QgsProcessingException(
                 self.tr('Could not create destination layer'))
 
         request = QgsFeatureRequest()
         request.setSubsetOfAttributes(attrs)
-        features = QgsProcessingUtils.getFeatures(layer, context, request)
+        features = layer.getFeatures(request)
         total = 100.0 / layer.featureCount() if layer.featureCount() else 0
         for current, f in enumerate(features):
+            if feedback.isCanceled():
+                break
+
             if kde.addFeature(f) != QgsKernelDensityEstimation.Success:
-                QgsMessageLog.logMessage(self.tr('Error adding feature with ID {} to heatmap').format(f.id()), self.tr('Processing'), QgsMessageLog.CRITICAL)
+                feedback.reportError(self.tr('Error adding feature with ID {} to heatmap').format(f.id()), self.tr('Processing'), QgsMessageLog.CRITICAL)
 
             feedback.setProgress(int(current * total))
 
         if kde.finalise() != QgsKernelDensityEstimation.Success:
-            raise GeoAlgorithmExecutionException(
+            raise QgsProcessingException(
                 self.tr('Could not save destination layer'))
+
+        return {self.OUTPUT_LAYER: outputFile}
diff --git a/python/plugins/processing/algs/qgis/QGISAlgorithmProvider.py b/python/plugins/processing/algs/qgis/QGISAlgorithmProvider.py
@@ -59,6 +59,7 @@
 from .ExtentFromLayer import ExtentFromLayer
 from .FixGeometry import FixGeometry
 from .GridPolygon import GridPolygon
+from .Heatmap import Heatmap
 from .ImportIntoPostGIS import ImportIntoPostGIS
 from .ImportIntoSpatialite import ImportIntoSpatialite
 from .Intersection import Intersection
@@ -159,7 +160,6 @@
 # from .GeometryByExpression import GeometryByExpression
 # from .PoleOfInaccessibility import PoleOfInaccessibility
 # from .RasterCalculator import RasterCalculator
-# from .Heatmap import Heatmap
 # from .Orthogonalize import Orthogonalize
 # from .ShortestPathPointToPoint import ShortestPathPointToPoint
 # from .ShortestPathPointToLayer import ShortestPathPointToLayer
@@ -277,7 +277,8 @@ def getAlgs(self):
                 Union(),
                 VectorSplit(),
                 VoronoiPolygons(),
-                ZonalStatistics()
+                ZonalStatistics(),
+                Heatmap()
                 ]
 
         if hasPlotly:
diff --git a/python/plugins/processing/tests/testdata/qgis_algorithm_tests.yaml b/python/plugins/processing/tests/testdata/qgis_algorithm_tests.yaml
@@ -2591,3 +2591,19 @@ tests:
         type: vector
         name: points_weighted.gml
         compare: false
+
+  - algorithm: qgis:heatmapkerneldensityestimation
+    name: Heatmap (Kernel density estimation)
+    params:
+      DECAY: 0.0
+      INPUT_LAYER:
+        name: points.gml
+        type: vector
+      KERNEL: '0'
+      OUTPUT_VALUE: '0'
+      PIXEL_SIZE: 0.1
+      RADIUS: 0.5
+    results:
+      OUTPUT_LAYER:
+        hash: a594d4fa994b81641bfadde54447dc548fa6e23f630c6a4c383487c5
+        type: rasterhash
