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

1import os 

2import shutil 

3import time 

4import copy 

5import numpy as np 

6import functools 

7 

8from casatools import image as _image, table as _table 

9from casatools import quanta, ms 

10from casatasks import casalog 

11from casatools import synthesisutils as su 

12from .imager_base import PySynthesisImager 

13from .input_parameters import ImagerParameters 

14from typing import Tuple, List, Union, Optional 

15import pdb 

16_ia = _image() 

17_tb = _table() 

18_qa = quanta() 

19_ms = ms() 

20_su = su() 

21 

22SW=True 

23 

24############################################# 

25def time_func(func): 

26 @functools.wraps(func) 

27 def wrap_time(*args, **kwargs): 

28 t0 = time.time() 

29 result = func(*args, **kwargs) 

30 t1 = time.time() 

31 ####print(f'#######Function {func.__name__!r} took {(t1-t0)}s') 

32 return result 

33 return wrap_time 

34 

35class PyMtmfsViaCubeSynthesisImager(PySynthesisImager): 

36 """A subclass of PySynthesisImager, for specmode='mvc' 

37 

38 The idea is to do the major cycle with cube imaging, then convert the cube images 

39 to taylor term ".ttN" images, then do the minor cycle, then convert back to cubes. 

40 """ 

41 @time_func 

42 def __init__(self, params: ImagerParameters) -> None: 

43 # Set up the mfs part for deconv 

44 mfsparams = copy.deepcopy(params) 

45 mfsparams.allimpars["0"]["specmode"] = "mfs" 

46 mfsparams.alldecpars["0"]["specmode"] = "mfs" 

47 # if there is a startmodel make the model 

48 if ( 

49 len(mfsparams.alldecpars["0"]["startmodel"]) 

50 == mfsparams.alldecpars["0"]["nterms"] 

51 ): 

52 self.copy_startmodel( 

53 mfsparams.alldecpars["0"], 

54 mfsparams.allimpars["0"], 

55 mfsparams.allnormpars["0"], 

56 ) 

57 mfsparams.alldecpars["0"]["startmodel"] = "" 

58 mfsparams.allimpars["0"]["startmodel"] = "" 

59 params.alldecpars["0"]["startmodel"] = "" 

60 params.allimpars["0"]["startmodel"] = "" 

61 

62 

63 ################################# 

64 super().__init__(params) 

65 ## print(f'self all impars {self.allimpars}, \n allvars={vars(self)}') 

66 if self.allimpars["0"]["specmode"] != "mvc": 

67 raise RuntimeError( 

68 f"Can't use specmode {self.allimpars['0']['specmode']} with imager helper {self.__class__.__name__}!" 

69 ) 

70 

71 ### Set up and check nchan and reffreq settings.  

72 nchan = self.allimpars["0"]["nchan"] 

73 freqbeg, freqwidth = self.determineFreqRange() 

74 if nchan < 1 : 

75 nchan=int((freqwidth)/(0.1*freqbeg)) #gives around 10 channel for 2:1 BW 

76 if nchan < 5: 

77 nchan=mfsparams.alldecpars["0"]["nterms"] + 1 

78 casalog.post('Calculating nchan from the data range to be '+str(nchan),'INFO') 

79 

80 ## If nchan < nterms, complain. 

81 in_nterms = mfsparams.alldecpars["0"]["nterms"] 

82 if nchan<in_nterms: 

83 raise RuntimeError( 

84 f"nchan (={nchan}) should be >= nterms ( {in_nterms} ) for valid polynomial fits to be feasible. " 

85 ) 

86 

87 if nchan>50: 

88 casalog.post('For mvc (mtmfs_via_cube), one usually needs only about 10 channels across the freq range, to fit Taylor polynomials of a low order','WARN') 

89 

90 in_reffreq = mfsparams.allimpars["0"]["reffreq"] ##### NEED TO MAKE THIS WORK FOR MULTIFIELD... 

91 if in_reffreq == "": ## User has not set it. Calculate default 

92 midfreq = freqbeg + 0.5*freqwidth # midpoint of the freq range found by "determineFreqRange()" in Hz 

93 

94 for k in mfsparams.allimpars: 

95 mfsparams.allimpars[k]["reffreq"] = str(midfreq) + "Hz" 

96 params.allimpars[k]["reffreq"] = str(midfreq)+"Hz" 

97 

98 

99 #print("params.impars", mfsparams.allimpars) 

100 

101 rfreq = _qa.convert(mfsparams.allimpars["0"]["reffreq"] , "Hz")["value"] 

