Repository: MarcoLandtHayen/climate_index_collection

Issue:
Time stamp problems
```python
from climate_index_collection import data_loading

import xarray as xr

FOCI_files = data_loading.find_data_files(
    data_path="../data/test_data/", data_source_name="FOCI"
)
CESM_files = data_loading.find_data_files(
    data_path="../data/test_data/", data_source_name="CESM"
)

FOCI_ds = {
    fname: xr.open_dataset(fname) for fname in FOCI_files
}

CESM_ds = {
    fname: xr.open_dataset(fname) for fname in CESM_files
}

FOCI_times = {
    k: str(ds.time.data[0]) for k, ds in FOCI_ds.items()
}

CESM_times = {
    k: str(ds.time.data[0]) for k, ds in CESM_ds.items()
}

print("FOCI\n")

for k, v in FOCI_times.items():
    print(f"{k}:\n\t{v}")

print("CESM\n")
    
for k, v in CESM_times.items():
    print(f"{k}:\n\t{v}")
```
yields
```
FOCI

../data/test_data/FOCI/FOCI1.3-SW038_1m_23500101_23541231_grid_T_atmos_grid.nc:
	2350-02-28 23:52:00
../data/test_data/FOCI/FOCI1.3-SW038_echam6_ATM_mm_2350-2359_geopoth_pl_monthly_50000.nc:
	2350-02-28 23:52:00
../data/test_data/FOCI/FOCI1.3-SW038_echam6_BOT_mm_2350-2359_precip_monthly_1.nc:
	2350-02-28 23:52:00
../data/test_data/FOCI/FOCI1.3-SW038_echam6_BOT_mm_2350-2359_slp_monthly_1.nc:
	2350-02-28 23:52:00
../data/test_data/FOCI/FOCI1.3-SW038_echam6_BOT_mm_2350-2359_temp2_monthly_1.nc:
	2350-02-15 00:00:00
../data/test_data/FOCI/FOCI1.3-SW038_echam6_BOT_mm_2350-2359_tsw_monthly_1.nc:
	2350-02-28 23:52:00
CESM

../data/test_data/CESM/B1850WCN_f19g16_1000y_v3.2_mod-S15-G16.cam2.h0.0001-0010.PRECT.nc:
	2350-02-15 00:00:00
../data/test_data/CESM/B1850WCN_f19g16_1000y_v3.2_mod-S15-G16.cam2.h0.0001-0010.PSL.nc:
	0001-03-01 00:00:00
../data/test_data/CESM/B1850WCN_f19g16_1000y_v3.2_mod-S15-G16.cam2.h0.0001-0010.SST.nc:
	0001-03-01 00:00:00
../data/test_data/CESM/B1850WCN_f19g16_1000y_v3.2_mod-S15-G16.cam2.h0.0001-0010.TS.nc:
	2350-02-15 00:00:00
../data/test_data/CESM/B1850WCN_f19g16_1000y_v3.2_mod-S15-G16.cam2.h0.0001-0010.Z500.nc:
	0001-03-01 00:00:00
../data/test_data/CESM/B1850WCN_f19g16_1000y_v3.2_mod-S15-G16.pop.h.0001-0005.SSS.atmos_grid.nc:
	0001-03-01 00:00:00
```

So:
- For FOCI, only `FOCI1.3-SW038_echam6_BOT_mm_2350-2359_temp2_monthly_1.nc` seems to have mid-month time labeling, while all others are end of month
- But for CESM, there's a wild mix:
  - mid-month labels from year 2350 (!) for two files
  - start-of-next month labels from year 0001 for all others

What to do here?
- For FOCI, we can just label everything to mid months in the loading function (or on disk)
- For CESM, we need to find out if the two files labelled as years 2350+ are wrong and maybe ask Sebastian for the correct data

Discussion:

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 43ac7b0..bad1f2f 100644
--- a/climate_index_collection/indices.py
+++ b/climate_index_collection/indices.py
@@ -135,6 +135,50 @@ def el_nino_southern_oscillation_34(data_set, sst_name="sea-surface-temperature"
     return ENSO34_index
 
 
+def north_atlantic_sea_surface_salinity(data_set, sss_name="sea-surface-salinity"):
+    """Calculate North-Atlantic Sea-Surface Salinity index.
+
+    Following https://doi.org/10.1126/sciadv.1501588
+    the index is derived from Atlantic sea-surface salinity (SSS) anomalies in a box
+    spanning 50W-15W, 25N-50N.
+
+    Note that in the original ref, there was a smaller box spanning 50W-40W, 37.5N-50N
+    cut away from the one used here. We skip this here, for now.
+
+    Computation is done as follows:
+    1. Calculate the weighted area average of SSS over 50W-15W, 25N-50N.
+    2. Standardize time series.
+
+    Parameters
+    ----------
+    data_set: xarray.DataSet
+        Dataset containing an SSS field.
+    sss_name: str
+        Name of the Sea-Surface Salinity field. Defaults to "sea-surface-salinity".
+
+    Returns
+    -------
+    xarray.DataArray
+        Time series containing the NASSS index.
+
+    """
+    sss = data_set[sss_name]
+
+    sss_box_ave = area_mean_weighted(
+        sss,
+        lat_south=25,
+        lat_north=50,
+        lon_west=-50,
+        lon_east=-15,
+    )
+
+    NASSS = (sss_box_ave - sss_box_ave.mean("time")) / sss_box_ave.std("time")
+
+    NASSS = NASSS.rename("NASSS")
+
+    return NASSS
+
+
 class ClimateIndexFunctions(Enum):
     """Enumeration of all index functions.
 
@@ -151,6 +195,7 @@ class ClimateIndexFunctions(Enum):
     southern_annular_mode = partial(southern_annular_mode)
     north_atlantic_oscillation = partial(north_atlantic_oscillation)
     el_nino_southern_oscillation_34 = partial(el_nino_southern_oscillation_34)
+    north_atlantic_sea_surface_salinity = partial(north_atlantic_sea_surface_salinity)
 
     @classmethod
     def get_all_names(cls):
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 ff1fb75..3e7697a 100644
--- a/tests/test_indices.py
+++ b/tests/test_indices.py
@@ -10,6 +10,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_sea_surface_salinity,
     southern_annular_mode,
 )
 
@@ -124,3 +125,46 @@ def test_ENSO34_naming(source_name):
     result = el_nino_southern_oscillation_34(data_set)
 
     assert result.name == "ENSO34"
+
+
+@pytest.mark.parametrize("source_name", list(VARNAME_MAPPING.keys()))
+def test_NASSS_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 NASSS index
+    NASSS = north_atlantic_sea_surface_salinity(data_set)
+
+    # Check, if calculated NASSS index only has one dimension: 'time'
+    assert NASSS.dims[0] == "time"
+    assert len(NASSS.dims) == 1
+
+
+@pytest.mark.parametrize("source_name", list(VARNAME_MAPPING.keys()))
+def test_NASSS_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 NASSS index
+    NASSS = north_atlantic_sea_surface_salinity(data_set)
+
+    assert NASSS.name == "NASSS"
+
+
+@pytest.mark.parametrize("source_name", list(VARNAME_MAPPING.keys()))
+def test_NASSS_standardisation(source_name):
+    """Ensure that standardisation works 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 NASSS index
+    NASSS = north_atlantic_sea_surface_salinity(data_set)
+
+    # Check, if calculated NASSS index has zero mean and unit std dev:
+    assert_almost_equal(actual=NASSS.mean("time").values[()], desired=0, decimal=3)
+    assert_almost_equal(actual=NASSS.std("time").values[()], desired=1, decimal=3)
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,
+    )