fix grid

src/tools/plots.py

@@ -21,23 +21,31 @@ def plot_solve_rate(df, task, rows=30, cols=38):
     keys = df["task_id"]
     values = df["solve_rate"]
     
-    values = np.array(values)
+    values = np.array(values, dtype=float)  # Ensure values are floats
         
     n = len(values)
-    values = np.pad(values, (0, rows * cols - n), 'constant', constant_values=np.nan).reshape((rows, cols))
-    keys = np.pad(keys, (0, rows * cols - n), 'constant', constant_values='').reshape((rows, cols))
+    pad_width = rows * cols - n
+    
+    # Use masked array to handle NaN values
+    masked_values = np.ma.array(values)
+    masked_values = np.ma.pad(masked_values, (0, pad_width), 'constant', constant_values=np.ma.masked)
+    masked_values = masked_values.reshape((rows, cols))
+
+    keys = np.pad(keys, (0, pad_width), 'constant', constant_values='').reshape((rows, cols))
 
-    hover_text = np.empty_like(values, dtype=object)
+    hover_text = np.empty_like(masked_values, dtype=object)
     for i in range(rows):
         for j in range(cols):
-            if not np.isnan(values[i, j]):
-                hover_text[i, j] = f"{keys[i, j]}<br>Solve Rate: {values[i, j]:.2f}"
+            if not masked_values.mask[i, j]:
+                hover_text[i, j] = f"{keys[i, j]}<br>Solve Rate: {masked_values[i, j]:.2f}"
             else:
                 hover_text[i, j] = "NaN"
 
-    upper_solve_rate = round(np.count_nonzero(values)/n*100, 2)
+    # Use compressed array to count non-masked (finite) values
+    upper_solve_rate = round(np.count_nonzero(~masked_values.mask) / n * 100, 2)
+    
     fig = go.Figure(data=go.Heatmap(
-        z=values,
+        z=masked_values,
         text=hover_text,
         hoverinfo='text',
         colorscale='teal',
@@ -52,8 +60,6 @@ def plot_solve_rate(df, task, rows=30, cols=38):
         xaxis=dict(showticklabels=False),
         yaxis=dict(showticklabels=False),
         autosize=True,
-        # width=760,
-        # height=600,
     )
     
     return fig