<commit_before>"""Wrapper functions around rasterio and useful raster functions."""

from affine import Affine
from collections import namedtuple
import itertools
import logging
import numpy as np
import numpy.ma as ma
from retry import retry
import rasterio
from rasterio.enums import Resampling
from rasterio.errors import RasterioIOError
from rasterio.io import MemoryFile
from rasterio.transform import from_bounds as affine_from_bounds
from rasterio.vrt import WarpedVRT
from rasterio.warp import reproject
from rasterio.windows import from_bounds
from tilematrix import clip_geometry_to_srs_bounds, Shape, Bounds
from types import GeneratorType
import warnings

from mapchete.tile import BufferedTile
from mapchete.io import path_is_remote, get_gdal_options


logger = logging.getLogger(__name__)

ReferencedRaster = namedtuple("ReferencedRaster", ("data", "affine", "bounds", "crs"))


def read_raster_window(
    input_files,
    tile,
    indexes=None,
    resampling="nearest",
    src_nodata=None,
    dst_nodata=None,
    gdal_opts=None
):
    """
    Return NumPy arrays from an input raster.

    NumPy arrays are reprojected and resampled to tile properties from input
    raster. If tile boundaries cross the antimeridian, data on the other side
    of the antimeridian will be read and concatenated to the numpy array
    accordingly.

    Parameters
    ----------
    input_files : string or list
        path to a raster file or list of paths to multiple raster files readable by
        rasterio.
    tile : Tile
        a Tile object
    indexes : list or int
        a list of band numbers; None will read all.
    resampling : string
        one of "nearest", "average", "bilinear" or "lanczos"
    src_nodata : int or float, optional
        if not set, the nodata value from the source dataset will be used
    dst_nodata : int or float, optional
        if not set, the nodata value from the source dataset will be used
    gdal_opts : dict
        GDAL options passed on to rasterio.Env()

    Returns
    -------
    raster : MaskedArray
    """
    with rasterio.Env(
        **get_gdal_options(
            gdal_opts,
            is_remote=path_is_remote(
                input_files[0] if isinstance(input_files, list) else input_files, s3=True
            )
        )
    ) as env:
        logger.debug("reading %s with GDAL options %s", input_files, env.options)
        return _read_raster_window(
            input_files,
            tile,
            indexes=indexes,
            resampling=resampling,
            src_nodata=src_nodata,
            dst_nodata=dst_nodata
        )


