<commit_before>#!/usr/bin/env python

"""
assorted utility functions needed by gridded
"""
from __future__ import absolute_import, division, print_function, unicode_literals
import collections
import numpy as np
import netCDF4 as nc4
import six

must_have = ['dtype', 'shape', 'ndim', '__len__', '__getitem__', '__getattribute__']


def gen_mask(mask_var, add_boundary=False):
    """
    Takes a mask variable (netCDF4.Variable) and creates a mask to apply to
    vertexes in the grid. This map is for GEOMETRY PURPOSES ONLY. It should not
    be applied to any data that may fit the grid. Such data should already have
    a mask of it's own that is to be treated as canon.

    For example ROMS mask_psi:
    <type 'netCDF4._netCDF4.Variable'>
    float64 mask_psi(eta_psi, xi_psi)
        long_name: mask on psi-points
        flag_values: [ 0.  1.]
        flag_meanings: land water
        grid: grid
        location: node
        coordinates: lat_psi lon_psi
    unlimited dimensions:
    current shape = (1043, 723)
    filling off

    The return of this function is a True/False array of the same shape as the variable passed in.
    If flag_values and flag_meanings are attributes on the provided netCDF4.Variable, it will treat
    any value not containing 'water', or 'lake' as True (masked). Otherwise, it will convert to a
    boolean array directly on the 'truthiness' of the values.
    Anything != 0 -> True
    0 = False

    This MAY result in an incorrect mask, if this function doesn't handle special flag
    meanings as above.

    Additionally, it will transform any True (masked) value that is adjacent to a False (unmasked)
    value to False (unmasked). This is to avoid boundary condition problems.
    """
    ret_mask = np.ones(mask_var.shape, dtype=bool)
    input_mask = mask_var[:]
    if isinstance(mask_var, nc4.Variable) and hasattr(mask_var, 'flag_values') and hasattr(mask_var, 'flag_meanings'):
        fm = mask_var.flag_meanings
        if isinstance(fm, six.string_types):
            fm = fm.split(' ')
        meaning_mask = [False if ('water' in s or 'lake' in s) else True for s in fm]
        tfmap = dict(zip(mask_var.flag_values, meaning_mask))
        for k, v in tfmap.items():
            ret_mask[input_mask == k] = v
    else:
        ret_mask[:] = input_mask
    #add in boundary pts
    if add_boundary:
        testarr = np.ones(ret_mask.shape, dtype=bool)
        np.logical_and(testarr[:, 0:-1], ~(ret_mask[:, 0:-1] & ~ret_mask[:, 1:]), testarr[:,0:-1]) #check 'right'
        np.logical_and(testarr[:, 1:], ~(ret_mask[:, 1:] & ~ret_mask[:, 0:-1]), testarr[:, 1:]) #check 'left'
        np.logical_and(testarr[1:, :], ~(ret_mask[1:, :] & ~ret_mask[0:-1, :]), testarr[1:, :]) #check 'down'
        np.logical_and(testarr[0:-1, :], ~(ret_mask[0:-1, :] & ~ret_mask[1:, :]), testarr[0:-1, :]) #check 'up'
        np.logical_and(ret_mask, testarr, ret_mask)
    return ret_mask


def regrid_variable(grid, o_var, location='node'):
    from gridded.variable import Variable
    from gridded.grids import Grid_S, Grid_U<commit_msg>code to handle special case for land masks<commit_after>#!/usr/bin/env python

"""
assorted utility functions needed by gridded
"""
from __future__ import absolute_import, division, print_function, unicode_literals
import collections
import numpy as np
import netCDF4 as nc4
import six

must_have = ['dtype', 'shape', 'ndim', '__len__', '__getitem__', '__getattribute__']


def gen_mask(mask_var, add_boundary=False):
    """
    Takes a mask variable (netCDF4.Variable) and creates a mask to apply to
    vertexes in the grid. This map is for GEOMETRY PURPOSES ONLY. It should not
    be applied to any data that may fit the grid. Such data should already have
    a mask of it's own that is to be treated as canon.

    For example ROMS mask_psi:
    <type 'netCDF4._netCDF4.Variable'>
    float64 mask_psi(eta_psi, xi_psi)
        long_name: mask on psi-points
        flag_values: [ 0.  1.]
        flag_meanings: land water
        grid: grid
        location: node
        coordinates: lat_psi lon_psi
    unlimited dimensions:
    current shape = (1043, 723)
    filling off

    The return of this function is a True/False array of the same shape as the variable passed in.
    If flag_values and flag_meanings are attributes on the provided netCDF4.Variable, it will treat
    any value not containing 'water', or 'lake' as True (masked). Otherwise, it will convert to a
    boolean array directly on the 'truthiness' of the values.
    Anything != 0 -> True
    0 = False

    This MAY result in an incorrect mask, if this function doesn't handle special flag
    meanings as above.

    Additionally, it will transform any True (masked) value that is adjacent to a False (unmasked)
    value to False (unmasked). This is to avoid boundary condition problems.
    """
    ret_mask = np.ones(mask_var.shape, dtype=bool)
    input_mask = mask_var[:]
    type1 = (isinstance(mask_var, nc4.Variable) and hasattr(mask_var, 'flag_values') and hasattr(mask_var, 'flag_meanings'))
    type2 = (isinstance(mask_var, nc4.Variable) and hasattr(mask_var, 'option_0'))
    if type1:
        fm = mask_var.flag_meanings
        if isinstance(fm, six.string_types):
            fm = fm.split(' ')
        meaning_mask = [False if ('water' in s or 'lake' in s) else True for s in fm]
        tfmap = dict(zip(mask_var.flag_values, meaning_mask))
        for k, v in tfmap.items():
            ret_mask[input_mask == k] = v
    elif type2: #special case where option_0 == land, option_1 == water, etc #TODO generalize this properly
        meaning_mask = [True, False]
        tfmap = dict(zip([0, 1], meaning_mask))
        for k, v in tfmap.items():
            ret_mask[input_mask == k] = v
    else:
        ret_mask[:] = input_mask
    #add in boundary pts
    if add_boundary:
        testarr = np.ones(ret_mask.shape, dtype=bool)
        np.logical_and(testarr[:, 0:-1], ~(ret_mask[:, 0:-1] & ~ret_mask[:, 1:]), testarr[:,0:-1]) #check 'right'
        np.logical_and(testarr[:, 1:], ~(ret_mask[:, 1:] & ~ret_mask[:, 0:-1]), testarr[:, 1:]) #check 'left'
        np.logical_and(testarr[1:, :], ~(ret_mask[1:, :] & ~ret_mask[0:-1, :]), testarr[1:, :]) #check 'down'
        np.logical_and(testarr[0:-1, :], ~(ret_mask[0:-1, :] & ~ret_mask[1:, :]), testarr[0:-1, :]) #check 'up'
        np.logical_and(ret_mask, testarr, ret_mask)
    return ret_mask


def regrid_variable(grid, o_var, location='node'):
    from gridded.variable import Variable
    from gridded.grids import Grid_S, Grid_U
{'fraction_non_alphanumeric': 0.11366402921485089, 'fraction_numerical': 0.014911746804625686, 'mean_word_length': 3.265412070084361, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 7, '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': '1960182', 'n_tokens_mistral': 2203, 'n_tokens_neox': 2148, 'n_words': 717}