<commit_before>#!/usr/bin/env python
"""
  analysis.py

  Methods to assist in the analysis of ROMS fields

  Written by <NAME> on 05/24/15
  Copyright (c)2017 University of Hawaii under the MIT-License.
"""

import numpy as np
from joblib import Parallel, delayed
import seapy
import netCDF4

# Create a function to do the multiprocessing vertical interpolation.
# This function is used because we don't want to construct an array of
# tuples from the oa routine.


def __dinterp(x, y, z, dat, fz, pmap):
    ndat, pm = seapy.oavol(x, y, z, dat, x, y, fz, pmap, 5, 1, 1)
    ndat = np.squeeze(ndat)
    return np.ma.masked_where(np.abs(ndat) > 9e10, ndat, copy=False)


def constant_depth(field, grid, depth, zeta=None, threads=-2):
    """
    Find the values of a 3-D field at a constant depth for all times given.

    Parameters
    ----------
    field : ndarray,
        ROMS 3-D field to interpolate onto a constant depth level. Can be
        two- or three-dimensional array (first dimension assumed to be time).
    grid : seapy.model.grid or string or list,
        Grid that defines the depths and stretching for the field given
    depth : float,
        Depth (in meters) to find all values
    zeta : ndarray, optional,
        ROMS zeta field corresponding to field if you wish to apply the SSH
        correction to the depth calculations.
    threads : int, optional,
        Number of threads to use for processing

    Returns
    -------
    nfield : ndarray,
        Values from ROMS field on the given constant depth
    """

    # Make sure our inputs are all valid
    grid = seapy.model.asgrid(grid)
    if np.ndim(field) == 3:
        field = seapy.adddim(field)
    if zeta is not None and np.ndim(zeta == 2):
        zeta = seapy.adddim(zeta)
    depth = depth if depth < 0 else -depth

    # Set up some arrays
    x, y = np.meshgrid(np.arange(field.shape[-1]), np.arange(field.shape[-2]))
    fz, pmap = seapy.oasurf(x, y, x, x, y, None, 5, 1, 1)
    fz = seapy.adddim(np.ones(x.shape)) * depth
    # Loop over all times, generate new field at depth
    nfield = np.ma.array(Parallel(n_jobs=threads, verbose=2)
                         (delayed(__dinterp)(x, y, grid.depth_rho,
                                             np.squeeze(field[i, :, :, :]), fz, pmap)
                          for i in range(field.shape[0])), copy=False)

    return nfield


def depth_average(field, grid, depth, top_depth=0, zeta=None):
    """
    Compute the depth-averaged field down to the depth specified. NOTE:
    This just finds the nearest layer, so at every grid cell, it may not be
    exactly the specified depth.

    Parameters
    ----------
    field : ndarray,
        ROMS 3-D field to integrate from a depth level. Must be
        three-dimensional array (single time).
    grid : seapy.model.grid or string or list,
        Grid that defines the depths and stretching for the field given
    depth : float,
        Depth (in meters) to integrate from
    top_depth : float,
        Depth (in meters) to integrate to
    zeta : ndarray, optional,
        ROMS zeta field corresponding to field if you wish to apply the SSH
        correction to the depth calculations.

    Returns
    -------
    ndarray,
        Values from depth integrated ROMS field
    """
    grid = seapy.model.asgrid(grid)
    depth = depth if depth < 0 else -depth
    top_depth = top_depth if top_depth < 0 else -top_depth
    if depth > top_depth:
        depth, top_depth = top_depth, depth
    drange = top_depth - depth

    # If we have zeta, we need to compute thickness
    if zeta is not None:
        s_w, cs_w = seapy.roms.stretching(grid.vstretching, grid.theta_s,
                                          grid.theta_b, grid.hc,
                                          grid.n, w_grid=True)
        depths = np.ma.masked_equal(seapy.roms.depth(
            grid.vtransform, grid.h, grid.hc, grid.s_rho, grid.cs_r) *
            grid.mask_rho, 0)

        thickness = np.ma.masked_array(seapy.roms.thickness(
            grid.vtransform, grid.h, grid.hc, s_w, cs_w, zeta) *
            grid.mask_rho, 0)

    else:
        depths = np.ma.masked_equal(grid.depth_rho * grid.mask_rho, 0)
        thickness = np.ma.masked_equal(grid.thick_rho * grid.mask_rho, 0)

    # If we are on u- or v-grid, transform
    if field.shape == grid.thick_u.shape:
        depths = seapy.model.rho2u(depths)
        thickness = seapy.model.rho2u(thickness)
    elif field.shape == grid.thick_v.shape:
        depths = seapy.model.rho2v(depths)
        thickness = seapy.model.rho2v(thickness)

    # 1. pick all of the points that are deeper and shallower than the limits
    k_ones = np.arange(grid.n, dtype=int)
    top_depth = depths[-1, :, :] if top_depth == 0 else top_depth
    upper = depths - top_depth
    upper[np.where(upper < 0)] = np.float('inf')
    lower = depths - depth
    lower[np.where(lower > 0)] = -np.float('inf')
    thickness *= np.ma.masked_equal(np.logical_and(
        k_ones[:, np.newaxis, np.newaxis] <= np.argmin(upper, axis=0),
        k_ones[:, np.newaxis, np.newaxis] >=
        np.argmax(lower, axis=0)).astype(int), 0)

    # Do the integration
    return np.sum(field * thickness, axis=0) / \
        np.sum(thickness, axis=0)
