fix: correct the if in

resgroup · Dec 5, 2024 · eff9dfd · eff9dfd
1 parent 308361f
commit eff9dfd
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 31 deletions.
diff --git a/examples/smarteole_example.ipynb b/examples/smarteole_example.ipynb
diff --git a/tests/test_ops_curve_shift.py b/tests/test_ops_curve_shift.py
@@ -1,4 +1,5 @@
 import logging
+from pathlib import Path
 from unittest.mock import Mock, patch
 
 import numpy as np

diff --git a/wind_up/ops_curve_shift.py b/wind_up/ops_curve_shift.py
@@ -226,7 +226,7 @@ def calculate_curve_shift(curve_shift_input: CurveShiftInput) -> CurveShiftOutpu
     post_df["expected_y"] = np.interp(post_df[conf.x_col], mean_curve["x_mean"], mean_curve["y_mean"])
     mean_df = post_df.mean()
 
-    if conf.y_col == CurveTypes.PITCH.value:
+    if conf.name in [CurveTypes.PITCH.value, CurveTypes.PITCH]:
         result = mean_df[conf.y_col] - mean_df["expected_y"]
     else:
         result = (mean_df[conf.y_col] / mean_df["expected_y"] - 1).clip(-1, 1)