102 #print("REFFREQ = ",rfreq, " ----" , mfsparams.allimpars["0"]["reffreq"]) 

103 if rfreq < freqbeg or rfreq > freqbeg+freqwidth: 

104 casalog.post('The reffreq of ' + mfsparams.allimpars["0"]["reffreq"] + ' is outside the selected frequency range of ' + str(freqbeg) + ' Hz - ' + str(freqbeg+freqwidth) + ' Hz','WARN') 

105 

106 self.mfsImager = PySynthesisImager(mfsparams) 

107 

108 

109 freqwidth = freqwidth / nchan 

110 #print(f"#####freqbeg={freqbeg}, freqwidth={freqwidth}, nchan={nchan} for cube") 

111 # Update some settings: 

112 # - specmode to cube so that we run a cube major cycle 

113 # - deconvolver to hogbom so that major cycle doesn't get confused TODO is this necessary? 

114 for k in self.allimpars: 

115 self.allimpars[k]["specmode"] = "cube" 

116 self.allimpars[k]["start"] = f"{freqbeg}Hz" 

117 self.allimpars[k]["width"] = f"{freqwidth}Hz" 

118 # self.alldecpars[k]['specmode']='cube' 

119 # this is needed for basic check...it is not used 

120 self.allimpars[k]["deconvolver"] = "hogbom" 

121 for k in self.allgridpars: 

122 self.allgridpars[k]["deconvolver"] = "hogbom" 

123 self.allgridpars[k]["interpolation"] = "nearest" 

124 self.allgridpars[k]["conjbeams"]=False 

125 self.allgridpars[k]["wbawp"]=True 

126 for k in self.allnormpars: 

127 self.allnormpars[k]["deconvolver"] = "hogbom" 

128 self.weightpars['usecubebriggs']=False 

129 self.fresh_images: List[str] = [] 

130 self.verify_dec_pars() 

131 ####################################### 

132 @time_func 

133 def determineFreqRange(self) -> Tuple[np.double, np.double]: 

134 minFreq = 1.0e20 

135 maxFreq = 0.0 

136 #pdb.set_trace() 

137 for msid in self.allselpars: 

138 msname = self.allselpars[msid]["msname"] 

139 spwsel = ( 

140 self.allselpars[msid]["spw"] if (self.allselpars[msid]["spw"]) else "*" 

141 ) 

142 fieldsel = ( 

143 self.allselpars[msid]["field"] 

144 if (self.allselpars[msid]["field"]) 

145 else "*" 

146 ) 

147 fieldid = _ms.msseltoindex(vis=msname, spw=spwsel, field=fieldsel)["field"][0] 

148 _tb.open(msname) 

149 fieldids=_tb.getcol("FIELD_ID") 

150 _tb.done() 

151 # have to do this because advisechansel does not work for fieldids not in main 

152 if fieldid not in fieldids: 

153 fieldid=fieldids[0] 

154 frange = _su.advisechansel( 

155 msname=msname, getfreqrange=True, fieldid=fieldid, spwselection=spwsel 

156 ) 

157 if minFreq > _qa.convert(frange["freqstart"], "Hz")["value"]: 

158 minFreq = _qa.convert(frange["freqstart"], "Hz")["value"] 

159 if maxFreq < _qa.convert(frange["freqend"], "Hz")["value"]: 

160 maxFreq = _qa.convert(frange["freqend"], "Hz")["value"] 

161 if(minFreq > maxFreq): 

162 raise Exception("Failed to determine the frequency range to build the cube from the field and spw selection") 

163 

164 #print(f"@@@@@@@MinFreq and MaxFreq = {minFreq}, {maxFreq}") 

165 freqwidth = maxFreq - minFreq 

166 return (minFreq, freqwidth) 

167 

168 ############################################# 

169 @time_func 

170 def verify_dec_pars(self) -> bool: 

171 for immod in range(0, self.NF): 

172 pars = self.alldecpars[str(immod)] 

173 if pars["specmode"] != "mvc": 

174 raise RuntimeError( 

175 f"Creating instance of class {type(self).__name__} with the wrong specmode! Expected 'mvc' but instead got '{pars['specmode']}'!" 

176 ) 

177 if pars["deconvolver"] != "mtmfs": 

178 raise RuntimeError( 

179 f"specmode {pars['specmode']} requires 'mtmfs' deconvolver but instead got '{pars['deconvolver']}'!" 

180 ) 

181 if pars["nterms"] < 2: 

182 raise RuntimeError( 

183 f"specmode {pars['specmode']} requires nterms >1 !" 

184 ) 

185 return True 

186 ############################################## 

187 def get_dec_pars_for_immod(self, immod: int) -> dict: 

