From 69b818f265f8b70915c54edf2d76ad68a0a72402 Mon Sep 17 00:00:00 2001
From: paulf81 <paul.fleming@nrel.gov>
Date: Mon, 21 Oct 2024 15:24:08 -0600
Subject: [PATCH] Fix line in time fix (#227)

---
 flasc/utilities/tuner_utilities.py |  2 +-
 tests/model_tuning_test.py         | 36 ++++++++++++++++++++++++++++++
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/flasc/utilities/tuner_utilities.py b/flasc/utilities/tuner_utilities.py
index 7bc0283c..bcb4100b 100644
--- a/flasc/utilities/tuner_utilities.py
+++ b/flasc/utilities/tuner_utilities.py
@@ -50,7 +50,7 @@ def replicate_nan_values(
     common_columns = df_1.columns.intersection(df_2.columns)
 
     # Remove the time column from the common columns if included
-    common_columns.drop("time", errors="ignore")
+    common_columns = common_columns.drop("time", errors="ignore")
 
     # Use assign to create a new DataFrame with NaN values replaced
     df_2_updated = df_2.assign(
diff --git a/tests/model_tuning_test.py b/tests/model_tuning_test.py
index 0bfe6318..cfd3b210 100644
--- a/tests/model_tuning_test.py
+++ b/tests/model_tuning_test.py
@@ -28,6 +28,42 @@ def test_replicate_nan_values(self):
         assert result_df.equals(expected_df_2)
         assert df_1.equals(expected_df_1)
 
+    def test_replicate_nan_values_with_time(self):
+        # Sample dataframes
+        data_1 = {"A": [1, 2, np.nan, 4], "B": [5, np.nan, 7, 8], "C": [np.nan, 1, 1, 1]}
+        data_2 = {"A": [10, 20, 30, 40], "B": [50, 60, 70, 80]}
+        df_1 = pd.DataFrame(data_1)
+        df_2 = pd.DataFrame(data_2)
+
+        df_1["time"] = pd.date_range("2021-01-01", periods=4)
+        df_2["time"] = df_1["time"]
+
+        # Call the function to replicate NaN values
+        result_df = replicate_nan_values(df_1, df_2)
+
+        print(result_df)
+
+        # Expected output
+        expected_df_1 = pd.DataFrame(
+            {
+                "A": [1, 2, np.nan, 4],
+                "B": [5, np.nan, 7, 8],
+                "C": [np.nan, 1, 1, 1],
+                "time": pd.date_range("2021-01-01", periods=4),
+            }
+        )
+        expected_df_2 = pd.DataFrame(
+            {
+                "A": [10, 20, np.nan, 40],
+                "B": [50, np.nan, 70, 80],
+                "time": pd.date_range("2021-01-01", periods=4),
+            }
+        )
+
+        # Check if the result matches the expected output
+        assert result_df.equals(expected_df_2)
+        assert df_1.equals(expected_df_1)
+
     def test_evaluate_overall_wake_loss(self):
         # Create a sample DataFrame for testing
         data = {"pow_ref": [5, 2, 3], "pow_test": [1, 2, 2]}