Coverage for /wheeldirectory/casa-6.7.0-12-py3.10.el8/lib/py/lib/python3.10/site-packages/casatasks/private/task_defintent.py: 84%
126 statements
« prev ^ index » next coverage.py v7.6.4, created at 2024-11-01 07:19 +0000
1import os
2import numpy as np
3import shutil
4import casatools
6def defintent(vis='', intent='', mode='append',
7 outputvis='', scan='', field='',
8 obsid=''):
9 """
10 Description:
11 Allows users to manually set the intents for a selection of scans, fields, or obsids.
13 Keyword arguments:
14 mode: 'set' or 'append'.
15 Set allows the user to fully define a new intent.
16 Append allows the user to add to the current intent for the current selection.
17 Accepts string values, no change if left undefined.
18 intent: new intent to add.
19 User provides the new intent to be set.
20 Accepts string values, no change if left undefined.
21 scan: Select the scans to modify.
22 Defaults to all scans selected.
23 field: Select the fields to modify
24 Defaults to all fields selected.
25 obsid: Select the obsids to modify
26 Defaults to all obsids selected
28 Return: none
29 """
31 tb = casatools.table()
32 ms = casatools.ms()
33 #changeList = []
35 # If no intent has been provided exit the task and print
36 if vis == '':
37 print('You must specify a MS')
38 return
40 if intent == '':
41 print('you must specify an Intent')
42 return
44 # if there is an outputvis make a copy
45 if outputvis != '' and outputvis != vis:
46 if os.path.exists(outputvis):
47 print("outputvis already exists! Exiting task...")
48 return
49 shutil.copytree(vis, outputvis)
50 # if the outputvis is the same as the vis. edit the vis table and don't make a copy
51 elif outputvis == vis:
52 outputvis = vis
53 print("outputvis and vis are the same. Editing provided vis...")
54 else:
55 print("No outputvis has been specified, please enter an outputvis name")
56 return
58 # Table tool query?
59 # ----- TABLE SELECTION -----
61 # Get field names
62 tb.open(vis+'/FIELD')
63 fieldnames = tb.getcol('NAME')
64 tb.close()
66 # Get field ids for names
67 fieldSplit = field.split(',')
68 fieldSplit = [x.strip() for x in fieldSplit]
70 # Link field names to ID
71 nameDict = {}
72 for i in range(len(fieldnames)):
73 if fieldnames[i] not in nameDict:
74 nameDict[fieldnames[i]] = str(i)
75 else:
76 nameDict[fieldnames[i]] += ',' + str(i)
78 # Replace names in selection with IDs
79 for i in range(len(fieldSplit)):
80 if fieldSplit[i] in fieldnames:
81 field = field.replace(fieldSplit[i], nameDict[fieldSplit[i]])
83 # Get exitsing intents
84 tb.open(vis+'/STATE')
85 intentcol = tb.getcol('OBS_MODE')
86 tb.close()
88 # Allowed Intents? / Intents reformatting?
90 # If selected field is found
91 foundField = False
92 if (type(scan) != list):
93 scan = str(scan)
95 #selectedRows = set()
96 selectedRows = []
97 selectedIntents = dict()
99 # NEW get query using ms tool selection
100 ms.open(vis)
101 ms.msselect({'field':field, 'scan':scan, 'observation':obsid}, onlyparse=True)
102 selectedIndex = ms.msselectedindices()
103 ms.close()
105 tb.open(vis)
106 fieldIds = tb.getcol('FIELD_ID')
107 scanNum = tb.getcol('SCAN_NUMBER')
108 stateIds = tb.getcol('STATE_ID')
109 obsIds = tb.getcol('OBSERVATION_ID')
111 '''# Dict to write to the outfile if it exists
112 outfileDict = {}
113 outfileDict["origin_state_ids"] = stateIds
114 paramDict = {'vis':vis, 'intent':intent, 'mode':mode, 'outfile':outfile,
115 'originfile':originfile, 'scan':scan, 'field':field,
116 'obsid':obsid}
117 outfileDict["task_parameters"] = paramDict
118 outfileDict["execution_time"] = date.today().strftime("%B %d, %Y")'''
121 # mstool query version
122 toJoin = []
123 if len(selectedIndex['field']) > 0:
124 toJoin.append(f"FIELD_ID in {selectedIndex['field'].tolist()}")
125 if len(selectedIndex['scan']) > 0:
126 toJoin.append(f"SCAN_NUMBER in {selectedIndex['scan'].tolist()}")
127 if len(selectedIndex['observationid']) > 0:
128 toJoin.append(f"OBSERVATION_ID in {selectedIndex['observationid'].tolist()}")
129 # join into query string
130 taskQuery = " && ".join(toJoin)
132 selectedData = tb.query(taskQuery)
133 #selectedRows = set(selectedData.rownumbers())
134 selectedRows = selectedData.rownumbers()
136 selectedStateIds = selectedData.getcol('STATE_ID')
137 for i in range(len(selectedRows)):
138 selectedIntents[selectedStateIds[i]] = selectedStateIds[i]
140 #tmpString = str(selectedRows[i]) + ':' + str(selectedStateIds[i])
141 #changeList.append(tmpString)
143 tb.close()
145 print("Number of matching rows found: ", len(selectedRows))
147 # for Set if intent not in state table
148 # then add a new row to the state table and change index (STATE_ID) in main table
149 if mode.lower() == 'set':
150 # Keep track of the new value to set the state_id to
151 newState = -1
152 # Adding to intents col
153 statetb = outputvis+'/STATE'
154 tb.open(statetb, nomodify=False)
155 intents = tb.getcol('OBS_MODE')
156 # Check if the intent already exists
157 if intent in intents:
158 print("Intent already exists")
159 newState = np.where(intents == intent)
160 # If it doesn't add a row with the new intent
161 else:
162 # If there are no intents to begin with add blank zero to prevent listobs segfault (?)
163 numIntents = len(list(tb.getcol('OBS_MODE')))
164 if numIntents == 0:
165 tb.addrows(2)
166 else:
167 tb.addrows(1)
168 intents = list(tb.getcol('OBS_MODE'))
169 intents[-1] = intent
170 if numIntents == 0:
171 print("No intents present. filling unselected with UNSPECIFIED(DEFINTENT)")
172 intents[0] = 'UNSPECIFIED(DEFINTENT)'
173 intents = np.asarray(intents)
174 tb.putcol('OBS_MODE', intents)
175 newState = len(intents) - 1
176 tb.close()
178 # For all selected rows replace with new state_id
179 tb.open(outputvis, nomodify=False)
180 stateCol = tb.getcol('STATE_ID')
182 for row in selectedRows:
183 stateCol[row] = newState
184 tb.putcol('STATE_ID', stateCol)
185 tb.close()
187 tb.close()
189 # For Append mode
190 elif mode.lower() == 'append':
191 statetb = outputvis+'/STATE'
192 # Find our selected intents
193 for i in selectedIntents:
194 newState = -1
195 tb.open(statetb, nomodify=False)
196 intents = tb.getcol('OBS_MODE')
197 # Add a row with old intent + new, if it was UNSPECIFIED just do the new intent (?)
198 if intents[i] == 'UNSPECIFIED(DEFINTENT)':
199 newIntent = intent
200 else:
201 newIntent = intents[i] + ',' + intent
202 # Check if thie intent already exists
203 if newIntent in intents:
204 print("Intent already exists")
205 newState = intents.index(newIntent)
206 else:
207 tb.addrows(1)
208 intents = list(tb.getcol('OBS_MODE'))
209 intents[-1] = newIntent
210 intents = np.asarray(intents)
211 tb.putcol('OBS_MODE', intents)
212 newState = len(intents) - 1
213 selectedIntents[i] = newState
214 tb.close()
215 tb.close()
217 # For all selected rows replace with new ID
218 tb.open(outputvis, nomodify=False)
219 stateCol = tb.getcol('STATE_ID')
220 for row in selectedRows:
221 if stateCol[row] in selectedIntents:
222 stateCol[row] = selectedIntents[stateCol[row]]
223 tb.putcol('STATE_ID', stateCol)
224 tb.close()
226 return