From c03a5ce3a04867290b0e46d3a065f507afa4356e Mon Sep 17 00:00:00 2001
From: tgalvin <tim.galvin@csiro.au>
Date: Wed, 25 Dec 2024 15:29:27 +1100
Subject: [PATCH] remove leakage attempts to use stokes .i.

---
 flint/coadd/linmos.py                         | 11 +++++--
 flint/prefect/flows/subtract_cube_pipeline.py | 29 +++++++++++++++++++
 2 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/flint/coadd/linmos.py b/flint/coadd/linmos.py
index 7e02dee9..423fd73c 100644
--- a/flint/coadd/linmos.py
+++ b/flint/coadd/linmos.py
@@ -404,7 +404,9 @@ def _get_alpha_linmos_option(pol_axis: Optional[float] = None) -> str:
 
 
 def _get_holography_linmos_options(
-    holofile: Optional[Path] = None, pol_axis: Optional[float] = None
+    holofile: Optional[Path] = None,
+    pol_axis: Optional[float] = None,
+    remove_leakage: bool = False,
 ) -> str:
     """Construct the appropriate set of linmos options that
     describe the use of the holography cube file to primary
@@ -414,6 +416,7 @@ def _get_holography_linmos_options(
     Args:
         holofile (Optional[Path], optional): Path to the holography cube file to primary beam correct with. Defaults to None.
         pol_axis (Optional[float], optional): The rotation of the third axis as described in an ASAKP MS. Defaults to None.
+        remove_leakage (bool, optional): Add the directive to remove leakage. Defaults to False.
 
     Returns:
         str: Set of linmos options to add to a parset file
@@ -431,7 +434,7 @@ def _get_holography_linmos_options(
     parset = (
         f"linmos.primarybeam      = ASKAP_PB\n"
         f"linmos.primarybeam.ASKAP_PB.image = {str(holofile.absolute())}\n"
-        f"linmos.removeleakage    = true\n"
+        f"linmos.removeleakage    = {'true' if remove_leakage else 'false'}\n"
     )
     parset += _get_alpha_linmos_option(pol_axis=pol_axis)
 
@@ -522,7 +525,9 @@ def generate_linmos_parameter_set(
         "linmos.imageaccess.axis  = 1\n"  # WSClean outputs frequency as second dimension (so 1 in zero indexing)
     )
     # Construct the holography section of the linmos parset
-    parset += _get_holography_linmos_options(holofile=holofile, pol_axis=pol_axis)
+    parset += _get_holography_linmos_options(
+        holofile=holofile, pol_axis=pol_axis, remove_leakage=".i." not in str(images[0])
+    )
 
     # Now write the file, me hearty
     logger.info(f"Writing parset to {str(parset_output_path)}.")
diff --git a/flint/prefect/flows/subtract_cube_pipeline.py b/flint/prefect/flows/subtract_cube_pipeline.py
index b63fe472..4142c73c 100644
--- a/flint/prefect/flows/subtract_cube_pipeline.py
+++ b/flint/prefect/flows/subtract_cube_pipeline.py
@@ -30,6 +30,7 @@
 )
 from flint.options import (
     AddModelSubtractFieldOptions,
+    BaseOptions,
     SubtractFieldOptions,
     add_options_to_parser,
     create_options_from_parser,
@@ -42,6 +43,15 @@
 from flint.naming import get_sbid_from_path
 
 
+class CrystalBallOptions(BaseOptions):
+    """Options related to running crystal ball"""
+
+    attempt_crystalball: bool = False
+    """Attempt to predict the model visibilities using ``crystalball``"""
+    wsclean_pol_mode: List[str] = ["i"]
+    """The polarisation of the wsclean model that was generated"""
+
+
 def _check_and_verify_options(
     options: Union[AddModelSubtractFieldOptions, SubtractFieldOptions],
 ) -> None:
@@ -174,6 +184,25 @@ def task_addmodel_to_ms(
     return ms.with_options(model_column="MODEL_DATA")
 
 
+def task_crystalball_to_ms(ms: MS, crystalball_options: CrystalBallOptions) -> MS:
+    from prefect_dask import get_dask_client
+    from flint.imager.wsclean import get_wsclean_output_source_list_path
+
+    logger.info(f"Searching for wsclean source list for {ms.path}")
+    for idx, pol in enumerate(crystalball_options.wsclean_pol_mode):
+        wsclean_source_list_path = get_wsclean_output_source_list_path(
+            name_path=ms.path, pol=pol
+        )
+        assert (
+            wsclean_source_list_path.exists()
+        ), f"{wsclean_source_list_path=} was requested, but does not exist"
+
+    with get_dask_client():
+        logger.info("Running crystalball in prefect dask client")
+
+    return ms
+
+
 @task
 def task_combine_all_linmos_images(
     linmos_commands: List[LinmosCommand],