188 pars = self.alldecpars[str(immod)] 

189 # Do not do these sneaky things here ...let the user change the parameters themselves 

190 # pars['specmode'] = 'mfs' 

191 return pars 

192 

193 @time_func 

194 def initializeNormalizers(self): 

195 super().initializeNormalizers() 

196 self.mfsImager.initializeNormalizers() 

197 @time_func 

198 def initializeDeconvolvers(self): 

199 # for immod in range(0,self.NF): 

200 # self.SDtools.append(synthesisdeconvolver()) 

201 # self.SDtools[immod].setupdeconvolution(decpars=self.get_dec_pars_for_immod(immod)) 

202 # should initialize mfs deconvolvers 

203 self.mfsImager.initializeDeconvolvers() 

204 

205 ############################################# 

206 @time_func 

207 def check_psf(self, immod): 

208 self.cube2tt(immod, suffixes=["psf", "sumwt"]) 

209 return super().check_psf(immod) 

210 

211 ################################################## 

212 @time_func 

213 def copy_startmodel(self, decpars: dict, impars: dict, normpars: dict) -> None: 

214 """ 

215 As tclean provides capacity for startmodel to be for field 0 only we don't 

216 need to deal with outlier fields 

217 """ 

218 # decpars = self.get_dec_pars_for_immod(0) 

219 # print(f"DECPARS {decpars}") 

220 imagename = decpars["imagename"] 

221 basemod = imagename + ".model" 

222 if len(decpars["startmodel"]) == decpars["nterms"]: 

223 if not os.path.exists(basemod + ".tt0"): 

224 refImage = imagename + ".psf.tt0" 

225 if not os.path.exists(refImage): 

226 self.copy_image(imagename + ".psf", refImage) 

227 for k in range(decpars["nterms"]): 

228 self.regrid_image( 

229 refImage, f"{basemod}.tt{k}", decpars["startmodel"][k] 

230 ) 

231 # As this is being called before __init__ temporarily assign alldecpars 

232 self.alldecpars = {"0": decpars} 

233 self.allimpars = {"0": impars} 

234 self.tt2cube(0) 

235 inpcube = self.get_image_name(0, "model") 

236 pbcube = self.get_image_name(0, "pb") 

237 pblimit = normpars["pblimit"] 

238 if SW==False: 

239 self.modify_cubemodel_with_pb( 

240 modcube=inpcube, pbcube=pbcube, pbtt0=pbcube + ".tt0", pblimit=np.fabs(pblimit) 

241 ) 

242 else: 

243 imname = self.get_dec_pars_for_immod(0)['imagename'] 

244 t0 = time.time() 

245 _su.apply_freq_dep_pb(cubename=imname,mtname=imname,pblimit=np.fabs(pblimit)) 

246 t1 = time.time() 

247 #print(f'#######---- Function apply_freq_dep_pb took {(t1-t0)}s') 

248 

249 del self.alldecpars 

250 del self.allimpars 

251 

252 ################################################## 

253 

254 def get_image_name(self, immod: int, suffix: str, ttN: Optional[int] = None): 

255 decpars = self.get_dec_pars_for_immod(immod) 

256 imagename = decpars["imagename"] 

257 basename = lambda img: f"{img}.{suffix}" 

258 

259 bn = basename(imagename) 

260 if ttN is not None: 

261 return f"{bn}.tt{ttN}" 

262 return bn 

263 

264 def hasConverged(self) -> bool: 

265 # create .ttN taylor term images for the mtmfs deconvolver 

266 # for immod in range(0,self.NF): 

267 # suffixes = ["residual", "sumwt"] 

268 # only need to create the psf taylor term images once (shouldn't change after check_psf) 

269 # self.cube2tt(immod, suffixes=suffixes) 

270 return self.mfsImager.hasConverged() 

271 

272 ############################################# 

273 

274 def initializeIterationControl(self) -> None: 

275 self.mfsImager.initializeIterationControl() 

276 

277 ############################################# 

278 @time_func 

279 def runMajorCycle(self, isCleanCycle: bool = True) -> None: 

280 # hopefully this carries all info for writing last model 

281 self.IBtool = self.mfsImager.IBtool 

282 #time0 = time.time() 

283 super().runMajorCycle(isCleanCycle) 

284 time1 = time.time() 

285 suffixes = ["residual"] #, "sumwt"] 

286 for immod in range(0, self.NF): 

287 inpcube = self.get_image_name(immod, "residual") 

288 pbcube = self.get_image_name(immod, "pb") 

289 cubewt = self.get_image_name(immod, "sumwt") 

