From 14870ae20e2e03ec49c456dc68aea83e68987498 Mon Sep 17 00:00:00 2001
From: Thomas Williams <32936518+thomaswilliamsastro@users.noreply.github.com>
Date: Wed, 15 Jul 2026 12:54:17 +0100
Subject: [PATCH] Fix crash with swapped Spectral/Stokes axes for large mosaics

This PR fixes a crash that can occur if the spectral/Stokes axes are swapped when creating a large mosaic.

- Fixed crash if spectral/Stokes axis is swapped when making large mosaics
---
 CHANGELOG.md                         |  1 +
 phangsPipeline/casaMosaicRoutines.py | 55 ++++++++++++++++++++++------
 2 files changed, 44 insertions(+), 12 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f28dfbe..3e780c4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -40,6 +40,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fix stokes passed as integer string to imsubimage in 4D mosaic path (#349).
 - Import recipe_phangs_flat_mask in handlerDerived (#357).
 - Fix crash with using cleanmasks combined with new sdintimaging implementation (#360)
+- Fixed crash if spectral/Stokes axis is swapped when making large mosaics (#366).
 
 ### Dependencies
 - Bump actions/upload-artifact from 6 to 7 (#313).
diff --git a/phangsPipeline/casaMosaicRoutines.py b/phangsPipeline/casaMosaicRoutines.py
index b3e6164..8400049 100644
--- a/phangsPipeline/casaMosaicRoutines.py
+++ b/phangsPipeline/casaMosaicRoutines.py
@@ -1134,6 +1134,8 @@ def mosaic_aligned_data(
     local_outfile = os.path.basename(outfile)
     local_maskfile = os.path.basename(mask_file)
 
+    ndim = len(cube_shape)
+
     if not has_memory_issue:
         # Original approach for smaller cubes
         casaStuff.immath(imagename = local_imlist, mode='evalexpr',
@@ -1158,8 +1160,6 @@ def mosaic_aligned_data(
         myia_sum.open(sum_file)
         myia_weight.open(weight_file)
         
-        ndim = len(cube_shape)
-        
         try:
             if ndim == 3:
                 nchan = cube_shape[2]
@@ -1207,15 +1207,21 @@ def mosaic_aligned_data(
                         os.system('rm -rf ' + temp)
                         
             elif ndim == 4:
-                nstokes = cube_shape[3]
-                nchan = cube_shape[2]
-                total = nstokes * nchan
-                logger.info(f'Processing {nstokes} Stokes x {nchan} channels = {total} planes')
 
+                # The spectral/Stokes axis may be flipped. Keep track of this
                 csys_tmp = myia_sum.coordsys()
+                specaxis = csys_tmp.axiscoordinatetypes().index('Spectral')
+                stokesaxis = csys_tmp.axiscoordinatetypes().index('Stokes')
                 stokes_labels = csys_tmp.stokes()
                 csys_tmp.done()
 
+                nchan = cube_shape[specaxis]
+                nstokes = cube_shape[stokesaxis]
+                blc = [0, 0, 0, 0]
+
+                total = nstokes * nchan
+                logger.info(f'Processing {nstokes} Stokes x {nchan} channels = {total} planes')
+
                 counter = 0
                 for istokes in range(nstokes):
                     for ichan in range(nchan):
@@ -1223,15 +1229,20 @@ def mosaic_aligned_data(
                         if counter % 10 == 0:
                             logger.info(f'Processing plane {counter}/{total}')
                         
-                        blc = [0, 0, ichan, istokes]
+                        blc[specaxis] = ichan
+                        blc[stokesaxis] = istokes
                         
                         # Extract channel/stokes slices
                         temp_chan_images = []
                         for idx, im in enumerate(local_imlist):
                             temp_chan = f'temp_ch{ichan}_st{istokes}_img{idx}'
-                            casaStuff.imsubimage(imagename=im, outfile=temp_chan,
-                                               chans=str(ichan), stokes=stokes_labels[istokes],
-                                               dropdeg=False)
+                            casaStuff.imsubimage(
+                                imagename=im,
+                                outfile=temp_chan,
+                                chans=str(ichan),
+                                stokes=stokes_labels[istokes],
+                                dropdeg=False,
+                            )
                             temp_chan_images.append(temp_chan)
                         
                         # Process sum
@@ -1327,10 +1338,30 @@ def mosaic_aligned_data(
                     myia_mask.putchunk(mask_slice, blc)
                     
             elif ndim == 4:
+
+                # The spectral/Stokes axis may be flipped. Keep track of this
+                myia_sum.open(sum_file)
+                csys_tmp = myia_sum.coordsys()
+                specaxis = csys_tmp.axiscoordinatetypes().index('Spectral')
+                stokesaxis = csys_tmp.axiscoordinatetypes().index('Stokes')
+                csys_tmp.done()
+
+                myia_sum.close()
+
+                nchan = cube_shape[specaxis]
+                nstokes = cube_shape[stokesaxis]
+
+                blc = [0, 0, 0, 0]
+                trc = [cube_shape[0]-1, cube_shape[1]-1, 0, 0]
+
                 for istokes in range(nstokes):
                     for ichan in range(nchan):
-                        blc = [0, 0, ichan, istokes]
-                        trc = [cube_shape[0]-1, cube_shape[1]-1, ichan, istokes]
+
+                        blc[specaxis] = ichan
+                        blc[stokesaxis] = istokes
+
+                        trc[specaxis] = ichan
+                        trc[stokesaxis] = istokes
                         
                         myia_sum.open(sum_file)
                         sum_slice = myia_sum.getchunk(blc, trc)