<commit_before>"""
Utility functions for atmospheric data wrangling / preparation.

- ndarrays
- netCDF files
- Lat-lon geophysical data
- Pressure level data and topography
"""

from __future__ import division
import numpy as np
import collections
import scipy.interpolate as interp
from mpl_toolkits import basemap
import xray
from xray import Dataset

from atmos.utils import print_if, disptime
import atmos.utils as utils
import atmos.xrhelper as xr
#from atmos.xrhelper import subset
from atmos.constants import const as constants

# ======================================================================
# NDARRAYS
# ======================================================================

# ----------------------------------------------------------------------
def biggify(small, big, tile=False):
    """Add dimensions or tile an array for broadcasting.

    Parameters
    ----------
    small : ndarray
        Array which singleton dimensions will be added to.  Its
        dimensions must be a subset of big's dimensions.
    big : ndarray
        Array whose shape will be used to determine the shape of
        the output.
    tile : bool, optional
        If True, tile the array along the additional dimensions.
        If False, add singleton dimensions.

    Returns
    -------
    biggified : ndarray
        Array of data from small, with dimensions added
        for any dimension that is in big but not in small.
    """

    debug = False
    dbig, dsmall = big.shape, small.shape

    # Check that all of the dimensions of small are contained within big
    check = [d in dbig or d == 1 for d in dsmall]
    if not np.all(check):
        msg = ('Dimensions of small ' + str(dsmall) +
            ' are not a subset of big ' + str(dbig))
        raise ValueError(msg)

    # Check that the dimensions appear in a compatible order
    inds = list()
    for d in dsmall:
        try:
            inds.append(dbig.index(d))
        except ValueError:
            inds.append(-1)
    if not utils.non_decreasing(inds):
        msg = ('Dimensions of small ' + str(dsmall) +
            ' are not in an order compatible with big ' + str(dbig))
        raise ValueError(msg)

    # Biggify the small array
    biggified = small
    ibig = big.ndim - 1
    ismall = small.ndim - 1
    n = -1

    # First add singleton dimensions
    while ismall >= 0 and ibig >= 0:
        print_if('ibig %d, ismall %d, n %d' % (ibig, ismall, n), debug)
        if dbig[ibig] == dsmall[ismall] or dsmall[ismall] == 1:
            print_if('  Same %d' % dbig[ibig], debug)
            ismall -= 1
        else:
            print_if('  Different.  Big %d, small %d' %
                (dbig[ibig], dsmall[ismall]), debug)
            biggified = np.expand_dims(biggified, n)
        n -= 1
        ibig -= 1

    # Expand with tiles if selected
    if tile:
        dims = list(biggified.shape)

        # First add any additional singleton dimensions needed to make
        # biggified of the same dimension as big\
        for i in range(len(dims), len(dbig)):
            dims.insert(0, 1)

        # Tile the array
        for i in range(-1, -1 - len(dims), -1):
            if dims[i] == dbig[i]:
                dims[i] = 1
            else:
                dims[i] = dbig[i]
        biggified = np.tile(biggified, dims)

    return biggified


# ----------------------------------------------------------------------
def nantrapz(y, x=None, axis=-1):
    """
    Integrate using the composite trapezoidal rule, ignoring NaNs

    Integrate `ym` (`x`) along given axis, where `ym` is a masked
    array of `y` with NaNs masked.

    Parameters
    ----------
    y : array_like
        Input array to integrate.
    x : array_like, optional
        If `x` is None, then spacing between all `y` elements is `dx`.
    axis : int, optional
        Specify the axis.

    Returns
    -------
    trapz : float
        Definite integral as approximated by trapezoidal rule.
    """

    ym = np.ma.masked_array(y, np.isnan(y))
    trapz = np.trapz(ym, x, axis=axis)

    # Convert from masked array back to regular ndarray
    if isinstance(trapz, np.ma.masked_array):
        trapz = trapz.filled(np.nan)

    return trapz


# ======================================================================
# NETCDF FILE I/O
# ======================================================================

# ----------------------------------------------------------------------
def ncdisp(filename, verbose=True, decode_cf=False, indent=2, width=None):
    """Display the attributes of data in a netcdf file."""
    with xray.open_dataset(filename, decode_cf=decode_cf) as ds:
        if verbose:
            xr.ds_print(ds, indent, width)
        else:
            print(ds)


# ----------------------------------------------------------------------
def ncload(filename, verbose=True, unpack=True, missing_name=u'missing_value',
           offset_name=u'add_offset', scale_name=u'scale_factor',
           decode_cf=False):
    """
    Read data from netcdf file into xray dataset.

    If options are selected, unpacks from compressed form and/or replaces
    missing values with NaN.  Returns data as an xray.Dataset object.
    """
    with xray.open_dataset(filename, decode_cf=decode_cf) as ds:
        print_if('****** Reading file: ' + filename + '********', verbose)
        print_if(ds, verbose, printfunc=xr.ds_print)
        if unpack:
            print_if('****** Unpacking data *********', verbose)
            ds = xr.ds_unpack(ds, verbose=verbose, missing_name=missing_name,
                offset_name=offset_name, scale_name=scale_name)

        # Use the load() function so that the dataset is available after
        # the file is closed
        ds.load()
        return ds


# ----------------------------------------------------------------------
def load_concat(paths, var, concat_dim=None, verbose=False):
    """Load a variable from multiple files and concatenate into one.

    Especially useful for extracting variables split among multiple
    OpenDAP files.
<commit_msg>Changed default concat_dim in load_concat<commit_after>"""
Utility functions for atmospheric data wrangling / preparation.

