Add option to error if there are really bad z^2

hera_notebook_templates/notebooks/lststack.ipynb

@@ -159,6 +159,7 @@
     "freq_max: float = 0.0\n",
     "history: str = \"\"\n",
     "plot_every: int = 1\n",
+    "exception_for_zsq_above: float = np.inf  # Raise an exception if any baseline has a zsq greater than this. Useful for flagging weird behaviour for validation.\n",
     "\n",
     "# In-painting config\n",
     "inpaint_horizon: float = 1.0\n",
@@ -1301,6 +1302,17 @@
     "    zsquare = do_flagging([direct_zscore_pruning, watershed], cross_stacks, auto_stats, max_iter=max_flagging_iterations)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2c58d0b1-6900-4e65-981d-68fb0ec8d30c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "if any(np.any(zsq.metrics > exception_for_zsq_above) for zsq in zsquare):\n",
+    "    raise RuntimeError(f\"The maximum zsquares before inpainting are {[np.nanmax(zsq.metrics) for zsq in zsquares]}, which is over the threshold for raising an error\")"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "ca98fc9f-a22e-4a2d-8ff1-8bccfdb453f5",