add more stats for MoV driver

pcmdi_metrics/stats/__init__.py

@@ -1,4 +1,4 @@
-from .compute_statistics import (
+from .compute_statistics_dataset import (
  annual_mean,
  bias_xy,
  bias_xyt,
@@ -15,3 +15,8 @@
  std_xyt,
  zonal_mean,
 )
+
+from .compute_statistics_dataArray import (
+ calculate_temporal_correlation,
+ calculate_spatial_correlation
+)

pcmdi_metrics/stats/compute_statistics_dataArray.py

@@ -0,0 +1,51 @@
+from pcmdi_metrics.io import get_time_key, get_latitude_key, get_longitude_key
+import xarray as xr
+from typing import Union
+
+
+def calculate_temporal_correlation(data_array1: xr.DataArray, data_array2: xr.DataArray, dim:Union [str, tuple] = None) -> float:
+ return _calculate_correlation(data_array1, data_array2, method='temporal', dim=None)
+
+
+def calculate_spatial_correlation(data_array1: xr.DataArray, data_array2: xr.DataArray, wtarray: xr.DataArray = None, dim:Union [str, tuple] = None) -> float:
+ return _calculate_correlation(data_array1, data_array2, method='spatial', wtarray=wtarray, dim=None)
+
+
+def _calculate_correlation(data_array1: xr.DataArray, data_array2: xr.DataArray, wtarray: xr.DataArray = None, method='temporal', dim:Union [str, tuple] = None) -> float:
+ """
+ Calculate correlation between two xarray DataArrays.
+
+ Parameters:
+ - data_array1 (xarray.DataArray): First data array.
+ - data_array2 (xarray.DataArray): Second data array.
+ - method (str): Method for correlation calculation ('temporal' or 'spatial').
+ - dim (tuple or str): Dimensions along which to compute the correlation.
+ For 'temporal' correlation, specify the time dimension.
+ For 'spatial' correlation, specify the spatial dimensions.
+
+ Returns:
+ - correlation (float): Correlation coefficient.
+ """
+
+ if method not in ['temporal', 'spatial']:
+ raise ValueError("Invalid method. Use 'temporal' or 'spatial'.")
+
+ if method == 'temporal':
+ # Check if the specified dimension exists in both DataArrays
+ if dim is None:
+ dim = get_time_key(data_array1)
+ if (dim not in data_array1.dims or dim not in data_array2.dims):
+ raise ValueError(f"Temporal dimension '{dim}' not found in both input DataArrays.")
+
+ elif method == 'spatial':
+ # Check if the specified dimensions exist in both DataArrays
+ if dim is None:
+ lat_key = get_latitude_key(data_array1)
+ lon_key = get_longitude_key(data_array1)
+ dim = (lat_key, lon_key)
+ if (not isinstance(dim, tuple) and dim not in data_array1.dims) or (isinstance(dim, tuple) and any(d not in data_array1.dims or d not in data_array2.dims for d in dim)):
+ raise ValueError(f"Spatial dimension '{dim}' not found in both input DataArrays.")
+
+ correlation = xr.corr(data_array1, data_array2, dim=dim).values.item()
+
+ return correlation

pcmdi_metrics/stats/compute_statistics.py → pcmdi_metrics/stats/compute_statistics_dataset.py

@@ -2,6 +2,21 @@
 
 import numpy as np
 import xcdat as xc
+import xarray as xr
+from typing import Union
+
+
+
+def check_data_convert_to_ds_if_needed(da: Union[xr.Dataset, xr.DataArray], var: str=None):
+ if isinstance(da, xr.Dataset):
+ return da
+ elif isinstance(da, xr.DataArray):
+ if var is None:
+ var = "variable" # assign arbitrary variable name
+ return da.to_dataset(name=var)
+ else:
+ raise TypeError("Input must be an instance of DataArray or Dataset")
+
 
 
 def annual_mean(dm, do, var=None):
@@ -13,6 +28,10 @@ def annual_mean(dm, do, var=None):
  "Contact": "[email protected]",
  "Comments": "Assumes input are 12 months climatology",
  }
+
+ dm = check_data_convert_to_ds_if_needed(dm, var)
+ do = check_data_convert_to_ds_if_needed(do, var)
+
  dm_am = dm.temporal.average(var)
  do_am = do.temporal.average(var)
  return dm_am, do_am # DataSets
@@ -64,6 +83,10 @@ def bias_xy(dm, do, var=None, weights=None):
  "Abstract": "Compute Full Average of Model - Observation",
  "Contact": "[email protected]",
  }
+
+ dm = check_data_convert_to_ds_if_needed(dm, var)
+ do = check_data_convert_to_ds_if_needed(do, var)
+
  dif = dm[var] - do[var]
  if weights is None:
  weights = dm.spatial.get_weights(axis=["X", "Y"])
@@ -79,6 +102,10 @@ def bias_xyt(dm, do, var=None):
  "Abstract": "Compute Full Average of Model - Observation",
  "Contact": "[email protected]",
  }
