fixes to get auto_selfcal to run cleanly with allow_cocal=True

auto_selfcal/auto_selfcal.py

@@ -8,13 +8,15 @@
 import sys
 import pickle
 import pprint
+import copy
 
 from .selfcal_helpers import *
 from .run_selfcal import run_selfcal
 from .image_analysis_helpers import *
 from .weblog_creation import *
 from .prepare_selfcal import prepare_selfcal, set_clean_thresholds, plan_selfcal_per_solint
 from .original_ms_helpers import applycal_to_orig_MSes, uvcontsub_orig_MSes
+from .prepare_cocal import prepare_cocal
 try:
     from .alignment_helpers import align_measurement_sets
 except:
@@ -630,6 +632,7 @@ def auto_selfcal(
      with open('selfcal_library.pickle', 'wb') as handle:
         pickle.dump(selfcal_library, handle, protocol=pickle.HIGHEST_PROTOCOL)
 
+
     import json
 
     class NpEncoder(json.JSONEncoder):
@@ -709,7 +712,8 @@ def default(self, obj):
     with open('selfcal_library.pickle', 'wb') as handle:
         pickle.dump(selfcal_library, handle, protocol=pickle.HIGHEST_PROTOCOL)
 
-
+    with open('selfcal_plan.pickle', 'wb') as handle:
+        pickle.dump(selfcal_plan, handle, protocol=pickle.HIGHEST_PROTOCOL)
 
     ##
     ## Begin Self-cal loops
@@ -729,109 +733,7 @@ def default(self, obj):
 
 
     if allow_cocal:
-        ##
-        ## Save the flags following the main iteration of self-calibration since we will need to revert to the beginning for the fallback mode.
-        ##
-        # PS: I don't need this anymore?
-        for vis in selfcal_library[target][band]['vislist']:
-           if not os.path.exists(vis+'.flagversions/flags.fb_selfcal_starting_flags'):
-              flagmanager(vis=vis,mode='save',versionname='fb_selfcal_starting_flags')
-           else:
-              flagmanager(vis=vis,mode='restore',versionname='fb_selfcal_starting_flags')
-
-        ##
-        ## For sources that self-calibration failed, try to use the inf_EB and the inf solutions from the sources that
-        ## were successful.
-
-        for target in selfcal_library.keys():
-            for band in selfcal_library[target].keys():
-                print(target, selfcal_library[target][band]["final_solint"])
-
-        inf_EB_fields = {}
-        inf_fields = {}
-        fallback_fields = {}
-        calibrators = {}
-        for band in bands:
-            # Initialize the lists for this band.
-            inf_EB_fields[band] = []
-            inf_fields[band] = []
-            fallback_fields[band] = []
-
-            # Loop through and identify which sources belong where.
-            for target in selfcal_library.keys():
-                if selfcal_library[target][band]['SC_success'] and 'fb' not in selfcal_library[target][band]['final_solint']:
-                    inf_EB_fields[band].append(target)
-                    if selfcal_library[target][band]['final_solint'] != 'inf_EB':
-                        inf_fields[band].append(target)
-                    elif 'inf' in selfcal_plan[target][band]['solints']:
-                        fallback_fields[band].append(target)
-                else:
-                    fallback_fields[band].append(target)
-
-                # Update the relevant lists if we are going to do a fallback mode.
-                if len(fallback_fields[band]) > 0:
-                    selfcal_plan[target][band]['solints'] += ["inf_EB_fb","inf_fb1","inf_fb2","inf_fb3"]
-                    selfcal_plan[target][band]['solmode'] += ["p","p","p","p"]
-                    selfcal_plan[target][band]['gaincal_combine'] += [selfcal_plan[target][band]['gaincal_combine'][0], selfcal_plan[target][band]['gaincal_combine'][1], selfcal_plan[target][band]['gaincal_combine'][1], selfcal_plan[target][band]['gaincal_combine'][1]]
-                    applycal_mode[band][target] += [applycal_mode[band][target][0], applycal_mode[band][target][1], applycal_mode[band][target][1], applycal_mode[band][target][1]]
-                    calibrators[band] = [inf_EB_fields[band], inf_fields[band], inf_fields[band], inf_fields[band]]
-                    selfcal_library[target][band]["nsigma"] = np.concatenate((selfcal_library[target][band]["nsigma"],[selfcal_library[target][band]["nsigma"][0], \
-                            selfcal_library[target][band]["nsigma"][1], selfcal_library[target][band]["nsigma"][1], selfcal_library[target][band]["nsigma"][1]]))
-
-        print(inf_EB_fields)
-        print(inf_fields)
-        print(fallback_fields)
-
-        ##
-        ## Reset the inf_EB informational dictionaries.
-        ##
-
-        for target in selfcal_library:
-         for band in solint_snr[target].keys():
-           # If the target had a successful inf_EB solution, no need to reset.
-           if target in inf_EB_fields[band]:
-               continue
-
-           for vis in selfcal_library[target][band]['vislist']:
-            selfcal_plan[target][band][vis]['inf_EB_gaincal_combine']=inf_EB_gaincal_combine #'scan'
-            if selfcal_library[target][band]['obstype']=='mosaic':
-               selfcal_plan[target][band][vis]['inf_EB_gaincal_combine']+=',field'   
-            selfcal_plan[target][band][vis]['inf_EB_gaintype']=inf_EB_gaintype #G
-            selfcal_plan[target][band][vis]['inf_EB_fallback_mode']='' #'scan'
-
-
-        calculate_inf_EB_fb_anyways = True
-        preapply_targets_own_inf_EB = False
-
-        ## The below sets the calibrations back to what they were prior to starting the fallback mode. It should not be needed
-        ## for the final version of the codue, but is used for testing.
-
-
-        for target in selfcal_library:
-         sani_target=sanitize_string(target)
-         for band in selfcal_library[target].keys():
-           if target not in fallback_fields[band]:
-               continue
-           if 'gaintable_final' in selfcal_library[target][band][vislist[0]]:
-              print('****************Reapplying previous solint solutions*************')
-              for vis in selfcal_library[target][band]['vislist']:
-                 print('****************Applying '+str(selfcal_library[target][band][vis]['gaintable_final'])+' to '+target+' '+band+'*************')
-                 ## NOTE: should this be selfcal_starting_flags instead of fb_selfcal_starting_flags ???
-                 flagmanager(vis=vis,mode='delete',versionname='fb_selfcal_starting_flags_'+sani_target)
-                 applycal(vis=vis,\
-                         gaintable=selfcal_library[target][band][vis]['gaintable_final'],\
-                         interp=selfcal_library[target][band][vis]['applycal_interpolate_final'],\
-                         calwt=True,spwmap=selfcal_library[target][band][vis]['spwmap_final'],\
-                         applymode=selfcal_library[target][band][vis]['applycal_mode_final'],\
-                         field=target,spw=selfcal_library[target][band][vis]['spws'])    
-           else:            
-              print('****************Removing all calibrations for '+target+' '+band+'**************')
-              for vis in selfcal_library[target][band]['vislist']:
-                 flagmanager(vis=vis,mode='delete',versionname='fb_selfcal_starting_flags_'+sani_target)
-                 clearcal(vis=vis,field=target,spw=selfcal_library[target][band][vis]['spws'])
-        ## END
-
-
+        fallback_fields, calibrators = prepare_cocal(selfcal_library, selfcal_plan, inf_EB_gaincal_combine, inf_EB_gaintype)
         ##
         ## Begin fallback self-cal loops
         ##
@@ -845,8 +747,7 @@ def default(self, obj):
                    inf_EB_gaincal_combine=inf_EB_gaincal_combine, inf_EB_gaintype=inf_EB_gaintype, unflag_only_lbants=unflag_only_lbants, \
                    unflag_only_lbants_onlyap=unflag_only_lbants_onlyap, calonly_max_flagged=calonly_max_flagged, \
                    second_iter_solmode=second_iter_solmode, unflag_fb_to_prev_solint=unflag_fb_to_prev_solint, rerank_refants=rerank_refants, \
-                   mode="cocal", calibrators=calibrators, calculate_inf_EB_fb_anyways=calculate_inf_EB_fb_anyways, \
-                   preapply_targets_own_inf_EB=preapply_targets_own_inf_EB, gaincalibrator_dict=gaincalibrator_dict, allow_gain_interpolation=True)
+                   mode="cocal", calibrators=calibrators, gaincalibrator_dict=gaincalibrator_dict, allow_gain_interpolation=True)
 
         if debug:
             print(json.dumps(selfcal_library, indent=4, cls=NpEncoder))