<commit_before>""" Classes for basic grid types """

import sys
import copy
from math import sqrt
import numpy as np
import tiffile          # Temporarily used for TIF IO
import _aai             # Contains the ascread driver
try:
    from fill_sinks import fill_sinks
except ImportError:
    sys.stderr.write("Compiled fill_sinks not available. Falling back to "
                     "Python version\n")
    from raster import fill_sinks

class Grid(object):
    """ Grid baseclass. Don't use this directly except to implement subclasses.
    """
    _hdr = {}
    data = np.empty((0,0))

    def __len__(self):
        # Number of bands
        return np.atleast_3d(self.data).shape[2]

    def _check_hdr(self, hdr):
        """ Check that hdr contains required fields. Intended to be overloaded.
        """
        return

    def get_hdr(self):
        """ Return the header dictionary. """
        return self._hdr

    def set_hdr(self, hdr):
        """ Set the header dictionary. """
        self._check_hdr(hdr)
        self._hdr = hdr
        return

    def clipz(self, bounds):
        """ Clip the z-range in place to bounds = [min, max]. """
        self.data = self.data.clip(bounds[0], bounds[1])
        return

    def max(self):
        """ Return the maximum non-nan in self.data. """
        return self.data[np.isnan(self.data)==False].max()

    def min(self):
        """ Return the minimum non-nan in self.data. """
        return self.data[np.isnan(self.data)==False].min()

    def minmax(self):
        """ Return the minimum and maximum value of data array. """
        A = self.data[np.isnan(self.data)==False]
        return (A.min(), A.max())

    def copy(self):
        """ Return a copy. """
        return copy.deepcopy(self)