- ndarrays
- netCDF files
- Lat-lon geophysical data
- Pressure level data and topography
"""

from __future__ import division
import numpy as np
import collections
import scipy.interpolate as interp
from mpl_toolkits import basemap
import xray
from xray import Dataset

from atmos.utils import print_if, disptime
import atmos.utils as utils
import atmos.xrhelper as xr
#from atmos.xrhelper import subset
from atmos.constants import const as constants

# ======================================================================
# NDARRAYS
# ======================================================================

# ----------------------------------------------------------------------
def biggify(small, big, tile=False):
    """Add dimensions or tile an array for broadcasting.

    Parameters
    ----------
    small : ndarray
        Array which singleton dimensions will be added to.  Its
        dimensions must be a subset of big's dimensions.
    big : ndarray
        Array whose shape will be used to determine the shape of
        the output.
    tile : bool, optional
        If True, tile the array along the additional dimensions.
        If False, add singleton dimensions.

    Returns
    -------
    biggified : ndarray
        Array of data from small, with dimensions added
        for any dimension that is in big but not in small.
    """

    debug = False
    dbig, dsmall = big.shape, small.shape

    # Check that all of the dimensions of small are contained within big
    check = [d in dbig or d == 1 for d in dsmall]
    if not np.all(check):
        msg = ('Dimensions of small ' + str(dsmall) +
            ' are not a subset of big ' + str(dbig))
        raise ValueError(msg)

    # Check that the dimensions appear in a compatible order
    inds = list()
    for d in dsmall:
        try:
            inds.append(dbig.index(d))
        except ValueError:
            inds.append(-1)
    if not utils.non_decreasing(inds):
        msg = ('Dimensions of small ' + str(dsmall) +
            ' are not in an order compatible with big ' + str(dbig))
        raise ValueError(msg)

    # Biggify the small array
    biggified = small
    ibig = big.ndim - 1
    ismall = small.ndim - 1
    n = -1

    # First add singleton dimensions
    while ismall >= 0 and ibig >= 0:
        print_if('ibig %d, ismall %d, n %d' % (ibig, ismall, n), debug)
        if dbig[ibig] == dsmall[ismall] or dsmall[ismall] == 1:
            print_if('  Same %d' % dbig[ibig], debug)
            ismall -= 1
        else:
            print_if('  Different.  Big %d, small %d' %
                (dbig[ibig], dsmall[ismall]), debug)
            biggified = np.expand_dims(biggified, n)
        n -= 1
        ibig -= 1

    # Expand with tiles if selected
    if tile:
        dims = list(biggified.shape)

        # First add any additional singleton dimensions needed to make
        # biggified of the same dimension as big\
        for i in range(len(dims), len(dbig)):
            dims.insert(0, 1)

        # Tile the array
        for i in range(-1, -1 - len(dims), -1):
            if dims[i] == dbig[i]:
                dims[i] = 1
            else:
                dims[i] = dbig[i]
        biggified = np.tile(biggified, dims)

    return biggified


# ----------------------------------------------------------------------
def nantrapz(y, x=None, axis=-1):
    """
    Integrate using the composite trapezoidal rule, ignoring NaNs

    Integrate `ym` (`x`) along given axis, where `ym` is a masked
    array of `y` with NaNs masked.

    Parameters
    ----------
    y : array_like
        Input array to integrate.
    x : array_like, optional
        If `x` is None, then spacing between all `y` elements is `dx`.
    axis : int, optional
        Specify the axis.

    Returns
    -------
    trapz : float
        Definite integral as approximated by trapezoidal rule.
    """

    ym = np.ma.masked_array(y, np.isnan(y))
    trapz = np.trapz(ym, x, axis=axis)

    # Convert from masked array back to regular ndarray
    if isinstance(trapz, np.ma.masked_array):
        trapz = trapz.filled(np.nan)

    return trapz


# ======================================================================
# NETCDF FILE I/O
# ======================================================================

# ----------------------------------------------------------------------
def ncdisp(filename, verbose=True, decode_cf=False, indent=2, width=None):
    """Display the attributes of data in a netcdf file."""
    with xray.open_dataset(filename, decode_cf=decode_cf) as ds:
        if verbose:
            xr.ds_print(ds, indent, width)
        else:
            print(ds)


# ----------------------------------------------------------------------
def ncload(filename, verbose=True, unpack=True, missing_name=u'missing_value',
           offset_name=u'add_offset', scale_name=u'scale_factor',
           decode_cf=False):
    """
    Read data from netcdf file into xray dataset.

    If options are selected, unpacks from compressed form and/or replaces
    missing values with NaN.  Returns data as an xray.Dataset object.
    """
    with xray.open_dataset(filename, decode_cf=decode_cf) as ds:
        print_if('****** Reading file: ' + filename + '********', verbose)
        print_if(ds, verbose, printfunc=xr.ds_print)
        if unpack:
            print_if('****** Unpacking data *********', verbose)
            ds = xr.ds_unpack(ds, verbose=verbose, missing_name=missing_name,
                offset_name=offset_name, scale_name=scale_name)

        # Use the load() function so that the dataset is available after
        # the file is closed
        ds.load()
        return ds


# ----------------------------------------------------------------------
def load_concat(paths, var, concat_dim='TIME', verbose=False):
    """Load a variable from multiple files and concatenate into one.

    Especially useful for extracting variables split among multiple
    OpenDAP files.

{'fraction_non_alphanumeric': 0.19003115264797507, 'fraction_numerical': 0.003115264797507788, 'mean_word_length': 2.0444222610431746, 'pattern_counts': {'":': 0, '<': 3, '<?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': '2585263', 'n_tokens_mistral': 3480, 'n_tokens_neox': 3313, 'n_words': 1164}