Merge 38beb7c into 8c82c75

.gitignore

@@ -103,3 +103,6 @@ venv.bak/
 # mypy
 .mypy_cache/
 *.DS_Store
+
+# VS Code
+.vscode/*

.pre-commit-config.yaml

@@ -22,3 +22,10 @@ repos:
       hooks:
       - id: flake8
         args: ["--max-line-length=88", "--exclude=__init__.py", "--ignore=W605,W503"]
+
+  -   repo: https://github.com/asottile/blacken-docs
+      rev: v1.4.0
+      hooks:
+      - id: blacken-docs
+        additional_dependencies: [black]
+        args: ["--line-length", "88"]

README.rst

@@ -22,7 +22,7 @@ a toolbox for Earth System Model analysis
 Installation
 ============
 
-You can install the latest release of ``esmtools`` using ``pip``:
+You can install the latest release of ``esmtools`` using ``pip``. We will also release ``esmtools`` on conda for the next release.
 
 .. code-block:: bash
 

ci/environment-dev-3.6.yml

@@ -3,32 +3,37 @@ channels:
   - conda-forge
 dependencies:
   - python=3.6
-  # development
-  - black
-  - coveralls
-  - flake8
-  - importlib_metadata
-  - jupyterlab
+  # Documentation
   - nbsphinx
-  - pytest
-  - pytest-cov
-  - sphinx
   - sphinx_rtd_theme
-  # input/output
+  - sphinx
+  - sphinxcontrib-napoleon
+  # IDE
+  - jupyterlab
+  # Input/Output
   - netcdf4
-  # numerics
+  # Miscellaneous
   - cftime
+  # Numerics
   - bottleneck
   - dask
   - numpy
   - xarray
-  # stats
+  # Package Management
+  - black
+  - coveralls
+  - flake8
+  - importlib_metadata
+  - pre-commit
+  - pytest
+  - pytest-cov
+  - pytest-sugar
+  # Statistics
   - climpred
   - scipy
   - statsmodels
-  # visualization
+  # Visualization
   - matplotlib
   - pip
   - pip:
-      - sphinxcontrib-napoleon
-      - pre-commit
+      - pytest-tldr

esmtools/__init__.py

@@ -10,6 +10,7 @@
     stats,
     temporal,
     testing,
+    utils,
 )
 from .accessor import GridAccessor
 from .versioning.print_versions import show_versions

esmtools/carbon.py

@@ -8,10 +8,10 @@
 import matplotlib.pyplot as plt
 
 from .stats import linear_regression, nanmean, rm_poly
-from .utils import check_xarray
+from .checks import is_xarray
 
 
-@check_xarray([0, 1])
+@is_xarray([0, 1])
 def co2_sol(t, s):
     """Compute CO2 solubility per the equation used in CESM.
 
@@ -58,7 +58,7 @@ def sol_calc(t, s):
     return ff
 
 
-@check_xarray(0)
+@is_xarray(0)
 def schmidt(t):
     """Computes the dimensionless Schmidt number.
 
@@ -94,7 +94,7 @@ def calc_schmidt(t):
     return Sc
 
 
-@check_xarray(0)
+@is_xarray(0)
 def temp_decomp_takahashi(ds, time_dim='time', temperature='tos', pco2='spco2'):
     """Decompose surface pCO2 into thermal and non-thermal components.
 
@@ -147,7 +147,7 @@ def temp_decomp_takahashi(ds, time_dim='time', temperature='tos', pco2='spco2'):
     return decomp
 
 
-@check_xarray([0, 1])
+@is_xarray([0, 1])
 def potential_pco2(t_insitu, pco2_insitu):
     """Calculate potential pCO2 in the interior ocean.
 
@@ -180,7 +180,7 @@ def potential_pco2(t_insitu, pco2_insitu):
     return pco2_potential
 
 
-@check_xarray(0)
+@is_xarray(0)
 def spco2_sensitivity(ds):
     """Compute sensitivity of surface pCO2 to changes in driver variables.
 
@@ -257,7 +257,7 @@ def _check_variables(ds):
 
 
 # TODO: adapt for CESM and MPI output.
-@check_xarray([0, 1])
+@is_xarray([0, 1])
 def spco2_decomposition_index(
     ds_terms,
     index,
@@ -356,7 +356,7 @@ def regression_against_index(ds, index, psig=None):
         return terms_in_pCO2_units
 
 
-@check_xarray(0)
+@is_xarray(0)
 def spco2_decomposition(ds_terms, detrend=True, order=1, deseasonalize=False):
     """Decompose oceanic surface pco2 in a first order Taylor-expansion.
 

esmtools/checks.py

@@ -1,9 +1,41 @@
 from .exceptions import DimensionError
+from functools import wraps
+from pandas.core.indexes.datetimes import DatetimeIndex
+import xarray as xr
+
+
+# https://stackoverflow.com/questions/10610824/
+# python-shortcut-for-writing-decorators-which-accept-arguments
+def dec_args_kwargs(wrapper):
+    return lambda *dec_args, **dec_kwargs: lambda func: wrapper(
+        func, *dec_args, **dec_kwargs
+    )
+
+
+def get_dims(da):
+    """
+    Simple function to retrieve dimensions from a given dataset/datarray.
+    Currently returns as a list, but can add keyword to select tuple or
+    list if desired for any reason.
+    """
+    return list(da.dims)
 
 
 def has_dims(xobj, dims, kind):
     """
     Checks that at the minimum, the object has provided dimensions.
+
+    Args:
+        xobj (xarray DataArray or Dataset): xarray object over which to check for
+        specified dims.
+        dims (str or list): Dimensions that object should have at the minimum.
+        kind (str): Relayed in error message (see below).
+
+    Returns:
+        True, if dimensions are contained in the xarray object.
+
+    Raises:
+        DimensionError: If xarray object does not have listed dimensions at the minimum.
     """
     if isinstance(dims, str):
         dims = [dims]
@@ -14,3 +46,73 @@ def has_dims(xobj, dims, kind):
             f"following dimensions at the minimum: {dims}"
         )
     return True
+
+
+def is_coordinate_of_applied_dimension(da, dim):
+    """Checks if xarray DataArray is the coordinate variable of the dimension the
+        function is being applied over."""
+    if isinstance(da, xr.DataArray):  # Must be a DataArray to be a coordinate.
+        try:
+            if da.name == dim:
+                return True
+            else:
+                return False
+        except KeyError:  # KeyError occurs when it is an unnamed DataArray.
+            return False
+    else:
+        return False
+
+
+def is_time_index(xobj, kind):
+    """
+    Checks that xobj coming through is a DatetimeIndex or CFTimeIndex.
+
+    This checks that `esmtools` is converting the DataArray to an index,
+    i.e. through .to_index()
+    """
+    if not (isinstance(xobj, xr.CFTimeIndex) or isinstance(xobj, DatetimeIndex)):
+        raise ValueError(
+            f"Your {kind} object must be either an xr.CFTimeIndex or "
+            f"pd.DatetimeIndex."
+        )
+    return True
+
+
+@dec_args_kwargs
+# Lifted from climpred. This was originally written by Andrew Huang.
+def is_xarray(func, *dec_args):
+    """
+    Decorate a function to ensure the first arg being submitted is
+    either a Dataset or DataArray.
+    """
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        try:
+            ds_da_locs = dec_args[0]
+            if not isinstance(ds_da_locs, list):
+                ds_da_locs = [ds_da_locs]
+
+            for loc in ds_da_locs:
+                if isinstance(loc, int):
+                    ds_da = args[loc]
+                elif isinstance(loc, str):
+                    ds_da = kwargs[loc]
+
+                is_ds_da = isinstance(ds_da, (xr.Dataset, xr.DataArray))
+                if not is_ds_da:
+                    typecheck = type(ds_da)
+                    raise IOError(
+                        f"""The input data is not an xarray DataArray or
+                        Dataset.
+
+                        Your input was of type: {typecheck}"""
+                    )
+        except IndexError:
+            pass
+        # this is outside of the try/except so that the traceback is relevant
+        # to the actual function call rather than showing a simple Exception
+        # (probably IndexError from trying to subselect an empty dec_args list)
+        return func(*args, **kwargs)
+
+    return wrapper

esmtools/composite.py

@@ -1,11 +1,10 @@
+from .checks import is_xarray
+from .stats import standardize
+from .testing import ttest_ind_from_stats
 import warnings
 
 import xarray as xr
 
-from .stats import standardize
-from .testing import ttest_ind_from_stats
-from .utils import check_xarray
-
 
 def _create_composites(anomaly_field, index, threshold=1, dim="time"):
     """Creates composite from some variable's anomaly field and a climate
@@ -17,7 +16,7 @@ def _create_composites(anomaly_field, index, threshold=1, dim="time"):
     return composite
 
 
-@check_xarray([0, 1])
+@is_xarray([0, 1])
 def composite_analysis(
     field, index, threshold=1, plot=False, ttest=False, psig=0.05, **plot_kwargs
 ):

esmtools/constants.py

@@ -1,3 +1,40 @@
+# Number of days per month for each `cftime` calendar.
+# Taken from xarray example:
+# http://xarray.pydata.org/en/stable/examples/monthly-means.html
+DAYS_PER_MONTH = {
+    "noleap": [0, 31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31],
+    "365_day": [0, 31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31],
+    "standard": [0, 31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31],
+    "gregorian": [0, 31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31],
+    "proleptic_gregorian": [0, 31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31],
+    "all_leap": [0, 31, 29, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31],
+    "366_day": [0, 31, 29, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31],
+    "360_day": [0, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30],
+    "julian": [0, 31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31],
+}
+# Not as simple as summing since there's nuance here in which years get a quarter year.
+DAYS_PER_YEAR = {
+    "noleap": 365,
+    "365_day": 365.25,
+    "standard": 365.25,
+    "gregorian": 365.25,
+    "proleptic_gregorian": 365.25,
+    "all_leap": 366,
+    "366_day": 366,
+    "360_day": 360,
+    "julian": 365,
+}
+CALENDARS = [k for k in DAYS_PER_MONTH]
+
+# Converts from `cftime` class name to netCDF convention for calendar
+CFTIME_TO_NETCDF = {
+    "DatetimeJulian": "julian",
+    "DatetimeProlepticGregorian": "proleptic_gregorian",
+    "DatetimeNoLeap": "noleap",
+    "DatetimeAllLeap": "all_leap",
+    "DatetimeGregorian": "gregorian",
+}
+
 MULTIPLE_TESTS = [
     "bonferroni",
     "sidak",
@@ -10,3 +47,5 @@
     "fdr_tsbh",
     "fdr_tsbky",
 ]
+
+NAN_METHODS = ["return", "skip", "interpolate"]

esmtools/conversions.py

@@ -1,7 +1,7 @@
-from .utils import check_xarray
+from .checks import is_xarray
 
 
-@check_xarray(0)
+@is_xarray(0)
 def convert_mpas_fgco2(mpas_fgco2):
     """Convert native MPAS CO2 flux (mmol m-3 m s-1) to (molC m-2 yr-1)
 

esmtools/grid.py

@@ -1,8 +1,8 @@
 from .exceptions import CoordinateError
-from .utils import check_xarray
+from .checks import is_xarray
 
 
-@check_xarray(0)
+@is_xarray(0)
 def _convert_lon_to_180to180(ds, coord="lon"):
     """Convert from 0 to 360 (degrees E) grid to -180 to 180 (W-E) grid.
 
@@ -27,7 +27,7 @@ def _convert_lon_to_180to180(ds, coord="lon"):
     return ds
 
 
-@check_xarray(0)
+@is_xarray(0)
 def _convert_lon_to_0to360(ds, coord="lon"):
     """Convert from -180 to 180 (W-E) to 0 to 360 (degrees E) grid.
 
@@ -54,7 +54,7 @@ def _convert_lon_to_0to360(ds, coord="lon"):
 
 # NOTE: Check weird POP grid that goes up to 240 or something. How do we deal with
 # that?
-@check_xarray(0)
+@is_xarray(0)
 def convert_lon(ds, coord="lon"):
     """Converts longitude grid from -180to180 to 0to360 and vice versa.
 

esmtools/spatial.py

@@ -1,5 +1,5 @@
 import numpy as np
-from .utils import check_xarray
+from .checks import is_xarray
 
 
 def find_indices(xgrid, ygrid, xpoint, ypoint):
@@ -43,7 +43,7 @@ def find_indices(xgrid, ygrid, xpoint, ypoint):
     return i, j
 
 
-@check_xarray(0)
+@is_xarray(0)
 def extract_region(ds, xgrid, ygrid, coords, lat_dim="lat", lon_dim="lon"):
     """Extract a subset of some larger spatial data.