Coverage for /wheeldirectory/casa-6.7.0-12-py3.10.el8/lib/py/lib/python3.10/site-packages/casatasks/private/task_sdfit.py: 88%

1import datetime 

2import os 

3import time 

4 

5from casatasks import casalog 

6from casatools import ms as mstool 

7from casatools import singledishms, table 

8 

9from . import casaxmlutil, sdutil 

10 

11ms = mstool() 

12sdms = singledishms() 

13tb = table() 

14 

15@casaxmlutil.xml_constraints_injector 

16@sdutil.sdtask_decorator 

17def sdfit(infile=None, datacolumn=None, antenna=None, field=None, spw=None, 

18 timerange=None, scan=None, pol=None, intent=None, 

19 timebin=None, timespan=None, 

20 polaverage=None, 

21 fitfunc=None, fitmode=None, nfit=None, thresh=None, avg_limit=None, 

22 minwidth=None, edge=None, outfile=None, overwrite=None): 

23 

24 casalog.origin('sdfit') 

25 try: 

26 if os.path.exists(outfile): 

27 if overwrite: 

28 os.system('rm -rf %s' % outfile) 

29 else: 

30 raise ValueError(outfile + ' exists.') 

31 if (fitmode not in ['list', 'auto']): 

32 raise ValueError("fitmode='%s' is not supported yet" % fitmode) 

33 if (spw == ''): 

34 spw = '*' 

35 

36 selection = ms.msseltoindex(vis=infile, spw=spw, field=field, 

37 baseline=antenna, time=timerange, 

38 scan=scan) 

39 

40 sdms.open(infile) 

41 sdms.set_selection(spw=sdutil.get_spwids(selection), field=field, 

42 antenna=antenna, timerange=timerange, 

43 scan=scan, polarization=pol, intent=intent) 

44 

45 tempfile = 'temp_sdfit_' + \ 

46 str(datetime.datetime.fromtimestamp(time.time())).replace( 

47 '-', '').replace(' ', '').replace(':', '') 

48 if os.path.exists(tempfile): 

49 tempfile += str(datetime.datetime.fromtimestamp(time.time()) 

50 ).replace('-', '').replace(' ', '').replace(':', '') 

51 if os.path.exists(tempfile): 

52 raise Exception('temporary file ' + tempfile + ' exists...') 

53 tempoutfile = tempfile + '_temp_output_ms' 

54 if os.path.exists(tempoutfile): 

55 tempoutfile += str(datetime.datetime.fromtimestamp(time.time()) 

56 ).replace('-', '').replace(' ', '').replace(':', '') 

57 if os.path.exists(tempoutfile): 

58 raise Exception('temporary ms file ' + tempoutfile + ' exists...') 

59 

60 if fitmode == 'auto': 

61 nfit = [-1] 

62 num_fit_str = str(',').join(map(str, nfit)) 

63 

64 sdms.fit_line(datacolumn=datacolumn, spw=spw, pol=pol, 

65 timebin=timebin, timespan=timespan, 

66 polaverage=polaverage, 

67 fitfunc=fitfunc, nfit=num_fit_str, 

68 linefinding=(fitmode == 'auto'), threshold=thresh, 

69 avg_limit=avg_limit, minwidth=minwidth, edge=edge, 

70 tempfile=tempfile, tempoutfile=tempoutfile) 

71 

72 if os.path.exists(tempfile): 

73 return _get_results(tempfile, fitfunc, nfit, outfile) 

74 else: 

75 raise Exception('temporary file was unexpectedly not created.') 

76 

77 finally: 

78 if 'tempfile' in locals() and os.path.exists(tempfile): 

79 os.system('rm -f %s' % tempfile) 

80 if 'tempoutfile' in locals() and os.path.exists(tempoutfile): 

81 os.system('rm -rf %s' % tempoutfile) 

82 

83 

84def _get_results(tempfile, fitfunc, nfit, outfile): 

85 res = {'cent': [], 'peak': [], 'fwhm': [], 'nfit': []} 

86 if (fitfunc == 'gaussian'): 

87 func = 'gauss' 

88 elif (fitfunc == 'lorentzian'): 

89 func = 'lorentz' 

90 ncomp = sum(nfit) 

91 iline = 0 

92 with open(tempfile, 'r') as f: 

93 outfile_exists = (outfile != '') 

94 if outfile_exists: 

95 fout = open(outfile, 'a') 

96 s = '#SCAN\tTIME\t\tANT\tBEAM\tSPW\tPOL\tFunction\tP0\t\tP1\t\tP2\n' 

97 fout.write(s) 

98 

99 for line in f: 

100 component = line.strip().split(':') # split into each component 

101 if (ncomp > 0): 

102 assert(len(component) == ncomp) 

103 res['cent'].append([]) 

104 res['peak'].append([]) 

105 res['fwhm'].append([]) 

106 res['nfit'].append(ncomp if ncomp >= 0 else len(component)) 

107 for icomp in range(res['nfit'][-1]): 

108 fit_result = component[icomp].split(',') # split into each parameter 

109 num_ids = 6 # scan, time, ant, beam, spw, pol 

110 assert(len(fit_result) == 2 * (len(res.keys()) - 1) + num_ids) 

111 res['cent'][iline].append([float(fit_result[num_ids + 0]), 

112 float(fit_result[num_ids + 1])]) 

113 res['peak'][iline].append([float(fit_result[num_ids + 2]), 

114 float(fit_result[num_ids + 3])]) 

115 res['fwhm'][iline].append([float(fit_result[num_ids + 4]), 

116 float(fit_result[num_ids + 5])]) 

117 

118 if outfile_exists: 

119 s = fit_result[0] + '\t' # scanID 

120 s += fit_result[1] + '\t' # time 

121 s += fit_result[2] + '\t' # antennaID 

122 s += fit_result[3] + '\t' # beamID 

123 s += fit_result[4] + '\t' # spwID 

124 s += fit_result[5] + '\t' # polID 

125 s += func + str(icomp) + '\t\t' # function 

126 s += fit_result[8] + '\t' # P0 (peak) 

127 s += fit_result[6] + '\t' # P1 (cent) 

128 s += fit_result[10] + '\t' # P2 (fwhm) 

129 s += '\n' 

130 fout.write(s) 

131 iline += 1 

132 

133 if outfile_exists: 

134 fout.close() 

135 

136 return res