290 pblimit = self.allnormpars[str(immod)]["pblimit"] 

291 

292 ##print("USING PB2TTPB from runMajor") 

293 ##self.cubePB2ttPB(pbcube, pbcube + ".tt0", cubewt, np.fabs(pblimit)) 

294 ##self.cube2tt(immod, suffixes=["pb"]) 

295 

296 # self.modify_with_pb(inpcube=inpcube, pbcube=pbcube, cubewt=cubewt, action='div', pblimit=pblimit, freqdep=True) 

297 # self.modify_with_pb(inpcube=inpcube, pbcube=pbcube, cubewt=cubewt, action='mult', pblimit=pblimit, freqdep=False) 

298 if SW==False: 

299 self.removePBSpectralIndex(inpcube, pbcube, pbcube + ".tt0", np.fabs(pblimit)) 

300 else: 

301 imname = self.get_dec_pars_for_immod(immod)['imagename'] 

302 t0 = time.time() 

303 _su.remove_freq_dep_pb(cubename=imname,mtname=imname,pblimit=np.fabs(pblimit)) 

304 t1 = time.time() 

305 #print(f'#######---- Function remove_freq_dep_pb took {(t1-t0)}s') 

306 

307 self.cube2tt(immod, suffixes=suffixes) 

308 #time2 = time.time() 

309 #print(f"MAKE RESidual time, core={time1-time0} s, cube2tt={time2-time1}") 

310 

311 ############################################## 

312 @time_func 

313 def makePSF(self) -> None: 

314 # pdb.set_trace() 

315 time0 = time.time() 

316 super().makePSFCore() 

317 time1 = time.time() 

318 #####have to ensure the pb is made 

319 super().makePB() 

320 pblimit = self.allnormpars['0']["pblimit"] 

321 cubewt = self.get_image_name(0, "sumwt") 

322 pbcube = self.get_image_name(0, "pb") 

323 

324 #print("USING PB2TTPB from makePSF") 

325 #self.cubePB2ttPB(pbcube, pbcube + ".tt0", cubewt, np.fabs(pblimit)) 

326 #suffixes = ["psf", "sumwt", "weight"] 

327 

328 suffixes = ["pb","psf", "sumwt"] 

329 for immod in range(0, self.NF): 

330 self.cube2tt(immod, suffixes=suffixes) 

331 #now that we donot call dividepsfbyweight which did the fitting 

332 #we have to do it now explicitly 

333 self.mfsImager.PStools[immod].makepsfbeamset() 

334 

335# for immod in range(0, self.NF): 

336# self.mfsImager.PStools[immod].gatherpsfweight() 

337# self.mfsImager.PStools[immod].dividepsfbyweight() 

338 time2 = time.time() 

339 #print(f"MAKE psf time, core={time1-time0} s, cube2tt={time2-time1}") 

340 

341 ############################################### 

342 @time_func 

343 def runMinorCycle(self) -> bool: 

344 # convert from cube to .ttN taylor term images for the mtmfs deconvolver 

345 time0 = time.time() 

346 # for immod in range(0,self.NF): 

347 # Before minorcycle : Divide out the frequency-dependent PB, multiply by a common PB. 

348 # inpcube = self.get_image_name(immod, "residual") 

349 # pbcube = self.get_image_name(immod, "pb") 

350 # cubewt = self.get_image_name(immod, "sumwt") 

351 # pblimit = self.allnormpars[str(immod)]['pblimit'] 

352 # self.modify_with_pb(inpcube=inpcube, pbcube=pbcube, cubewt=cubewt, action='div', pblimit=pblimit, freqdep=True) 

353 # self.modify_with_pb(inpcube=inpcube, pbcube=pbcube, cubewt=cubewt, action='mult', pblimit=pblimit, freqdep=False) 

354 

355 # suffixes = ["residual", "psf", "sumwt", "model"] 

356 # only need to create the psf taylor term images once (shouldn't change after check_psf) 

357 # suffixes.remove('psf') 

358 # suffixes=['model'] 

359 # self.cube2tt(immod, suffixes=suffixes) 

360 

361 # run the mtmfs deconvolver 

362 ret = self.mfsImager.runMinorCycle() 

363 time1 = time.time() 

364 # convert back to cube images for the cube major cycle 

365 for immod in range(0, self.NF): 

366 self.tt2cube(immod) 

367 

368 # After minorcycle : Divide out the common PB, Multiply by frequency-dependent PB. 

369 inpcube = self.get_image_name(immod, "model") 

370 pbcube = self.get_image_name(immod, "pb") 

