Stress test early-kernel operator predictor

1 files changed, 17 insertions, 2 deletions

diff --git a/scripts/plot_compressed_operator_results.py b/scripts/plot_compressed_operator_results.py
index c23fa14..ee2f472 100644
--- a/scripts/plot_compressed_operator_results.py
+++ b/scripts/plot_compressed_operator_results.py
@@ -53,8 +53,23 @@ def read_rows(pattern: str, input_dim: int, output_dim: int, width: int) -> pd.D
     if not frames:
         raise FileNotFoundError(f"No rows found for pattern: {pattern}")
     df = pd.concat(frames, ignore_index=True)
-    parameter_count = input_dim * width + width * width + width * output_dim
-    constraint_rank = max(width * width - 1, 0) + max(width * output_dim - 1, 0)
+    if "width" in df.columns:
+        row_width = df["width"].astype(int)
+    else:
+        row_width = pd.Series(width, index=df.index)
+    if "hidden_layers" in df.columns:
+        hidden_layers = df["hidden_layers"].astype(int)
+    else:
+        hidden_layers = pd.Series(2, index=df.index)
+    parameter_count = (
+        input_dim * row_width
+        + (hidden_layers - 1).clip(lower=0) * row_width * row_width
+        + row_width * output_dim
+    )
+    constraint_rank = (
+        (hidden_layers - 1).clip(lower=0) * np.maximum(row_width * row_width - 1, 0)
+        + np.maximum(row_width * output_dim - 1, 0)
+    )
     df["hard_fa_capacity_margin"] = (
         parameter_count - constraint_rank - output_dim * df["train_samples"]
     )