+
+ dm = check_data_convert_to_ds_if_needed(dm, var)
+ do = check_data_convert_to_ds_if_needed(do, var)
+
  ds = dm.copy(deep=True)
  ds["dif"] = dm[var] - do[var]
  stat = (
@@ -95,6 +122,10 @@ def cor_xy(dm, do, var=None, weights=None):
  "Abstract": "Compute Spatial Correlation",
  "Contact": "[email protected]",
  }
+
+ dm = check_data_convert_to_ds_if_needed(dm, var)
+ do = check_data_convert_to_ds_if_needed(do, var)
+
  if weights is None:
  weights = dm.spatial.get_weights(axis=["X", "Y"])
 
@@ -123,6 +154,8 @@ def mean_xy(d, var=None, weights=None):
  "Contact": "[email protected]",
  }
 
+ d = check_data_convert_to_ds_if_needed(d, var)
+
  if weights is None:
  weights = d.spatial.get_weights(axis=["X", "Y"])
  stat = float(d[var].weighted(weights).mean(("lon", "lat")))
@@ -140,6 +173,10 @@ def meanabs_xy(dm, do, var=None, weights=None):
  }
  if weights is None:
  weights = dm.spatial.get_weights(axis=["X", "Y"])
+
+ dm = check_data_convert_to_ds_if_needed(dm, var)
+ do = check_data_convert_to_ds_if_needed(do, var)
+
  dif = abs(dm[var] - do[var])
  stat = dif.weighted(weights).mean(("lon", "lat"))
  return float(stat)
@@ -154,6 +191,10 @@ def meanabs_xyt(dm, do, var=None):
  + "Absolute Difference Between Model And Observation",
  "Contact": "[email protected]",
  }
+
+ dm = check_data_convert_to_ds_if_needed(dm, var)
+ do = check_data_convert_to_ds_if_needed(do, var)
+
  ds = dm.copy(deep=True)
  ds["absdif"] = abs(dm[var] - do[var])
  stat = (
@@ -172,6 +213,10 @@ def rms_0(dm, do, var=None, weighted=True):
  "Abstract": "Compute Root Mean Square over the first axis",
  "Contact": "[email protected]",
  }
+
+ dm = check_data_convert_to_ds_if_needed(dm, var)
+ do = check_data_convert_to_ds_if_needed(do, var)
+
  dif_square = (dm[var] - do[var]) ** 2
  if weighted:
  weights = dm.spatial.get_weights(axis=["Y"])
@@ -189,6 +234,10 @@ def rms_xy(dm, do, var=None, weights=None):
  "Abstract": "Compute Spatial Root Mean Square",
  "Contact": "[email protected]",
  }
+
+ dm = check_data_convert_to_ds_if_needed(dm, var)
+ do = check_data_convert_to_ds_if_needed(do, var)
+
  dif_square = (dm[var] - do[var]) ** 2
  if weights is None:
  weights = dm.spatial.get_weights(axis=["X", "Y"])
@@ -204,6 +253,10 @@ def rms_xyt(dm, do, var=None):
  "Abstract": "Compute Spatial and Temporal Root Mean Square",
  "Contact": "[email protected]",
  }
+
+ dm = check_data_convert_to_ds_if_needed(dm, var)
+ do = check_data_convert_to_ds_if_needed(do, var)
+
  ds = dm.copy(deep=True)
  ds["diff_square"] = (dm[var] - do[var]) ** 2
  ds["diff_square_sqrt"] = np.sqrt(
@@ -221,6 +274,10 @@ def rmsc_xy(dm, do, var=None, weights=None):
  "Abstract": "Compute Centered Spatial Root Mean Square",
  "Contact": "[email protected]",
  }
+
+ dm = check_data_convert_to_ds_if_needed(dm, var)
+ do = check_data_convert_to_ds_if_needed(do, var)
+
  if weights is None:
  weights = dm.spatial.get_weights(axis=["X", "Y"])
 
@@ -259,6 +316,8 @@ def std_xy(ds, var=None, weights=None):
  if weights is None:
  weights = ds.spatial.get_weights(axis=["X", "Y"])
 
+ ds = check_data_convert_to_ds_if_needed(ds, var)
+
  lat_key = xc.axis.get_dim_keys(ds, axis="Y")
  lon_key = xc.axis.get_dim_keys(ds, axis="X")
 
@@ -278,6 +337,7 @@ def std_xyt(d, var=None):
  "Contact": "[email protected]",
  }
  ds = d.copy(deep=True)
+ ds = check_data_convert_to_ds_if_needed(ds, var)
  average = d.spatial.average(var, axis=["X", "Y"]).temporal.average(var)[var]
  ds["anomaly"] = (d[var] - average) ** 2
  variance = (
@@ -296,6 +356,9 @@ def zonal_mean(dm, do, var=None):
  "Contact": "[email protected]",
  "Comments": "",
  }
+ dm = check_data_convert_to_ds_if_needed(dm, var)
+ do = check_data_convert_to_ds_if_needed(do, var)
+
  dm_zm = dm.spatial.average(var, axis=["X"])
  do_zm = do.spatial.average(var, axis=["X"])
  return dm_zm, do_zm # DataSets