rasterio merge mean

merge_mean.py

from contextlib import ExitStack
import numpy as np
from pathlib import Path
from rasterio.merge import merge
from rasterio.io import DatasetWriter, DatasetReader
import rasterio as rio

def merge_mean(
        sources: list[Path | str | DatasetReader],
        dst_path: Path | str | DatasetWriter | None = None,
        dst_kwds: dict | None = None,
        **kwargs: dict | None
) -> tuple[np.ndarray, rio.Affine]:
    """
    Copy valid pixels from input files to an output file.

    Input files are merged in their listed order using the mean value.
    If the output file exists, its values will be overwritten by input values.

    All files must have the same number of bands, data type, and
    coordinate reference system. Rotated, flipped, or upside-down
    rasters cannot be merged.

    Parameters

    sources : list
        A sequence of dataset objects opened in 'r' mode or Path-like
        objects.
    dst_path : str or PathLike, optional
        Path of output dataset
    dst_kwds : dict, optional
        Dictionary of creation options and other parameters that will be
        overlaid on the profile of the output dataset.
    ** kwargs : optional
        Other keyword arguments to pass to [rasterio.merge.merge](https://rasterio.readthedocs.io/en/stable/api/rasterio.merge.html#rasterio.merge.merge)
        Note method is not supported
    Returns
    -------
    tuple
        Two elements:
            dest: numpy.ndarray
                Contents of all input rasters in single array
            out_transform: affine.Affine()
                Information for mapping pixel coordinates in `dest` to
                another coordinate system

    Raises
    ------
    MergeError
        When sources cannot be merged due to incompatibility between
        them or limitations of the tool.
        
    Example
    -------
    # Merge rasters. Tell rasterio to use float32 when merging otherwise we'll get truncated / overflowed
    # data back as my rasters are byte dtype
    merged, meta = merge_mean(
        ["raster1.vrt", "raster2.vrt", "raster3.vrt"], dtype="float32",
        dst_path="mean.tif", dst_kwds={"driver": "GTIFF"}, indexes=(1,2)

    Credits
    -------
    Some code/docstring text is copied/modified from:
     - rasterio.merge.merge - https://github.com/rasterio/rasterio/blob/main/rasterio/merge.py
       Which is licensed under the 3-Clause BSD License https://github.com/rasterio/rasterio/blob/main/LICENSE.txt
     - https://gis.stackexchange.com/a/490079
       Which is licensed under CC BY-SA 4.0
    """

    if "method" in kwargs:
        raise TypeError("The `method` argument is not supported.")

    # calculate merged sum and count
    sum_array, merged_transform = merge(sources, method="sum", **kwargs)
    count_array, merged_transform = merge(sources, method="count", **kwargs)

    # calculate mean
    with np.errstate(divide='ignore', invalid='ignore'): # Hide RuntimeWarning: invalid value encountered in divide
        mean_array = sum_array/count_array

    # Get rid of NaN caused by divide by 0
    mean_array[count_array == 0] = 0

    with ExitStack() as stack:

        if isinstance(sources[0], DatasetReader):
            src = sources[0]
        else:
            src = stack.enter_context(rio.open(sources[0]))

        if dst_path is not None:
            out_profile = src.profile.copy()
            out_profile.update(**(dst_kwds or {}))
            out_profile["transform"] = merged_transform
            out_profile["height"] = mean_array.shape[-2]
            out_profile["width"] = mean_array.shape[-1]
            out_profile["count"] = 1 if len(mean_array.shape) == 2 else mean_array.shape[0]
            out_profile["dtype"] = mean_array.dtype
            out_profile["nodata"] = kwargs.get("nodata")

            if isinstance(dst_path, DatasetWriter):
                dst = dst_path
            else:
                dst = stack.enter_context((rio.open(dst_path, "w", **out_profile)))

            dst.write(mean_array)

        return mean_array, merged_transform