371 # cubewt = self.get_image_name(immod, "sumwt") 

372 pblimit = self.allnormpars[str(immod)]["pblimit"] 

373 # self.modify_with_pb(inpcube=inpcube, pbcube=pbcube, cubewt=cubewt, action='div', pblimit=pblimit, freqdep=False) 

374 # self.modify_with_pb(inpcube=inpcube, pbcube=pbcube, cubewt=cubewt, action='mult', pblimit=pblimit, freqdep=True) 

375 if SW==False: 

376 self.modify_cubemodel_with_pb( 

377 modcube=inpcube, pbcube=pbcube, pbtt0=pbcube + ".tt0", pblimit=np.fabs(pblimit) 

378 ) 

379 else: 

380 imname = self.get_dec_pars_for_immod(immod)['imagename'] 

381 t0 = time.time() 

382 _su.apply_freq_dep_pb(cubename=imname,mtname=imname,pblimit=np.fabs(pblimit)) 

383 t1 = time.time() 

384 #print(f'#######---- Function apply_freq_dep_pb took {(t1-t0)}s') 

385 

386 time2 = time.time() 

387 #print(f"Minorcycle time, minor={time1-time0} s, tt2cube={time2-time1}") 

388 return ret 

389 

390 ############################################################# 

391 @time_func 

392 def updateMask(self) -> None: 

393 return self.mfsImager.updateMask() 

394 

395 ########################################################### 

396 def getSummary(self, fignum : int =1) -> dict: 

397 return self.mfsImager.getSummary(fignum) 

398 

399 ########################################################## 

400 @time_func 

401 def restoreImages(self) -> None: 

402 self.mfsImager.restoreImages() 

403 

404 #################################################### 

405 @time_func 

406 def pbcorImages(self) -> None: 

407 self.mfsImager.pbcorImages() 

408 

409 #################################### 

410 @time_func 

411 def cube2tt(self, immod: int = 0, suffixes: Optional[List[str]] = None) -> None: 

412 """Creates the necessary taylor term images. 

413 

414 Args: 

415 immod: which image facet/outlier field to convert 

416 suffixes: list of images to convert, can include any of ["residual", "psf", "sumwt"] 

417 

418 Outputs: 

419 pb.tt0 

420 residual.tt0..residual.tt(N-1), model.tt0..model.tt(N-1) 

421 psf.tt0..psf.tt(2*N-2) 

422 

423 If incompatible images already exist with the same name, replace them.""" 

424 time0 = time.time() 

425 if suffixes is None: 

426 suffixes = ["pb", "residual", "psf", "sumwt"] 

427 decpars = self.get_dec_pars_for_immod(immod) 

428 nterms = decpars["nterms"] 

429 pblimit = self.allnormpars[str(immod)]["pblimit"] 

430 # determine which images are being converted 

431 imgs = [ 

432 ("pb",1), 

433 ("residual", nterms), 

434 ("psf", nterms * 2 - 1), 

435 ("sumwt", nterms * 2 - 1) 

436 # ("weight", 1), 

437 ] # , ('model',nterms)] 

438 tmp_imgs = [] 

439 for suffix, num_terms in imgs: 

440 if suffix not in suffixes: 

441 continue 

442 if os.path.exists(self.get_image_name(immod, suffix)): 

443 tmp_imgs.append((suffix, num_terms)) 

444 imgs = tmp_imgs 

445 

446 # create the .ttN images 

447 for suffix, num_terms in imgs: 

448 basename = self.get_image_name(immod, suffix) 

449 for N in range(num_terms): 

450 ttname = self.get_image_name(immod, suffix, ttN=N) 

451 

452 # remove the existing image, if any 

453 if os.path.exists(ttname): 

454 # only create the images once per execution 

455 if ttname in self.fresh_images: 

456 continue 

457 # TODO possible optimization where all the data is set to 0 instead 

458 shutil.rmtree(ttname) 

459 

460 # create a new, blank image based off the template baseimage 

461 self.copy_image(template_img=basename, output_img=ttname) 

462 self.copy_nonexistant_keywords(template_img=basename, output_img=ttname) 

463 dopsf = suffix == "psf" or suffix == "sumwt" 

464 if not dopsf and pblimit>0.0: 

465 self.add_mask( 

466 ttname, self.get_image_name(immod, "pb", ttN=0), np.fabs(pblimit) 

467 ) 

468 self.fresh_images.append(ttname) 

469 

470 # convert them images! 

471 cubewt = self.get_image_name(immod, "sumwt") 

472 chanweight = None 

473 if not os.path.exists(cubewt): 

474 cubewt = "" 

475 else: 

