Fix crash with swapped Spectral/Stokes axes for large mosaics

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).

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,31 +1207,42 @@ 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):
                         counter += 1
                         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)