class RegularGrid(Grid):
    """ Regular (structured) grid class. A RegularGrid contains a fixed number
    of rows and columns with a regular spacing and a scalar or vector field
    defined as `Z`.
    """
    def __init__(self, hdr, Z=None):
        """
        Parameters:
        -----------
        hdr : dictionary of header fields, which must contain
            - xllcorner (float)
            - yllcorner (float)
            - nx (int)
            - ny (int)
            - dx (float)
            - dy (float)
            - nbands (int)
        Z : dependent m-dimensional quantity (nrows x ncols x m)
        """
        if ('nbands' not in hdr) and (Z is not None):
            if Z.ndim >= 3:
                hdr['nbands'] = Z.ndim[2]
            else:
                hdr['nbands'] = 1
        self.set_hdr(hdr)

        shape = (self._hdr['ny'], self._hdr['nx'])
        if Z is None:
            Z = np.empty(shape)
        else:
            if Z.shape[:2] != shape:
                raise GridError('Data array `Z` has an invalid shape')
        self.data = Z
        return

    def _check_hdr(self, hdr):
        """ Check that the header contains the required fields. """
        required_fields = ('xllcorner', 'yllcorner', 'nx', 'ny', 'dx', 'dy',
                           'nbands')
        for key in required_fields:
            if key not in hdr:
                raise GridError('Header missing {0}'.format(key))
        assert isinstance(hdr['nx'], int)
        assert isinstance(hdr['ny'], int)

    def coordmesh(self, *args, **kwargs):
        """ Alias for center_coords() """
        return self.center_coords(*args, **kwargs)

    def center_llref(self):
        """ Return the 'lower-left' reference in terms of a center coordinate.
        """
        return (self._hdr['xllcorner'] + 0.5*self._hdr['dx'],
                self._hdr['yllcorner'] + 0.5*self._hdr['dy'])

    def center_coords(self):
        """ Return the cell-center coordinates. """
        xllcenter, yllcenter = self.center_llref()
        xurcenter = xllcenter + self._hdr['dx'] * (self._hdr['nx'] - 1)
        yurcenter = yllcenter + self._hdr['dy'] * (self._hdr['ny'] - 1)
        return np.meshgrid(np.linspace(xllcenter, xurcenter, self._hdr['nx']),
                           np.linspace(yurcenter, yllcenter, self._hdr['ny']))

    def vertex_coords(self):
        """ Return the coordinates of vertices. """
        xllcorner = self._hdr['xllcorner']
        yllcorner = self._hdr['yllcorner']
        xurcorner = xllcorner + self._hdr['nx'] * self._hdr['dx']
        yurcorner = yllcorner + self._hdr['ny'] * self._hdr['dy']
        nx = self._hdr['nx']
        ny = self._hdr['ny']
        return np.meshgrid(np.linspace(xllcorner, xurcorner, nx + 1),
                           np.linspace(yurcorner, yllcorner, ny + 1))

    def get_region(self, reference='center'):
        """ Return the region characteristics as a tuple, relative to either
        cell centers or the grid edges. """
        if reference == 'center':
            x0, y0  = self.center_llref()
            n = 0
        elif reference == 'edge':
            x0, y0  = self._hdr['xllcorner'], self._hdr['yllcorner']
            n = 1
        else:
            raise GridError("`reference` must be 'center' or 'edge'")
        return (x0, x0 + self._hdr['dx'] * (self._hdr['nx'] + n),
                y0, y0 + self._hdr['dy'] * (self._hdr['ny'] + n))

    def resample(self, dx, dy, method='nearest'):
        """ Resample array to have spacing (`dx`, `dy'). Modify header
        in-place.

        Parameters
        ----------

        cellsize : cell dimension, float

        method : interpolation method, string ('nearest', 'linear')
        """
        xllcenter, yllcenter = self.center_llref()
        xurcenter = xllcenter + self._hdr['dx'] * self._hdr['nx']
        yurcenter = yllcenter + self._hdr['dy'] * self._hdr['ny']
        nx = int((xurcenter - xllcenter) // dx)
        ny = int((yurcenter - yllcenter) // dy)
        dimratio = (dx / self._hdr['dx'], dy / self._hdr['dy'])

        if method == 'nearest':
            JJ, II = np.meshgrid(np.arange(nx), np.arange(ny))
            srcII = np.around(II * dimratio[1]) \
                            .astype(int) \
                            .clip(0, self._hdr['ny'] - 1)
            srcJJ = np.around(JJ * dimratio[0]) \
                            .astype(int).\
                            clip(0, self._hdr['nx'] - 1)
            self.data = self.data[srcII, srcJJ]
        else:
            raise NotImplementedError('method "{0}" not '
                                      'implemented'.format(method))

        hdr = self.get_hdr()
        hdr['dx'] = dx
        hdr['dy'] = dy
        hdr['nx'] = nx
        hdr['ny'] = ny
        self.set_hdr(hdr)
        return

    def resize(self, te):
        """ Resize array to fit within extents given by te. If the new
        dimensions are smaller, the data is clipped. If they are larger,
        nan padding is added.

        *te*        :   tuple of center coordinates in the form
                        (xmin, xmax, ymin, ymax).

        Returns None.
        """
        if self.data is not None:
            xmin1, ymin1 = self.center_llref()
            xmax1 = xmin1 + self._hdr['dx'] * (self._hdr['nx'] - 1)
            ymax1 = ymin1 + self._hdr['dy'] * (self._hdr['ny'] - 1)

            xmin2 = te[0]
            xmax2 = te[1]
            ymin2 = te[2]
            ymax2 = te[3]<commit_msg>use scipy.interpolate.griddata for resampling when available<commit_after>""" Classes for basic grid types """

import sys
import copy
from math import sqrt
import numpy as np
import tiffile          # Temporarily used for TIF IO
import _aai             # Contains the ascread driver
try:
    from fill_sinks import fill_sinks
except ImportError:
    sys.stderr.write("Compiled fill_sinks not available. Falling back to "
                     "Python version\n")
    from raster import fill_sinks
try:
    from scipy.interpolate import griddata
    HAS_SCIPY = True
except ImportError:
    HAS_SCIPY = False

class Grid(object):
    """ Grid baseclass. Don't use this directly except to implement subclasses.
    """
    _hdr = {}
    data = np.empty((0,0))

    def __len__(self):
        # Number of bands
        return np.atleast_3d(self.data).shape[2]

    def _check_hdr(self, hdr):
        """ Check that hdr contains required fields. Intended to be overloaded.
        """
        return

    def get_hdr(self):
        """ Return the header dictionary. """
        return self._hdr

    def set_hdr(self, hdr):
        """ Set the header dictionary. """
        self._check_hdr(hdr)
        self._hdr = hdr
        return

    def clipz(self, bounds):
        """ Clip the z-range in place to bounds = [min, max]. """
        self.data = self.data.clip(bounds[0], bounds[1])
        return

    def max(self):
        """ Return the maximum non-nan in self.data. """
        return self.data[np.isnan(self.data)==False].max()

    def min(self):
        """ Return the minimum non-nan in self.data. """
        return self.data[np.isnan(self.data)==False].min()

    def minmax(self):
        """ Return the minimum and maximum value of data array. """
        A = self.data[np.isnan(self.data)==False]
        return (A.min(), A.max())

    def copy(self):
        """ Return a copy. """
        return copy.deepcopy(self)


class RegularGrid(Grid):
    """ Regular (structured) grid class. A RegularGrid contains a fixed number
    of rows and columns with a regular spacing and a scalar or vector field
    defined as `Z`.
    """
    def __init__(self, hdr, Z=None):
        """
        Parameters:
        -----------
        hdr : dictionary of header fields, which must contain
            - xllcorner (float)
            - yllcorner (float)
            - nx (int)
            - ny (int)
            - dx (float)
            - dy (float)
            - nbands (int)
        Z : dependent m-dimensional quantity (nrows x ncols x m)
        """
        if ('nbands' not in hdr) and (Z is not None):
            if Z.ndim >= 3:
                hdr['nbands'] = Z.ndim[2]
            else:
                hdr['nbands'] = 1
        self.set_hdr(hdr)

        shape = (self._hdr['ny'], self._hdr['nx'])
        if Z is None:
            Z = np.empty(shape)
        else:
            if Z.shape[:2] != shape:
                raise GridError('Data array `Z` has an invalid shape')
        self.data = Z
        return

    def _check_hdr(self, hdr):
        """ Check that the header contains the required fields. """
        required_fields = ('xllcorner', 'yllcorner', 'nx', 'ny', 'dx', 'dy',
                           'nbands')
        for key in required_fields:
            if key not in hdr:
                raise GridError('Header missing {0}'.format(key))
        assert isinstance(hdr['nx'], int)
        assert isinstance(hdr['ny'], int)

    def coordmesh(self, *args, **kwargs):
        """ Alias for center_coords() """
        return self.center_coords(*args, **kwargs)

    def center_llref(self):
        """ Return the 'lower-left' reference in terms of a center coordinate.
        """
        return (self._hdr['xllcorner'] + 0.5*self._hdr['dx'],
                self._hdr['yllcorner'] + 0.5*self._hdr['dy'])

    def center_coords(self):
        """ Return the cell-center coordinates. """
        xllcenter, yllcenter = self.center_llref()
        xurcenter = xllcenter + self._hdr['dx'] * (self._hdr['nx'] - 1)
        yurcenter = yllcenter + self._hdr['dy'] * (self._hdr['ny'] - 1)
        return np.meshgrid(np.linspace(xllcenter, xurcenter, self._hdr['nx']),
                           np.linspace(yurcenter, yllcenter, self._hdr['ny']))

    def vertex_coords(self):
        """ Return the coordinates of vertices. """
        xllcorner = self._hdr['xllcorner']
        yllcorner = self._hdr['yllcorner']
        xurcorner = xllcorner + self._hdr['nx'] * self._hdr['dx']
        yurcorner = yllcorner + self._hdr['ny'] * self._hdr['dy']
        nx = self._hdr['nx']
        ny = self._hdr['ny']
        return np.meshgrid(np.linspace(xllcorner, xurcorner, nx + 1),
                           np.linspace(yurcorner, yllcorner, ny + 1))

    def get_region(self, reference='center'):
        """ Return the region characteristics as a tuple, relative to either
        cell centers or the grid edges. """
        if reference == 'center':
            x0, y0  = self.center_llref()
            n = 0
        elif reference == 'edge':
            x0, y0  = self._hdr['xllcorner'], self._hdr['yllcorner']
            n = 1
        else:
            raise GridError("`reference` must be 'center' or 'edge'")
        return (x0, x0 + self._hdr['dx'] * (self._hdr['nx'] + n),
                y0, y0 + self._hdr['dy'] * (self._hdr['ny'] + n))

    def resample(self, dx, dy, method='nearest'):
        """ Resample array in-place to have spacing `dx`, `dy'.

        Parameters
        ----------

        dx : cell dimension, float

        dy : cell dimension, float

        method : interpolation method, string ('nearest', 'linear')
        """
        xllcenter, yllcenter = self.center_llref()
        xurcenter = xllcenter + self._hdr['dx'] * self._hdr['nx']
        yurcenter = yllcenter + self._hdr['dy'] * self._hdr['ny']
        nx = int((xurcenter - xllcenter) // dx)
        ny = int((yurcenter - yllcenter) // dy)

        if HAS_SCIPY:

            xx, yy = self.coordmesh()
            xxi, yyi = np.meshgrid(np.linspace(xllcenter, xurcenter, nx),
                                   np.linspace(yllcenter, yurcenter, ny))
            idata = griddata((xx.flatten(), yy.flatten()), self.data.flatten(),
                             (xxi.flatten(), yyi.flatten()), method=method)
            self.data = idata.reshape(ny, nx)[::-1]

        else:
            dimratio = (dx / self._hdr['dx'], dy / self._hdr['dy'])

            if method == 'nearest':
                JJ, II = np.meshgrid(np.arange(nx), np.arange(ny))
                srcII = np.around(II * dimratio[1]) \
                                .astype(int) \
                                .clip(0, self._hdr['ny'] - 1)
                srcJJ = np.around(JJ * dimratio[0]) \
                                .astype(int).\
                                clip(0, self._hdr['nx'] - 1)
                self.data = self.data[srcII, srcJJ]
            else:
                raise NotImplementedError('method "{0}" not '
                                          'implemented'.format(method))

        hdr = self.get_hdr()
        hdr['dx'] = dx
        hdr['dy'] = dy
        hdr['nx'] = nx
        hdr['ny'] = ny
        self.set_hdr(hdr)
        return

    def resize(self, te):
        """ Resize array to fit within extents given by te. If the new
        dimensions are smaller, the data is clipped. If they are larger,
        nan padding is added.

        *te*        :   tuple of center coordinates in the form
                        (xmin, xmax, ymin, ymax).

        Returns None.
        """
        if self.data is not None:
            xmin1, ymin1 = self.center_llref()
            xmax1 = xmin1 + self._hdr['dx'] * (self._hdr['nx'] - 1)
            ymax1 = ymin1 + self._hdr['dy'] * (self._hdr['ny'] - 1)

            xmin2 = te[0]
            xmax2 = te[1]
            ymin2 = te[2]
            ymax2 = te[3]
{'fraction_non_alphanumeric': 0.1357415449401237, 'fraction_numerical': 0.006908803789972365, 'mean_word_length': 1.7862511457378552, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7436411', 'n_tokens_mistral': 4720, 'n_tokens_neox': 4528, 'n_words': 1267}