476 _ia.open(cubewt) 

477 # nchan = _ia.shape()[3] 

478 chanweight = _ia.getchunk()[0, 0, 0, :] 

479 _ia.done() 

480 # print(f'IMGS={imgs}') 

481 for suffix, num_terms in imgs: 

482 basename = self.get_image_name(immod, suffix) 

483 reffreq = self.allimpars[str(immod)]["reffreq"] 

484 dopsf = suffix == "psf" or suffix == "sumwt" 

485 chwgt = None if suffix != "weight" else chanweight 

486 if SW==False: 

487 self.cube_to_taylor_sum( 

488 cubename=basename, 

489 cubewt=cubewt, 

490 chanwt=chwgt, 

491 mtname=basename, 

492 reffreq=reffreq, 

493 nterms=num_terms, 

494 dopsf=dopsf, 

495 ) 

496 else: 

497 imname = self.get_dec_pars_for_immod(immod)['imagename'] 

498 if suffix == "psf": 

499 imtype=0 ## num_terms should be 2*nterms-1 

500 if suffix == "residual": 

501 imtype=1 ## num_terms should be nterms 

502 if suffix == 'pb': ## may not be used.....  

503 imtype=2 ## num_terms is 1 (for now) 

504 if suffix == "sumwt": 

505 imtype=3 

506 t0=time.time() 

507 _su.cube_to_taylor_sum(cubename=imname,mtname=imname,nterms=nterms,reffreq=reffreq,imtype=imtype,pblimit=np.fabs(pblimit)) 

508 #print(f"Imname : {imname} , suffix : {suffix} and type : {imtype} with pblimit : {pblimit}") 

509 t1 = time.time() 

510 #print(f'#######---- Function cube_to_taylor_sum took {(t1-t0)}s') 

511 

512 

513 if not dopsf and pblimit>0.0: 

514 for theTerm in range(num_terms): 

515 ttname = self.get_image_name(immod, suffix, ttN=theTerm) 

516 # print(f'Adding masks to {ttname}') 

517 self.add_mask( 

518 ttname, self.get_image_name(immod, "pb", ttN=0), np.fabs(pblimit) 

519 ) 

520 time1 = time.time() 

521 #print(f"Time taken in cube2tt={time1-time0}") 

522 # special case: just copy pb 

523 # basename = self.get_image_name(immod, "pb") 

524 # ttname = self.get_image_name(immod, "pb", ttN=0) 

525 # if ((not os.path.exists(ttname)) and os.path.exists(basename)): 

526 # self.copy_image(template_img=basename, output_img=ttname) 

527 # self.copy_nonexistant_keywords(template_img=basename, output_img=ttname) 

528 # #shutil.rmtree(ttname) 

529 # #shutil.copytree(basename, ttname) 

530 # self.cube_to_taylor_sum(cubename=basename, cubewt=cubewt, mtname=basename, reffreq=reffreq, nterms=1, dopsf=False) 

531 # _ia.open(ttname) 

532 # _ia.calc(pixels="'"+ttname+"'/max('"+ttname+"')") 

533 # _ia.done() 

534 

535 # end 

536 @time_func 

537 def tt2cube(self, immod: int =0) -> None: 

538 """Creates or updates the .model image with all new data obtained 

539 from the .model.ttN images. 

540 """ 

541 time0 = time.time() 

542 decpars = self.get_dec_pars_for_immod(immod) 

543 nterms = decpars["nterms"] 

544 imagename = decpars["imagename"] 

545 reffreq = self.allimpars[str(immod)]["reffreq"] 

546 

547 # run the conversion 

548 if SW==False: 

549 self.taylor_model_to_cube( 

550 cubename=imagename, mtname=imagename, reffreq=reffreq, nterms=nterms 

551 ) 

552 else: 

553 t0=time.time() 

554 _su.taylor_coeffs_to_cube(cubename=imagename,mtname=imagename,reffreq=reffreq,nterms=nterms) 

555 t1 = time.time() 

556 #print(f'#######---- Function taylor_coeffs_to_cube took {(t1-t0)}s') 

557 

558 time1 = time.time() 

559 #print(f"Time taken in tt2cube {time1-time0}") 

560 

561 ######################################################################### 

562 @time_func 

563 def regrid_image(self, template_image: str ="", output_image: str ="", input_image: str ="") -> None: 

564 """ 

565 template_image provides the coordinatesystem and shape onto which the 

566 input_image is regridded to and named output_image. 

567 if output_image exists on disk it gets overwritten 

568 """ 

569 _ia.open(template_image) 

570 csys = _ia.coordsys() 

