From 9da0cc670493e3e86f3f4aa28fcbdc2d5d116c40 Mon Sep 17 00:00:00 2001
From: "Thomson, Alec (CASS, Kensington)" <alec.thomson@csiro.au>
Date: Tue, 2 Apr 2024 12:33:42 +1100
Subject: [PATCH] Add imager options

---
 arrakis/imager.py        | 13 +++++++++++++
 arrakis/process_spice.py | 13 +++++++------
 2 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/arrakis/imager.py b/arrakis/imager.py
index 4134286b..17137327 100644
--- a/arrakis/imager.py
+++ b/arrakis/imager.py
@@ -157,6 +157,7 @@ def image_beam(
     local_rms_window: Optional[float] = None,
     multiscale: bool = False,
     multiscale_scale_bias: Optional[float] = None,
+    multiscale_scales: Optional[str] = "0,2,4,8,16,32,64,128",
     data_column: str = "CORRECTED_DATA",
     no_mf_weighting: bool = False,
     no_update_model_required: bool = True,
@@ -212,6 +213,7 @@ def image_beam(
             local_rms_window=local_rms_window,
             multiscale_scale_bias=multiscale_scale_bias,
             multiscale=multiscale,
+            multiscale_scales=multiscale_scales,
             data_column=data_column,
             no_mf_weighting=no_mf_weighting,
             no_update_model_required=no_update_model_required,
@@ -267,6 +269,7 @@ def image_beam(
             local_rms_window=local_rms_window,
             multiscale=multiscale if not squared_channel_joining else False,
             multiscale_scale_bias=multiscale_scale_bias,
+            multiscale_scales=multiscale_scales,
             data_column=data_column,
             no_mf_weighting=no_mf_weighting,
             no_update_model_required=no_update_model_required,
@@ -626,6 +629,7 @@ def main(
     wsclean_path: Union[Path, str] = "docker://alecthomson/wsclean:latest",
     multiscale: Optional[bool] = None,
     multiscale_scale_bias: Optional[float] = None,
+    multiscale_scales: Optional[str] = "0,2,4,8,16,32,64,128",
     absmem: Optional[float] = None,
     make_residual_cubes: Optional[bool] = False,
     ms_glob_pattern: str = "scienceData*_averaged_cal.leakage.ms",
@@ -697,6 +701,7 @@ def main(
             local_rms_window=local_rms_window,
             multiscale=multiscale,
             multiscale_scale_bias=multiscale_scale_bias,
+            multiscale_scales=multiscale_scales,
             absmem=absmem,
             data_column=data_column,
             no_mf_weighting=no_mf_weighting,
@@ -911,6 +916,12 @@ def imager_parser(parent_parser: bool = False) -> argparse.ArgumentParser:
         default=None,
         help="The multiscale scale bias term provided to wsclean. ",
     )
+    parser.add_argument(
+        "--multiscale_scales",
+        type=str,
+        default="0,2,4,8,16,32,64,128",
+        help="The scales used in the multiscale clean. ",
+    )
     parser.add_argument(
         "--absmem",
         type=float,
@@ -991,6 +1002,8 @@ def cli():
             Path(args.local_wsclean) if args.local_wsclean else args.hosted_wsclean
         ),
         multiscale=args.multiscale,
+        multiscale_scale_bias=args.multiscale_scale_bias,
+        multiscale_scales=args.multiscale_scales,
         ms_glob_pattern=args.ms_glob_pattern,
         data_column=args.data_column,
         skip_fix_ms=args.skip_fix_ms,
diff --git a/arrakis/process_spice.py b/arrakis/process_spice.py
index 05fa53ee..7b275f2d 100644
--- a/arrakis/process_spice.py
+++ b/arrakis/process_spice.py
@@ -250,12 +250,13 @@ def main(args: configargparse.Namespace) -> None:
     output_args_path = save_args(args)
     logger.info(f"Saved arguments to {output_args_path}.")
 
-    # Test the mongoDB
-    test_db(
-        host=host,
-        username=args.username,
-        password=args.password,
-    )
+    if not args.imager_only:
+        # Test the mongoDB
+        test_db(
+            host=host,
+            username=args.username,
+            password=args.password,
+        )
 
     if not args.skip_imager:
         # This is the client for the imager component of the arrakis