<commit_msg>Add transect method to analysis to produce slices<commit_after>#!/usr/bin/env python
"""
  analysis.py

  Methods to assist in the analysis of ROMS fields

  Written by <NAME> on 05/24/15
  Copyright (c)2017 University of Hawaii under the MIT-License.
"""

import numpy as np
from joblib import Parallel, delayed
import seapy
import netCDF4

# Create a function to do the multiprocessing vertical interpolation.
# This function is used because we don't want to construct an array of
# tuples from the oa routine.


def __dinterp(x, y, z, dat, fz, pmap):
    ndat, pm = seapy.oavol(x, y, z, dat, x, y, fz, pmap, 5, 1, 1)
    ndat = np.squeeze(ndat)
    return np.ma.masked_where(np.abs(ndat) > 9e10, ndat, copy=False)


def constant_depth(field, grid, depth, zeta=None, threads=-2):
    """
    Find the values of a 3-D field at a constant depth for all times given.

    Parameters
    ----------
    field : ndarray,
        ROMS 3-D field to interpolate onto a constant depth level. Can be
        two- or three-dimensional array (first dimension assumed to be time).
    grid : seapy.model.grid or string or list,
        Grid that defines the depths and stretching for the field given
    depth : float,
        Depth (in meters) to find all values
    zeta : ndarray, optional,
        ROMS zeta field corresponding to field if you wish to apply the SSH
        correction to the depth calculations.
    threads : int, optional,
        Number of threads to use for processing

    Returns
    -------
    nfield : ndarray,
        Values from ROMS field on the given constant depth
    """

    # Make sure our inputs are all valid
    grid = seapy.model.asgrid(grid)
    if np.ndim(field) == 3:
        field = seapy.adddim(field)
    if zeta is not None and np.ndim(zeta == 2):
        zeta = seapy.adddim(zeta)
    depth = depth if depth < 0 else -depth

    # Set up some arrays
    x, y = np.meshgrid(np.arange(field.shape[-1]), np.arange(field.shape[-2]))
    fz, pmap = seapy.oasurf(x, y, x, x, y, None, 5, 1, 1)
    fz = seapy.adddim(np.ones(x.shape)) * depth
    # Loop over all times, generate new field at depth
    nfield = np.ma.array(Parallel(n_jobs=threads, verbose=2)
                         (delayed(__dinterp)(x, y, grid.depth_rho,
                                             np.squeeze(field[i, :, :, :]), fz, pmap)
                          for i in range(field.shape[0])), copy=False)

    return nfield


