def predictKernelPolClusterCatalog(self,
                                       DicoData,
                                       SM,
                                       iDirection=None,
                                       ApplyJones=None,
                                       ApplyTimeJones=None,
                                       Noise=None):
        self.DicoData = DicoData
        self.SourceCat = SM.SourceCat
        self.SM = SM

        freq = DicoData["freqs_full"]
        times = DicoData["times"]
        nf = freq.size
        na = DicoData["infos"][0]

        nrows = DicoData["A0"].size
        DataOut = np.zeros((nrows, nf, 4), self.CType)
        if nrows == 0: return DataOut

        self.freqs = freq
        self.wave = 299792458. / self.freqs

        if iDirection != None:
            ListDirection = [iDirection]
        else:
            ListDirection = SM.Dirs  #range(SM.NDir)

        A0 = DicoData["A0"]
        A1 = DicoData["A1"]
        if ApplyJones is not None:
            #print "!!!!!",ApplyJones.shape
            #print "!!!!!",ApplyJones.shape
            #print "!!!!!",ApplyJones.shape
            na, NDir, _ = ApplyJones.shape
            Jones = np.swapaxes(ApplyJones, 0, 1)
            Jones = Jones.reshape((NDir, na, 4))
            JonesH = ModLinAlg.BatchH(Jones)

        TSmear = 0.
        FSmear = 0.
        DT = DicoData["infos"][1]
        UVW_dt = DicoData["uvw"]
        if self.DoSmearing:
            if "T" in self.DoSmearing:
                TSmear = 1.
                UVW_dt = DicoData["UVW_dt"]
            if "F" in self.DoSmearing:
                FSmear = 1.

        # self.SourceCat.m[:]=0
        # self.SourceCat.l[:]=0.1
        # self.SourceCat.I[:]=10
        # self.SourceCat.alpha[:]=0

        # DataOut=DataOut[1:2]
        # self.DicoData["uvw"]=self.DicoData["uvw"][1:2]
        # self.DicoData["A0"]=self.DicoData["A0"][1:2]
        # self.DicoData["A1"]=self.DicoData["A1"][1:2]
        # self.DicoData["IndexTimesThisChunk"]=self.DicoData["IndexTimesThisChunk"][1:2]
        # self.SourceCat=self.SourceCat[0:1]

        ColOutDir = np.zeros(DataOut.shape, np.complex64)

        for iCluster in ListDirection:
            ColOutDir.fill(0)

            indSources = np.where(self.SourceCat.Cluster == iCluster)[0]
            T = ClassTimeIt("predict")
            T.disable()
            ### new
            SourceCat = self.SourceCat[indSources].copy()
            #l=np.ones((1,),dtype=np.float64)#,float64(SourceCat.l).copy()
            l = np.require(SourceCat.l,
                           dtype=np.float64,
                           requirements=["A", "C"])
            m = np.require(SourceCat.m,
                           dtype=np.float64,
                           requirements=["A", "C"])

            #m=SourceCat.m#np.float64(SourceCat.m).copy()
            I = np.float32(SourceCat.I)
            Gmaj = np.float32(SourceCat.Gmaj)
            Gmin = np.float32(SourceCat.Gmin)
            GPA = np.float32(SourceCat.Gangle)
            alpha = np.float32(SourceCat.alpha)
            WaveL = np.float64(299792458. / self.freqs)
            WaveL = np.require(WaveL,
                               dtype=np.float64,
                               requirements=["A", "C"])

            flux = np.float32(SourceCat.I)
            alpha = SourceCat.alpha
            dnu = np.float32(self.DicoData["dfreqs_full"])
            f0 = (self.freqs / SourceCat.RefFreq[0])
            fluxFreq = np.float32(
                flux.reshape((flux.size, 1)) * (f0.reshape(
                    (1, f0.size)))**(alpha.reshape((alpha.size, 1))))

            LSM = [l, m, fluxFreq, Gmin, Gmaj, GPA]
            LFreqs = [WaveL, np.float32(self.freqs), dnu]
            LUVWSpeed = [UVW_dt, DT]

            LSmearMode = [FSmear, TSmear]
            T.timeit("init")

            AllowEqualiseChan = IsChanEquidistant(DicoData["freqs_full"])

            if ApplyTimeJones != None:

                #predict.predictJones(ColOutDir,(DicoData["uvw"]),LFreqs,LSM,LUVWSpeed,LSmearMode,ParamJonesList)

                #ColOutDir.fill(0)
                #predict.predictJones(ColOutDir,(DicoData["uvw"]),LFreqs,LSM,LUVWSpeed,LSmearMode,ParamJonesList,0)
                #d0=ColOutDir.copy()
                #ColOutDir.fill(0)

                # predict.predictJones(ColOutDir,(DicoData["uvw"]),LFreqs,LSM,LUVWSpeed,LSmearMode,ParamJonesList,AllowEqualiseChan)
                # ColOutDir0=ColOutDir.copy()
                # ColOutDir.fill(0)

                #predict.predictJones2(ColOutDir,(DicoData["uvw"]),LFreqs,LSM,LUVWSpeed,LSmearMode,ParamJonesList,AllowEqualiseChan)
                #print LSmearMode
                predict.predictJones2_Gauss(ColOutDir, (DicoData["uvw"]),
                                            LFreqs, LSM, LUVWSpeed, LSmearMode,
                                            AllowEqualiseChan, self.LExp,
                                            self.LSinc)

                T.timeit("predict")

                ParamJonesList = self.GiveParamJonesList(
                    ApplyTimeJones, A0, A1)
                ParamJonesList = ParamJonesList + [iCluster]

                predict.ApplyJones(ColOutDir, ParamJonesList)
                T.timeit("apply")

                # print ColOutDir

                #d1=ColOutDir.copy()
                #ind=np.where(d0!=0)
                #print np.max((d0-d1)[ind]/(d0[ind]))
                #stop
            else:
                #predict.predict(ColOutDir,(DicoData["uvw"]),LFreqs,LSM,LUVWSpeed,LSmearMode)
                #AllowEqualiseChan=0
                #predict.predict(ColOutDir,(DicoData["uvw"]),LFreqs,LSM,LUVWSpeed,LSmearMode,AllowEqualiseChan)
                #d0=ColOutDir.copy()
                #ColOutDir.fill(0)

                predict.predictJones2_Gauss(ColOutDir, (DicoData["uvw"]),
                                            LFreqs, LSM, LUVWSpeed, LSmearMode,
                                            AllowEqualiseChan, self.LExp,
                                            self.LSinc)
                #print ColOutDir
                #predict.predict(ColOutDir,(DicoData["uvw"]),LFreqs,LSM,LUVWSpeed,LSmearMode,AllowEqualiseChan)
                T.timeit("predict")
                # d0=ColOutDir.copy()
                # ColOutDir.fill(0)

                # predict_np19.predict(ColOutDir,(DicoData["uvw"]),LFreqs,LSM,LUVWSpeed,LSmearMode,AllowEqualiseChan)
                # print ColOutDir,d0
                # d1=ColOutDir.copy()
                # ind=np.where(d0!=0)
                # print np.max((d0-d1)[ind]/(d0[ind]))
                # stop

            del (l, m, I, SourceCat, alpha, WaveL, flux, dnu, f0, fluxFreq,
                 LSM, LFreqs)
            T.timeit("del")

            # d1=ColOutDir
            # ind=np.where(d0!=0)
            # print np.max((d0-d1)[ind]/(d0[ind]))
            # stop

            if Noise != None:
                ColOutDir += Noise * (np.random.randn(*ColOutDir.shape) +
                                      1j * np.random.randn(*ColOutDir.shape))
                stop
            DataOut += ColOutDir
            T.timeit("add")
        #del(LFreqs,LSM,LUVWSpeed,LSmearMode)
        #del(ColOutDir)

        return DataOut
    def predictKernelPolCluster(self,
                                DicoData,
                                SM,
                                iDirection=None,
                                ApplyJones=None,
                                ApplyTimeJones=None,
                                Noise=None,
                                VariableFunc=None):
        T = ClassTimeIt("predictKernelPolCluster")
        T.disable()
        self.DicoData = DicoData
        self.SourceCat = SM.SourceCat

        freq = DicoData["freqs"]
        times = DicoData["times"]
        nf = freq.size
        na = int(DicoData["infos"][0])

        nrows = DicoData["A0"].size
        DataOut = np.zeros((nrows, nf, 4), self.CType)
        if nrows == 0: return DataOut

        self.freqs = freq
        self.wave = 299792458. / self.freqs

        if iDirection != None:
            ListDirection = [iDirection]
        else:
            ListDirection = SM.Dirs  #range(SM.NDir)
        T.timeit("0")
        A0 = DicoData["A0"]
        A1 = DicoData["A1"]
        if ApplyJones != None:
            na, NDir, _ = ApplyJones.shape
            Jones = np.swapaxes(ApplyJones, 0, 1)
            Jones = Jones.reshape((NDir, na, 4))
            JonesH = ModLinAlg.BatchH(Jones)
        T.timeit("1")

        for iCluster in ListDirection:
            print("IIIIIIIIIIIIIIIIII", iCluster)
            ColOutDir = self.PredictDirSPW(iCluster)
            T.timeit("2")
            if ColOutDir is None: continue

            # print(iCluster,ListDirection)
            # print(ColOutDir.shape)
            # ColOutDir.fill(0)
            # print(ColOutDir.shape)
            # ColOutDir[:,:,0]=1
            # print(ColOutDir.shape)
            # ColOutDir[:,:,3]=1
            # print(ColOutDir.shape)

            # Apply Jones
            if ApplyJones != None:

                J = Jones[iCluster]
                JH = JonesH[iCluster]
                for ichan in range(nf):
                    ColOutDir[:, ichan, :] = ModLinAlg.BatchDot(
                        J[A0, :], ColOutDir[:, ichan, :])
                    ColOutDir[:, ichan, :] = ModLinAlg.BatchDot(
                        ColOutDir[:, ichan, :], JH[A1, :])
            T.timeit("3")

            if VariableFunc is not None:  #"DicoBeam" in DicoData.keys():
                tt = np.unique(times)
                lt0, lt1 = tt[0:-1], tt[1::]
                for it in range(lt0.size):
                    t0, t1 = lt0[it], lt1[it]
                    ind = np.where((times >= t0) & (times < t1))[0]
                    if ind.size == 0: continue
                    data = ColOutDir[ind]

                    if "ChanMap" in ApplyTimeJones.keys():
                        ChanMap = ApplyTimeJones["ChanMap"]
                    else:
                        ChanMap = range(nf)

                    for ichan in range(len(ChanMap)):
                        tc = (t0 + t1) / 2.
                        nuc = freq[ichan]
                        ColOutDir[ind, ichan, :] *= VariableFunc(tc, nuc)
                        # c0=ColOutDir[ind,ichan,:].copy()
                        # ColOutDir[ind,ichan,:]*=VariableFunc(tc,nuc)
                        # print(c0-ColOutDir[ind,ichan,:])
                        #print(it,ichan,VariableFunc(tc,nuc))

            if ApplyTimeJones is not None:  #"DicoBeam" in DicoData.keys():
                D = ApplyTimeJones  #DicoData["DicoBeam"]
                Beam = D["Beam"]
                BeamH = D["BeamH"]

                lt0, lt1 = D["t0"], D["t1"]

                for it in range(lt0.size):
                    t0, t1 = lt0[it], lt1[it]
                    ind = np.where((times >= t0) & (times < t1))[0]
                    if ind.size == 0: continue
                    data = ColOutDir[ind]
                    A0sel = A0[ind]
                    A1sel = A1[ind]

                    if "ChanMap" in ApplyTimeJones.keys():
                        ChanMap = ApplyTimeJones["ChanMap"]
                    else:
                        ChanMap = range(nf)

                    #print("ChanMap:",ChanMap)

                    for ichan in range(len(ChanMap)):
                        JChan = ChanMap[ichan]
                        J = Beam[it, iCluster, :, JChan, :, :].reshape((na, 4))
                        JH = BeamH[it, iCluster, :, JChan, :, :].reshape(
                            (na, 4))
                        data[:, ichan, :] = ModLinAlg.BatchDot(
                            J[A0sel, :], data[:, ichan, :])
                        data[:, ichan, :] = ModLinAlg.BatchDot(
                            data[:, ichan, :], JH[A1sel, :])

                    ColOutDir[ind] = data[:]
            T.timeit("4")

            DataOut += ColOutDir
            T.timeit("5")

        if Noise is not None:
            DataOut += Noise / np.sqrt(2.) * (np.random.randn(
                *ColOutDir.shape) + 1j * np.random.randn(*ColOutDir.shape))

        return DataOut
