Update visualization.py

visualization.py

@@ -16,37 +16,50 @@ def plot_mse(df, mse_values, title, color='navy', time_threshold=3, anomaly_thre
 
     # Ensure df and mse_values have the same length and remove NaN values
     min_length = min(len(df), len(mse_values))
-    df = df.iloc[:min_length]
+    df = df.iloc[:min_length].copy()
     mse_values = mse_values[:min_length]
 
-    # Remove NaN values
-    mask = ~np.isnan(mse_values)
-    df = df[mask]
-    mse_values = mse_values[mask]
-
-    # Calculate rolling mean and std
-    mean = pd.Series(mse_values).rolling(window=10, min_periods=1).mean()
-    std = pd.Series(mse_values).rolling(window=10, min_periods=1).std()
-
-    # Plot scatter points
-    ax.scatter(df['Seconds'], mse_values, color=color, alpha=0.3, s=5)
-
-    # Plot mean line and std fill only for continuous valid segments
+    # Remove NaN values and create a mask for valid data
     valid_mask = ~np.isnan(mse_values)
-    segments = np.split(np.arange(len(df)), np.where(~valid_mask)[0])
+    df = df[valid_mask]
+    mse_values = mse_values[valid_mask]
+
+    # Function to identify continuous segments
+    def get_continuous_segments(seconds, values, max_gap=1):
+        segments = []
+        current_segment = []
+        for i, (sec, val) in enumerate(zip(seconds, values)):
+            if not current_segment or (sec - current_segment[-1][0] <= max_gap):
+                current_segment.append((sec, val))
+            else:
+                segments.append(current_segment)
+                current_segment = [(sec, val)]
+        if current_segment:
+            segments.append(current_segment)
+        return segments
+
+    # Get continuous segments
+    segments = get_continuous_segments(df['Seconds'], mse_values)
+
+    # Plot each segment separately
     for segment in segments:
-        if len(segment) > 0 and valid_mask[segment[0]]:
-            ax.plot(df['Seconds'].iloc[segment], mean.iloc[segment], color=color, linewidth=0.5)
-            ax.fill_between(df['Seconds'].iloc[segment], 
-                            mean.iloc[segment] - std.iloc[segment],
-                            mean.iloc[segment] + std.iloc[segment],
-                            color=color, alpha=0.1)
-
-    # Add median line
+        segment_seconds, segment_mse = zip(*segment)
+        ax.scatter(segment_seconds, segment_mse, color=color, alpha=0.3, s=5)
+        
+        # Calculate and plot rolling mean and std for this segment
+        if len(segment) > 1:  # Only if there's more than one point in the segment
+            segment_df = pd.DataFrame({'Seconds': segment_seconds, 'MSE': segment_mse})
+            segment_df = segment_df.sort_values('Seconds')
+            mean = segment_df['MSE'].rolling(window=min(10, len(segment)), min_periods=1, center=True).mean()
+            std = segment_df['MSE'].rolling(window=min(10, len(segment)), min_periods=1, center=True).std()
+            
+            ax.plot(segment_df['Seconds'], mean, color=color, linewidth=0.5)
+            ax.fill_between(segment_df['Seconds'], mean - std, mean + std, color=color, alpha=0.1)
+
+    # Rest of the function remains the same
     median = np.median(mse_values)
     ax.axhline(y=median, color='black', linestyle='--', label='Median Baseline')
 
-    # Add threshold line
     threshold = np.mean(mse_values) + anomaly_threshold * np.std(mse_values)
     ax.axhline(y=threshold, color='red', linestyle='--', label=f'Threshold: {anomaly_threshold:.1f}')
     ax.text(ax.get_xlim()[1], threshold, f'Threshold: {anomaly_threshold:.1f}', verticalalignment='center', horizontalalignment='left', color='red')