def depth_average(field, grid, depth, top_depth=0, zeta=None):
    """
    Compute the depth-averaged field down to the depth specified. NOTE:
    This just finds the nearest layer, so at every grid cell, it may not be
    exactly the specified depth.

    Parameters
    ----------
    field : ndarray,
        ROMS 3-D field to integrate from a depth level. Must be
        three-dimensional array (single time).
    grid : seapy.model.grid or string or list,
        Grid that defines the depths and stretching for the field given
    depth : float,
        Depth (in meters) to integrate from
    top_depth : float,
        Depth (in meters) to integrate to
    zeta : ndarray, optional,
        ROMS zeta field corresponding to field if you wish to apply the SSH
        correction to the depth calculations.

    Returns
    -------
    ndarray,
        Values from depth integrated ROMS field
    """
    grid = seapy.model.asgrid(grid)
    depth = depth if depth < 0 else -depth
    top_depth = top_depth if top_depth < 0 else -top_depth
    if depth > top_depth:
        depth, top_depth = top_depth, depth
    drange = top_depth - depth

    # If we have zeta, we need to compute thickness
    if zeta is not None:
        s_w, cs_w = seapy.roms.stretching(grid.vstretching, grid.theta_s,
                                          grid.theta_b, grid.hc,
                                          grid.n, w_grid=True)
        depths = np.ma.masked_equal(seapy.roms.depth(
            grid.vtransform, grid.h, grid.hc, grid.s_rho, grid.cs_r) *
            grid.mask_rho, 0)

        thickness = np.ma.masked_array(seapy.roms.thickness(
            grid.vtransform, grid.h, grid.hc, s_w, cs_w, zeta) *
            grid.mask_rho, 0)

    else:
        depths = np.ma.masked_equal(grid.depth_rho * grid.mask_rho, 0)
        thickness = np.ma.masked_equal(grid.thick_rho * grid.mask_rho, 0)

    # If we are on u- or v-grid, transform
    if field.shape == grid.thick_u.shape:
        depths = seapy.model.rho2u(depths)
        thickness = seapy.model.rho2u(thickness)
    elif field.shape == grid.thick_v.shape:
        depths = seapy.model.rho2v(depths)
        thickness = seapy.model.rho2v(thickness)

    # 1. pick all of the points that are deeper and shallower than the limits
    k_ones = np.arange(grid.n, dtype=int)
    top_depth = depths[-1, :, :] if top_depth == 0 else top_depth
    upper = depths - top_depth
    upper[np.where(upper < 0)] = np.float('inf')
    lower = depths - depth
    lower[np.where(lower > 0)] = -np.float('inf')
    thickness *= np.ma.masked_equal(np.logical_and(
        k_ones[:, np.newaxis, np.newaxis] <= np.argmin(upper, axis=0),
        k_ones[:, np.newaxis, np.newaxis] >=
        np.argmax(lower, axis=0)).astype(int), 0)

    # Do the integration
    return np.sum(field * thickness, axis=0) / \
        np.sum(thickness, axis=0)


def transect(lon, lat, depth, data, nx=200, nz=40, z=None):
    """
    Generate an equidistant transect from data at varying depths. Can be
    used to plot a slice of model or observational data.

    Parameters
    ----------
    lat: array
        n-dimensional array with latitude of points
    lon: array
        n-dimensional array with longitude of points
    depth: array
        [k,n] dimensional array of the depths for all k-layers at each n point
    data: array
        [k,n] dimensional array of values
    nx: int, optional
        number of horizontal points desired in the transect
    nz: int, optional
        number of vertical points desired in the transect
    z: array, optional
        list of depths to use if you do not want equidistant depths

    Returns
    -------
    x: array
        x-location values of the new transect
    z: array
        z-location values of the new transect
    vals: array
        data values of the new transect
    """
    from scipy.interpolate import griddata
    depth = np.atleast_2d(depth)
    data = np.atleast_2d(data)
    lon = np.atleast_1d(lon)
    lat = np.atleast_1d(lat)

    # Generate the depths
    depth[depth > 0] *= -1
    if z is None:
        z = np.linspace(depth.min() - 2, depth.max(), nz)
    else:
        z[z > 0] *= -1
        nz = len(z)
    dz = np.abs(np.diff(z).mean())

    # Determine the distance between points and the weighting to apply
    dist = np.hstack(([0], seapy.earth_distance(
        lon[0], lat[0], lon[1:], lat[1:])))
    dx = np.diff(dist).mean()
    zscale = np.maximum(1, 10**int(np.log10(dx / dz)))
    dx /= zscale
    x = np.linspace(0, dist.max(), nx)

    # All arrays have to be the same size
    xx, zz = np.meshgrid(x / zscale, z)

    # For the source data, we don't want to extrpolate,
    # so make the data go from the surface to twice its
    # depth.
    zl = np.argsort(depth[:, 0])
    dep = np.vstack((np.ones((1, depth.shape[1])) * 2 * depth.min(),
                     depth[zl, :],
                     np.zeros((1, depth.shape[1]))))

    # repeat the same data at the top and bottom
    dat = np.vstack((data[zl[0], :], data[zl], data[zl[-1], :]))
    dist = np.tile(dist.T, [dep.shape[0], 1]) / zscale

    # Find the bottom indices to create a mask for nodata/land
    idx = np.interp(xx[0, :], dist[0, :],
                    np.interp(depth.min(axis=0), z,
                              np.arange(nz))).astype(int)
    mask = np.arange(nz)[:, np.newaxis] <= idx

    # Interpolate
    ndat = np.ma.array(griddata(
        (dist.ravel(), dep.ravel()), dat.ravel(), (xx.ravel(), zz.ravel()),
        method='cubic').reshape(xx.shape), mask=mask)

    # Return everything
    return x, z, ndat

{'fraction_non_alphanumeric': 0.10845091564094866, 'fraction_numerical': 0.010657460222155508, 'mean_word_length': 2.2901234567901234, 'pattern_counts': {'":': 0, '<': 16, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1167460', 'n_tokens_mistral': 4449, 'n_tokens_neox': 4285, 'n_words': 1412}