Regenerating primary beams after switching to radio velocity results in the wrong primary beam
The worker assumes the 3rd axis to be frequency not vrad. It should check first. Easy to fix, will do it.
diff --git a/caracal/sample_configurations/mosaic_basic_config.yml b/caracal/sample_configurations/mosaic_basic_config.yml
index d675bb93..6bee3654 100644
--- a/caracal/sample_configurations/mosaic_basic_config.yml
+++ b/caracal/sample_configurations/mosaic_basic_config.yml
@@ -1,4 +1,4 @@
-schema_version: 1.0.3
+schema_version: 1.1.0
 
 getdata:
   dataid: ['']
diff --git a/caracal/schema/flag_schema.yml b/caracal/schema/flag_schema.yml
index 10465845..0755febb 100644
--- a/caracal/schema/flag_schema.yml
+++ b/caracal/schema/flag_schema.yml
@@ -10,7 +10,7 @@ mapping:
         required: True
         example: 'False'
       field:
-        desc: Fields that should be flagged. It can be set to either 'target' or 'calibrators' (i.e., all calibrators) as defined in the obsconf worker. Note that this selection is ignored -- i.e., all fields in the selected .MS file(s) are flagged -- in the flagging step 'flag_mask' (see below). If a user wants to only flag a subset of the calibrators the selection can be further refined using 'calibrator_fields' below. The value of 'field' is also used to compose the name of the .MS file(s) that should be flagged, as explained in 'label_in' below.
+        desc: Fields that should be flagged. It can be set to either 'target' or 'calibrators' (i.e., all calibrators) as defined in the obsconf worker. Note that this selection is ignored -- i.e., all fields in the selected .MS file(s) are flagged -- in the flagging step 'flag_mask' (see below). If a user wants to only flag a subset of the calibrators the selection can be further refined using 'calfields' below. The value of 'field' is also used to compose the name of the .MS file(s) that should be flagged, as explained in 'label_in' below.
         type: str
         enum: ["target", "calibrators"]
         required: False
diff --git a/caracal/schema/mosaic_schema.yml b/caracal/schema/mosaic_schema.yml
index 257a30d0..c8228e2f 100644
--- a/caracal/schema/mosaic_schema.yml
+++ b/caracal/schema/mosaic_schema.yml
@@ -52,17 +52,23 @@ mapping:
             required: false
             example: 'True'
       cutoff:
-         desc: The cutoff in the primary beam. It should be a number between 0 and 1.
-         type: float
-         required: false
-         example: '0.1'
+        desc: The cutoff in the primary beam. It should be a number between 0 and 1.
+        type: float
+        required: false
+        example: '0.1'
+      pb_type:
+        desc: If no continuum pb.fits are already in place, user needs to choose whether a rudimentary primary beam is created ('gaussian') or one that follows the model of Mauch et al. (2020), relevant for MeerKAT data ('mauchian').   
+        type: str
+        enum: ['gaussian', 'mauchian']
+        required: false
+        example: 'gaussian' 
       dish_diameter:
-        desc: If no continuum pb.fits are already in place, user needs to specify the dish diameter (in units of m) so that rudimentary primary beams can be created.
+        desc: If 'pb_type' has been set to 'gaussian', user needs to specify the dish diameter (in units of m).
         type: float
         required: false
         example: '13.5'
       ref_frequency:
-        desc: If no continuum pb.fits are already in place, user needs to specify the reference frequency (in units of Hz) so that rudimentary primary beams can be created.
+        desc: If no continuum pb.fits are already in place, user needs to specify the reference frequency (in units of Hz) so that primary beams can be created.
         type: float
         required: false
         example: '1383685546.875'
diff --git a/caracal/workers/line_worker.py b/caracal/workers/line_worker.py
index 19320688..23c4450f 100644
--- a/caracal/workers/line_worker.py
+++ b/caracal/workers/line_worker.py
@@ -128,6 +128,9 @@ def fix_specsys(filename, specframe):
             headcube['specsys3'] = specsys3
 
 def make_pb_cube(filename, apply_corr, typ, dish_size):
