<commit_before># -*- coding: utf-8 -*-
# -----------------------------------------------------------------------------
# (C) British Crown Copyright 2017 Met Office.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# * Neither the name of the copyright holder nor the names of its
#   contributors may be used to endorse or promote products derived from
#   this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

"""Gridded data extraction for the Improver site specific process chain."""

import warnings
import copy

import numpy as np
from numpy.linalg import lstsq
import iris
from iris.coords import AuxCoord
from iris.cube import Cube
from iris.exceptions import CoordinateNotFoundError

from improver.spotdata.common_functions import (nearest_n_neighbours,
                                                node_edge_check)
from improver.constants import (R_DRY_AIR,
                                CP_DRY_AIR)
from improver.utilities.cube_manipulation import build_coordinate


class ExtractData(object):
    """
    A series of methods for extracting data from grid points to derive
    diagnostic values at off grid positions.

    """

    def __init__(self, method='use_nearest'):
        """
        The class is called with the desired method to be used in extracting/
        interpolating data to the site of interest from gridded data. If no
        method is defined in the call to the class it defaults to the
        'use_nearest' method.

        """
        self.method = method

    def process(self, cube, sites, neighbours, ancillary_data, additional_data,
                no_neighbours=9, lower_level=1, upper_level=3,
                dz_tolerance=2., dthetadz_threshold=0.02,
                dz_max_adjustment=70.):
        """
        Call the correct function to enact the method of data extraction
        specified. This function also handles multiple timesteps, consolidating
        the resulting data cubes into an Iris.CubeList.

        Args:
            cube (iris.cube.Cube):
                Cube of diagnostic data from which to extract spotdata.

            sites (OrderedDict):
                A dictionary containing the properties of spotdata sites.

            neighbours (numpy.array):
                Array of neighbouring grid points that are associated with
                sites in the OrderedDict of sites.

            additional_data (dict):
                A dictionary containing any supplmentary time varying
                diagnostics that are needed for the selected extraction method.

            ancillary_data (dict):
                A dictionary containing additional model data that is needed.
                e.g. {'orography': <cube of orography>}

            no_neighbours (int):
                Number of grid points about the site to consider when
                calculating a local temperature lapse rate using orography
                variation.
                e.g. consider a 5x5 grid of points -> no_neighbours = 25.

            upper/lower_level (ints):
                Define the hybrid height model levels to use when calculating
                potential temperature gradients for use in lapse rate
                temperature adjustment.

            dz_tolerance (float (units: m)):
                Vertical displacement between spotdata site and neighbouring
                grid point below which there is no need to perform a lapse rate
                adjustment. Defaults to value of 2m.

            dthetadz_threshold (float (units: K/m)):
                Potential temperature gradient threshold, with gradients above
                this value deemed to have been calculated across an inversion.
                Defaults to UKPP value of 0.02 K/m.

            dz_max_adjustment (float (units: m)):
                Maximum vertical distance over which a temperature will be
                adjusted using the lapse rate. If the spotdata site is more
                than this distance above or below its neighbouring grid point,
                the adjustment will be made using dz = dz_max_adjustment.
                Defaults to UKPP value of 70m.

        Returns:
            cube (iris.cube.Cube):
                An irregular (i.e. non-gridded) cube of diagnostic data
                extracted at the spotdata sites.

        """
        if self.method == 'use_nearest':
            return self.use_nearest(cube, sites, neighbours)
        elif self.method == 'orography_derived_temperature_lapse_rate':
            orography = ancillary_data['orography'].data
            return self.orography_derived_temperature_lapse_rate(
                cube, sites, neighbours, orography,
                no_neighbours=no_neighbours)
        elif self.method == 'model_level_temperature_lapse_rate':
            pressure_on_height_levels = (
                additional_data['pressure_on_height_levels'])

            # Ensure that pressure units are in Pa.
            pressure_on_height_levels.convert_units('Pa')

            lower_pressure = pressure_on_height_levels[lower_level, ...]
            upper_pressure = pressure_on_height_levels[upper_level, ...]

            surface_pressure = additional_data['surface_pressure']
            surface_pressure.convert_units('Pa')

            temperature_on_height_levels = (
                additional_data['temperature_on_height_levels'])
            lower_temperature = temperature_on_height_levels[lower_level, ...]
            upper_temperature = temperature_on_height_levels[upper_level, ...]

            return self.model_level_temperature_lapse_rate(
                cube, sites, neighbours, surface_pressure, lower_pressure,
                upper_pressure, lower_temperature, upper_temperature,
                dz_tolerance, dthetadz_threshold, dz_max_adjustment)

        raise AttributeError('Unknown method "{}" passed to {}.'.format(
            self.method, self.__class__.__name__))

    @staticmethod
    def _aux_coords_to_make():
        """
        Define coordinates that need to be made for the cube to be produced.

        Returns:
            dict:
                Dictionary of coordinates that are required for a spotdata
                cube.

        """
        return {'latitude': {'units': 'degrees', 'data_type': float,
                             'coord_type': AuxCoord},
                'longitude': {'units': 'degrees', 'data_type': float,
                              'coord_type': AuxCoord},
                'altitude': {'units': 'm', 'data_type': float,
                             'coord_type': AuxCoord,
                             'custom_function': np.nan_to_num},
                'wmo_site': {'data_type': int, 'coord_type': AuxCoord},
                'utc_offset': {'units': 'hours', 'data_type': float,
                               'coord_type': AuxCoord}
               }

    @staticmethod
    def make_stat_coordinate_first(cube):
        """
        Reorder cube dimension coordinates to ensure the statistical
        coordinate is first.

        Args:
            cube (iris.cube.Cube):
                The cube to be reordered.

        Returns:
            cube (iris.cube.Cube):
                Cube with the statistical coordinate moved to be first.

        Raises:
            Warning if more than one statistical dimension is found. Then
            promotes the first found to become the leading dimension.

        """
        stat_coordinates = ['realization', 'percentile_over']
        cube_dimension_order = {
            coord.name(): cube.coord_dims(coord.name())[0]
            for coord in cube.dim_coords}

        stat_coord = []
        for crd in stat_coordinates:
            stat_coord += [coord for coord in cube_dimension_order.keys()<commit_msg>Fix pep8 bug<commit_after># -*- coding: utf-8 -*-
# -----------------------------------------------------------------------------
# (C) British Crown Copyright 2017 Met Office.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# * Neither the name of the copyright holder nor the names of its
#   contributors may be used to endorse or promote products derived from
#   this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

"""Gridded data extraction for the Improver site specific process chain."""

import warnings
import copy

import numpy as np
from numpy.linalg import lstsq
import iris
from iris.coords import AuxCoord
from iris.cube import Cube
from iris.exceptions import CoordinateNotFoundError

from improver.spotdata.common_functions import (nearest_n_neighbours,
                                                node_edge_check)
from improver.constants import (R_DRY_AIR,
                                CP_DRY_AIR)
from improver.utilities.cube_manipulation import build_coordinate


class ExtractData(object):
    """
    A series of methods for extracting data from grid points to derive
    diagnostic values at off grid positions.

    """

    def __init__(self, method='use_nearest'):
        """
        The class is called with the desired method to be used in extracting/
        interpolating data to the site of interest from gridded data. If no
        method is defined in the call to the class it defaults to the
        'use_nearest' method.

        """
        self.method = method

    def process(self, cube, sites, neighbours, ancillary_data, additional_data,
                no_neighbours=9, lower_level=1, upper_level=3,
                dz_tolerance=2., dthetadz_threshold=0.02,
                dz_max_adjustment=70.):
        """
        Call the correct function to enact the method of data extraction
        specified. This function also handles multiple timesteps, consolidating
        the resulting data cubes into an Iris.CubeList.

        Args:
            cube (iris.cube.Cube):
                Cube of diagnostic data from which to extract spotdata.

            sites (OrderedDict):
                A dictionary containing the properties of spotdata sites.

            neighbours (numpy.array):
                Array of neighbouring grid points that are associated with
                sites in the OrderedDict of sites.

            additional_data (dict):
                A dictionary containing any supplmentary time varying
                diagnostics that are needed for the selected extraction method.

            ancillary_data (dict):
                A dictionary containing additional model data that is needed.
                e.g. {'orography': <cube of orography>}

            no_neighbours (int):
                Number of grid points about the site to consider when
                calculating a local temperature lapse rate using orography
                variation.
                e.g. consider a 5x5 grid of points -> no_neighbours = 25.

            upper/lower_level (ints):
                Define the hybrid height model levels to use when calculating
                potential temperature gradients for use in lapse rate
                temperature adjustment.

            dz_tolerance (float (units: m)):
                Vertical displacement between spotdata site and neighbouring
                grid point below which there is no need to perform a lapse rate
                adjustment. Defaults to value of 2m.

            dthetadz_threshold (float (units: K/m)):
                Potential temperature gradient threshold, with gradients above
                this value deemed to have been calculated across an inversion.
                Defaults to UKPP value of 0.02 K/m.

            dz_max_adjustment (float (units: m)):
                Maximum vertical distance over which a temperature will be
                adjusted using the lapse rate. If the spotdata site is more
                than this distance above or below its neighbouring grid point,
                the adjustment will be made using dz = dz_max_adjustment.
                Defaults to UKPP value of 70m.

        Returns:
            cube (iris.cube.Cube):
                An irregular (i.e. non-gridded) cube of diagnostic data
                extracted at the spotdata sites.

        """
        if self.method == 'use_nearest':
            return self.use_nearest(cube, sites, neighbours)
        elif self.method == 'orography_derived_temperature_lapse_rate':
            orography = ancillary_data['orography'].data
            return self.orography_derived_temperature_lapse_rate(
                cube, sites, neighbours, orography,
                no_neighbours=no_neighbours)
        elif self.method == 'model_level_temperature_lapse_rate':
            pressure_on_height_levels = (
                additional_data['pressure_on_height_levels'])

            # Ensure that pressure units are in Pa.
            pressure_on_height_levels.convert_units('Pa')

            lower_pressure = pressure_on_height_levels[lower_level, ...]
            upper_pressure = pressure_on_height_levels[upper_level, ...]

            surface_pressure = additional_data['surface_pressure']
            surface_pressure.convert_units('Pa')

            temperature_on_height_levels = (
                additional_data['temperature_on_height_levels'])
            lower_temperature = temperature_on_height_levels[lower_level, ...]
            upper_temperature = temperature_on_height_levels[upper_level, ...]

            return self.model_level_temperature_lapse_rate(
                cube, sites, neighbours, surface_pressure, lower_pressure,
                upper_pressure, lower_temperature, upper_temperature,
                dz_tolerance, dthetadz_threshold, dz_max_adjustment)

        raise AttributeError('Unknown method "{}" passed to {}.'.format(
            self.method, self.__class__.__name__))

    @staticmethod
    def _aux_coords_to_make():
        """
        Define coordinates that need to be made for the cube to be produced.

        Returns:
            dict:
                Dictionary of coordinates that are required for a spotdata
                cube.

        """
        return {'latitude': {'units': 'degrees', 'data_type': float,
                             'coord_type': AuxCoord},
                'longitude': {'units': 'degrees', 'data_type': float,
                              'coord_type': AuxCoord},
                'altitude': {'units': 'm', 'data_type': float,
                             'coord_type': AuxCoord,
                             'custom_function': np.nan_to_num},
                'wmo_site': {'data_type': int, 'coord_type': AuxCoord},
                'utc_offset': {'units': 'hours', 'data_type': float,
                               'coord_type': AuxCoord}}

    @staticmethod
    def make_stat_coordinate_first(cube):
        """
        Reorder cube dimension coordinates to ensure the statistical
        coordinate is first.

        Args:
            cube (iris.cube.Cube):
                The cube to be reordered.

        Returns:
            cube (iris.cube.Cube):
                Cube with the statistical coordinate moved to be first.

        Raises:
            Warning if more than one statistical dimension is found. Then
            promotes the first found to become the leading dimension.

        """
        stat_coordinates = ['realization', 'percentile_over']
        cube_dimension_order = {
            coord.name(): cube.coord_dims(coord.name())[0]
            for coord in cube.dim_coords}

        stat_coord = []
        for crd in stat_coordinates:
            stat_coord += [coord for coord in cube_dimension_order.keys()
{'fraction_non_alphanumeric': 0.07028435230863034, 'fraction_numerical': 0.0028268942963250374, 'mean_word_length': 2.1836950767601904, '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': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5920873', 'n_tokens_mistral': 4718, 'n_tokens_neox': 4283, 'n_words': 1665}