571 shp = _ia.shape() 

572 _ia.done() 

573 _ia.open(input_image) 

574 _ia.regrid( 

575 outfile=output_image, 

576 shape=shp, 

577 csys=csys.torecord(), 

578 axes=[0, 1], 

579 overwrite=True, 

580 force=False, 

581 ) 

582 _ia.done() 

583 

584 ###################################################################### 

585 @time_func 

586 def copy_image(self, template_img: str ="try.psf", output_img: str ="try.zeros.psf") -> None: 

587 # get the shape 

588 _ia.open(template_img) 

589 shape = _ia.shape() 

590 nchan = _ia.shape()[3] 

591 csys = _ia.coordsys() 

592 pixeltype = _ia.pixeltype() 

593 _ia.close() 

594 _ia.done() 

595 

596 # get the data type 

597 dtype = np.single 

598 pixelprefix = pixeltype[0] # for 'f'loat, 'd'ouble, or 'c'omplex 

599 if pixeltype == "double": 

600 dtype = np.double 

601 elif pixeltype == "complex": 

602 dtype = np.csingle 

603 elif pixeltype == "dcomplex": 

604 dtype = np.cdouble 

605 pixelprefix = "cd" 

606 # Get the frequency and BW correct 

607 minfreq = csys.toworld([0, 0, 0, 0])["numeric"][3] 

608 maxfreq = csys.toworld([0, 0, 0, nchan - 1])["numeric"][3] 

609 a = csys.increment(type="spectral") 

610 a["numeric"] *= nchan 

611 csys.setincrement(a, type="spectral") 

612 b = csys.referencevalue(type="spectral") 

613 b["numeric"] = (minfreq + maxfreq) / 2.0 

614 csys.setreferencevalue(b, type="spectral") 

615 csys.setreferencepixel(0.0, "spectral") 

616 # populate some pixels 

617 shape[3] = 1 # taylor term images don't use channels 

618 pixels = np.zeros(shape, dtype=dtype) 

619 

620 # create the new outputmask 

621 _ia.fromarray(output_img, csys=csys.torecord(), pixels=pixels, type=pixelprefix) 

622 _ia.close() 

623 _ia.done() 

624 

625 ################################################ 

626 @time_func 

627 def copy_nonexistant_keywords( 

628 self, template_img: str ="try.psf", output_img: str ="try.zeros.psf" 

629 ) -> None : 

630 _tb.open(template_img) 

631 new_ii = _tb.getkeyword("imageinfo") 

632 _tb.close() 

633 if "perplanebeams" in new_ii: 

634 del new_ii["perplanebeams"] 

635 old_ii = {} 

636 _tb.open(output_img, nomodify=False) 

637 kws = _tb.keywordnames() 

638 if "imageinfo" in kws: 

639 old_ii = _tb.getkeyword("imageinfo") 

640 old_ii.update(new_ii) 

641 # for kw in old_kws: 

642 # del new_kws[kw] 

643 # for kw in new_kws: 

644 # old_kws[kw]=new_kws[kw] 

645 # casalog.post(f"{template_img} to {output_img} imageinfo: {old_ii}\n\n\n", "WARN") 

646 _tb.putkeyword("imageinfo", old_ii) 

647 _tb.close() 

648 _ia.open(template_img) 

649 miscinf = _ia.miscinfo() 

650 _ia.done() 

651 _ia.open(output_img) 

652 newmiscinf = _ia.miscinfo() 

653 newmiscinf.update(miscinf) 

654 _ia.setmiscinfo(newmiscinf) 

655 _ia.done() 

656 

657 ################################################ 

658 def get_freq_list(self, imname=""): 

659 """Get the list of frequencies for the given image, one for each channel. 

660 

661 Returns: 

662 list[float] The frequencies for each channel in the image, in Hz. 

663 

664 From: 

665 sdint_helper.py 

666 """ 

667 

668 _ia.open(imname) 

669 csys = _ia.coordsys() 

670 shp = _ia.shape() 

671 _ia.close() 

672 

673 if csys.axiscoordinatetypes()[3] == "Spectral": 

674 restfreq = csys.referencevalue()["numeric"][ 

675 3 

676 ] # /1.0e+09; # convert more generally.. 

677 freqincrement = csys.increment()["numeric"][3] # /1.0e+09; 

678 freqlist = [] 

679 for chan in range(0, shp[3]): 

680 freqlist.append(restfreq + chan * freqincrement) 

681 elif csys.axiscoordinatetypes()[3] == "Tabular": 

682 freqlist = csys.torecord()["tabular2"]["worldvalues"] # /1.0e+09; 