+    C = 2.99792458e+8       # m/s
+    HI = 1.4204057517667e+9  # Hz
+    
     if not os.path.exists(filename):
         caracal.log.warn(
             'Skipping primary beam cube for {0:s}. File does not exist.'.format(filename))
@@ -140,11 +143,26 @@ def make_pb_cube(filename, apply_corr, typ, dish_size):
             datacube[1] -= (headcube['crpix1'] - 1)
             datacube = np.sqrt((datacube**2).sum(axis=0))
             datacube.resize((1, datacube.shape[0], datacube.shape[1]))
+
             datacube = np.repeat(datacube,
                                  headcube['naxis3'],
                                  axis=0) * np.abs(headcube['cdelt1'])
-            freq = (headcube['crval3'] + headcube['cdelt3'] * (
-                np.arange(headcube['naxis3']) - headcube['crpix3'] + 1))
+            
+            cdelt3 = float(headcube['cdelt3'])
+            crval3 = float(headcube['crval3'])
+            
+            # Convert radio velocity to frequency if required
+            if 'VRAD' in headcube['ctype3']:
+                if 'restfreq' in headcube:
+                    restfreq = float(headcube['restfreq'])
+                else:
+                    restfreq = HI
+                cdelt3 = - restfreq*cdelt3/C
+                crval3 = restfreq*(1-crval3/C)
+                
+            freq = (crval3 + cdelt3 * (np.arange(headcube['naxis3']) -
+                                       headcube['crpix3'] + 1))
+            
             if typ == 'gauss':
                sigma_pb = 17.52 / (freq / 1e+9) / dish_size / 2.355
                sigma_pb.resize((sigma_pb.shape[0], 1, 1))
diff --git a/caracal/workers/mosaic_worker.py b/caracal/workers/mosaic_worker.py
index 7f2c60f3..9457142b 100644
--- a/caracal/workers/mosaic_worker.py
+++ b/caracal/workers/mosaic_worker.py
@@ -10,7 +10,7 @@ from astropy.io import fits
 from astropy import wcs
 from caracal.dispatch_crew import utils
 
-NAME = "Mosaic Images and Cubes"
+NAME = "Mosaic 2D-images or cubes"
 LABEL = 'mosaic'
 
 
@@ -65,7 +65,7 @@ def worker(pipeline, recipe, config):
         last_subdirectory = subdirectory_prefix + str(max_num)
         return max_num, last_subdirectory
 
-    # Copied from masking_worker.py and edited
+    # Copied from masking_worker.py and edited. This is to get a Gaussian beam.
     def build_beam(obs_freq, centre, cell, imsize, out_beam):
 
         # if copy_head == True:
@@ -112,6 +112,32 @@ def worker(pipeline, recipe, config):
 
         fits.writeto(out_beam, gaussian, hdr, overwrite=True)
 
+    # Copied from line_worker.py and edited. This is to get a Mauchian beam.
+    # The original version makes the build_beam function above redundant but I do not want to change too many things at once.
+    def make_mauchian_pb(filename, freq): # pbtype):
+        with fits.open(filename) as image:
+            headimage = image[0].header
+            ang_offset = np.indices(
+                (headimage['naxis2'], headimage['naxis1']), dtype=np.float32)
+            ang_offset[0] -= (headimage['crpix2'] - 1)
+            ang_offset[1] -= (headimage['crpix1'] - 1)
+            ang_offset = np.sqrt( (ang_offset**2).sum(axis=0) )  # Using offset in x and y direction to calculate the total offset from the pointing centre
+            ang_offset = ang_offset * np.abs(headimage['cdelt1'])  # Now offset is in units of deg
+            #if pbtype == 'gaussian':
+            #    sigma_pb = 17.52 / (freq / 1e+9) / dish_size / 2.355
+            #    sigma_pb.resize((sigma_pb.shape[0], 1, 1))
+            #    datacube = np.exp(-datacube**2 / 2 / sigma_pb**2)
+            #elif pbtype == 'mauchian':
+            FWHM_pb = (57.5/60) * (freq / 1.5e9)**-1  # Eqn 4 of Mauch et al. (2020), but in deg   # freq is just a float for the 2D case 
+            pb_image = (np.cos(1.189 * np.pi * (ang_offset / FWHM_pb)) / (
+                           1 - 4 * (1.189 * ang_offset / FWHM_pb)**2))**2  # Eqn 3 of Mauch et al. (2020)
+            fits.writeto(filename.replace('image.fits','pb.fits'),
+                pb_image, header=headimage, overwrite=True)
+            caracal.log.info('Created Mauchian primary-beam  FITS {0:s}'.format(
+                filename.replace('image.fits', 'pb.fits')))
+
+
+
     ##########################################
     # Main part of the worker
     ##########################################
