Detect and retry wsclean divergence error

CHANGELOG.md

@@ -9,6 +9,7 @@
     - the strategy format now has a `operations` set of keywords, including `stokesv` to drawn options from
     - naming format of output linmos files could contain the pol field
     - `stokesv` imaging will not linmos the cleaning residuals together, even if the `--linmos-residuals` CLI is provided
+- Capture `CleanDivergenceError` from `wsclean` and rerun with larger image size and lower gain
 
 # 0.2.5
 - added in skip rounds for masking and selfcal

flint/prefect/common/imaging.py

@@ -276,6 +276,8 @@ def task_wsclean_imager(
     Returns:
         WSCleanCommand: A resulting wsclean command and resulting meta-data
     """
+    from flint.exceptions import CleanDivergenceError
+
     ms = in_ms if isinstance(in_ms, MS) else in_ms.ms
 
     update_wsclean_options = (
@@ -286,11 +288,41 @@ def task_wsclean_imager(
         update_wsclean_options["fits_mask"] = fits_mask.mask_fits
 
     logger.info(f"wsclean inager {ms=}")
-    return wsclean_imager(
-        ms=ms,
-        wsclean_container=wsclean_container,
-        update_wsclean_options=update_wsclean_options,
-    )
+    try:
+        return wsclean_imager(
+            ms=ms,
+            wsclean_container=wsclean_container,
+            update_wsclean_options=update_wsclean_options,
+        )
+    except CleanDivergenceError:
+        # NOTE: If the cleaning failed retry with some larger images
+        # and slower cleaning. Perhaps this should be moved closer
+        # to the wscleaning functionality
+        size = (
+            update_wsclean_options["size"] + 1024
+            if "size" in update_wsclean_options
+            else 8196
+        )
+        mgain = (
+            max(0, update_wsclean_options["mgain"] - 0.1)
+            if "mgain" in update_wsclean_options
+            else 0.6
+        )
+        convergence_wsclean_options = dict(size=size, mgain=mgain)
+        # dicts are mutable. Don't want to change for everything. Unclear to me
+        # how prefect would behave here.
+        update_wsclean_options = update_wsclean_options.copy().update(
+            **convergence_wsclean_options
+        )
+        logger.warn(
+            f"Clean divergence dertected. Reruning. Updated options {convergence_wsclean_options=}"
+        )
+
+        return wsclean_imager(
+            ms=ms,
+            wsclean_container=wsclean_container,
+            update_wsclean_options=update_wsclean_options,
+        )
 
 
 @task