Repository: MarcoLandtHayen/climate_index_collection

Issue:
Add North Atlantic Oscillation (NAO) index EOF-based (https://climatedataguide.ucar.edu/climate-data/hurrell-north-atlantic-oscillation-nao-index-pc-based).

Discussion:
willirath: I can confirm that this works with a local merge of #75.

Patch:
diff --git a/climate_index_collection/data_loading.py b/climate_index_collection/data_loading.py
index d93de76..8ed93e8 100644
--- a/climate_index_collection/data_loading.py
+++ b/climate_index_collection/data_loading.py
@@ -1,7 +1,11 @@
 from pathlib import Path
 
+import numpy as np
+import pandas as pd
 import xarray as xr
 
+from .timestamp_handling import fix_monthly_time_stamps, fix_year_in_CESM_files
+
 
 def find_data_files(data_path="data/test_data/", data_source_name="FOCI"):
     """Find all files for given data source.
@@ -28,7 +32,8 @@ def load_and_preprocess_single_data_file(file_name, **kwargs):
 
     Current pre-processing steps:
     - squeeze (to get rid of, e.g., some of the vertical degenerate dims)
-    - ...
+    - fixing monthly timestamps
+    - fixing start years
 
     Parameters
     ----------
@@ -48,6 +53,12 @@ def load_and_preprocess_single_data_file(file_name, **kwargs):
     # get rid of singleton dims
     ds = ds.squeeze()
 
+    # fix time stamps to always be mid month
+    ds = fix_monthly_time_stamps(ds)
+
+    # fix start years
+    ds = fix_year_in_CESM_files(dataset=ds, filename=file_name)
+
     return ds
 
 
diff --git a/climate_index_collection/indices.py b/climate_index_collection/indices.py
index bad1f2f..fa45840 100644
--- a/climate_index_collection/indices.py
+++ b/climate_index_collection/indices.py
@@ -2,6 +2,8 @@ from enum import Enum
 from functools import partial
 
 import numpy as np
+import scipy as sp
+import xarray as xr
 
 from .reductions import area_mean_weighted, mean_unweighted
 
@@ -86,6 +88,59 @@ def north_atlantic_oscillation(data_set, slp_name="sea-level-pressure"):
     return NAO_index
 
 
+def north_atlantic_oscillation_pc(data_set, slp_name="sea-level-pressure"):
+    """Calculate the principal component based North Atlantic Oscillation (NAO) index
+
+    Following
+    https://climatedataguide.ucar.edu/climate-data/hurrell-north-atlantic-oscillation-nao-index-pc-based
+    this index is obtained as Principle Component (PC) time series of the leading Empirical Orthogonal Function (EOF)
+    of monthly sea-level pressure anomalies over the Atlantic sector, 20°-80°N, 90°W-40°E.
+
+    Computation is done as follows:
+    1. Compute sea level pressure anomalies over Atlantic sector.
+    2. Flatten spatial dimensions and subtract mean in time.
+    3. Perform Singular Value Decomposition.
+    4. Normalize Principal Components.
+    5. Obtain NAO index as PC time series related to leading EOF.
+
+    Parameters
+    ----------
+    data_set: xarray.DataSet
+        Dataset containing an SLP field.
+    slp_name: str
+        Name of the Sea-Level Pressure field. Defaults to "sea-level-pressure".
+
+    Returns
+    -------
+    xarray.DataArray
+        Time series containing the NAO index.
+
+    """
+
+    mask = (
+        (data_set.coords["lat"] >= 20)
+        & (data_set.coords["lat"] <= 80)
+        & ((data_set.coords["lon"] >= 270) | (data_set.coords["lon"] <= 40))
+    )
+
+    slp = data_set[slp_name].where(mask)
+    climatology = slp.groupby("time.month").mean("time")
+    slp = (slp.groupby("time.month") - climatology).drop("month")
+
+    slp_flat = slp.stack(tmp_space=("lat", "lon")).dropna(dim="tmp_space")
+
+    pc, s, eof = sp.linalg.svd(slp_flat - slp_flat.mean(axis=0), full_matrices=False)
+
+    pc_std = pc.std(axis=0)
+    pc /= pc_std
+
+    NAO_index = xr.DataArray(pc[:, 0], dims=("time"), coords={"time": slp_flat["time"]})
+
+    NAO_index = NAO_index.rename("NAO_PC")
+
+    return NAO_index
+
+
 def el_nino_southern_oscillation_34(data_set, sst_name="sea-surface-temperature"):
     """Calculate the El Nino Southern Oscillation 3.4 index (ENSO 3.4)
 
@@ -194,6 +249,7 @@ class ClimateIndexFunctions(Enum):
 
     southern_annular_mode = partial(southern_annular_mode)
     north_atlantic_oscillation = partial(north_atlantic_oscillation)
+    north_atlantic_oscillation_pc = partial(north_atlantic_oscillation_pc)
     el_nino_southern_oscillation_34 = partial(el_nino_southern_oscillation_34)
     north_atlantic_sea_surface_salinity = partial(north_atlantic_sea_surface_salinity)
 
diff --git a/climate_index_collection/timestamp_handling.py b/climate_index_collection/timestamp_handling.py
new file mode 100644
index 0000000..402ff4d
--- /dev/null
+++ b/climate_index_collection/timestamp_handling.py
@@ -0,0 +1,163 @@
+from datetime import datetime
+from pathlib import Path
+from tracemalloc import start
+
+import cftime
+import numpy as np
+import pandas as pd
+import xarray as xr
+
+
+def _is_start_of_next_month(timestamp):
+    return (timestamp.day == 1) & (timestamp.hour == 0)
+
+
+def _get_prev_month(year=None, month=None):
+    if month == 1:
+        return year - 1, 12
+    else:
+        return year, month - 1
+
+
+def _get_next_month(year=None, month=None):
+    if month == 12:
+        return year + 1, 1
+    else:
+        return year, month + 1
+
+
+def _get_fixed_year_month(timestamp):
+    # if possible, convert to datetime with Pandas
+    # (this won't be necessary for cftime types)
+    try:
+        timestamp = pd.to_datetime(timestamp)
+    except Exception as e:
+        pass
+    if _is_start_of_next_month(timestamp):
+        return _get_prev_month(year=timestamp.year, month=timestamp.month)
+    else:
+        return timestamp.year, timestamp.month
+
+
+def _get_mid_of_month(
+    year=None, month=None, datetime_type=cftime.DatetimeProlepticGregorian
+):
+    year_upper, month_upper = _get_next_month(year=year, month=month)
+    lower = datetime_type(year, month, 1, 0, 0, 0)
+    upper = datetime_type(year_upper, month_upper, 1, 0, 0, 0)
+    return lower + (upper - lower) / 2
+
+
+def fix_monthly_time_stamps(dobj, time_name="time"):
+    """Fix monthly time stamps to be exactly the middle of months.
+
+    There's a few input data sets which label monthly data at the turn of months.
+    Let's unify the convention to mid month everywhere.
+
+    Currently, only cftime Datetimes are supported.
+
+    Parameters
+    ----------
+    dobj: xarray.Dataset or xarray.DataArray
+        Should contain a time dimension.
+    time_name: str
+        Defaults to "time".
+
+    Returns
+    -------
+    xarray.Dataset or xarray.DataArray
+        Same as dobj but with fixed time axis.
+
+    """
+    # extract time axis
+    orig_time = dobj.coords[time_name]
+    datetime_type = type(dobj.coords[time_name].data[0])
+
+    # check for not implemented numpy datetimes
+    if datetime_type is np.datetime64:
+        raise NotImplementedError("Currently only works for cftime Datetimes.")
+
+    # construct new time axis
+    fixed_time = [
+        _get_mid_of_month(*_get_fixed_year_month(ts), datetime_type)
+        for ts in orig_time.data
+    ]
+
+    # assign coord
+    fixed_dobj = dobj.copy()
+    fixed_dobj = fixed_dobj.assign_coords(
+        {time_name: xr.DataArray(fixed_time, dims=(time_name,))}
+    )
+
+    return fixed_dobj
+
+
+def fix_year_in_CESM_files(dataset=None, filename=None):
+    """Makes sure the time axis starts in the year given in the file name.
+
+    Note that this will be obsolete once we've fixed the data at the source.
+
+    Parameters
+    ----------
+    dataset: xr.Dataset
+        Contains a potentially broken time axis.
+    filename: str or pathlike
+        Contains info on the desired time stamp.
+
+    Returns
+    -------
+    dataset
+        With fixed time axis.
+
+    """
+    # example filenames
+    # - B1850WCN_f19g16_1000y_v3.2_mod-S15-G16.cam2.h0.0001-0010.PRECT.nc
+    # - B1850WCN_f19g16_1000y_v3.2_mod-S15-G16.cam2.h0.0001-0010.PSL.nc
+    # - B1850WCN_f19g16_1000y_v3.2_mod-S15-G16.cam2.h0.0001-0010.SST.nc
+    # - B1850WCN_f19g16_1000y_v3.2_mod-S15-G16.cam2.h0.0001-0010.TS.nc
+    # - B1850WCN_f19g16_1000y_v3.2_mod-S15-G16.cam2.h0.0001-0010.Z500.nc
+    # - B1850WCN_f19g16_1000y_v3.2_mod-S15-G16.pop.h.0001-0005.SSS.atmos_grid.nc
+    #
+    # SO we
+    # - check that the filename starts with "B1850WCN"
+    # - extract the 5th (index 4) part separated by a dot and then get the first part
+
+    # make sure we have a string
+    try:
+        filename = str(Path(filename).name)
+    except Exception as e:
+        pass
+
+    # skip if filename not from CESM
+    if not filename.startswith("B1850WCN"):
+        return dataset
+
+    # if we're still here
+    start_year = int(filename.split(".")[4].split("-")[0])
+
+    # fix offsets
+    #
+    # Note that we're taking a winded route here which in the end substracts all the same
+    # offset from all the time steps, because we want to avoid to completely re-create the
+    # time axis Data Array from scratch.
+    #
+    old_time_axis = dataset.coords["time"].data
+    year_offset = old_time_axis[0].year - start_year  # calc offset from first item only
+    datetime_type = type(old_time_axis[0])
+    new_time_axis = np.array(
+        [
+            datetime_type(
+                year=odt.year - year_offset,
+                month=odt.month,
+                day=odt.day,
+                hour=odt.hour,
+                minute=odt.minute,
+                second=odt.second,
+            )
+            for odt in old_time_axis
+        ]
+    )
+    all_offsets = old_time_axis - new_time_axis
+    dataset.coords["time"] = dataset.coords["time"] - all_offsets
+
+    return dataset
diff --git a/pyproject.toml b/pyproject.toml
index ac2d294..8b79841 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -19,7 +19,7 @@ quiet = false
 color = true
 
 [tool.isort]
-known_third_party = ["click", "numpy", "pandas", "pkg_resources", "pytest", "setuptools", "xarray", "xhistogram"]
+known_third_party = ["cftime", "click", "numpy", "pandas", "pkg_resources", "pytest", "setuptools", "xarray", "xhistogram"]
 
 [tool.pytest.ini_options]
 minversion = "6.0"


Test patch:
diff --git a/tests/test_data_loading.py b/tests/test_data_loading.py
index 0e3373a..b969125 100644
--- a/tests/test_data_loading.py
+++ b/tests/test_data_loading.py
@@ -1,8 +1,16 @@
+from datetime import datetime
 from pathlib import Path
 
+import cftime
+import numpy as np
 import pytest
+import xarray as xr
 
-from climate_index_collection.data_loading import VARNAME_MAPPING, load_data_set
+from climate_index_collection.data_loading import (
+    VARNAME_MAPPING,
+    fix_monthly_time_stamps,
+    load_data_set,
+)
 
 
 TEST_DATA_PATH = Path(__file__).parent / "../data/test_data/"
diff --git a/tests/test_indices.py b/tests/test_indices.py
index 3e7697a..4f806dc 100644
--- a/tests/test_indices.py
+++ b/tests/test_indices.py
@@ -2,6 +2,7 @@ from pathlib import Path
 
 import numpy as np
 import pytest
+import scipy as sp
 import xarray as xr
 
 from numpy.testing import assert_allclose, assert_almost_equal
@@ -10,6 +11,7 @@ from climate_index_collection.data_loading import VARNAME_MAPPING, load_data_set
 from climate_index_collection.indices import (
     el_nino_southern_oscillation_34,
     north_atlantic_oscillation,
+    north_atlantic_oscillation_pc,
     north_atlantic_sea_surface_salinity,
     southern_annular_mode,
 )
@@ -100,6 +102,49 @@ def test_NAO_naming(source_name):
     assert NAO.name == "NAO"
 
 
+@pytest.mark.parametrize("source_name", list(VARNAME_MAPPING.keys()))
+def test_NAO_PC_metadata(source_name):
+    """Ensure that index only contains time dimension."""
+    # Load test data
+    TEST_DATA_PATH = Path(__file__).parent / "../data/test_data/"
+    data_set = load_data_set(data_path=TEST_DATA_PATH, data_source_name=source_name)
+
+    # Calculate NAO index PC-based
+    NAO_PC = north_atlantic_oscillation_pc(data_set)
+
+    # Check, if calculated NAO index only has one dimension: 'time'
+    assert NAO_PC.dims[0] == "time"
+    assert len(NAO_PC.dims) == 1
+
+
+@pytest.mark.parametrize("source_name", list(VARNAME_MAPPING.keys()))
+def test_NAO_PC_standardisationn(source_name):
+    """Ensure that NAO has zero mean and unit std dev."""
+    # Load test data
+    TEST_DATA_PATH = Path(__file__).parent / "../data/test_data/"
+    data_set = load_data_set(data_path=TEST_DATA_PATH, data_source_name=source_name)
+
+    # Calculate NAO index PC-based
+    NAO_PC = north_atlantic_oscillation_pc(data_set)
+
+    # Check, if calculated NAO index has zero mean:
+    assert_almost_equal(actual=NAO_PC.mean("time").values[()], desired=0, decimal=3)
+    assert_almost_equal(actual=NAO_PC.std("time").values[()], desired=1, decimal=3)
+
+
+@pytest.mark.parametrize("source_name", list(VARNAME_MAPPING.keys()))
+def test_NAO_PC_naming(source_name):
+    """Ensure that the index is named correctly."""
+    # Load test data
+    TEST_DATA_PATH = Path(__file__).parent / "../data/test_data/"
+    data_set = load_data_set(data_path=TEST_DATA_PATH, data_source_name=source_name)
+
+    # Calculate NAO index PC-based
+    NAO_PC = north_atlantic_oscillation_pc(data_set)
+
+    assert NAO_PC.name == "NAO_PC"
+
+
 @pytest.mark.parametrize("source_name", list(VARNAME_MAPPING.keys()))
 def test_enso34_zeromean(source_name):
     """Ensure that ENSO 3.4 has zero mean."""
diff --git a/tests/test_timestamp_handling.py b/tests/test_timestamp_handling.py
new file mode 100644
index 0000000..e35b72e
--- /dev/null
+++ b/tests/test_timestamp_handling.py
@@ -0,0 +1,54 @@
+import datetime
+
+import cftime
+import numpy as np
+import xarray as xr
+
+from climate_index_collection.timestamp_handling import fix_monthly_time_stamps
+
+
+def test_time_stamp_fixing_to_mid_month():
+    """Test that end-of-month labels are moved to mid months"""
+
+    # construct time axes
+    wrong_time_axis = [
+        cftime.DatetimeProlepticGregorian(2022, 1, 31, 23, 58, 0),
+        cftime.DatetimeProlepticGregorian(2022, 2, 15, 12, 0, 0),  # that's not mid Feb!
+        cftime.DatetimeProlepticGregorian(2022, 4, 1, 0, 0, 0),  # that's end of March!
+    ]
+    desired_time_axis = [
+        # note we can't do `(date1 + date2) / 2` but `date1 + (date2 - date1) / 2` works
+        cftime.DatetimeProlepticGregorian(2022, 1, 1, 0, 0, 0)
+        + (
+            cftime.DatetimeProlepticGregorian(2022, 2, 1, 0, 0, 0)
+            - cftime.DatetimeProlepticGregorian(2022, 1, 1, 0, 0, 0)
+        )
+        / 2,
+        cftime.DatetimeProlepticGregorian(2022, 2, 1, 0, 0, 0)
+        + (
+            cftime.DatetimeProlepticGregorian(2022, 3, 1, 0, 0, 0)
+            - cftime.DatetimeProlepticGregorian(2022, 2, 1, 0, 0, 0)
+        )
+        / 2,
+        cftime.DatetimeProlepticGregorian(2022, 3, 1, 0, 0, 0)
+        + (
+            cftime.DatetimeProlepticGregorian(2022, 4, 1, 0, 0, 0)
+            - cftime.DatetimeProlepticGregorian(2022, 3, 1, 0, 0, 0)
+        )
+        / 2,
+    ]
+
+    # put into data array
+    dobj = xr.DataArray(
+        [1, 2, 3], name="dummy", dims=("time",), coords={"time": wrong_time_axis}
+    )
+
+    # get fixed array
+    fixed_dobj = fix_monthly_time_stamps(dobj, time_name="time")
+
+    # check closeness
+    np.testing.assert_almost_equal(
+        (fixed_dobj.time.data - desired_time_axis) / datetime.timedelta(seconds=1),
+        desired=0,
+        decimal=3,
+    )