Commits

Python 3 compat and imports for Python 2/3 - CASA 5/6, CAS-12282
No tags

gcwrap/python/scripts/partitionhelper.py

Modified
1 1 import os
2 2 import sys
3 3 import shutil
4 4 import pprint as pp
5 5 import traceback
6 6 import time
7 -import commands
8 7 import numpy as np
9 -import matplotlib.pyplot as plt
10 -from __main__ import *
11 -from taskinit import *
8 +from matplotlib import pyplot as plt
9 +
10 +try:
11 + # CASA 6
12 + from casatasks import casalog
13 + from casatools.platform import bytes2str
14 + from casatools import table, ms, msmetadata
15 +
16 + import subprocess
17 + use_old_casa5_commands = False
18 +except ImportError:
19 + # CASA 5
20 + from __main__ import *
21 + from taskinit import *
22 +
23 + import commands
24 + use_old_casa5_commands = True
12 25
13 26
14 27 class convertToMMS():
15 28 def __init__(self,\
16 29 inpdir=None, \
17 30 mmsdir=None, \
18 31 axis='auto', \
19 32 numsubms=4,
20 33 # createmslink=False, \
21 34 cleanup=False):
104 117 sys.exit(2)
105 118
106 119 # Verify later if this is still needed
107 120 time.sleep(10)
108 121
109 122 casalog.origin('convertToMMS')
110 123 casalog.post('--------------- Successfully created MMS -----------------')
111 124
112 125
113 126 # Copy non-MS files to MMS directory
114 - for ff in nonmslist:
115 - bfile = os.path.basename(ff)
127 + for nfile in nonmslist:
128 + bfile = os.path.basename(nfile)
116 129 lfile = os.path.join(self.mmsdir, bfile)
117 130 casalog.post('Copying non-MS file '+bfile)
118 131 # os.symlink(file, lfile)
119 -# shutil.copytree(ff, lfile, symlinks=False)
120 - os.system("cp -RL "+ff+" "+lfile)
132 +# shutil.copytree(nfile, lfile, symlinks=False)
133 + os.system("cp -RL {0} {1}".format(nfile, lfile))
121 134
122 135
123 136 def getMSlist(self, files):
124 137 '''Get a list of MSs from a directory.
125 138 files -> a tuple that is returned by the following call:
126 139 files = os.walk(self.inpdir,followlinks=True).next()
127 140
128 141 It will test if a directory is an MS and will only return
129 142 true MSs, that have Type:Measurement Set in table.info. It will skip
130 143 directories that start with . and those that do not end with
137 150 # Loop through list of directories
138 151 for d in files[1]:
139 152 # Skip . entries
140 153 if d.startswith('.'):
141 154 continue
142 155
143 156 # if not d.endswith('.ms'):
144 157 # continue
145 158
146 159 # Full path for directory
147 - dir = os.path.join(topdir,d)
160 + mydir = os.path.join(topdir,d)
148 161
149 162 # It is probably an MS
150 - if self.isItMS(dir) == 1:
151 - mslist.append(dir)
163 + if self.isItMS(mydir) == 1:
164 + mslist.append(mydir)
152 165
153 166 return mslist
154 167
155 168 def isItMS(self, mydir):
156 169 '''Check the type of a directory.
157 170 mydir --> full path of a directory.
158 171 Returns 1 for an MS, 2 for a cal table and 3 for a MMS.
159 172 If 0 is returned, it means any other type or an error.'''
160 173
161 174 ret = 0
162 175
163 176 # Listing of this directory
164 177 ldir = os.listdir(mydir)
165 178
166 179 if not ldir.__contains__('table.info'):
167 180 return ret
168 181
169 182 cmd1 = 'grep Type '+mydir+'/table.info'
170 - mytype = commands.getoutput(cmd1)
171 183 cmd2 = 'grep SubType '+mydir+'/table.info'
172 - stype = commands.getoutput(cmd2)
173 -
184 + if use_old_casa5_commands:
185 + mytype = commands.getoutput(cmd1)
186 + stype = commands.getoutput(cmd2)
187 + else:
188 + mytype = bytes2str(subprocess.check_output(cmd1)).rstrip("\n")
189 + stype = bytes2str(subprocess.check_output(cmd2)).rstrip("\n")
190 +
174 191 # It is a cal table
175 192 if mytype.__contains__('Calibration'):
176 193 ret = 2
177 194
178 195 elif mytype.__contains__('Measurement'):
179 196 # It is a Multi-MS
180 197 if stype.__contains__('CONCATENATED'):
181 198 # Further check
182 199 if ldir.__contains__('SUBMSS'):
183 200 ret = 3
233 250
234 251
235 252 def runPartition(self, ms, mmsdir, axis, subms):
236 253 '''Run partition with default values to create an MMS.
237 254 ms --> full pathname of the MS
238 255 mmsdir --> directory to save the MMS to
239 256 axis --> separationaxis to use (spw, scan, auto)
240 257 subms --> number of subMss to create
241 258
242 259 '''
243 - from tasks import partition
244 - from __main__ import default
260 + try:
261 + # CASA 6
262 + from casatasks import partition
263 + except ImportError:
264 + # CASA 5
265 + from tasks import partition
245 266
246 267 if not os.path.lexists(ms):
247 268 return False
248 269
249 270 # Create MMS name
250 271 # bname = os.path.basename(ms)
251 272 # if bname.endswith('.ms'):
252 273 # mmsname = bname.replace('.ms','.mms')
253 274 # else:
254 275 # mmsname = bname+'.mms'
268 289 # return False
269 290
270 291 # Check for remainings of corrupted mms
271 292 # corrupted = mms.replace('.mms','.data')
272 293 corrupted = MMSFullName + '.data'
273 294 if os.path.exists(corrupted):
274 295 casalog.post('Cleaning up left overs','WARN')
275 296 shutil.rmtree(corrupted)
276 297
277 298 # Run partition
278 - default('partition')
279 299 partition(vis=ms, outputvis=MMSFullName, createmms=True, datacolumn='all', flagbackup=False,
280 300 separationaxis=axis, numsubms=subms)
281 301 casalog.origin('convertToMMS')
282 302
283 303 # Check if MMS was created
284 304 if not os.path.exists(MMSFullName):
285 305 casalog.post('Cannot create MMS ->'+MMSFullName, 'ERROR')
286 306 return False
287 307
288 308 # If requested, create a link to this MMS with the original MS name
292 312 # mmsname = os.path.basename(mms)
293 313 ## lms = mmsname.replace('.mms', '.ms')
294 314 # casalog.post('Creating symbolic link to MMS')
295 315 ## os.symlink(mmsname, lms)
296 316 # os.symlink(mmsname, bname)
297 317 # os.chdir(here)
298 318
299 319 return True
300 320
301 321 def usage(self):
302 - print '========================================================================='
303 - print ' convertToMMS will create a directory with multi-MSs.'
304 - print 'Usage:\n'
305 - print ' import partitionhelper as ph'
306 - print ' ph.convertToMMS(inpdir=\'dir\') \n'
307 - print 'Options:'
308 - print ' inpdir <dir> directory with input MS.'
309 - print ' mmsdir <dir> directory to save output MMS. If not given, it will save '
310 - print ' the MMS in a directory called mmsdir in the current directory.'
311 - print " axis='auto' separationaxis parameter of partition (spw,scan,auto)."
312 - print " numsubms=4 number of subMSs to create in output MMS"
313 - print ' cleanup=False if True it will remove the output directory before starting.\n'
322 + print('=========================================================================')
323 + print(' convertToMMS will create a directory with multi-MSs.')
324 + print('Usage:\n')
325 + print(' import partitionhelper as ph')
326 + print(' ph.convertToMMS(inpdir=\'dir\') \n')
327 + print('Options:')
328 + print(' inpdir <dir> directory with input MS.')
329 + print(' mmsdir <dir> directory to save output MMS. If not given, it will save ')
330 + print(' the MMS in a directory called mmsdir in the current directory.')
331 + print(" axis='auto' separationaxis parameter of partition (spw,scan,auto).")
332 + print(" numsubms=4 number of subMSs to create in output MMS")
333 + print(' cleanup=False if True it will remove the output directory before starting.\n')
314 334
315 - print ' NOTE: this script will run using the default values of partition. It will try to '
316 - print ' create an MMS for every MS in the input directory. It will skip non-MS directories '
317 - print ' such as cal tables. If partition succeeds, the script will create a link to every '
318 - print ' other directory or file in the output directory. '
319 - print ' The script will not walk through sub-directories of inpdir. It will also skip '
320 - print ' files or directories that start with a .'
321 - print '=========================================================================='
335 + print(' NOTE: this script will run using the default values of partition. It will try to ')
336 + print(' create an MMS for every MS in the input directory. It will skip non-MS directories ')
337 + print(' such as cal tables. If partition succeeds, the script will create a link to every ')
338 + print(' other directory or file in the output directory. ')
339 + print(' The script will not walk through sub-directories of inpdir. It will also skip ')
340 + print(' files or directories that start with a .')
341 + print('==========================================================================')
322 342 return
323 343
324 344 #
325 345 # -------------- HELPER functions for dealing with an MMS --------------
326 346 #
327 347 # getMMSScans 'Get the list of scans of an MMS dictionary'
328 348 # getScanList 'Get the list of scans of an MS or MMS'
329 349 # getScanNrows 'Get the number of rows of a scan in a MS. It will add the
330 350 # nrows of all sub-scans.'
331 351 # getMMSScanNrows 'Get the number of rows of a scan in an MMS dictionary.'
336 356
337 357 # def getNumberOf(msfile, item='row'):
338 358 # '''Using the msmd tool, it gets the number of
339 359 # scan, spw, antenna, baseline, field, state,
340 360 # channel, row in a MS or MMS'''
341 361 #
342 362 # md = msmdtool()
343 363 # try:
344 364 # md.open(msfile)
345 365 # except:
346 -# print 'Cannot open the msfile'
366 +# print('Cannot open the msfile')
347 367 # return 0
348 368 #
349 369 # if item == 'row':
350 370 # numof = md.nrows()
351 371 # elif item == 'scan':
352 372 # numof = md.nscans()
353 373 # elif item == 'spw':
354 374 # numof = md.nspw()
355 375 # elif item == 'antenna':
356 376 # numof = md.nantennas()
371 391
372 392 # NOTE
373 393 # There is a bug in ms.getscansummary() that does not give the scans for all
374 394 # observation Ids, but only for the last one. See CAS-4409
375 395 def getMMSScans(mmsdict):
376 396 '''Get the list of scans of an MMS dictionary.
377 397 mmsdict --> output dictionary from listpartition(MMS,createdict=true)
378 398 Return a list of the scans in this MMS. '''
379 399
380 400 if not isinstance(mmsdict, dict):
381 - print 'ERROR: Input is not a dictionary'
401 + print('ERROR: Input is not a dictionary')
382 402 return []
383 403
384 404 tkeys = mmsdict.keys()
385 405 scanlist = []
386 406 slist = set(scanlist)
387 407 for k in tkeys:
388 408 skeys = mmsdict[k]['scanId'].keys()
389 409 for j in skeys:
390 410 slist.add(j)
391 411
431 451 '''
432 452 msTool=mstool()
433 453 msTool.open(msfile)
434 454 if isinstance(selection, dict) and selection != {}:
435 455 msTool.msselect(items=selection)
436 456
437 457 scand = msTool.getscansummary()
438 458 msTool.close()
439 459
440 460 Nrows = 0
441 - if not scand.has_key(str(myscan)):
461 + if not str(myscan) in scand:
442 462 return Nrows
443 463
444 464 subscans = scand[str(myscan)]
445 465 for ii in subscans.keys():
446 466 Nrows += scand[str(myscan)][ii]['nRow']
447 467
448 468 return Nrows
449 469
450 470
451 471 def getMMSScanNrows(thisdict, myscan):
452 472 '''Get the number of rows of a scan in an MMS dictionary.
453 473 thisdict --> output dictionary from listpartition(MMS,createdict=true)
454 474 myscan --> scan ID (int)
455 475 Return the number of rows in the given scan. '''
456 476
457 477 if not isinstance(thisdict, dict):
458 - print 'ERROR: Input is not a dictionary'
478 + print('ERROR: Input is not a dictionary')
459 479 return -1
460 480
461 481 tkeys = thisdict.keys()
462 482 scanrows = 0
463 483 for k in tkeys:
464 - if thisdict[k]['scanId'].has_key(myscan):
484 + if myscan in thisdict[k]['scanId']:
465 485 scanrows += thisdict[k]['scanId'][myscan]['nrows']
466 486
467 487 return scanrows
468 488
469 489
470 490 def getSpwIds(msfile, myscan, selection={}):
471 491 '''Get the Spw IDs of a scan.
472 492 msfile --> name of the MS or MMS
473 493 myscan --> scan Id (int)
474 494 selection --> dictionary with data selection
481 501 msTool=mstool()
482 502 msTool.open(msfile)
483 503 if isinstance(selection, dict) and selection != {}:
484 504 msTool.msselect(items=selection)
485 505
486 506 scand = msTool.getscansummary()
487 507 msTool.close()
488 508
489 509 spwlist = []
490 510
491 - if not scand.has_key(str(myscan)):
511 + if not str(myscan) in scand:
492 512 return spwlist
493 513
494 514 subscans = scand[str(myscan)]
495 515 aspws = np.array([],dtype=int)
496 516
497 517 for ii in subscans.keys():
498 518 sscanid = ii
499 519 spwids = scand[str(myscan)][sscanid]['SpwIds']
500 520 aspws = np.append(aspws,spwids)
501 521
549 569 for subms in mslist:
550 570 try:
551 571 mslocal1.open(subms)
552 572 scans = mslocal1.getscansummary()
553 573 msscanlist.append(scans)
554 574 spws = mslocal1.getspectralwindowinfo()
555 575 msspwlist.append(spws)
556 576 mslocal1.close()
557 577 except:
558 578 mslocal1.close()
559 - raise Exception, 'Cannot get scan/spw information from subMS'
579 + raise Exception('Cannot get scan/spw information from subMS')
560 580
561 581 # Get the data volume in bytes per sub-MS
562 582 sizelist.append(getDiskUsage(subms))
563 583
564 584 # Get the information to list in output
565 585 # Dictionary to return
566 586 outdict = {}
567 587
568 588 for ims in range(mslist.__len__()):
569 589 # Create temp dictionary for each sub-MS
629 649
630 650
631 651 def getMMSSpwIds(thisdict):
632 652 '''Get the list of spws from an MMS dictionary.
633 653 thisdict --> output dictionary from listpartition(MMS,createdict=true)
634 654 Return a list of the spw Ids in the dictionary. '''
635 655
636 656 import numpy as np
637 657
638 658 if not isinstance(thisdict, dict):
639 - print 'ERROR: Input is not a dictionary'
659 + print('ERROR: Input is not a dictionary')
640 660 return []
641 661
642 662 tkeys = thisdict.keys()
643 663
644 664 aspws = np.array([],dtype='int32')
645 665 for k in tkeys:
646 666 scanlist = thisdict[k]['scanId'].keys()
647 667 for s in scanlist:
648 668 spwids = thisdict[k]['scanId'][s]['spwIds']
649 669 aspws = np.append(aspws, spwids)
685 705
686 706 Keyword arguments:
687 707 msfile --> name of the MS
688 708 This function will return a value given by the command du -hs
689 709 """
690 710
691 711 from subprocess import Popen, PIPE, STDOUT
692 712
693 713 # Command line to run
694 714 ducmd = 'du -hs '+msfile
695 -
696 - p = Popen(ducmd, shell=True, stdin=PIPE, stdout=PIPE, stderr=STDOUT, close_fds=True)
697 -
698 - sizeline = p.stdout.read()
699 -
715 +
716 + if use_old_casa5_commands:
717 + p = Popen(ducmd, shell=True, stdin=PIPE, stdout=PIPE, stderr=STDOUT, close_fds=True)
718 + sizeline = p.stdout.read()
719 + _out, _err = p.communicate()
720 + else:
721 + p = Popen(ducmd, shell=True, stdin=None, stdout=PIPE, stderr=STDOUT, close_fds=True)
722 + o, e = p.communicate() ### previously 'sizeline = p.stdout.read()' here
723 + ### left process running...
724 + sizeline = bytes2str(o.split( )[0])
725 +
700 726 # Create a list of the output string, which looks like this:
701 727 # ' 75M\tuidScan23.data/uidScan23.0000.ms\n'
702 728 # This will create a list with [size,sub-ms]
703 729 mssize = sizeline.split()
704 730
705 731 return mssize[0]
706 732
707 733
708 734 def getSubtables(vis):
709 735 tbTool = tbtool()
729 755 copysubtables -- True will copy the sub-tables from the first subMS to the others in the
730 756 output MMS. Default to False.
731 757 omitsubtables -- List of sub-tables to omit when copying to output MMS. They will be linked instead
732 758 parallelasxis -- Optionally, set the value to be written to AxisType in table.info of the output MMS
733 759 Usually this value comes from the separationaxis keyword of partition or mstransform.
734 760
735 761 Be AWARE that this function will remove the tables listed in submslist.
736 762 """
737 763
738 764 if os.path.exists(outputvis):
739 - raise ValueError, "Output MS already exists"
765 + raise ValueError('Output MS already exists')
740 766
741 767 if len(submslist)==0:
742 - raise ValueError, "No SubMSs given"
768 + raise ValueError('No SubMSs given')
743 769
744 770 ## make an MMS with all sub-MSs contained in a SUBMSS subdirectory
745 771 origpath = os.getcwd()
746 772 mymstool = mstool()
747 773 mytbtool = tbtool()
748 774
749 775 try:
750 776 try:
751 777 mymstool.createmultims(outputvis,
752 778 submslist,
801 827 os.system('rm -rf '+s)
802 828 os.symlink('../'+mastersubms+'/'+s, s)
803 829
804 830 # Write the AxisType info in the MMS
805 831 if parallelaxis != '':
806 832 setAxisType(outputvis, parallelaxis)
807 833
808 834 except:
809 835 theproblem = str(sys.exc_info())
810 836 os.chdir(origpath)
811 - raise ValueError, "Problem in MMS creation: "+theproblem
837 + raise ValueError('Problem in MMS creation: {0}'.format(theproblem))
812 838
813 839 os.chdir(origpath)
814 840
815 841 return True
816 842
817 843 def axisType(mmsname):
818 844 """Get the axisType information from a Multi-MS. The AxisType information
819 845 is usually added for Multi-MS with the axis which data is parallelized across.
820 846
821 847 Keyword arguments:
822 848 mmsname -- name of the Multi-MS
823 849
824 850 It returns the value of AxisType or an empty string if it doesn't exist.
825 851 """
826 852 tblocal = tbtool()
827 853
828 854 axis = ''
829 855
830 856 try:
831 857 tblocal.open(mmsname, nomodify=True)
832 858 except:
833 - raise ValueError, "Unable to open table %s" % mmsname
859 + raise ValueError('Unable to open table {0}'.format(mmsname))
834 860
835 861 tbinfo = tblocal.info()
836 862 tblocal.close()
837 863
838 - if tbinfo.has_key('readme'):
864 + if 'readme' in tbinfo:
839 865 readme = tbinfo['readme']
840 866 readlist = readme.splitlines()
841 867 for val in readlist:
842 868 if val.__contains__('AxisType'):
843 869 a,b,axis = val.partition('=')
844 870
845 871
846 872 return axis.strip()
847 873
848 874 def setAxisType(mmsname, axis=''):
849 875 """Set the AxisType keyword in a Multi-MS info. If AxisType already
850 876 exists, it will be overwritten.
851 877
852 878 Keyword arguments:
853 879 mmsname -- name of the Multi-MS
854 880 axis -- parallel axis of the Multi-MS. Options: scan; spw or scan,spw
855 881
856 882 Return True on success, False otherwise.
857 883 """
858 884
859 885 if axis == '':
860 - raise ValueError, "Axis value cannot be empty"
886 + raise ValueError('Axis value cannot be empty')
861 887
862 888 tblocal = tbtool()
863 889 try:
864 890 tblocal.open(mmsname, nomodify=False)
865 891 except:
866 - raise ValueError, "Unable to open table %s" % mmsname
892 + raise ValueError('Unable to open table {0}'.format(mmsname))
867 893
868 894 import copy
869 895
870 896 tbinfo = tblocal.info()
871 897 readme = ''
872 898 # Save original readme
873 - if tbinfo.has_key('readme'):
899 + if 'readme' in tbinfo:
874 900 readme = tbinfo['readme']
875 901
876 902 # Check if AxisType already exist and remove it
877 903 if axisType(mmsname) != '':
878 - print 'WARN: Will overwrite the existing AxisType value'
904 + print('WARN: Will overwrite the existing AxisType value')
879 905 readlist = readme.splitlines()
880 906 newlist = copy.deepcopy(readlist)
881 907 for val in newlist:
882 908 if val.__contains__('AxisType'):
883 909 readlist.remove(val)
884 910
885 911 # Recreate the string
886 912 nr = ''
887 913 for val in readlist:
888 914 nr = nr + val + '\n'
933 959 # Initialize scan sub-map
934 960 scanDdiMap[scan] = {}
935 961 # Iterate over timestamps for this scan
936 962 for timestamp in scanSummary[scan]:
937 963 # Get list of ddis for this timestamp
938 964 DDIds = scanSummary[scan][timestamp]['DDIds']
939 965 fieldId = str(scanSummary[scan][timestamp]['FieldId'])
940 966 # Get number of rows per ddi (assume all DDIs have the same number of rows)
941 967 # In ALMA data WVR DDI has only one row per antenna but it is separated from the other DDIs
942 968 nrowsPerDDI = scanSummary[scan][timestamp]['nRow'] / len(DDIds)
943 -))
944 969 # Iterate over DDIs for this timestamp
945 970 for ddi in DDIds:
946 971 # Convert to string to be used as a map key
947 972 ddi = str(ddi)
948 973 # Check if DDI entry is already present for this scan, otherwise initialize it
949 - if not scanDdiMap[scan].has_key(ddi):
974 + if ddi not in scanDdiMap[scan]:
950 975 scanDdiMap[scan][ddi] = {}
951 976 scanDdiMap[scan][ddi]['nVis'] = 0
952 977 scanDdiMap[scan][ddi]['fieldId'] = fieldId
953 978 scanDdiMap[scan][ddi]['isWVR'] = ddIspectralWindowInfo[ddi]['isWVR']
954 979 # Calculate number of visibilities
955 980 nvis = nrowsPerDDI*ddIspectralWindowInfo[ddi]['NumChan']*ddIspectralWindowInfo[ddi]['NumCorr']
956 981 # Add number of rows and vis from this timestamp
957 982 scanDdiMap[scan][ddi]['nVis'] = scanDdiMap[scan][ddi]['nVis'] + nvis
958 983 # Update ddi nvis
959 - if not nVisPerDDI.has_key(ddi):
984 + if ddi not in nVisPerDDI:
960 985 nVisPerDDI[ddi] = nvis
961 986 else:
962 987 nVisPerDDI[ddi] = nVisPerDDI[ddi] + nvis
963 988 # Update scan nvis
964 - if not nVisPerScan.has_key(scan):
989 + if scan not in nVisPerScan:
965 990 nVisPerScan[scan] = nvis
966 991 else:
967 992 nVisPerScan[scan] = nVisPerScan[scan] + nvis
968 993 # Update field nvis
969 - if not nVisPerField.has_key(fieldId):
994 + if fieldId not in nVisPerField:
970 995 nVisPerField[fieldId] = nvis
971 996 else:
972 997 nVisPerField[fieldId] = nVisPerField[fieldId] + nvis
973 998
974 999 return scanDdiMap, nVisPerDDI, nVisPerScan, nVisPerField
975 1000
976 1001 def getPartitionMap(msfilename, nsubms, selection={}, axis=['field','spw','scan'],plotMode=0):
977 1002 """Generates a partition scan/spw map to obtain optimal load balancing with the following criteria:
978 1003
979 1004 1st - Maximize the scan/spw/field distribution across sub-MSs
1052 1077 if nsubms > nScanDDIPairs:
1053 1078 casalog.post("Number of subMSs (%i) is greater than available scan,ddi pairs (%i), setting nsubms to %i"
1054 1079 % (nsubms,nScanDDIPairs,nScanDDIPairs),"WARN","getPartitionMap")
1055 1080 nsubms = nScanDDIPairs
1056 1081
1057 1082 ddiArray = np.array(ddiList)
1058 1083 scanArray = np.array(scanList)
1059 1084 nVisArray = np.array(nVisList)
1060 1085
1061 1086 nVisSortIndex = np.lexsort((ddiArray, scanArray, nVisArray))
1087 + # argsort/lexsort return indices by increasing value. This reverses the indices by
1088 + # decreasing value
1062 1089 nVisSortIndex[:] = nVisSortIndex[::-1]
1063 1090
1064 1091 ddiArray = ddiArray[nVisSortIndex]
1065 1092 scanArray = scanArray[nVisSortIndex]
1066 1093 nVisArray = nVisArray[nVisSortIndex]
1067 1094
1068 1095 # Make a map for the contribution of each subMS to each scan
1069 1096 scanNvisDistributionPerSubMs = {}
1070 1097 for scan in scanSummary:
1071 1098 scanNvisDistributionPerSubMs[scan] = np.zeros(nsubms)
1104 1131
1105 1132 ddi = ddiArray[pair]
1106 1133 scan = scanArray[pair]
1107 1134 field = scanDdiMap[scan][ddi]['fieldId']
1108 1135
1109 1136 # Select the subMS that with bigger (scan/ddi/field gap)
1110 1137 # We use the average as a refLevel to include global structure information
1111 1138 # But we also take into account the actual max value in case we are distributing large uneven chunks
1112 1139 jointNvisGap = np.zeros(nsubms)
1113 1140 if 'scan' in axis:
1114 - refLevel = max(nVisPerScan[scan]/nsubms,scanNvisDistributionPerSubMs[scan].max())
1141 + refLevel = max(nVisPerScan[scan] //
1142 + nsubms,scanNvisDistributionPerSubMs[scan].max())
1115 1143 jointNvisGap = jointNvisGap + refLevel - scanNvisDistributionPerSubMs[scan]
1116 1144 if 'spw' in axis:
1117 - refLevel = max(nVisPerDDI[ddi]/nsubms,ddiNvisDistributionPerSubMs[ddi].max())
1145 + refLevel = max(nVisPerDDI[ddi] //
1146 + nsubms,ddiNvisDistributionPerSubMs[ddi].max())
1118 1147 jointNvisGap = jointNvisGap + refLevel - ddiNvisDistributionPerSubMs[ddi]
1119 1148 if 'field' in axis:
1120 - refLevel = max(nVisPerField[field]/nsubms,fieldNvisDistributionPerSubMs[field].max())
1121 - jointNvisGap = jointNvisGap + refLevel - fieldNvisDistributionPerSubMs[field]
1149 + refLevel = max(nVisPerField[field] //
1150 + nsubms,fieldNvisDistributionPerSubMs[field].max())
1151 + jointNvisGap = jointNvisGap + refLevel - fieldNvisDistributionPerSubMs[field]
1122 1152
1123 1153 optimalSubMs = np.where(jointNvisGap == jointNvisGap.max())
1124 1154 optimalSubMs = optimalSubMs[0] # np.where returns a tuple
1125 1155
1126 1156 # In case of multiple candidates select the subms with minum number of total visibilities
1127 1157 if len(optimalSubMs) > 1:
1128 1158 subIdx = np.argmin(nvisPerSubMs[optimalSubMs])
1129 1159 optimalSubMs = optimalSubMs[subIdx]
1130 1160 else:
1131 1161 optimalSubMs = optimalSubMs[0]

Everything looks good. We'll let you know here if there's anything you should know about.

Add shortcut