Updated test data functions

Author: AgentOxygen

hdp/tests/test_utils.py

@@ -1,46 +1,57 @@
-import hdp.utils
-import xarray
+from hdp.utils import *
 import numpy as np
-
+from math import isclose
+from xarray import DataArray
 
 def test_get_time_stamp():
-    assert type(hdp.utils.get_time_stamp()) is str
+    assert type(get_time_stamp()) is str
 
 
 def test_get_version():
-    assert type(hdp.utils.get_version()) is str
-
-
-def test_test_data_functions():
-    var = "test"
-    center_val = 10
-    amplitude_val = 1
-    units = "test_units"
-    
-    ds = hdp.utils.generate_test_dataset(name=var, units=units, center=center_val, amplitude=amplitude_val)
-    assert type(ds) is xarray.Dataset
-    assert len(ds.data_vars) == 1
-    assert var in ds
-    
-    assert "units" in ds[var].attrs
-    assert ds[var].attrs["units"] == units
-    assert ds[var].dims == ("lat", "lon", "time")
-    assert ds[var].dtype == float
-    assert ds[var].time.values[0].calendar == "noleap"
-    assert ds[var].lat.size > 1
-    assert ds[var].lon.size > 1
-    assert ds[var].time.size >= 2*365
-
-    data = ds[var].compute()
-    
-    assert np.isclose(data.mean(), center_val)
-    assert np.isclose(data.max(), center_val + amplitude_val)
-    assert np.isclose(data.values, data.values[0]).all()
-
-    exceed_data = hdp.utils.generate_exceedance_dataarray(ds[var], exceedance_pattern=[1, 0, 1, 0]).compute()
-
-    assert np.isclose(exceed_data.mean(), center_val + 0.5)
-    assert exceed_data.attrs == ds[var].attrs
-    assert exceed_data.dims == ds[var].dims
-    assert exceed_data.shape == ds[var].shape
-    assert exceed_data.dtype == ds[var].dtype
+    assert type(get_version()) is str
+
+
+def test_control_dataarray():
+    control_da = generate_test_control_dataarray(grid_shape=(1,1), start_date="2000", end_date="2100")
+    assert type(control_da) is DataArray
+    var = control_da.name
+    assert control_da.attrs["units"] == "degC"
+    assert control_da.dims == ("lon", "lat", "time")
+    assert control_da.dtype == float
+    assert control_da.time.values[0].calendar == "noleap"
+    assert control_da.time.values.size >= 365
+    assert np.sum(np.isnan(control_da)) == 0
+
+    control_slope = np.polyfit(np.arange(control_da["time"].size), control_da.mean(dim=["lat", "lon"]).values, 1)[0]
+    assert np.abs(control_slope) < 0.01
+
+
+def test_warming_dataarray():
+    warming_da = generate_test_warming_dataarray(grid_shape=(1,1), start_date="2000", end_date="2100")
+    avg_da = warming_da.mean(dim=["lat", "lon"]).values
+    warm_slope = np.polyfit(np.arange(avg_da.size), avg_da, 1)[0]
+
+    assert warm_slope > 0
+    assert not isclose(warm_slope, 0)
+
+
+def test_rh_dataarray():
+    rh_da = generate_test_rh_dataarray()
+
+    assert rh_da.max() <= 1
+    assert rh_da.min() >= 0
+
+
+def test_defaults_compatibility():
+    control_da = generate_test_control_dataarray()
+    warming_da = generate_test_warming_dataarray()
+    rh_da = generate_test_rh_dataarray()
+
+    assert control_da.shape == warming_da.shape
+    assert warming_da.shape == rh_da.shape
+
+    assert control_da.dims == warming_da.dims
+    assert warming_da.dims == rh_da.dims
+
+    assert control_da.name == warming_da.name
+    assert control_da.units == warming_da.units

hdp/tests/test_workflow.py

@@ -4,8 +4,8 @@
 
 
 def test_full_data_workflow():
-    baseline_temp = hdp.utils.generate_test_dataset(name="temp")["temp"]
-    baseline_rh = hdp.utils.generate_test_dataset(name="rh", units="%", center=90, amplitude=15)["rh"]
+    baseline_temp = hdp.utils.generate_test_control_dataarray().rename("temp")
+    baseline_rh = hdp.utils.generate_test_rh_dataarray().rename("rh")
     baseline_measures = hdp.measure.format_standard_measures([baseline_temp], rh=baseline_rh)
 
     percentiles = np.arange(0.9, 1, 0.01)
@@ -14,9 +14,8 @@ def test_full_data_workflow():
 
     exceedance_pattern = [1, 0, 1, 1, 1, 0, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 0, 1, 0, 1, 0, 1]
 
-    test_temp = hdp.utils.generate_test_dataset(name="temp")["temp"]
-    test_temp = hdp.utils.generate_exceedance_dataarray(test_temp, exceedance_pattern)
-    test_rh = hdp.utils.generate_test_dataset(name="rh", units="%", center=90, amplitude=15)["rh"]
+    test_temp = hdp.utils.generate_test_warming_dataarray().rename("temp")
+    test_rh = baseline_rh
 
     hw_definitions = [[3,0,0], [3,1,1], [4,2,0], [4,1,3], [5,0,1], [5,1,4]]
 
