ZonalStats.py

import os
from datetime import date, timedelta
import processing
from processing.core.GeoAlgorithmExecutionException import GeoAlgorithmExecutionException
from processing.tools import dataobjects, system
import numpy as np


def ZonalStats(start_date, end_date, vec_name, sb_column, file_list, subcatchmap_res,
               corr_by_num=None, corr_by_fact=None, progress = None):

    if not os.path.isfile(vec_name):
        raise GeoAlgorithmExecutionException('No shapefile: \"' + vec_name + '\" ')
    if not file_list:
        raise GeoAlgorithmExecutionException('List of files is empty')

    # Creating a list of dates (year + julian day)
    dates = []
    for n in range((end_date-start_date).days + 1):
        d = start_date + timedelta(days=n)
        dates.append(d.strftime("%Y%j"))

    # Initialising array for results
    layer = dataobjects.getObjectFromUri(vec_name)
    if layer.crs().toProj4() != "+proj=longlat +datum=WGS84 +no_defs":
        raise IOError('Vector file must be in geographic WGS 1984 projection')
    result_ts = np.ones([len(dates), layer.featureCount()]) * -99.0

    # Extracting data and saving in array
    for file_name in file_list:
        if (file_name.endswith('.tif')) or (file_name.endswith('.tiff')):
            f = os.path.split(file_name)[1]
            file_date = date(int(f[0:4]), int(f[4:6]), int(f[6:8]))
            ind = dates.index(file_date.strftime("%Y%j"))  # index in dates list
        else:
            continue
        if progress is not None:
            progress.setConsoleInfo("Processing %s..." % file_name)

        # Check the raster projection
        raster = dataobjects.getObjectFromUri(file_name)
        if raster.crs().toProj4() != "+proj=longlat +datum=WGS84 +no_defs":
            raise IOError('Datafiles must be in geographic WGS 1984 projection')

        # Apply correction factors if given
        formula = ""
        if corr_by_num is not None:
            formula = "A+%f" % (corr_by_num)
        elif corr_by_fact is not None:
            formula = "A*%f" % (corr_by_fact)
        if formula:
            temp_rast_name = system.getTempFilename("tif")
            calculator_params = {"INPUT_A": file_name, "OUTPUT": temp_rast_name,
                                 "FORMULA": formula}
            processing.runalg("gdalogr:rastercalculator", calculator_params)
        else:
            temp_rast_name = file_name

        # Resample raster to given resolution
        out_rast_name = system.getTempFilename("tif")
        warp_params = {"INPUT": temp_rast_name, "OUTPUT": out_rast_name, "DEST_SRS": "EPSG:4326",
                       "METHOD": 1, "TR": subcatchmap_res, "USE_RASTER_EXTENT": True,
                       "RASTER_EXTENT": dataobjects.extent([layer]), "EXTENT_CRS": "EPSG:4326"}
        processing.runalg("gdalogr:warpreproject", warp_params)

        # Run zonal stats
        zs_output_name = system.getTempFilename("shp")
        zonal_statistics_params = {"INPUT_RASTER": out_rast_name, "RASTER_BAND": 1,
                                   "INPUT_VECTOR": vec_name, "COLUMN_PREFIX": "stats_",
                                   "GLOBAL_EXTENT": False, "OUTPUT_LAYER": zs_output_name}
        processing.runalg("qgis:zonalstatistics", zonal_statistics_params)

        # Save result
        zs_output = dataobjects.getObjectFromUri(zs_output_name)
        for feature in zs_output.getFeatures():
            subbasin_id = feature[sb_column]
            subbasin_value = feature["stats_mean"]
            try:
                result_ts[ind, subbasin_id-1] = subbasin_value
            except TypeError:
                result_ts[ind, subbasin_id-1] = -99.0

    # Return results
    return dates, result_ts