@@ -123,13 +149,14 @@ def worker(pipeline, recipe, config):
     specified_images = config['target_images']
     label = config['label_in']
     line_name = config['line_name']
+    pb_type = config['pb_type']
 
     # Parameters that depend on the mosaictype
     if specified_mosaictype == 'spectral':
-        pb_origin = 'generated by the image_line_worker'
+        pb_origin = 'generated by the line_worker'
         input_directory = pipeline.cubes
     else:
-        pb_origin = 'that are already in place (generated by the self_cal_worker, or during a previous run of the mosaic_worker)'
+        pb_origin = 'that are already in place (generated by the selfcal_worker, or during a previous run of the mosaic_worker)'
         input_directory = pipeline.continuum
         if specified_mosaictype == 'continuum':
             pass
@@ -181,7 +208,7 @@ def worker(pipeline, recipe, config):
             field = utils.filter_name(target)
 
             # Use the mosaictype to infer the filenames of the images
-            if specified_mosaictype == 'continuum':  # Add name of 2D image output by selfcal
+            if specified_mosaictype == 'continuum':  # Add name of 2D image output by selfcal_worker
 
                 path_to_image = pipeline.continuum
                 pathnames.append(path_to_image)
@@ -191,7 +218,7 @@ def worker(pipeline, recipe, config):
                 # Note that the path is not included in image_name
                 specified_images.append(image_name)
 
-            else:  # i.e. mosaictype = 'spectral', so add name of cube output by image_line
+            else:  # i.e. mosaictype = 'spectral', so add name of cube output by line_worker
 
                 path_to_cube = pipeline.cubes
                 pathnames.append(path_to_cube)
@@ -199,7 +226,7 @@ def worker(pipeline, recipe, config):
                 image_name = '{0:s}/{1:s}_{2:s}_{3:s}{4:s}-image.fits'.format(
                     last_subdirectory, prefix, field, line_name, mfsprefix)  # max_num not needed for finalcubename
                 if mfsprefix == '':
-                    # Following the naming in image_line_worker
+                    # Following the naming in line_worker
                     image_name = image_name.replace('-image', '.image')
                 # Note that the path is not included in image_name
                 specified_images.append(image_name)
@@ -214,7 +241,7 @@ def worker(pipeline, recipe, config):
     caracal.log.info(pathnames)
 
     # Although montage_mosaic checks whether pb.fits files are present, we need to do this earlier in the worker,
-    # so that we can create simple Gaussian primary beams if need be
+    # so that we can create simple Gaussian (or Mauchian) primary beams if need be
     for image_name in specified_images:
 
         pb_name = image_name.replace('image.fits', 'pb.fits')
@@ -233,42 +260,64 @@ def worker(pipeline, recipe, config):
                 caracal.log.error(
                     '{0:s}/{1:s} does not exist. Please make sure that it is in place before proceeding.'.format(pathname, pb_name))
                 caracal.log.error(
-                    'You may need to re-run the image_line worker with pb_cube enabled. EXITING.')
+                    'You may need to re-run the line_worker with pb_cube enabled. EXITING.')
                 raise caracal.ConfigurationError("missing primary beam file {}/{}".format(pathname, pb_name))
 
