Commits
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] |