@@ -37,11 +36,6 @@ def test_full_data_workflow():
     assert metrics.definition.values[5] == "5-1-4"
     assert (metrics.percentile.values == percentiles).all()
 
-    assert (metrics["temp.temp_threshold.HWF"] == metrics["temp_hi.temp_hi_threshold.HWF"]).all()
-    assert (metrics["temp.temp_threshold.HWD"] == metrics["temp_hi.temp_hi_threshold.HWD"]).all()
-    assert (metrics["temp.temp_threshold.HWA"] == metrics["temp_hi.temp_hi_threshold.HWA"]).all()
-    assert (metrics["temp.temp_threshold.HWN"] == metrics["temp_hi.temp_hi_threshold.HWN"]).all()
-
     metric_means = metrics.mean()
 
     assert metric_means["temp.temp_threshold.HWF"] >= metric_means["temp.temp_threshold.HWD"]

hdp/utils.py

@@ -36,45 +36,46 @@ def get_func_description(func):
     return desc
 
 
-def generate_test_dataset(center=25, amplitude=10, name="temperature", units="degC"):
+def generate_test_warming_dataarray(start_date="2000-01-01", end_date="2009-12-31", grid_shape=(10, 10), warming_period=10):
+    base_data = generate_test_control_dataarray(start_date=start_date, end_date=end_date, grid_shape=grid_shape)
+    base_data += xarray.DataArray(np.arange(base_data["time"].size) / (365*warming_period), dims=["time"], coords={"time": base_data["time"]})
+    return base_data
+
+
+def generate_test_rh_dataarray(start_date="2000-01-01", end_date="2009-12-31", grid_shape=(10, 10)):
+    base_data = generate_test_control_dataarray(start_date=start_date, end_date=end_date, grid_shape=grid_shape)
+    base_data = abs(base_data / base_data.max() - 0.3)
+    base_data = base_data.rename("test_rh_data")
+    base_data.attrs["units"] = 'g/g'
+    return base_data
+
+
+def generate_test_control_dataarray(start_date="2000-01-01", end_date="2009-12-31", grid_shape=(10, 10)):
     time_values = xarray.date_range(
-        start=cftime.DatetimeNoLeap(2000, 1, 1),
-        end=cftime.DatetimeNoLeap(2009, 12, 31),
+        start=start_date,
+        end=end_date,
         freq="D",
         calendar="noleap",
         use_cftime=True
     )
+    temperature_seasonal_ts = 20 + 15*np.sin(np.pi*np.arange(time_values.size, dtype=float) / 365)
+    temperature_seasonal_vals = np.broadcast_to(temperature_seasonal_ts, (grid_shape[0], grid_shape[1], temperature_seasonal_ts.size))
 
-    temp_timeseries = center + amplitude*np.sin(np.pi*np.arange(time_values.size, dtype=float) / 365)
-    temp_values = np.broadcast_to(temp_timeseries, (3, 3, temp_timeseries.size))
-    
-    temp_da = xarray.DataArray(
-        data=da.from_array(temp_values),
-        dims=["lat", "lon", "time"],
+    lat_vals = np.linspace(-90, 90, grid_shape[1], dtype=float)
+
+    lat_grad = np.broadcast_to(np.abs(lat_vals) / 90 * 15, grid_shape).T
+    temperature_seasonal_vals = temperature_seasonal_vals - lat_grad[:, :, None]
+
+    return xarray.DataArray(
+        data=temperature_seasonal_vals,
+        dims=["lon", "lat", "time"],
         coords={
-            "lat": np.array([-90, 0, 90], dtype=float),
-            "lon": np.array([-180, 0, 180], dtype=float),
+            "lon": np.linspace(-180, 180, grid_shape[0], dtype=float),
+            "lat": lat_vals,
             "time": time_values
         },
-        name=name,
+        name="test_temperature_data",
         attrs={
-            "units": units
+            "units": "degC"
         }
-    ).chunk(dict(lat=1, lon=1))
-    return xarray.Dataset({name: temp_da})
-
-
-def generate_exceedance_dataarray(measure, exceedance_pattern, multiplier=1.0):
-    tiles = np.ceil(measure.time.size / len(exceedance_pattern)).astype(int)
-    pattern_data = np.broadcast_to(np.tile(exceedance_pattern, tiles)[:measure.time.size], measure.shape)*multiplier
-
-    ret_da = None
-    with xarray.set_options(keep_attrs=True): 
-        ret_da = measure + xarray.DataArray(
-            data=da.from_array(pattern_data),
-            dims=measure.dims,
-            coords=measure.coords,
-            name=measure.name,
-            attrs=measure.attrs
-        ).chunk(measure.chunksizes)
-    return ret_da
+    ).chunk('auto')