Add katbeam S band option (#20)

* add S band to katbeam options

* debug

* add residual after fit as output

* fix uninitialised beam image when using katbeam

* np.float64 -> float in spimple-spifit parser

---------

Co-authored-by: landmanbester <lbester@ska.ac.za>
Co-authored-by: Athanaseus Javas Ramaila <aramaila@ska.ac.za>

spimple/apps/image_convolver.py

@@ -36,7 +36,7 @@ def image_convolver():
                         "Use power_beam_maker to make power beam "
                         "corresponding to image. ")
     parser.add_argument('-band', "--band", type=str, default='l',
-                        help="Band to use with JimBeam. L or UHF")
+                        help="Band to use with JimBeam. L, UHF or S")
     parser.add_argument('-pb-min', '--pb-min', type=float, default=0.05,
                         help="Set image to zero where pb falls below this value")
     parser.add_argument('-pf', '--padding-frac', type=float, default=0.5,
@@ -186,8 +186,13 @@ def image_convolver():
             from katbeam import JimBeam
             if opts.band.lower() == 'l':
                 beam = JimBeam('MKAT-AA-L-JIM-2020')
-            else:
+            elif opts.band.lower() == 'uhf':
                 beam = JimBeam('MKAT-AA-UHF-JIM-2020')
+            elif opts.band.lower() == 's':
+                beam = JimBeam('MKAT-AA-S-JIM-2020')
+            else:
+                raise ValueError(f"Unknown beam model for katbeam in band {opts.band}")
+
             beam_image = np.zeros(image.shape, dtype=opts.out_dtype)
             for v in range(freqs.size):
                 beam_image[v] = beam.I(xx, yy, freqs[v]/1e6)  # freqs in MHz

spimple/apps/spi_fitter.py

@@ -43,7 +43,7 @@ def spi_fitter():
                         "Default of zero means use all threads")
     parser.add_argument('-pb-min', '--pb-min', type=float, default=0.15,
                         help="Set image to zero where pb falls below this value")
-    parser.add_argument('-products', '--products', default='aeikIcmrb', type=str,
+    parser.add_argument('-products', '--products', default='aeikIcmrbd', type=str,
                         help="Outputs to write. Letter correspond to: \n"
                         "a - alpha map \n"
                         "e - alpha error map \n"
@@ -54,6 +54,7 @@ def spi_fitter():
                         "m - convolved model \n"
                         "r - convolved residual \n"
                         "b - average power beam \n"
+                        "d - difference between data and fitted model \n"
                         "Default is to write all of them")
     parser.add_argument('-pf', "--padding-frac", default=0.5, type=float,
                         help="Padding factor for FFT's.")
@@ -67,7 +68,7 @@ def spi_fitter():
                         help="Per-channel freqs to use during fit to frequency axis\n"+
                         "These are automatically generated from the FITS header.\n"+
                         "Only specify if you have a non-standard FITS header.")
-    parser.add_argument('-rf', '--ref-freq', default=None, type=np.float64,
+    parser.add_argument('-rf', '--ref-freq', default=None, type=float,
                         help='Reference frequency where the I0 map is sought. \n'
                         "Will overwrite in fits headers of output.")
     parser.add_argument('-otype', '--out_dtype', default='f4', type=str,
@@ -94,7 +95,7 @@ def spi_fitter():
     parser.add_argument('-ct', '--corr-type', type=str, default='linear',
                         help="Correlation typ i.e. linear or circular. ")
     parser.add_argument('-band', "--band", type=str, default='l',
-                        help="Band to use with JimBeam. L or UHF")
+                        help="Band to use with JimBeam. L, UHF or S")
     parser.add_argument('-db', '--deselect-bands', default=None, nargs='+', type=int,
                         help="Indices of subbands to exclude from the fitting \n"
                         "By default, all the sub-bands are used for the residual image. \n"
@@ -250,13 +251,19 @@ def spi_fitter():
             beam_image = load_fits(opts.beam_model, dtype=opts.out_dtype).squeeze()
         elif opts.beam_model == "JimBeam":
             from katbeam import JimBeam
+            beam_image = []
             if opts.band.lower() == 'l':
                 beam = JimBeam('MKAT-AA-L-JIM-2020')
-            else:
+            elif opts.band.lower() == 'uhf':
                 beam = JimBeam('MKAT-AA-UHF-JIM-2020')
-            beam_image = np.zeros(model.shape, dtype=opts.out_dtype)
+            elif opts.band.lower() == 's':
+                beam = JimBeam('MKAT-AA-S-JIM-2020')
+            else:
+                raise ValueError(f"Unknown beam model for katbeam in band {opts.band}")
+            beam_image = np.zeros_like(model)
             for v in range(freqs.size):
-                beam_image[v] = beam.I(xx, yy, freqs[v]/1e6)  # freqs in MHz
+                beam_image.append(beam.I(xx, yy, freqs[v]/1e6))  # freqs in MHz
+            beam_image = np.stack(beam_image)
 
         else:
             beam_image = interpolate_beam(xx, yy, freqs, opts)
@@ -449,15 +456,26 @@ def spi_fitter():
     alpha_err_map[maskindices[:, 0], maskindices[:, 1]] = alpha_err
     i0map[maskindices[:, 0], maskindices[:, 1]] = Iref
     i0_err_map[maskindices[:, 0], maskindices[:, 1]] = i0_err
+    Irec_cube = i0map[None, :, :] * (freqs[:, None, None]/ref_freq)**alphamap[None, :, :]
+    fit_diff = np.zeros_like(model)
+    fit_diff[...] = np.nan
+    ix = maskindices[:, 0]
+    iy = maskindices[:, 1]
+    fit_diff[:, ix, iy] = model[:, ix, iy]/beam_image[:, ix, iy]
+    fit_diff[:, ix, iy] -= Irec_cube[:, ix, iy]
 
     if 'I' in opts.products:
         # get the reconstructed cube
-        Irec_cube = i0map[None, :, :] * \
-            (freqs[:, None, None]/ref_freq)**alphamap[None, :, :]
         name = outfile + '.Irec_cube.fits'
         save_fits(name, np.expand_dims(Irec_cube, axis=4 - stokes_axis), mhdr, dtype=opts.out_dtype)
         print(f"Wrote reconstructed cube to {name}", file=log)
 
+    if 'd' in opts.products:
+        # get the reconstructed cube
+        name = outfile + '.fit_diff.fits'
+        save_fits(name, np.expand_dims(fit_diff, axis=4 - stokes_axis), mhdr, dtype=opts.out_dtype)
+        print(f"Wrote reconstructed cube to {name}", file=log)
+
     # save alpha map
     if 'a' in opts.products:
         name = outfile + '.alpha.fits'