From 7bfcf73732279bf92b9268f1d639389744dcbf0f Mon Sep 17 00:00:00 2001
From: "Thomson, Alec (CASS, Kensington)" <alec.thomson@csiro.au>
Date: Wed, 15 Jan 2025 20:02:43 +1100
Subject: [PATCH] Add cube rotation for linmos

---
 flint/coadd/linmos.py                        |  2 --
 flint/imager/wsclean.py                      | 29 ++++++++++++++++++++
 flint/prefect/flows/polarisation_pipeline.py | 21 +++++++-------
 3 files changed, 40 insertions(+), 12 deletions(-)

diff --git a/flint/coadd/linmos.py b/flint/coadd/linmos.py
index e0d0095a..2c3caef9 100644
--- a/flint/coadd/linmos.py
+++ b/flint/coadd/linmos.py
@@ -541,8 +541,6 @@ def generate_linmos_parameter_set(
         f"linmos.weightstate      = Inherent\n"
         f"linmos.cutoff           = {cutoff}\n"
         f"linmos.finalcutoff           = 0.01\n"
-        "linmos.imageaccess.axis  = 3\n"  # WSClean outputs frequency as second dimension (so 3 in zero indexing)
-        # Assuming linmos uses FITS/fortran indexing - which is weird...
     )
     # Construct the holography section of the linmos parset
     remove_leakage = (holofile is not None) and (".i." not in str(next(iter(images))))
diff --git a/flint/imager/wsclean.py b/flint/imager/wsclean.py
index 588a8131..ab956568 100644
--- a/flint/imager/wsclean.py
+++ b/flint/imager/wsclean.py
@@ -24,6 +24,7 @@
 from typing import Any, Collection, NamedTuple
 
 import numpy as np
+from astropy.io import fits
 from fitscube.combine_fits import combine_fits
 from prefect import task
 
@@ -849,6 +850,31 @@ def create_wsclean_cmd(
     )
 
 
+def _rotate_cube(output_cube_name) -> None:
+    logger.info(f"Rotating {output_cube_name=}")
+    # This changes the output cube to a shape of (chan, pol, dec, ra)
+    # which is what yandasoft linmos tasks like
+    with fits.open(output_cube_name, mode="update", memmap=True) as hdulist:
+        hdu = hdulist[0]
+        new_header = hdu[0].header
+        data_cube = hdu[0].data
+
+        tmp_header = new_header.copy()
+        # Need to swap NAXIS 3 and 4 to make LINMOS happy - booo
+        for a, b in ((3, 4), (4, 3)):
+            new_header[f"CTYPE{a}"] = tmp_header[f"CTYPE{b}"]
+            new_header[f"CRPIX{a}"] = tmp_header[f"CRPIX{b}"]
+            new_header[f"CRVAL{a}"] = tmp_header[f"CRVAL{b}"]
+            new_header[f"CDELT{a}"] = tmp_header[f"CDELT{b}"]
+            new_header[f"CUNIT{a}"] = tmp_header[f"CUNIT{b}"]
+
+        # Cube is currently STOKES, FREQ, RA, DEC - needs to be FREQ, STOKES, RA, DEC
+        data_cube = np.moveaxis(data_cube, 1, 0)
+        hdu[0].data = data_cube
+        hdu[0].header = new_header
+        hdulist.flush()
+
+
 def combine_image_set_to_cube(
     imageset: ImageSet,
     remove_original_images: bool = False,
@@ -895,6 +921,8 @@ def combine_image_set_to_cube(
         logger.info(f"Combining {len(subband_images)} images. {subband_images=}")
         freqs = combine_fits(file_list=subband_images, out_cube=output_cube_name)
 
+        _rotate_cube(output_cube_name)
+
         # Write out the hdu to preserve the beam table constructed in fitscube
         logger.info(f"Writing {output_cube_name=}")
 
@@ -936,6 +964,7 @@ def combine_images_to_cube(
 
     logger.info(f"Combining {len(images)} images. {images=}")
     freqs = combine_fits(file_list=images, out_cube=output_cube_name)
+    _rotate_cube(output_cube_name)
 
     # Write out the hdu to preserve the beam table constructed in fitscube
     logger.info(f"Writing {output_cube_name=}")
diff --git a/flint/prefect/flows/polarisation_pipeline.py b/flint/prefect/flows/polarisation_pipeline.py
index caa31f25..16f45f61 100644
--- a/flint/prefect/flows/polarisation_pipeline.py
+++ b/flint/prefect/flows/polarisation_pipeline.py
@@ -181,16 +181,17 @@ def process_science_fields_pol(
         force_remove_leakage = False
 
     for stokes, beam_cubes in stokes_beam_cubes.items():
-        linmos_result = task_linmos_images.submit(
-            images=beam_cubes,
-            container=pol_field_options.yandasoft_container,
-            holofile=pol_field_options.holofile,
-            cutoff=pol_field_options.pb_cutoff,
-            field_summary=field_summary,
-            stokesi_images=stokes_beam_cubes.get("i"),
-            force_remove_leakage=force_remove_leakage,
-        )
-        linmos_result_list.append(linmos_result)
+        with tags(f"stokes-{stokes}"):
+            linmos_result = task_linmos_images.submit(
+                images=beam_cubes,
+                container=pol_field_options.yandasoft_container,
+                holofile=pol_field_options.holofile,
+                cutoff=pol_field_options.pb_cutoff,
+                field_summary=field_summary,
+                stokesi_images=stokes_beam_cubes.get("i"),
+                force_remove_leakage=force_remove_leakage,
+            )
+            linmos_result_list.append(linmos_result)
 
     # wait for all linmos results to be completed
     _ = [linmos_result.result() for linmos_result in linmos_result_list]