Esempio n. 3
0
    def GiveJones(self):
        if self.Sols is None:
            Sols = self.GiveSols()
        else:
            Sols = self.Sols

        MS = self.MS
        SM = self.SM
        VS = self.VS
        ApplyBeam = self.ApplyBeam
        na = MS.na
        nd = SM.NDir

        if self.BeamAt == "facet":
            NDir = SM.SourceCat.shape[0]
            SM.SourceCat.Cluster = np.arange(NDir)
            SM.Dirs = SM.SourceCat.Cluster
            SM.NDir = NDir
            SM.ClusterCat = np.zeros((NDir, ), SM.ClusterCat.dtype)
            SM.ClusterCat = SM.ClusterCat.view(np.recarray)
            SM.ClusterCat.ra = SM.SourceCat.ra
            SM.ClusterCat.dec = SM.SourceCat.dec

        Jones = {}
        Jones["t0"] = Sols.t0
        Jones["t1"] = Sols.t1

        nt, nch, na, nd, _, _ = Sols.G.shape
        G = np.swapaxes(Sols.G, 1, 3).reshape((nt, nd, na, nch, 2, 2))

        # G[:,:,:,:,0,0]/=np.abs(G[:,:,:,:,0,0])
        # G[:,:,:,:,1,1]=G[:,:,:,:,0,0]
        # G.fill(0)
        # G[:,:,:,:,0,0]=1
        # G[:,:,:,:,1,1]=1

        nt, nd, na, nch, _, _ = G.shape

        # G=np.random.randn(*G.shape)+1j*np.random.randn(*G.shape)

        useArrayFactor = True
        useElementBeam = False
        if ApplyBeam:
            print(ModColor.Str("Apply Beam"))
            MS.LoadSR(useElementBeam=useElementBeam,
                      useArrayFactor=useArrayFactor)
            RA = SM.ClusterCat.ra
            DEC = SM.ClusterCat.dec

            NDir = RA.size
            Tm = Sols.tm
            T0s = Sols.t0
            T1s = Sols.t1
            DicoBeam = {}
            DicoBeam["Jones"] = np.zeros(
                (Tm.size, NDir, MS.na, MS.NSPWChan, 2, 2), dtype=np.complex64)
            DicoBeam["t0"] = np.zeros((Tm.size, ), np.float64)
            DicoBeam["t1"] = np.zeros((Tm.size, ), np.float64)
            DicoBeam["tm"] = np.zeros((Tm.size, ), np.float64)

            rac, decc = MS.OriginalRadec

            def GB(time, ra, dec):
                Beam = np.zeros(
                    (ra.shape[0], self.MS.na, self.MS.NSPWChan, 2, 2),
                    dtype=np.complex)
                # Beam[...,0,0]=1
                # Beam[...,1,1]=1
                # return Beam

                for i in range(ra.shape[0]):
                    self.MS.SR.setDirection(ra[i], dec[i])
                    Beam[i] = self.MS.SR.evaluate(time)
                return Beam

            for itime in range(Tm.size):
                print(itime)
                DicoBeam["t0"][itime] = T0s[itime]
                DicoBeam["t1"][itime] = T1s[itime]
                DicoBeam["tm"][itime] = Tm[itime]
                ThisTime = Tm[itime]
                Beam = GB(ThisTime, RA, DEC)

                ###### Normalise
                Beam0 = GB(ThisTime, np.array([rac]), np.array([decc]))
                Beam0inv = ModLinAlg.BatchInverse(Beam0)
                nd, _, _, _, _ = Beam.shape
                Ones = np.ones((nd, 1, 1, 1, 1), np.float32)
                Beam0inv = Beam0inv * Ones
                nd_, na_, nf_, _, _ = Beam.shape
                # Beam_=np.ones((nd_,na_,nf_),np.float32)*(1+np.arange(nd_).reshape((-1,1,1)))
                # Beam.fill(0)
                # Beam[:,:,:,0,0]=Beam_[:,:,:]
                # Beam[:,:,:,1,1]=Beam_[:,:,:]
                Beam = ModLinAlg.BatchDot(Beam0inv, Beam)
                ######

                DicoBeam["Jones"][itime] = Beam

            nt, nd, na, nch, _, _ = DicoBeam["Jones"].shape

            #m=np.mean(np.abs(DicoBeam["Jones"][:,1,:,:,:,:]))
            # m=np.mean(np.abs(DicoBeam["Jones"][:,1,:,:,:,:]))
            # DicoBeam["Jones"][:,1,0:6,:,:,:]*=2
            # DicoBeam["Jones"][:,1,:,:,:,:]/=np.mean(np.abs(DicoBeam["Jones"][:,1,:,:,:,:]))
            # DicoBeam["Jones"][:,1,:,:,:,:]*=m

            # #################"
            # # Single Channel
            # DicoBeam["Jones"]=np.mean(DicoBeam["Jones"],axis=3).reshape((nt,nd,na,1,2,2))

            # G=ModLinAlg.BatchDot(G,DicoBeam["Jones"])

            # #################"
            # Multiple Channel
            Ones = np.ones((1, 1, 1, nch, 1, 1), np.float32)
            G = G * Ones
            G = ModLinAlg.BatchDot(G, DicoBeam["Jones"])

            # #################"

            # G[:,:,:,:,0,0]=1
            # G[:,:,:,:,0,1]=0.5
            # G[:,:,:,:,1,0]=2.
            # G[:,:,:,:,1,1]=1

            print("Done")

        # #################
        # Multiple Channel
        self.ChanMap = range(nch)
        # #################

        Jones["Beam"] = G
        Jones["BeamH"] = ModLinAlg.BatchH(G)
        if self.ChanMap is None:
            self.ChanMap = np.zeros((VS.MS.NSPWChan, ), np.int32).tolist()

        Jones["ChanMap"] = self.ChanMap

        # ###### for PM5
        # Jones["Map_VisToJones_Freq"]=self.ChanMap
        # Jones["Jones"]=Jones["Beam"]
        # nt=VS.MS.times_all.size
        # ntJones=DicoBeam["tm"].size
        # d=VS.MS.times_all.reshape((nt,1))-DicoBeam["tm"].reshape((1,ntJones))
        # Jones["Map_VisToJones_Time"]=np.argmin(np.abs(d),axis=1)

        return Jones