updates pearson and r_squared to consider epsilon

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "teehr"
-version = "0.5.1dev6"
+version = "0.5.1dev9"
 description = "Tools for Exploratory Evaluation in Hydrologic Research"
 authors = [
     "RTI International",

src/teehr/__init__.py

@@ -1,7 +1,7 @@
 """Initialize the TEEHR package."""
 import warnings
 
-__version__ = "0.5.1dev8"
+__version__ = "0.5.1dev9"
 
 with warnings.catch_warnings():
     warnings.simplefilter("ignore", UserWarning)

src/teehr/metrics/deterministic_funcs.py

@@ -2,6 +2,7 @@
 import numpy as np
 import numpy.typing as npt
 import pandas as pd
+from scipy.stats import rankdata
 
 from teehr.models.metrics.basemodels import MetricsBasemodel
 from teehr.models.metrics.basemodels import TransformEnum
@@ -201,7 +202,21 @@ def pearson_correlation(model: MetricsBasemodel) -> Callable:
     def pearson_correlation_inner(p: pd.Series, s: pd.Series) -> float:
         """Pearson Correlation Coefficient."""
         p, s = _transform(p, s, model)
-        return np.corrcoef(s, p)[0][1]
+
+        if model.add_epsilon:
+            # Calculate covariance between p and s
+            numerator = np.cov(p, s)[0, 1]
+
+            # Calculate standard deviations and multiply them
+            denominator = np.nanstd(p) * np.nanstd(s) + EPSILON
+
+            # Calculate correlation coefficient
+            result = numerator / denominator
+
+        else:
+            result = np.corrcoef(s, p)[0][1]
+
+        return result
 
     return pearson_correlation_inner
 
@@ -236,8 +251,23 @@ def r_squared(model: MetricsBasemodel) -> Callable:
     def r_squared_inner(p: pd.Series, s: pd.Series) -> float:
         """R-squared."""
         p, s = _transform(p, s, model)
-        pearson_correlation_coefficient = np.corrcoef(s, p)[0][1]
-        return np.power(pearson_correlation_coefficient, 2)
+
+        if model.add_epsilon:
+            # Calculate covariance between p and s
+            numerator = np.cov(p, s)[0, 1]
+
+            # Calculate standard deviations and multiply them
+            denominator = np.nanstd(p) * np.nanstd(s) + EPSILON
+
+            # Calculate correlation coefficient and square it
+            pearson_correlation_coefficient = numerator / denominator
+            result = np.power(pearson_correlation_coefficient, 2)
+
+        else:
+            pearson_correlation_coefficient = np.corrcoef(s, p)[0][1]
+            result = np.power(pearson_correlation_coefficient, 2)
+
+        return result
 
     return r_squared_inner
 
@@ -311,19 +341,21 @@ def spearman_correlation_inner(p: pd.Series, s: pd.Series) -> float:
         """Spearman Rank Correlation Coefficient."""
         p, s = _transform(p, s, model)
 
-        primary_rank = p.rank()
-        secondary_rank = s.rank()
-        count = len(p)
+        # calculate ranks (average method for ties)
+        primary_ranks = rankdata(p, method='average')
+        secondary_ranks = rankdata(s, method='average')
+
+        # calculate covariance between p_rank and s_rank
+        covariance = np.cov(primary_ranks, secondary_ranks)[0, 1]
+
+        # calculate standard deviations of ranks
+        std_primary = np.std(primary_ranks)
+        std_secondary = np.std(secondary_ranks)
+
         if model.add_epsilon:
-            result = 1 - (
-                6 * np.sum(np.abs(primary_rank - secondary_rank)**2)
-                / (count * (count**2 - 1)) + EPSILON
-                )
+            result = covariance / (std_primary * std_secondary + EPSILON)
         else:
-            result = 1 - (
-                6 * np.sum(np.abs(primary_rank - secondary_rank)**2)
-                / (count * (count**2 - 1))
-                )
+            result = covariance / (std_primary * std_secondary)
 
         return result
 

tests/query/test_get_metrics_query.py

@@ -574,7 +574,7 @@ def test_ensemble_metrics(tmpdir):
 def test_metrics_transforms(tmpdir):
     """Test applying metric transforms (non-bootstrap)."""
     # Define the evaluation object.
-    eval = setup_v0_3_study(tmpdir)
+    test_eval = setup_v0_3_study(tmpdir)
 
     # define metric requiring p,s
     kge = DeterministicMetrics.KlingGuptaEfficiency()
@@ -592,21 +592,21 @@ def test_metrics_transforms(tmpdir):
     mvtd_t.transform = 'log'
 
     # get metrics_df
-    metrics_df_tansformed_e = eval.metrics.query(
+    metrics_df_tansformed_e = test_eval.metrics.query(
         group_by=["primary_location_id", "configuration_name"],
         include_metrics=[
             kge_t_e,
             mvtd_t
         ]
     ).to_pandas()
-    metrics_df_transformed = eval.metrics.query(
+    metrics_df_transformed = test_eval.metrics.query(
         group_by=["primary_location_id", "configuration_name"],
         include_metrics=[
             kge_t,
             mvtd_t
         ]
     ).to_pandas()
-    metrics_df = eval.metrics.query(
+    metrics_df = test_eval.metrics.query(
         group_by=["primary_location_id", "configuration_name"],
         include_metrics=[
             kge,
@@ -628,6 +628,42 @@ def test_metrics_transforms(tmpdir):
     assert result_kge != result_kge_t
     assert result_mvtd == result_mvtd_t
 
+    # test epsilon on R2 and Pearson
+    r2 = DeterministicMetrics.Rsquared()
+    r2_e = DeterministicMetrics.Rsquared()
+    r2_e.add_epsilon = True
+    pearson = DeterministicMetrics.PearsonCorrelation()
+    pearson_e = DeterministicMetrics.PearsonCorrelation()
+    pearson_e.add_epsilon = True
+
+    # ensure we can obtain a divide by zero error
+    sdf = test_eval.joined_timeseries.to_sdf()
+    from pyspark.sql.functions import lit
+    sdf = sdf.withColumn("primary_value", lit(100.0))
+    test_eval.joined_timeseries._write_spark_df(sdf, write_mode="overwrite")
+
+    # get metrics df control and assert divide by zero occurs
+    metrics_df_e_control = test_eval.metrics.query(
+        group_by=["primary_location_id", "configuration_name"],
+        include_metrics=[
+            r2,
+            pearson
+        ]
+    ).to_pandas()
+    assert np.isnan(metrics_df_e_control.r_squared.values).all()
+    assert np.isnan(metrics_df_e_control.pearson_correlation.values).all()
+
+    # get metrics df test and ensure no divide by zero occurs
+    metrics_df_e_test = test_eval.metrics.query(
+        group_by=["primary_location_id", "configuration_name"],
+        include_metrics=[
+            r2_e,
+            pearson_e
+        ]
+    ).to_pandas()
+    assert np.isfinite(metrics_df_e_test.r_squared.values).all()
+    assert np.isfinite(metrics_df_e_test.pearson_correlation.values).all()
+
 
 def test_bootstrapping_transforms(tmpdir):
     """Test applying metric transforms (bootstrap)."""