-            else:  # i.e. mosaictype = 'continuum'
+            else:  # i.e. mosaictype == 'continuum'
 
                 caracal.log.info(
-                    '{0:s}/{1:s} does not exist, so going to create a rudimentary pb.fits file instead.'.format(pathname, pb_name))
-
-                # Create rudimentary primary-beam, which is assumed to be a Gaussian with FWMH = 1.02*lambda/D
-                image_hdu = fits.open(pathname + '/' + image_name)
-                image_header = image_hdu[0].header
-                # i.e. [ RA, Dec ]. Assuming that these are in units of deg.
-                image_centre = [image_header['CRVAL1'], image_header['CRVAL2']]
-                # Again assuming that these are in units of deg.
-                image_cell = image_header['CDELT2']
-                image_imsize = image_header['NAXIS1']
-
-                recipe.add(build_beam, 'build_pb',
-                           {
-                               # Units of Hz. The default assumes that MeerKAT data is being processed
-                               'obs_freq': config['ref_frequency'],
-                               'centre': image_centre,
-                               'cell': image_cell,
-                               'imsize': image_imsize,
-                               'out_beam': pathname + '/' + pb_name,
-                           },
-                           input=pipeline.input,
-                           # Was pipeline=pipeline.output before the restructure of the output directory
-                           output=pathname,
-                           label='build_pb:: Generating {0:s}'.format(pathname + '/' + pb_name))
+                    '{0:s}/{1:s} does not exist, so going to create a pb.fits file instead.'.format(pathname, pb_name))
+
+                if pb_type == 'gaussian':
+
+                    # Create rudimentary primary-beam, which is assumed to be a Gaussian with FWMH = 1.02*lambda/D
+                    image_hdu = fits.open(pathname + '/' + image_name)
+                    image_header = image_hdu[0].header
+                    # i.e. [ RA, Dec ]. Assuming that these are in units of deg.
+                    image_centre = [image_header['CRVAL1'], image_header['CRVAL2']]
+                    # Again assuming that these are in units of deg.
+                    image_cell = image_header['CDELT2']
+                    image_imsize = image_header['NAXIS1']
+
+                    recipe.add(build_beam, 'build_gaussian_pb',
+                               {
+                                   # Units of Hz. The default assumes that MeerKAT data is being processed
+                                   'obs_freq': config['ref_frequency'],
+                                   'centre': image_centre,
+                                   'cell': image_cell,
+                                   'imsize': image_imsize,
+                                   'out_beam': pathname + '/' + pb_name,
+                               },
+                               input=pipeline.input,
+                               # Was pipeline=pipeline.output before the restructure of the output directory
+                               output=pathname,
+                               label='build_gaussian_pb:: Generating {0:s}'.format(pathname + '/' + pb_name))
+
+                    # Confirming freq and dish_diameter values being used for the primary beam
+                    caracal.log.info('Observing frequency = {0:f} Hz, dish diameter = {1:f} m'.format(
+                        config['ref_frequency'], config['dish_diameter']))
+                    caracal.log.info('If these are not the values that you were expecting to be used for primary-beam creation, then '
+                            'please delete the newly-created beams and re-run the mosaic worker with ref_frequency and dish_diameter '
+                            'set in the config file.')
+
+
+                else:  # i.e. pb_type == 'mauchian'
+
+                    filename = pathname + '/' + image_name
+                    freq = config['ref_frequency'] # Units of Hz. The default assumes that MeerKAT data is being processed
+                    make_mauchian_pb(filename, freq)
+
+                    # Confirming freq value being used for the primary beam
+                    caracal.log.info('Observing frequency = {0:f} Hz'.format(freq))
+                    if freq == 1383685546.875:  # i.e. if the default value was used 
+                        caracal.log.info('If you did not want this value (i.e. the default) to be used for primary-beam creation, then '
+                            'please delete the newly-created beams and re-run the mosaic worker with ref_frequency set in the config file.')
+                    else:
+                        caracal.log.info('as set via ref_frequency in the config file, and used for primary-beam creation.')
+
 
                 pb_origin = 'generated by the mosaic_worker'
 
-                # Confirming freq and dish_diameter values being used for the primary beam
-                caracal.log.info('Observing frequency = {0:f} Hz, dish diameter = {1:f} m'.format(
-                    config['ref_frequency'], config['dish_diameter']))
 
     caracal.log.info('Checking for *pb.fits files now complete.')
 
['caracal/workers/line_worker.py', 'caracal/workers/mosaic_worker.py', 'caracal/sample_configurations/mosaic_basic_config.yml', 'caracal/schema/mosaic_schema.yml', 'caracal/schema/flag_schema.yml']