def _read_raster_window(
    input_files,
    tile,
    indexes=None,
    resampling="nearest",
    src_nodata=None,
    dst_nodata=None
):
    if isinstance(input_files, list):
        # in case multiple input files are given, merge output into one array
        # read first file
        dst_array = _read_raster_window(
            input_files[0],
            tile=tile,
            indexes=indexes,
            resampling=resampling,
            src_nodata=src_nodata,
            dst_nodata=dst_nodata
        )
        # read subsequent files and merge
        for f in input_files[1:]:
            f_array = _read_raster_window(
                f,
                tile=tile,
                indexes=indexes,
                resampling=resampling,
                src_nodata=src_nodata,
                dst_nodata=dst_nodata
            )
            dst_array = ma.MaskedArray(
                data=np.where(f_array.mask, dst_array, f_array).astype(dst_array.dtype),
                mask=np.where(f_array.mask, dst_array.mask, f_array.mask).astype(np.bool)
            )
        return dst_array
    else:
        input_file = input_files
        dst_shape = tile.shape

        if not isinstance(indexes, int):
            if indexes is None:
                dst_shape = (None,) + dst_shape
            elif len(indexes) == 1:
                indexes = indexes[0]
            else:
                dst_shape = (len(indexes),) + dst_shape
        # Check if potentially tile boundaries exceed tile matrix boundaries on
        # the antimeridian, the northern or the southern boundary.
        if tile.pixelbuffer and tile.is_on_edge():
            return _get_warped_edge_array(
                tile=tile,
                input_file=input_file,
                indexes=indexes,
                dst_shape=dst_shape,
                resampling=resampling,
                src_nodata=src_nodata,
                dst_nodata=dst_nodata
            )

        # If tile boundaries don't exceed pyramid boundaries, simply read window
        # once.
        else:
            return _get_warped_array(
                input_file=input_file,
                indexes=indexes,
                dst_bounds=tile.bounds,
                dst_shape=dst_shape,
                dst_crs=tile.crs,<commit_msg>read window only once if tile is on edge and tilepyramid is not global<commit_after>"""Wrapper functions around rasterio and useful raster functions."""

from affine import Affine
from collections import namedtuple
import itertools
import logging
import numpy as np
import numpy.ma as ma
from retry import retry
import rasterio
from rasterio.enums import Resampling
from rasterio.errors import RasterioIOError
from rasterio.io import MemoryFile
from rasterio.transform import from_bounds as affine_from_bounds
from rasterio.vrt import WarpedVRT
from rasterio.warp import reproject
from rasterio.windows import from_bounds
from tilematrix import clip_geometry_to_srs_bounds, Shape, Bounds
from types import GeneratorType
import warnings

from mapchete.tile import BufferedTile
from mapchete.io import path_is_remote, get_gdal_options


logger = logging.getLogger(__name__)

ReferencedRaster = namedtuple("ReferencedRaster", ("data", "affine", "bounds", "crs"))


def read_raster_window(
    input_files,
    tile,
    indexes=None,
    resampling="nearest",
    src_nodata=None,
    dst_nodata=None,
    gdal_opts=None
):
    """
    Return NumPy arrays from an input raster.

    NumPy arrays are reprojected and resampled to tile properties from input
    raster. If tile boundaries cross the antimeridian, data on the other side
    of the antimeridian will be read and concatenated to the numpy array
    accordingly.

    Parameters
    ----------
    input_files : string or list
        path to a raster file or list of paths to multiple raster files readable by
        rasterio.
    tile : Tile
        a Tile object
    indexes : list or int
        a list of band numbers; None will read all.
    resampling : string
        one of "nearest", "average", "bilinear" or "lanczos"
    src_nodata : int or float, optional
        if not set, the nodata value from the source dataset will be used
    dst_nodata : int or float, optional
        if not set, the nodata value from the source dataset will be used
    gdal_opts : dict
        GDAL options passed on to rasterio.Env()

    Returns
    -------
    raster : MaskedArray
    """
    with rasterio.Env(
        **get_gdal_options(
            gdal_opts,
            is_remote=path_is_remote(
                input_files[0] if isinstance(input_files, list) else input_files, s3=True
            )
        )
    ) as env:
        logger.debug("reading %s with GDAL options %s", input_files, env.options)
        return _read_raster_window(
            input_files,
            tile,
            indexes=indexes,
            resampling=resampling,
            src_nodata=src_nodata,
            dst_nodata=dst_nodata
        )


def _read_raster_window(
    input_files,
    tile,
    indexes=None,
    resampling="nearest",
    src_nodata=None,
    dst_nodata=None
):
    if isinstance(input_files, list):
        # in case multiple input files are given, merge output into one array
        # read first file
        dst_array = _read_raster_window(
            input_files[0],
            tile=tile,
            indexes=indexes,
            resampling=resampling,
            src_nodata=src_nodata,
            dst_nodata=dst_nodata
        )
        # read subsequent files and merge
        for f in input_files[1:]:
            f_array = _read_raster_window(
                f,
                tile=tile,
                indexes=indexes,
                resampling=resampling,
                src_nodata=src_nodata,
                dst_nodata=dst_nodata
            )
            dst_array = ma.MaskedArray(
                data=np.where(f_array.mask, dst_array, f_array).astype(dst_array.dtype),
                mask=np.where(f_array.mask, dst_array.mask, f_array.mask).astype(np.bool)
            )
        return dst_array
    else:
        input_file = input_files
        dst_shape = tile.shape

        if not isinstance(indexes, int):
            if indexes is None:
                dst_shape = (None,) + dst_shape
            elif len(indexes) == 1:
                indexes = indexes[0]
            else:
                dst_shape = (len(indexes),) + dst_shape
        # Check if potentially tile boundaries exceed tile matrix boundaries on
        # the antimeridian, the northern or the southern boundary.
        if tile.tp.is_global and tile.pixelbuffer and tile.is_on_edge():
            return _get_warped_edge_array(
                tile=tile,
                input_file=input_file,
                indexes=indexes,
                dst_shape=dst_shape,
                resampling=resampling,
                src_nodata=src_nodata,
                dst_nodata=dst_nodata
            )

        # If tile boundaries don't exceed pyramid boundaries, simply read window
        # once.
        else:
            return _get_warped_array(
                input_file=input_file,
                indexes=indexes,
                dst_bounds=tile.bounds,
                dst_shape=dst_shape,
                dst_crs=tile.crs,
{'fraction_non_alphanumeric': 0.06333837619768028, 'fraction_numerical': 0.001210287443267776, 'mean_word_length': 2.1261034047919294, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2059092', 'n_tokens_mistral': 2915, 'n_tokens_neox': 2816, 'n_words': 826}