LON_STR = 'lon' LAT_STR = 'lat' PFULL_STR = 'pfull' _RADEARTH = 6370997. def to_radians(arr, is_delta=False):
    """Force data with units either degrees or radians to be radians."""
    # Infer the units from embedded metadata, if it's there.
    arr_out = arr
    try:
        units = arr_out.units
    except AttributeError:
        pass
    else:
        if units.lower().startswith('degrees'):
            warn_msg = ("Conversion applied: degrees -> radians to arr_outay: "
                        "{}".format(arr_out))
            warnings.warn(warn_msg, UserWarning)
            return np.deg2rad(arr_out)
    # Otherwise, assume degrees if the values are sufficiently large.
    threshold = 0.1*np.pi if is_delta else 4*np.pi
    if np.max(np.abs(arr_out)) > threshold:
        warn_msg = ("Conversion applied: degrees -> radians to arr_outay: "
                    "{}".format(arr_out))
        warnings.warn(warn_msg, UserWarning)
        return np.deg2rad(arr_out)
    return arr_out def wraparound(arr, dim, left_to_right=0, right_to_left=0,
               circumf=360., spacing=1):
    """Append wrap-around point(s) to the DataArray or Dataset coord."""
    # TODO: Bugfix: if points on both sides are added, as written this
    # causes the points first added to the right to then be added to the left,
    # rather than the original points on the right.
    new = add_cyclic_to_right(arr, dim, left_to_right, circumf)
    return add_cyclic_to_left(new, dim, right_to_left, circumf) class Coord(object):
    """Generic base class for physical coordinates."""
    _POSSIBLY_CYCLIC = True

    def _prep_dim(self, dim):
        msg = ("Specified dim '{}' does not match any dim in self.arr.  "
               "self.arr.dims: {}").format(dim, self.arr.dims)
        if dim in self.arr.dims:
            return dim
        if dim is None:
            if self.arr.shape == 1:
                return self.arr.dims[0]
        raise ValueError(msg)

    def __init__(self, arr, dim=None, cyclic=False):
        self.arr = arr
        self.dim = self._prep_dim(dim)
        if cyclic and not self._POSSIBLY_CYCLIC:
            raise ValueError("The coordinate cannot be cyclic in a physically "
                             "meaningful way.")
        self.cyclic = cyclic

    def __getitem__(self, key):
        return self.arr[key]

    def deriv_prefactor(self, *args, **kwargs):
        raise NotImplementedError

    def deriv_factor(self, *args, **kwargs):
        raise NotImplementedError class Lon(XCoord):
    """Longitude spherical horizontal coordinate."""
    def __init__(self, lon, dim=None, cyclic=True, radius=_RADEARTH):
        super(Lon, self).__init__(lon, dim=dim, cyclic=cyclic)
        self.radius = radius

    def deriv_prefactor(self, lat):
        return 1. / (self.radius * np.cos(to_radians(lat))) class Lat(YCoord):
    """Latitude spherical horizontal coordinate."""
    _POSSIBLY_CYCLIC = False

    def __init__(self, lat, dim=None, radius=_RADEARTH):
        super(Lat, self).__init__(lat, dim=dim, cyclic=False)
        self.radius = radius
        self._lat_rad = to_radians(lat)

    def deriv_prefactor(self, oper='grad'):
        if oper == 'grad':
            return 1. / self.radius
        if oper == 'divg':
            return 1. / (self.radius * np.cos(self._lat_rad))
        msg = "'oper' must be 'grad' or 'divg': value was '{}'".format(oper)
        raise ValueError(msg)

    def deriv_factor(self, oper='grad'):
        if oper == 'grad':
            return 1.
        if oper == 'divg':
            return np.cos(self._lat_rad)
        msg = "'oper' must be 'grad' or 'divg': value was '{}'".format(oper)
        raise ValueError(msg) class Eta(VertCoord):
    """Hybrid sigma-pressure vertical coordinates."""
    def __init__(self, pk, bk, pfull, dim=None):
        self.pk = pk
        self.bk = bk
        self.pfull = pfull
        self.phalf = self.pk[PHALF_STR]
        self.arr = self.phalf
        self.dim = dim if dim is not None else self.pk.dims[0]

    def phalf_from_ps(self, ps):
        """Compute pressure at level edges from surface pressure."""
        return self.pk + self.bk*ps

    def to_pfull_from_phalf(self, arr):
        """Compute data at full pressure levels from values at half levels."""
        arr_top = arr.copy()[{PHALF_STR: slice(1, None)}]
        arr_top = replace_coord(arr_top, PHALF_STR, PFULL_STR, self.pfull)

        arr_bot = arr.copy()[{PHALF_STR: slice(None, -1)}]
        arr_bot = replace_coord(arr_bot, PHALF_STR, PFULL_STR, self.pfull)
        return 0.5*(arr_bot + arr_top)

    def pfull_from_ps(self, ps):
        """Compute pressure at full levels from surface pressure."""
        return self.to_pfull_from_phalf(self.phalf_from_ps(ps))

    def d_deta_from_phalf(self, arr):
        """Compute pressure level thickness from half level pressures."""
        d_deta = arr.diff(dim=PHALF_STR, n=1)
        return replace_coord(d_deta, PHALF_STR, PFULL_STR, self.pfull)

    def d_deta_from_pfull(self, arr):
        """Compute $\partial/\partial\eta$ of the array on full hybrid levels.

        $\eta$ is the model vertical coordinate, and its value is assumed to
        simply increment by 1 from 0 at the surface upwards.  The data to be
        differenced is assumed to be defined at full pressure levels.
        """
        deriv = CenDiff(arr, PFULL_STR, spacing=1, fill_edge=True).diff() / 2.
        # Edges use 1-sided differencing, so only spanning one level, not two.
        deriv[{PFULL_STR: 0}] = deriv[{PFULL_STR: 0}] * 2.
        deriv[{PFULL_STR: -1}] = deriv[{PFULL_STR: -1}] * 2.
        return deriv

    def dp_from_ps(self, ps):
        """Compute pressure level thickness from surface pressure"""
        return self.d_deta_from_phalf(self.phalf_from_ps(ps)) class FiniteDeriv(object):
    """Base class for finite-diff based derivative classes."""
    _DIFF_CLS = FiniteDiff
    _VALID_ORDERS = range(1, 5)

    def _arr_coord(self, coord):
        """Get the coord to be used as the denominator for a derivative."""
        if coord is None:
            return self.arr[self.dim]
        return coord

    def __init__(self, arr, dim, coord=None, spacing=1, order=1,
                 fill_edge=True):
        """
        :param arr: Field to take derivative of.
        :param str dim: Name of dimension over which to take the derivative.
        :param xarray.DataArray coord: Coordinate array to use for the
            denominator.  If not given, arr[dim] is used.
        """
        self.arr = arr
        self.dim = dim
        self.coord = self._arr_coord(coord)
        self.spacing = spacing
        assert order in self._VALID_ORDERS
        self.order = order
        self.fill_edge = fill_edge

        self._arr_diff_obj = self._DIFF_CLS(self.arr, self.dim,
                                            spacing=self.spacing)
        self._coord_diff_obj = self._DIFF_CLS(self.coord, self.dim,
                                              spacing=self.spacing)

    def _deriv(self):
        """Core finite-differencing derivative; no edge handling."""
        return self._arr_diff_obj.diff() / self._coord_diff.diff()

    def _slice_edge(self, arr):
        raise NotImplementedError

    def _concat(self):
        raise NotImplementedError

    def deriv(self):
        raise NotImplementedError class FwdDeriv(OneSidedDeriv):
    """Derivatives using forward differencing."""
    _DIFF_CLS = FwdDiff
    _DIFF_REV_CLS = BwdDiff

    def __init__(self, arr, dim, coord=None, spacing=1, order=1,
                 fill_edge=True):
        super(FwdDeriv, self).__init__(arr, dim, coord=coord, spacing=spacing,
                                       order=order, fill_edge=fill_edge)

    def _slice_edge(self, arr):
        return arr[{self.dim: slice(-self.spacing*self.order, None)}]

    def _concat(self, interior, edge):
        return xr.concat([interior, edge], dim=self.dim) class BwdDeriv(OneSidedDeriv):
    """Derivatives using backward differencing."""
    _DIFF_CLS = BwdDiff
    _DIFF_REV_CLS = FwdDiff

    def __init__(self, arr, dim, coord=None, spacing=1, order=1,
                 fill_edge=True):
        super(BwdDeriv, self).__init__(arr, dim, coord=coord, spacing=spacing,
                                       order=order, fill_edge=fill_edge)

    def _slice_edge(self, arr):
        return arr[{self.dim: slice(None, self.spacing*self.order)}]

    def _concat(self, interior, edge):
        return xr.concat([edge, interior], dim=self.dim) class CenDeriv(FiniteDeriv):
    _DIFF_CLS = CenDiff
    _VALID_ORDERS = [2, 4]

    """Derivatives computed via centered finite differencing."""
    def __init__(self, arr, dim, coord=None, spacing=1, order=2,
                 fill_edge=True):
        """
        :param arr: Data to be center-differenced.
        :type arr: `xarray.DataArray` or `xarray.Dataset`
        :param str dim: Dimension over which to compute.
        :param xarray.DataArray coord: Coordinate array to use for the
            denominator.  If not given or None, arr[dim] is used.
        """
        super(CenDeriv, self).__init__(arr, dim, coord=coord, spacing=spacing,
                                       order=order, fill_edge=fill_edge)
        self._deriv_fwd_obj = FwdDeriv(arr, dim, coord=self.coord,
                                       spacing=self.spacing, order=2,
                                       fill_edge=self.fill_edge)
        self._deriv_bwd_obj = BwdDeriv(arr, dim, coord=self.coord,
                                       spacing=self.spacing, order=2,
                                       fill_edge=self.fill_edge)

    def _edge_deriv(self):
        left = self._deriv_bwd_obj._edge_deriv_rev()
        right = self._deriv_fwd_obj._edge_deriv_rev()
        pad = self.spacing*(self.order // 2)
        return (left[{self.dim: slice(None, pad)}],
                right[{self.dim: slice(-pad, None)}])

    def _concat(self, left, interior, right):
        return xr.concat([left, interior, right], dim=self.dim)

    def _deriv(self):
        """Lowest possible order derivative with this scheme."""
        interior = self._arr_diff_obj.diff() / self._coord_diff_obj.diff()
        if not self.fill_edge:
            return interior
        left, right = self._edge_deriv()
        return self._concat(left, interior, right)

    def deriv(self):
        """
        Centered differencing approximation of 1st derivative.

        :param int order: Order of accuracy to use.  Default 2.
        :param fill_edge: Whether or not to fill in the edge cells
            that don't have the needed neighbor cells for the stencil.  If
            `True`, use one-sided differencing with the same order of accuracy
            as `order`, and the outputted array is the same shape as `arr`.

            If `False`, the outputted array has a length in the computed axis
            reduced by `order`.
        """
        if self.order == 2:
            return self._deriv()
        if self.order == 4:
            single_space = self.__class__(self.arr, self.dim, coord=self.coord,
                                          spacing=self.spacing, order=2,
                                          fill_edge=self.fill_edge)._deriv()
            double_space = self.__class__(self.arr, self.dim, coord=self.coord,
                                          spacing=2*self.spacing, order=2,
                                          fill_edge=False)._deriv()
            interior = (4*single_space - double_space) / 3
            if not self.fill_edge:
                return interior
            left = single_space[{self.dim: slice(0, self.spacing*2)}]
            right = single_space[{self.dim: slice(-self.spacing*2, None)}]
            return self._concat(left, interior, right)
        raise NotImplementedError("Centered differencing only "
                                  "supported for 2nd and 4th order.")
import warnings
from .._constants import LON_STR, LAT_STR, PFULL_STR, _RADEARTH
from ..utils import to_radians, wraparound
from ..coord import Coord, Lon, Lat, Eta
from . import FiniteDeriv, FwdDeriv, BwdDeriv, CenDeriv
class LatCenDeriv(LatDeriv):
    _DERIV_CLS = CenDeriv


class EtaDeriv(object):
    _DERIV_CLS = FiniteDeriv
    _COORD_CLS = Eta

    def __init__(self, arr, pk, bk, ps, spacing=1, order=2, fill_edge=True,
                 **coord_kwargs):
        self.arr = arr.copy(deep=True)
        self.dim = PFULL_STR
        self.ps = ps
        self.spacing = spacing
        self.order = order
        self.fill_edge = fill_edge

        self._coord_obj = self._COORD_CLS(pk, bk, self.arr[self.dim],
                                          **coord_kwargs)
        self.pfull = self._coord_obj.pfull
        for method in ['phalf_from_ps',
                       'to_pfull_from_phalf',
                       'pfull_from_ps',
                       'd_deta_from_phalf',
                       'd_deta_from_pfull',
                       'dp_from_ps']:
            setattr(self, method, getattr(self._coord_obj, method))

    def deriv(self):
        pfull = self.pfull_from_ps(self.ps)
        return self._DERIV_CLS(self.arr.copy(deep=True), self.dim, coord=pfull,
                               spacing=self.spacing, order=self.order,
                               fill_edge=self.fill_edge).deriv()


class EtaFwdDeriv(EtaDeriv):
    _DERIV_CLS = FwdDeriv


class EtaBwdDeriv(EtaDeriv):
    _DERIV_CLS = BwdDeriv


class EtaCenDeriv(EtaDeriv):
    _DERIV_CLS = CenDeriv


class HorizPhysDeriv(object):
    """Horizontal derivatives."""
    _X_DERIV_CLS = PhysDeriv
    _Y_DERIV_CLS = PhysDeriv

    def __init__(self, arr, x_dim, y_dim, x_coord=None, y_coord=None,
                 **kwargs):
        self._x_deriv_obj = self._X_DERIV_CLS(arr.copy(deep=True), x_dim,
                                              coord=x_coord, **kwargs)
        self._y_deriv_obj = self._Y_DERIV_CLS(arr.copy(deep=True), y_dim,
                                              coord=y_coord, **kwargs)

    def d_dx(self, *args, **kwargs):
        return self._x_deriv_obj.deriv(*args, **kwargs)

    def d_dy(self, *args, **kwargs):
        return self._y_deriv_obj.deriv(*args, **kwargs)

    def horiz_grad(self, *args, **kwargs):
        return self.d_dx(*args, **kwargs) + self.d_dy(*args, **kwargs)


class SphereDeriv(HorizPhysDeriv):
    """Derivatives for data on a sphere."""
    _X_DERIV_CLS = LonDeriv
    _Y_DERIV_CLS = LatDeriv

    def __init__(self, arr, x_coord=None, y_coord=None, cyclic_lon=True,
                 fill_edge_lon=False, fill_edge_lat=True, **kwargs):
        self.arr = arr.copy(deep=True)
        self.cyclic_lon = cyclic_lon
        self._x_deriv_obj = self._X_DERIV_CLS(
            arr.copy(deep=True), LON_STR, coord=x_coord, cyclic=cyclic_lon,
            fill_edge=fill_edge_lon, **kwargs
        )
        self._y_deriv_obj = self._Y_DERIV_CLS(
            arr.copy(deep=True), LAT_STR, coord=y_coord,
            fill_edge=fill_edge_lat, **kwargs
        )
        self.d_dy = self._y_deriv_obj.deriv

    def d_dx(self):
        return self._x_deriv_obj.deriv(self.arr.copy(deep=True))

    def horiz_grad(self):
        return self.d_dx() + self.d_dy(oper='grad')


class SphereFwdDeriv(SphereDeriv):
    """Derivatives for data on a sphere."""
    _X_DERIV_CLS = LonFwdDeriv
    _Y_DERIV_CLS = LatFwdDeriv


class SphereBwdDeriv(SphereDeriv):
    """Derivatives for data on a sphere."""
    _X_DERIV_CLS = LonBwdDeriv
    _Y_DERIV_CLS = LatBwdDeriv


class SphereCenDeriv(SphereDeriv):
    """Derivatives for data on a sphere."""
    _X_DERIV_CLS = LonCenDeriv
    _Y_DERIV_CLS = LatCenDeriv


class SphereEtaDeriv(object):
    """Derivatives on the sphere with hybrid sigma-pressure in the vertical."""
    _HORIZ_DERIV_CLS = SphereDeriv
    _VERT_DERIV_CLS = EtaDeriv

    def __init__(self, arr, pk, bk, ps, spacing=1, order=2, cyclic_lon=True,
                 fill_edge_lon=False, fill_edge_lat=True, fill_edge_vert=True,

                 radius=_RADEARTH):