def SubStep(self, xc, yc, LocalSM, Residual): """ Sub-minor loop subtraction """ N0 = Residual.shape[-1] N1 = LocalSM.shape[-1] # Get overlap indices where psf should be subtracted Aedge, Bedge = GiveEdges(xc, yc, N0, N1 // 2, N1 // 2, N1) x0d, x1d, y0d, y1d = Aedge x0p, x1p, y0p, y1p = Bedge # Subtract from each channel/band Residual[:, :, x0d:x1d, y0d:y1d] -= LocalSM[:, :, x0p:x1p, y0p:y1p]
def AdaptArrayShape(self, A, Nout): nch, npol, Nin, _ = A.shape if Nin == Nout: return A elif Nin > Nout: # dx=Nout/2 # B=np.zeros((nch,npol,Nout,Nout),A.dtype) # print>>log," Adapt shapes: %s -> %s"%(str(A.shape),str(B.shape)) # B[:]=A[...,Nin/2-dx:Nin/2+dx+1,Nin/2-dx:Nin/2+dx+1] N0 = A.shape[-1] xc0 = yc0 = N0 / 2 N1 = Nout xc1 = yc1 = N1 / 2 Aedge, Bedge = GiveEdges((xc0, yc0), N0, (xc1, yc1), N1) x0d, x1d, y0d, y1d = Aedge x0p, x1p, y0p, y1p = Bedge B = A[..., x0d:x1d, y0d:y1d] return B else: return A
def SubStep(self, dx, dy, LocalSM): """ This is where subtraction in the image domain happens """ xc, yc = dx, dy N1 = LocalSM.shape[-1] # Get overlap indices where psf should be subtracted Aedge, Bedge = GiveEdges(xc, yc, self.Npix, N1 // 2, N1 // 2, N1) x0d, x1d, y0d, y1d = Aedge x0p, x1p, y0p, y1p = Bedge self._Dirty[:, :, x0d:x1d, y0d:y1d] -= LocalSM[:, :, x0p:x1p, y0p:y1p] # Subtract from the average if self.MultiFreqMode: # If multiple frequencies are present construct the weighted mean self._MeanDirty[:, 0, x0d:x1d, y0d:y1d] -= np.sum( LocalSM[:, :, x0p:x1p, y0p:y1p] * self.WeightsChansImages, axis=0) # Sum over freq else: self._MeanDirty = self._Dirty
def GiveGridFader(self, Image, DicoImager, iFacet, NormIm): nch, npol, NPixOut, _ = Image.shape _, _, N1, _ = self.GridShape xc, yc = DicoImager[iFacet]["pixCentral"] #x0,x1,y0,y1=DicoImager[iFacet]["pixExtent"] #xc,yc=(x0+x1)//2,(y0+y1)//2 Aedge, Bedge = GiveEdges(xc, yc, NPixOut, N1 // 2, N1 // 2, N1) #Bedge,Aedge=GiveEdges(N1//2,N1//2,N1,yc,xc,NPixOut) x0d, x1d, y0d, y1d = Aedge x0p, x1p, y0p, y1p = Bedge #print "xxA:",x0d,x1d #print "xxB:",x0p,x1p ModelIm = np.zeros((nch, npol, N1, N1), dtype=np.float32) for ch in range(nch): for pol in range(npol): #ModelIm[ch,pol][x0p:x1p,y0p:y1p]=Image[ch,pol].T[::-1,:].real[x0d:x1d,y0d:y1d] #ModelIm[ch,pol][x0p:x1p,y0p:y1p]=Image[ch,pol].real[x0d:x1d,y0d:y1d] ModelIm[ch, pol][x0p:x1p, y0p:y1p] = Image[ch, pol][x0d:x1d, y0d:y1d].real ModelIm[ch, pol][x0p:x1p, y0p:y1p] /= NormIm[x0d:x1d, y0d:y1d].real #ModelIm[ch,pol][x0p:x1p,y0p:y1p]/=NormIm[x0d:x1d,y0d:y1d].real ModelIm[ch, pol] = ModelIm[ch, pol].T[::-1, :] SumFlux = np.sum(ModelIm) #print iFacet,np.max(ModelIm) #return ModelIm, None ModelIm *= (self.OverS * N1)**2 Grid = np.complex64(self.FFTWMachine.fft(np.complex64(ModelIm))) return Grid, SumFlux
def Deconvolve(self): if self._Dirty.shape[-1] != self._Dirty.shape[-2]: # print "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!" # print self._Dirty.shape # print "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!" return "MaxIter", True, True dirty = self._Dirty nch, npol, nx, ny = dirty.shape Model = np.zeros_like(dirty) _, _, xp, yp = np.where(self._MeanDirty == np.max(self._MeanDirty)) self.PSFServer.setLocation(xp, yp) self.iFacet = self.PSFServer.iFacet psf, _ = self.PSFServer.GivePSF() nxPSF = psf.shape[-1] nxDirty = dirty.shape[-1] Nout = np.min([dirty.shape[-1], psf.shape[-1]]) dirty = self.AdaptArrayShape(dirty, Nout) SliceDirty = slice(0, None) if dirty.shape[-1] % 2 != 0: SliceDirty = slice(0, -1) d = dirty[:, :, SliceDirty, SliceDirty] psf = self.AdaptArrayShape(psf, d.shape[-1]) SlicePSF = slice(0, None) if psf.shape[-1] % 2 != 0: SlicePSF = slice(0, -1) p = psf[:, :, SlicePSF, SlicePSF] dirty_MUFFIN = np.squeeze(d[:, 0, :, :]) dirty_MUFFIN = dirty_MUFFIN.transpose((2, 1, 0)) psf_MUFFIN = np.squeeze(p[:, 0, :, :]) psf_MUFFIN = psf_MUFFIN.transpose((2, 1, 0)) EM = EasyMuffin(mu_s=self.GD['MUFFIN']['mu_s'], mu_l=self.GD['MUFFIN']['mu_l'], nb=self.GD['MUFFIN']['nb'], truesky=dirty_MUFFIN, psf=psf_MUFFIN, dirty=dirty_MUFFIN) EM.loop(nitermax=self.GD['MUFFIN']['NMinorIter']) nxModel = dirty_MUFFIN.shape[0] Aedge, Bedge = GiveEdges(nxModel // 2, nxModel // 2, nxModel, nxDirty // 2, nxDirty // 2, nxDirty) x0, x1, y0, y1 = Bedge Model = np.zeros((nxDirty, nxDirty, nch)) Model[x0:x1, y0:y1, :] = EM.x self.ModelMachine.setMUFFINModel(Model) # if self._Dirty.shape[-1]!=self._Dirty.shape[-2]: # # print "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!" # # print self._Dirty.shape # # print "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!" # return "MaxIter", True, True # dirty=self._Dirty # nch,npol,_,_=dirty.shape # Model=np.zeros_like(dirty) # _,_,xp,yp=np.where(self._MeanDirty==np.max(self._MeanDirty)) # self.PSFServer.setLocation(xp,yp) # self.iFacet=self.PSFServer.iFacet # psf,_=self.PSFServer.GivePSF() # Nout=np.min([dirty.shape[-1],psf.shape[-1]]) # dirty=self.AdaptArrayShape(dirty,Nout) # SliceDirty=slice(0,None) # if dirty.shape[-1]%2!=0: # SliceDirty=slice(0,-1) # d=dirty[:,:,SliceDirty,SliceDirty] # psf=self.AdaptArrayShape(psf,d.shape[-1]*2) # SlicePSF=slice(0,None) # if psf.shape[-1]%2!=0: # SlicePSF=slice(0,-1) # p=psf[:,:,SlicePSF,SlicePSF] # if p.shape[-1]!=2*d.shape[-1]: # print "!!!!!!!!!!!!!!!!!!!!!!!!!" # print "Could not adapt psf shape to 2*dirty shape!!!!!!!!!!!!!!!!!!!!!!!!!" # print p.shape[-1],d.shape[-1] # print "!!!!!!!!!!!!!!!!!!!!!!!!!" # psf=self.AdaptArrayShape(psf,d.shape[-1]) # SlicePSF=SliceDirty # for ch in range(nch): # CM=ClassMoresaneSingleSlice(dirty[ch,0,SliceDirty,SliceDirty],psf[ch,0,SlicePSF,SlicePSF],mask=None,GD=None) # model,resid=CM.giveModelResid(major_loop_miter=self.GD["MORESANE"]["NMajorIter"], # minor_loop_miter=self.GD["MORESANE"]["NMinorIter"], # loop_gain=self.GD["MORESANE"]["Gain"], # sigma_level=self.GD["MORESANE"]["SigmaCutLevel"],# tolerance=1., # enforce_positivity=self.GD["MORESANE"]["ForcePositive"]) # Model[ch,0,SliceDirty,SliceDirty]=model[:,:] # import pylab # pylab.clf() # pylab.subplot(2,2,1) # pylab.imshow(dirty[ch,0,SliceDirty,SliceDirty],interpolation="nearest") # pylab.colorbar() # pylab.subplot(2,2,2) # pylab.imshow(psf[ch,0,SlicePSF,SlicePSF],interpolation="nearest") # pylab.colorbar() # pylab.subplot(2,2,3) # pylab.imshow(model,interpolation="nearest") # pylab.colorbar() # pylab.subplot(2,2,4) # pylab.imshow(resid,interpolation="nearest") # pylab.colorbar() # pylab.draw() # pylab.show() # print # print np.max(np.max(Model,axis=-1),axis=-1) # print # print #_,_,nx,ny=Model.shape #Model=np.mean(Model,axis=0).reshape((1,1,nx,ny)) #Model.fill(0) #Model[:,:,xp,yp]=self._Dirty[:,:,xp,yp] return "MaxIter", True, True # stop deconvolution but do update model
def setSubDirty(self, ListPixParms): T = ClassTimeIt.ClassTimeIt("InitSSD.setSubDirty") T.disable() x, y = np.array(ListPixParms).T x0, x1 = x.min(), x.max() + 1 y0, y1 = y.min(), y.max() + 1 dx = x1 - x0 + self.Margin dy = y1 - y0 + self.Margin Size = np.max([dx, dy]) if Size % 2 == 0: Size += 1 _, _, N0, _ = self.Dirty.shape xc0, yc0 = int((x1 + x0) / 2.), int((y1 + y0) / 2.) self.xy0 = xc0, yc0 self.DeconvMachine.PSFServer.setLocation(*self.xy0) N1 = Size xc1 = yc1 = N1 / 2 Aedge, Bedge = GiveEdges((xc0, yc0), N0, (xc1, yc1), N1) x0d, x1d, y0d, y1d = Aedge x0p, x1p, y0p, y1p = Bedge self.SubDirty = self.Dirty[:, :, x0d:x1d, y0d:y1d].copy() T.timeit("0") self.blc = (x0d, y0d) self.DeconvMachine.PSFServer.setBLC(self.blc) _, _, nx, ny = self.SubDirty.shape ArrayPixParms = np.array(ListPixParms) ArrayPixParms[:, 0] -= x0d ArrayPixParms[:, 1] -= y0d self.ArrayPixParms = ArrayPixParms self.DicoSubDirty = {} for key in self.DicoDirty.keys(): if key in ["ImageCube", "MeanImage", 'FacetNorm', "JonesNorm"]: self.DicoSubDirty[key] = self.DicoDirty[key][..., x0d:x1d, y0d:y1d].copy() else: self.DicoSubDirty[key] = self.DicoDirty[key] T.timeit("1") # ModelImage=np.zeros_like(self.Dirty) # ModelImage[:,:,N0/2,N0/2]=10 # ModelImage[:,:,N0/2+3,N0/2]=10 # ModelImage[:,:,N0/2-2,N0/2-1]=10 # self.setSSDModelImage(ModelImage) # Mask=np.zeros((nx,ny),np.bool8) # Mask[x,y]=1 # self.SubMask=Mask x, y = ArrayPixParms.T Mask = np.zeros(self.DicoSubDirty["ImageCube"].shape[-2::], np.bool8) Mask[x, y] = 1 self.SubMask = Mask if self.SSDModelImage is not None: self.SubSSDModelImage = self.SSDModelImage[:, :, x0d:x1d, y0d:y1d].copy() for ch in range(self.NFreqBands): self.SubSSDModelImage[ch, 0][np.logical_not(self.SubMask)] = 0 self.addSubModelToSubDirty() T.timeit("2")
def setSubDirty(self, ListPixParms): T = ClassTimeIt.ClassTimeIt("InitSSD.setSubDirty") T.disable() x, y = np.array(ListPixParms).T x0, x1 = x.min(), x.max() + 1 y0, y1 = y.min(), y.max() + 1 dx = x1 - x0 + self.Margin dy = y1 - y0 + self.Margin Size = np.max([dx, dy]) if Size % 2 == 0: Size += 1 _, _, N0, _ = self.Dirty.shape xc0, yc0 = int((x1 + x0) / 2.), int((y1 + y0) / 2.) self.xy0 = xc0, yc0 self.DeconvMachine.PSFServer.setLocation(*self.xy0) N1 = Size xc1 = yc1 = N1 / 2 Aedge, Bedge = GiveEdges((xc0, yc0), N0, (xc1, yc1), N1) x0d, x1d, y0d, y1d = Aedge x0p, x1p, y0p, y1p = Bedge self.SubDirty = self.Dirty[:, :, x0d:x1d, y0d:y1d].copy() T.timeit("0") self.blc = (x0d, y0d) self.DeconvMachine.PSFServer.setBLC(self.blc) #self.DeconvMachine.PSFServer.iFacet=118 _, _, nx, ny = self.SubDirty.shape ArrayPixParms = np.array(ListPixParms) ArrayPixParms[:, 0] -= x0d ArrayPixParms[:, 1] -= y0d self.ArrayPixParms = ArrayPixParms self.DicoSubDirty = {} for key in self.DicoDirty.keys(): if key in ["ImageCube", "MeanImage", 'FacetNorm', "JonesNorm"]: self.DicoSubDirty[key] = self.DicoDirty[key][..., x0d:x1d, y0d:y1d].copy() else: self.DicoSubDirty[key] = self.DicoDirty[key] T.timeit("1") # ModelImage=np.zeros_like(self.Dirty) # ModelImage[:,:,N0/2,N0/2]=10 # ModelImage[:,:,N0/2+3,N0/2]=10 # ModelImage[:,:,N0/2-2,N0/2-1]=10 # self.setSSDModelImage(ModelImage) # Mask=np.zeros((nx,ny),np.bool8) # Mask[x,y]=1 # self.SubMask=Mask x, y = ArrayPixParms.T Mask = np.zeros(self.DicoSubDirty["ImageCube"].shape[-2::], np.bool8) Mask[x, y] = 1 self.SubMask = Mask # PSF,MeanPSF=self.DeconvMachine.PSFServer.GivePSF() # import pylab # pylab.clf() # ax=pylab.subplot(1,3,1) # N=self.DicoSubDirty["MeanImage"].shape[-1] # pylab.imshow(self.DicoSubDirty["MeanImage"][0,0], # interpolation="nearest",extent=(-N/2.,N/2.,-N/2.,N/2.),vmin=-0.1,vmax=1.) # pylab.colorbar() # pylab.subplot(1,3,2,sharex=ax,sharey=ax) # N=MeanPSF.shape[-1] # pylab.imshow(MeanPSF[0,0],interpolation="nearest",extent=(-N/2.,N/2.,-N/2.,N/2.),vmin=-0.1,vmax=1.) # pylab.colorbar() # pylab.draw() # pylab.show() if self.SSDModelImage is not None: self.SubSSDModelImage = self.SSDModelImage[:, :, x0d:x1d, y0d:y1d].copy() for ch in range(self.NFreqBands): self.SubSSDModelImage[ch, 0][np.logical_not(self.SubMask)] = 0 self.addSubModelToSubDirty() T.timeit("2")
class ClassImageDeconvMachine(): """ These methods may be called from ClassDeconvMachine Init(**kwargs) - contains minor cycle specific initialisations which are only used once Input: currently kwargs are minor cycle specific and should be set from ClassDeconvMachine but a ideally a generic interface has these set in the parset somehow. Deconvolve() - does joint deconvolution over all the channels/bands. Output: return_code - "MaxIter"???? continue - whether to continue the deconvolution updated - whether the model has been updated GiveModelImage(freq) - returns current model at freq Input: freq - tuple of frequencies at which to return the model Output: Mod - the current model at freq Update(DicoDirty,**kwargs) - updates to minor cycle at the end of each major cycle Input: DicoDirty - updated image dict at start of each major cycle Use kwargs to pass any other minor cycle specific options ToFile(fname) - saves dico model to file Input: fname - the name of the file to write the dico image to FromFile(fname) - reads model dict from file Input: fname - the name of the file to write the dico image to """ def __init__( self, Gain=0.1, MaxMinorIter=50000, NCPU=0, CycleFactor=2.5, FluxThreshold=None, RMSFactor=3, PeakFactor=0, GD=None, SearchMaxAbs=1, CleanMaskImage=None, ImagePolDescriptor=["I"], ModelMachine=None, MainCache=None, CacheFileName='WSCMS', **kw # absorb any unknown keywords arguments here ): self.SearchMaxAbs = SearchMaxAbs self.ModelImage = None self.MaxMinorIter = MaxMinorIter self.NCPU = NCPU self.MaskArray = None self.GD = GD self.MultiFreqMode = (self.GD["Freq"]["NBand"] > 1) self.NFreqBand = self.GD["Freq"]["NBand"] self.FluxThreshold = FluxThreshold self.CycleFactor = CycleFactor self.RMSFactor = RMSFactor self.PeakFactor = PeakFactor if ModelMachine is None: # raise RuntimeError("You need to supply ImageDeconvMachine with a instantiated ModelMachine") import ClassModelMachineWSCMS as ClassModelMachine self.ModelMachine = ClassModelMachine.ClassModelMachine( self.GD, GainMachine=ClassGainMachine.get_instance()) else: self.ModelMachine = ModelMachine self.GainMachine = self.ModelMachine.GainMachine self._niter = 0 # cache options self.maincache = MainCache self.CacheFileName = CacheFileName self.PSFHasChanged = False self.LastScale = 99999 # TODO - use MaskMachine for this CleanMaskImage = self.GD["Mask"]["External"] if CleanMaskImage is not None: print >> log, "Reading mask image: %s" % CleanMaskImage MaskArray = image(CleanMaskImage).getdata() nch, npol, nxmask, nymask = MaskArray.shape # if (nch > 1) or (npol > 1): # print>>log, "Warning - only single channel and pol mask supported. Will use mask for ch 0 pol 0" # MaskArray = MaskArray[0,0] # _, _, nxmod, nymod = self.ModelMachine.ModelShape # if (nxmod != nxmask) or (nymod !=nymask): # print>>log, "Warning - shape of mask != shape of your model. Will pad/trncate to match model shape" # nxdiff = nxmod - nxmask # nydiff = nymod - nymask # if nxdiff < 0: # MaskArray = MaskArray self._MaskArray = np.zeros(MaskArray.shape, np.bool8) for ch in range(nch): for pol in range(npol): self._MaskArray[ch, pol, :, :] = np.bool8( 1 - MaskArray[ch, pol].T[::-1].copy())[:, :] self.MaskArray = np.ascontiguousarray(self._MaskArray) # import matplotlib.pyplot as plt # plt.imshow(self.MaskArray[0,0]) # plt.colorbar() # plt.show() # # import sys # sys.exit(0) self._peakMode = "normal" self.CacheFileName = CacheFileName self.CurrentNegMask = None self._NoiseMap = None self._PNRStop = None # in _peakMode "sigma", provides addiitonal stopping criterion # # this is so that the relevant functions are registered as job handlers with APP # # pass to ModelMachine.setScaleMachine to set workers # self.FTMachine = FFTW_Scale_Manager(wisdom_file=self.GD["Cache"]["DirWisdomFFTW"]) # # APP.registerJobHandlers(self) def Init(self, cache=None, facetcache=None, **kwargs): # check for valid cache cachehash = dict([ (section, self.GD[section]) for section in ("Data", "Beam", "Selection", "Freq", "Image", "Facets", "Weight", "RIME", "Comp", "CF", "WSCMS") ]) cachepath, valid = self.maincache.checkCache(self.CacheFileName, cachehash, directory=True, reset=cache or self.PSFHasChanged) # export the hash self.maincache.saveCache(name='WSCMS') self.Freqs = kwargs["GridFreqs"] AllDegridFreqs = [] for i in kwargs["DegridFreqs"].keys(): AllDegridFreqs.append(kwargs["DegridFreqs"][i]) self.Freqs_degrid = np.asarray(AllDegridFreqs).flatten() self.SetPSF(kwargs["PSFVar"]) self.setSideLobeLevel(kwargs["PSFAve"][0], kwargs["PSFAve"][1]) self.ModelMachine.setPSFServer(self.PSFServer) self.ModelMachine.setFreqMachine( self.Freqs, self.Freqs_degrid, weights=kwargs["PSFVar"]["WeightChansImages"], PSFServer=self.PSFServer) from africanus.constants import c as lightspeed minlambda = lightspeed / self.Freqs.min() # LB - note MaskArray might be modified by ScaleMachine if GD{"WSCMS"]["AutoMask"] is True # so we should avoid keeping it as None # if self.MaskArray is None: # self.MaskArray = np.zeros([1, 1, self.Npix, self.Npix], dtype=np.bool8) self.ModelMachine.setScaleMachine(self.PSFServer, NCPU=self.NCPU, MaskArray=self.MaskArray, cachepath=cachepath, MaxBaseline=kwargs["MaxBaseline"] / minlambda) def Reset(self): pass def setMaskMachine(self, MaskMachine): self.MaskMachine = MaskMachine def SetModelRefFreq(self, RefFreq): """ Sets ref freq in ModelMachine. """ self.ModelMachine.setRefFreq(RefFreq) def SetModelShape(self): """ Sets the shape params of model, call in every update step """ self.ModelMachine.setModelShape(self._Dirty.shape) self.Nchan, self.Npol, self.Npix, _ = self._Dirty.shape self.NpixFacet = self.Npix // self.GD["Facets"]["NFacets"] def GiveModelImage(self, *args): return self.ModelMachine.GiveModelImage(*args) def setSideLobeLevel(self, SideLobeLevel, OffsetSideLobe): self.SideLobeLevel = SideLobeLevel self.OffsetSideLobe = OffsetSideLobe def SetPSF(self, DicoVariablePSF): """ The keys in DicoVariablePSF and what they mean: 'MeanFacetPSF' - 'MeanImage' - 'ImageCube' - 'CellSizeRad' - 'ChanMappingGrid' - 'ChanMappingGridChan' - 'CubeMeanVariablePSF' - 'CubeVariablePSF' - 'SumWeights' - 'MeanJonesBand' - 'PeakNormed_CubeMeanVariablePSF' 'PeakNormed_CubeVariablePSF' 'OutImShape' 'JonesNorm' 'Facets' 'PSFSidelobes' 'ImageInfo' 'CentralFacet' 'freqs' 'SumJonesChan' 'SumJonesChanWeightSq' 'EstimatesAvgPSF' 'WeightChansImages' 'FacetNorm' 'PSFGaussPars' 'FWHMBeam' """ self.PSFServer = ClassPSFServer(self.GD) # NormalisePSF must be true here for the beam to be applied correctly self.PSFServer.setDicoVariablePSF(DicoVariablePSF, NormalisePSF=True) self.DicoVariablePSF = DicoVariablePSF def setNoiseMap(self, NoiseMap, PNRStop=10): """ Sets the noise map. The mean dirty will be divided by the noise map before peak finding. If PNRStop is set, an additional stopping criterion (peak-to-noisemap) will be applied. Peaks are reported in units of sigmas. If PNRStop is not set, NoiseMap is treated as simply an (inverse) weighting that will bias peak selection in the minor cycle. In this mode, peaks are reported in units of flux. """ self._NoiseMap = NoiseMap self._PNRStop = PNRStop self._peakMode = "sigma" def SetDirty(self, DicoDirty): """ The keys in DicoDirty and what they mean (see also FacetMachine.FacetsToIm docs) 'JonesNorm' - array containing norm of Jones terms as an image 'ImageInfo' - dictionary containing 'CellSizeRad' and 'OutImShape' 'ImageCube' - array containing residual 'MeanImage' - array containing mean of the residual 'freqs' - dictionary keyed by band number containing the actual frequencies that got binned into that band 'SumWeights' - sum of visibility weights used in normalizing the gridded correlations 'FacetMeanResidual' - ??? 'WeightChansImages' - Weights corresponding to imaging bands (how is this computed?) 'FacetNorm' - self.FacetImage (grid-correcting map) see FacetMachine """ self.DicoDirty = DicoDirty self._Dirty = self.DicoDirty["ImageCube"] self._MeanDirty = self.DicoDirty["MeanImage"] self._JonesNorm = self.DicoDirty["JonesNorm"] self.WeightsChansImages = np.mean(np.float32( self.DicoDirty["WeightChansImages"]), axis=1)[:, None, None, None] # if self._peakMode is "sigma": # print>> log, "Will search for the peak in the SNR-weighted dirty map" # a, b = self._MeanDirty, self._NoiseMap.reshape(self._MeanDirty.shape) # self._PeakSearchImage = numexpr.evaluate("a/b") # else: # print>> log, "Will search for the peak in the unweighted dirty map" # self._PeakSearchImage = self._MeanDirty # # if self.ModelImage is None: # self._ModelImage = np.zeros_like(self._Dirty) # if self.MaskArray is None: # self._MaskArray = np.zeros(self._Dirty.shape, dtype=np.bool8) def SubStep(self, (dx, dy), LocalSM): """ This is where subtraction in the image domain happens """ xc, yc = dx, dy N1 = LocalSM.shape[-1] # Get overlap indices where psf should be subtracted Aedge, Bedge = GiveEdges((xc, yc), self.Npix, (N1 // 2, N1 // 2), N1) x0d, x1d, y0d, y1d = Aedge x0p, x1p, y0p, y1p = Bedge self._Dirty[:, :, x0d:x1d, y0d:y1d] -= LocalSM[:, :, x0p:x1p, y0p:y1p] # Subtract from the average if self.MultiFreqMode: # If multiple frequencies are present construct the weighted mean self._MeanDirty[:, 0, x0d:x1d, y0d:y1d] -= np.sum( LocalSM[:, :, x0p:x1p, y0p:y1p] * self.WeightsChansImages, axis=0) # Sum over freq else: self._MeanDirty = self._Dirty
class ClassModelMachine(ClassModelMachinebase.ClassModelMachine): def __init__(self, *args, **kwargs): ClassModelMachinebase.ClassModelMachine.__init__(self, *args, **kwargs) self.DicoSMStacked = {} self.DicoSMStacked["Type"] = "WSCMS" def setRefFreq(self, RefFreq, Force=False): if self.RefFreq is not None and not Force: print >> log, ModColor.Str( "Reference frequency already set to %f MHz" % (self.RefFreq / 1e6)) return self.RefFreq = RefFreq self.DicoSMStacked["RefFreq"] = RefFreq def setPSFServer(self, PSFServer): self.PSFServer = PSFServer _, _, self.Npix, _ = self.PSFServer.ImageShape self.NpixPadded = int(np.ceil(self.GD["Facets"]["Padding"] * self.Npix)) # make sure it is odd numbered if self.NpixPadded % 2 == 0: self.NpixPadded += 1 self.Npad = (self.NpixPadded - self.Npix) // 2 def setFreqMachine(self, GridFreqs, DegridFreqs, weights=None, PSFServer=None): self.PSFServer = PSFServer # Initiaise the Frequency Machine self.DegridFreqs = DegridFreqs self.GridFreqs = GridFreqs self.FreqMachine = ClassFrequencyMachine.ClassFrequencyMachine( GridFreqs, DegridFreqs, self.DicoSMStacked["RefFreq"], self.GD, weights=weights, PSFServer=self.PSFServer) self.FreqMachine.set_Method() if (self.GD["Freq"]["NBand"] > 1): self.Coeffs = np.zeros(self.GD["WSCMS"]["NumFreqBasisFuncs"]) else: self.Coeffs = np.zeros([1]) self.Nchan = self.FreqMachine.nchan self.Npol = 1 # self.DicoSMStacked["Eval_Degrid"] = self.FreqMachine.Eval_Degrid def setScaleMachine(self, PSFServer, NCPU=None, MaskArray=None, cachepath=None, MaxBaseline=None): # if self.GD["WSCMS"]["MultiScale"]: if NCPU is None: self.NCPU = self.GD['Parallel'][NCPU] if self.NCPU == 0: import multiprocessing self.NCPU = multiprocessing.cpu_count() else: self.NCPU = NCPU self.DoAbs = self.GD["Deconv"]["AllowNegative"] self.ScaleMachine = ClassScaleMachine.ClassScaleMachine( GD=self.GD, NCPU=NCPU, MaskArray=MaskArray) self.ScaleMachine.Init(PSFServer, self.FreqMachine, cachepath=cachepath, MaxBaseline=MaxBaseline) self.NpixPSF = self.ScaleMachine.NpixPSF self.halfNpixPSF = self.NpixPSF // 2 self.Nscales = self.ScaleMachine.Nscales # Initialise CurrentScale variable self.CurrentScale = 999999 # Initialise current facet variable self.CurrentFacet = 999999 self.DicoSMStacked["Scale_Info"] = {} for iScale, sigma in enumerate(self.ScaleMachine.sigmas): if iScale not in self.DicoSMStacked["Scale_Info"].keys(): self.DicoSMStacked["Scale_Info"][iScale] = {} self.DicoSMStacked["Scale_Info"][iScale][ "sigma"] = self.ScaleMachine.sigmas[iScale] self.DicoSMStacked["Scale_Info"][iScale][ "kernel"] = self.ScaleMachine.kernels[iScale] self.DicoSMStacked["Scale_Info"][iScale][ "extent"] = self.ScaleMachine.extents[iScale] # else: # # we need to keep track of what the sigma value of the delta scale corresponds to # # even if we don't do multiscale (because we need it in GiveModelImage) # (self.FWHMBeamAvg, _, _) = PSFServer.DicoVariablePSF["EstimatesAvgPSF"] # self.ListScales = [1.0/np.sqrt(2)*((self.FWHMBeamAvg[0] + self.FWHMBeamAvg[1])*np.pi / 180) / \ # (2.0 * self.GD['Image']['Cell'] * np.pi / 648000)] def ToFile(self, FileName, DicoIn=None): print >> log, "Saving dico model to %s" % FileName if DicoIn is None: D = self.DicoSMStacked else: D = DicoIn D["GD"] = self.GD D["Type"] = "WSCMS" # try: # D["ListScales"] = list(self.ScaleMachine.sigmas) # list containing std of Gaussian components # except: # D["ListScales"] = self.ListScales D["ModelShape"] = self.ModelShape MyPickle.Save(D, FileName) def FromFile(self, FileName): print >> log, "Reading dico model from %s" % FileName self.DicoSMStacked = MyPickle.Load(FileName) self.FromDico(self.DicoSMStacked) def FromDico(self, DicoSMStacked): self.DicoSMStacked = DicoSMStacked self.RefFreq = self.DicoSMStacked["RefFreq"] # self.ListScales = self.DicoSMStacked["ListScales"] self.ModelShape = self.DicoSMStacked["ModelShape"] def setModelShape(self, ModelShape): self.ModelShape = ModelShape self.Npix = self.ModelShape[-1] def AppendComponentToDictStacked(self, key, Sols, iScale, Gain): """ Adds component to model dictionary at a scale specified by Scale. The dictionary corresponding to each scale is keyed on pixel values (l,m location tupple). Each model component is therefore represented parametrically by a pixel value a scale and a set of coefficients describing the spectral axis. Currently only Stokes I is supported. Args: key: the (l,m) centre of the component in pixels Sols: Nd array of coeffs with length equal to the number of basis functions representing the component. iScale: the scale index Gain: clean loop gain Added component list to dictionary for particular scale. This dictionary is stored in self.DicoSMStacked["Comp"][iScale] and has keys: "SolsArray": solutions ndArray with shape [#basis_functions,#stokes_terms] "NumComps": scalar keeps tracl of the number of components found at a particular scale """ DicoComp = self.DicoSMStacked.setdefault("Comp", {}) if iScale not in DicoComp.keys(): DicoComp[iScale] = {} DicoComp[iScale]["NumComps"] = np.zeros( 1, np.int16) # keeps track of number of components at this scale if key not in DicoComp[iScale].keys(): DicoComp[iScale][key] = {} DicoComp[iScale][key]["SolsArray"] = np.zeros( Sols.size, np.float32) DicoComp[iScale]["NumComps"] += 1 DicoComp[iScale][key]["SolsArray"] += Sols.ravel() * Gain def GiveModelImage(self, FreqIn=None, out=None): RefFreq = self.DicoSMStacked["RefFreq"] # Default to reference frequency if no input given if FreqIn is None: FreqIn = np.array([RefFreq], dtype=np.float32) FreqIn = np.array([FreqIn.ravel()], dtype=np.float32).flatten() DicoComp = self.DicoSMStacked.setdefault("Comp", {}) _, npol, nx, ny = self.ModelShape # The model shape has nchan = len(GridFreqs) nchan = FreqIn.size if out is not None: # LB - is this for appending components to an existing model? if out.shape != (nchan, npol, nx, ny) or out.dtype != np.float32: raise RuntimeError( "supplied image has incorrect type (%s) or shape (%s)" % (out.dtype, out.shape)) ModelImage = out else: ModelImage = np.zeros((nchan, npol, nx, ny), dtype=np.float32) for iScale in DicoComp.keys(): ScaleModel = np.zeros((nchan, npol, nx, ny), dtype=np.float32) # get scale kernel if self.GD["WSCMS"]["MultiScale"]: sigma = self.DicoSMStacked["Scale_Info"][iScale]["sigma"] kernel = self.DicoSMStacked["Scale_Info"][iScale]["kernel"] extent = self.DicoSMStacked["Scale_Info"][iScale]["extent"] for key in DicoComp[iScale].keys(): if key != "NumComps": # LB - dirty dirty hack needs to die!!! Sol = DicoComp[iScale][key]["SolsArray"] # TODO - try soft thresholding components x, y = key try: # LB - Should we drop support for anything other than polynomials maybe? interp = self.FreqMachine.Eval_Degrid(Sol, FreqIn) except: interp = np.polyval(Sol[::-1], FreqIn / RefFreq) if interp is None: raise RuntimeError( "Could not interpolate model onto degridding bands. Inspect your data, check " "'WSCMS-NumFreqBasisFuncs' or if you think this is a bug report it." ) if self.GD["WSCMS"]["MultiScale"] and iScale != 0: Aedge, Bedge = GiveEdges( (x, y), nx, (extent // 2, extent // 2), extent) x0d, x1d, y0d, y1d = Aedge x0p, x1p, y0p, y1p = Bedge out = np.atleast_1d(interp)[:, None, None, None] * kernel ScaleModel[:, :, x0d:x1d, y0d:y1d] += out[:, :, x0p:x1p, y0p:y1p] else: ScaleModel[:, 0, x, y] += interp ModelImage += ScaleModel # print "Model - ", ModelImage.max(), ModelImage.min() return ModelImage def GiveSpectralIndexMap(self, GaussPars=[(1, 1, 0)], ResidCube=None, GiveComponents=False, ChannelWeights=None): # convert to radians ex, ey, pa = GaussPars ex *= np.pi / 180 ey *= np.pi / 180 pa *= np.pi / 180 # get in terms of number of cells CellSizeRad = self.GD['Image']['Cell'] * np.pi / 648000 # ex /= self.GD['Image']['Cell'] * np.pi / 648000 # ey /= self.GD['Image']['Cell'] * np.pi / 648000 # get Gaussian kernel GaussKern = ModFFTW.GiveGauss(self.Npix, CellSizeRad=CellSizeRad, GaussPars=(ex, ey, pa), parallel=False) # take FT Fs = np.fft.fftshift iFs = np.fft.ifftshift npad = self.Npad FTarray = self.ScaleMachine.FTMachine.xhatim.view() FTarray[...] = iFs(np.pad(GaussKern[None, None], ((0, 0), (0, 0), (npad, npad), (npad, npad)), mode='constant'), axes=(2, 3)) # this puts the FT in FTarray self.ScaleMachine.FTMachine.FFTim() # need to copy since FTarray and FTcube are views to the same array FTkernel = FTarray.copy() # evaluate model ModelImage = self.GiveModelImage(self.GridFreqs) # pad and take FT FTcube = self.ScaleMachine.FTMachine.Chatim.view() FTcube[...] = iFs(np.pad(ModelImage, ((0, 0), (0, 0), (npad, npad), (npad, npad)), mode='constant'), axes=(2, 3)) self.ScaleMachine.FTMachine.CFFTim() # multiply by kernel FTcube *= FTkernel # take iFT self.ScaleMachine.FTMachine.iCFFTim() I = slice(npad, -npad) ConvModelImage = Fs(FTcube, axes=(2, 3))[:, :, I, I].real ConvModelMean = np.mean(ConvModelImage.squeeze(), axis=0) ConvModelLow = ConvModelImage[-1, 0] ConvModelHigh = ConvModelImage[0, 0] if ResidCube is not None: ConvModelImage += ResidCube ConvModelImage = ConvModelImage.squeeze() RMS = np.std(ResidCube.flatten()) Threshold = self.GD["SPIMaps"]["AlphaThreshold"] * RMS # get minimum along any freq axis MinImage = np.amin(ConvModelImage, axis=0) MaskIndices = np.argwhere(MinImage > Threshold) FitCube = ConvModelImage[:, MaskIndices[:, 0], MaskIndices[:, 1]] # Initial guess for I0 I0i = ConvModelMean[MaskIndices[:, 0], MaskIndices[:, 1]] # initial guess for alphas Ilow = ConvModelLow[MaskIndices[:, 0], MaskIndices[:, 1]] / I0i Ihigh = ConvModelHigh[MaskIndices[:, 0], MaskIndices[:, 1]] / I0i alphai = (np.log(Ihigh) - np.log(Ilow)) / (np.log(self.GridFreqs[0] / self.RefFreq) - np.log(self.GridFreqs[-1] / self.RefFreq)) # import matplotlib.pyplot as plt # # for i in xrange(self.Nchan): # plt.imshow(np.where(ConvModelImage[i] > Threshold, ConvModelImage[i], 0.0)) # plt.show() if ChannelWeights is None: weights = np.ones(self.Nchan, dtype=np.float32) else: weights = ChannelWeights.astype(np.float32) if ChannelWeights.size != self.Nchan: import warnings warnings.warn( "The provided channel weights are of incorrect length. Ignoring weights.", RuntimeWarning) weights = np.ones(self.Nchan, dtype=np.float32) try: import traceback from africanus.model.spi.dask import fit_spi_components NCPU = self.GD["Parallel"]["NCPU"] if NCPU: from multiprocessing.pool import ThreadPool import dask dask.config.set(pool=ThreadPool(NCPU)) else: import multiprocessing NCPU = multiprocessing.cpu_count() import dask.array as da _, ncomps = FitCube.shape FitCubeDask = da.from_array(FitCube.T.astype(np.float64), chunks=(ncomps // NCPU, self.Nchan)) weightsDask = da.from_array(weights.astype(np.float64), chunks=(self.Nchan)) freqsDask = da.from_array(self.GridFreqs.astype(np.float64), chunks=(self.Nchan)) alpha, varalpha, Iref, varIref = fit_spi_components( FitCubeDask, weightsDask, freqsDask, self.RefFreq, dtype=np.float64, I0i=I0i, alphai=alphai).compute() # from africanus.model.spi import fit_spi_components # # alpha, varalpha, Iref, varIref = fit_spi_components(FitCube.T.astype(np.float64), weights.astype(np.float64), # self.GridFreqs.astype(np.float64), self.RefFreq.astype(np.float64), # dtype=np.float64, I0i=I0i, alphai=alphai) except Exception as e: traceback_str = traceback.format_exc(e) print>>log, "Warning - Failed at importing africanus spi fitter. This could be an issue with the dask " \ "version. Falling back to (slow) scipy version" print >> log, "Original traceback - ", traceback_str alpha, varalpha, Iref, varIref = self.FreqMachine.FitSPIComponents( FitCube, self.GridFreqs, self.RefFreq) _, _, nx, ny = ModelImage.shape alphamap = np.zeros([nx, ny]) Irefmap = np.zeros([nx, ny]) alphastdmap = np.zeros([nx, ny]) Irefstdmap = np.zeros([nx, ny]) alphamap[MaskIndices[:, 0], MaskIndices[:, 1]] = alpha Irefmap[MaskIndices[:, 0], MaskIndices[:, 1]] = Iref alphastdmap[MaskIndices[:, 0], MaskIndices[:, 1]] = np.sqrt(varalpha) Irefstdmap[MaskIndices[:, 0], MaskIndices[:, 1]] = np.sqrt(varIref) if GiveComponents: return alphamap[None, None], alphastdmap[None, None], alpha else: return alphamap[None, None], alphastdmap[None, None] def SubStep(self, (xc, yc), LocalSM, Residual): """ Sub-minor loop subtraction """ N0 = Residual.shape[-1] N1 = LocalSM.shape[-1] # Get overlap indices where psf should be subtracted Aedge, Bedge = GiveEdges((xc, yc), N0, (N1 // 2, N1 // 2), N1) x0d, x1d, y0d, y1d = Aedge x0p, x1p, y0p, y1p = Bedge # Subtract from each channel/band Residual[:, :, x0d:x1d, y0d:y1d] -= LocalSM[:, :, x0p:x1p, y0p:y1p]
def GiveModelImage(self, FreqIn=None, out=None): RefFreq = self.DicoSMStacked["RefFreq"] # Default to reference frequency if no input given if FreqIn is None: FreqIn = np.array([RefFreq], dtype=np.float32) FreqIn = np.array([FreqIn.ravel()], dtype=np.float32).flatten() DicoComp = self.DicoSMStacked.setdefault("Comp", {}) _, npol, nx, ny = self.ModelShape # The model shape has nchan = len(GridFreqs) nchan = FreqIn.size if out is not None: # LB - is this for appending components to an existing model? if out.shape != (nchan, npol, nx, ny) or out.dtype != np.float32: raise RuntimeError( "supplied image has incorrect type (%s) or shape (%s)" % (out.dtype, out.shape)) ModelImage = out else: ModelImage = np.zeros((nchan, npol, nx, ny), dtype=np.float32) for iScale in DicoComp.keys(): ScaleModel = np.zeros((nchan, npol, nx, ny), dtype=np.float32) # get scale kernel if self.GD["WSCMS"]["MultiScale"]: sigma = self.DicoSMStacked["Scale_Info"][iScale]["sigma"] kernel = self.DicoSMStacked["Scale_Info"][iScale]["kernel"] extent = self.DicoSMStacked["Scale_Info"][iScale]["extent"] for key in DicoComp[iScale].keys(): if key != "NumComps": # LB - dirty dirty hack needs to die!!! Sol = DicoComp[iScale][key]["SolsArray"] # TODO - try soft thresholding components x, y = key try: # LB - Should we drop support for anything other than polynomials maybe? interp = self.FreqMachine.Eval_Degrid(Sol, FreqIn) except: interp = np.polyval(Sol[::-1], FreqIn / RefFreq) if interp is None: raise RuntimeError( "Could not interpolate model onto degridding bands. Inspect your data, check " "'WSCMS-NumFreqBasisFuncs' or if you think this is a bug report it." ) if self.GD["WSCMS"]["MultiScale"] and iScale != 0: Aedge, Bedge = GiveEdges( (x, y), nx, (extent // 2, extent // 2), extent) x0d, x1d, y0d, y1d = Aedge x0p, x1p, y0p, y1p = Bedge out = np.atleast_1d(interp)[:, None, None, None] * kernel ScaleModel[:, :, x0d:x1d, y0d:y1d] += out[:, :, x0p:x1p, y0p:y1p] else: ScaleModel[:, 0, x, y] += interp ModelImage += ScaleModel # print "Model - ", ModelImage.max(), ModelImage.min() return ModelImage
def GiveModelTessel(self, Image, DicoImager, iFacet, NormIm, Sphe, SpacialWeight, ToGrid=False, ChanSel=None, ApplyNorm=True): nch, npol, NPixOut, _ = Image.shape N1 = DicoImager[iFacet]["NpixFacetPadded"] N1NonPadded = DicoImager[iFacet]["NpixFacetPadded"] dx = (N1 - N1NonPadded) // 2 xc, yc = DicoImager[iFacet]["pixCentral"] #x0,x1,y0,y1=DicoImager[iFacet]["pixExtent"] #xc,yc=(x0+x1)//2,(y0+y1)//2 Aedge, Bedge = GiveEdges(xc, yc, NPixOut, N1 // 2, N1 // 2, N1) #Bedge,Aedge=GiveEdges(N1//2,N1//2,N1,yc,xc,NPixOut) x0d, x1d, y0d, y1d = Aedge x0p, x1p, y0p, y1p = Bedge #print "xxA:",x0d,x1d #print "xxB:",x0p,x1p SumFlux = 1. ModelIm = np.zeros((nch, npol, N1, N1), dtype=np.float32) T = ClassTimeIt.ClassTimeIt("ClassImToGrid") T.disable() if ChanSel is None: CSel = range(nch) else: CSel = ChanSel SumFlux = 0 for ch in CSel: for pol in range(npol): #ModelIm[ch,pol][x0p:x1p,y0p:y1p]=Image[ch,pol].T[::-1,:].real[x0d:x1d,y0d:y1d] #ModelIm[ch,pol][x0p:x1p,y0p:y1p]=Image[ch,pol].real[x0d:x1d,y0d:y1d] ModelIm[ch, pol][x0p:x1p, y0p:y1p] = Image[ch, pol][x0d:x1d, y0d:y1d].real if (ModelIm[ch, pol] == 0).all(): continue T.timeit("0") M = ModelIm[ch, pol][dx:dx + N1NonPadded + 1, dx:dx + N1NonPadded + 1].copy() T.timeit("1") ModelIm[ch, pol].fill(0) T.timeit("2") ModelIm[ch, pol][dx:dx + N1NonPadded + 1, dx:dx + N1NonPadded + 1] = M[:, :] #ModelCutOrig=ModelIm[ch,pol].copy() T.timeit("3") #ind =np.where(np.abs(ModelIm)==np.max(np.abs(ModelIm))) ##print "!!!!!!!!!!!!!!!!!!!!!!" if ApplyNorm: # #print NormIm.max() # #print np.count_nonzero(np.isnan(NormIm)) # #print np.count_nonzero(np.isinf(NormIm)) # print NormIm.min() # np.save("NormIm",NormIm) # stop ModelIm[ch, pol][x0p:x1p, y0p:y1p] /= NormIm[x0d:x1d, y0d:y1d].real #ModelCutOrig_GNorm=NormIm[x0d:x1d,y0d:y1d].real.copy() T.timeit("4") if ApplyNorm: ModelIm[ch, pol][x0p:x1p, y0p:y1p] *= SpacialWeight[x0p:x1p, y0p:y1p] indPos = np.where(ModelIm[ch, pol] > 0) SumFlux += np.sum(ModelIm[ch, pol][indPos]) ModelCutOrig_SW = SpacialWeight[x0p:x1p, y0p:y1p].copy() #ModelCutOrig_GNorm_SW_Sphe_CorrT=ModelIm[ch,pol].copy() T.timeit("5") #SumFlux=np.sum(ModelIm) if ApplyNorm: ModelIm[ch, pol][x0p:x1p, y0p:y1p] /= Sphe[x0p:x1p, y0p:y1p].real # ModelIm[ch, pol][x0p:x1p, y0p:y1p] *= ModelCutOrig_SW / Sphe[x0p:x1p, # y0p:y1p].real # LB - added *SW #ModelCutOrig_Sphe=Sphe[x0p:x1p,y0p:y1p].real.copy() T.timeit("6") ModelIm[ch, pol][Sphe < 1e-3] = 0 T.timeit("7") ModelIm[ch, pol] = ModelIm[ch, pol].T[::-1, :] T.timeit("8") #ModelCutOrig_GNorm_SW_Sphe_CorrT=ModelIm[ch,pol].copy() #return True, ModelCutOrig, ModelCutOrig_GNorm, ModelCutOrig_SW, ModelCutOrig_Sphe, ModelCutOrig_GNorm_SW_Sphe_CorrT #print iFacet,DicoImager[iFacet]["l0m0"],DicoImager[iFacet]["NpixFacet"],DicoImager[iFacet]["NpixFacetPadded"],SumFlux # if np.max(np.abs(ModelIm))>1: # print ind #if np.abs(SumFlux)>1: stop # #print iFacet,np.max(ModelIm) # #return ModelIm, None # #Padding=self.GD["Image"]["Padding"] T.timeit("9") SumFlux /= nch if ToGrid: ModelIm *= (self.OverS * N1)**2 if SumFlux != 0: Grid = np.complex64( self.FFTWMachine.fft(np.complex64(ModelIm), ChanList=CSel)) else: Grid = np.complex64(ModelIm) return Grid, SumFlux elif ApplyNorm: ModelIm *= (self.OverS * N1)**2 return ModelIm, SumFlux else: return ModelIm, SumFlux