683 else: 

684 casalog.post("Unknown frequency axis. Exiting.", "SEVERE") 

685 return False 

686 

687 csys.done() 

688 return freqlist 

689 

690 ####################################################### 

691 @time_func 

692 def cubePB2ttPB(self, cubePB="", ttPB="", sumwt="", pblimit=0.2): 

693 """ 

694 convert the cube PB to an average PB 

695 """ 

696 time0 = time.time() 

697 # special case: just copy pb 

698 if (not os.path.exists(ttPB)) and os.path.exists(cubePB): 

699 self.copy_image(template_img=cubePB, output_img=ttPB) 

700 self.copy_nonexistant_keywords(template_img=cubePB, output_img=ttPB) 

701 # shutil.rmtree(ttname) 

702 # shutil.copytree(basename, ttname) 

703 # self.cube_to_taylor_sum(cubename=basename, cubewt=cubewt, mtname=basename, reffreq=reffreq, nterms=1, dopsf=False) 

704 _ia.open(ttPB) 

705 pix = np.zeros(_ia.shape(), dtype=np.float64) 

706 _ia.done() 

707 _ia.open(cubePB) 

708 #print(f"STATS of cube pb {_ia.statistics()}") 

709 shp = _ia.shape() 

710 cwt = np.ones((shp[3])) 

711 if os.path.exists(sumwt): 

712 _ib = _image() 

713 _ib.open(sumwt) 

714 if _ib.shape()[2] == shp[3]: 

715 cwt = _ib.getchunk(blc=[0, 0, 0, 0], trc=[0, 0, 0, shp[3] - 1]) 

716 _ib.done() 

717 

718 for k in range(shp[3]): 

719 # print(f'chan {k}, max {np.max(pix)}, {np.min(pix)}') 

720 pix += ( 

721 _ia.getchunk( 

722 blc=[0, 0, 0, k], trc=[shp[0] - 1, shp[1] - 1, shp[2] - 1, k] 

723 ) 

724 * cwt[k] 

725 ) 

726 _ia.done() 

727 

728 _ia.open(ttPB) 

729 _ia.putchunk(pix) 

730 _ia.done() 

731 _ia.open(ttPB) 

732 _ia.calc(pixels="'" + ttPB + "'/max('" + ttPB + "')") 

733 _ia.calcmask(mask="'" + ttPB + "' > " + str(pblimit)) 

734 _ia.done() 

735 time1 = time.time() 

736 #print(f"Time taken in cubePB2ttPB={time1-time0}") 

737 

738 ################################################ 

739 @time_func 

740 def cube_to_taylor_sum( 

741 self, 

742 cubename="", 

743 cubewt="", 

744 chanwt=None, 

745 mtname="", 

746 reffreq="1.5GHz", 

747 nterms=2, 

748 dopsf=False, 

749 ): 

750 """ 

751 Convert Cubes (output of major cycle) to Taylor weighted averages (inputs to the minor cycle) 

752 Input : Cube image <cubename>, with channels weighted by image <cubewt> 

753 Output : Set of images : <mtname>.tt0, <mtname>.tt1, etc... 

754 Algorithm: I_ttN = sum([ I_v * ((f-ref)/ref)**N for f in freqs ]) 

755 

756 Args: 

757 cubename: Name of a cube image to interpret into a set of taylor term .ttN images, eg "try.residual", "joint.cube.psf". 

758 cubewt: Name of a .sumwt image that contains the per-channel weighting for the interferometer image. 

759 chanwt: List of 0s and 1s, one per channel, to effectively disable the effect of a channel on the resulting images. 

760 mtname: The prefix output name, to be concatenated with ".ttN" strings, eg "try_mt.residual", "joint.multiterm.psf" 

761 These images should already exist by the time this function is called. 

762 It's suggested that this have the same suffix as cubename. 

763 reffreq: reference frequency, like for tclean 

764 nterms: number of taylor terms to fit the spectral index to 

765 dopsf: Signals that cubename represents a point source function, should be true if cubename ends with ".psf" or ".sumwt". 

766 If true, then output 2*nterms-1 ttN images. 

767 

768 From: 

769 sdint_helper.py 

770 """ 

771# if dopsf is True: 

772# nterms = 2 * nterms - 1 

773 

774 pix = [] 

775 for tt in range(0, nterms): 

776 _ia.open(mtname + ".tt" + str(tt)) 

777 pix.append(_ia.getchunk()) 

778 _ia.close() 

779 pix[tt].fill(0.0) 

780 

781 _ia.open(cubename) 

782 shp = _ia.shape() 

783 _ia.close()